feat: use end user credentials for bigframes.bigquery.ai functions when connection_id is not present (#2272)

Fixes #460856043 🦕

Author: sycai

bigframes/bigquery/_operations/ai.py

@@ -88,7 +88,7 @@ def generate(
             or pandas Series.
         connection_id (str, optional):
             Specifies the connection to use to communicate with the model. For example, `myproject.us.myconnection`.
-            If not provided, the connection from the current session will be used.
+            If not provided, the query uses your end-user credential.
         endpoint (str, optional):
             Specifies the Vertex AI endpoint to use for the model. For example `"gemini-2.5-flash"`. You can specify any
             generally available or preview Gemini model. If you specify the model name, BigQuery ML automatically identifies and
@@ -131,7 +131,7 @@ def generate(
 
     operator = ai_ops.AIGenerate(
         prompt_context=tuple(prompt_context),
-        connection_id=_resolve_connection_id(series_list[0], connection_id),
+        connection_id=connection_id,
         endpoint=endpoint,
         request_type=request_type,
         model_params=json.dumps(model_params) if model_params else None,
@@ -186,7 +186,7 @@ def generate_bool(
             or pandas Series.
         connection_id (str, optional):
             Specifies the connection to use to communicate with the model. For example, `myproject.us.myconnection`.
-            If not provided, the connection from the current session will be used.
+            If not provided, the query uses your end-user credential.
         endpoint (str, optional):
             Specifies the Vertex AI endpoint to use for the model. For example `"gemini-2.5-flash"`. You can specify any
             generally available or preview Gemini model. If you specify the model name, BigQuery ML automatically identifies and
@@ -216,7 +216,7 @@ def generate_bool(
 
     operator = ai_ops.AIGenerateBool(
         prompt_context=tuple(prompt_context),
-        connection_id=_resolve_connection_id(series_list[0], connection_id),
+        connection_id=connection_id,
         endpoint=endpoint,
         request_type=request_type,
         model_params=json.dumps(model_params) if model_params else None,
@@ -267,7 +267,7 @@ def generate_int(
             or pandas Series.
         connection_id (str, optional):
             Specifies the connection to use to communicate with the model. For example, `myproject.us.myconnection`.
-            If not provided, the connection from the current session will be used.
+            If not provided, the query uses your end-user credential.
         endpoint (str, optional):
             Specifies the Vertex AI endpoint to use for the model. For example `"gemini-2.5-flash"`. You can specify any
             generally available or preview Gemini model. If you specify the model name, BigQuery ML automatically identifies and
@@ -297,7 +297,7 @@ def generate_int(
 
     operator = ai_ops.AIGenerateInt(
         prompt_context=tuple(prompt_context),
-        connection_id=_resolve_connection_id(series_list[0], connection_id),
+        connection_id=connection_id,
         endpoint=endpoint,
         request_type=request_type,
         model_params=json.dumps(model_params) if model_params else None,
@@ -348,7 +348,7 @@ def generate_double(
             or pandas Series.
         connection_id (str, optional):
             Specifies the connection to use to communicate with the model. For example, `myproject.us.myconnection`.
-            If not provided, the connection from the current session will be used.
+            If not provided, the query uses your end-user credential.
         endpoint (str, optional):
             Specifies the Vertex AI endpoint to use for the model. For example `"gemini-2.5-flash"`. You can specify any
             generally available or preview Gemini model. If you specify the model name, BigQuery ML automatically identifies and
@@ -378,7 +378,7 @@ def generate_double(
 
     operator = ai_ops.AIGenerateDouble(
         prompt_context=tuple(prompt_context),
-        connection_id=_resolve_connection_id(series_list[0], connection_id),
+        connection_id=connection_id,
         endpoint=endpoint,
         request_type=request_type,
         model_params=json.dumps(model_params) if model_params else None,

bigframes/core/compile/sqlglot/expressions/ai_ops.py

@@ -104,10 +104,13 @@ def _construct_named_args(op: ops.NaryOp) -> list[sge.Kwarg]:
 
     op_args = asdict(op)
 
-    connection_id = op_args["connection_id"]
-    args.append(
-        sge.Kwarg(this="connection_id", expression=sge.Literal.string(connection_id))
-    )
+    connection_id = op_args.get("connection_id", None)
+    if connection_id is not None:
+        args.append(
+            sge.Kwarg(
+                this="connection_id", expression=sge.Literal.string(connection_id)
+            )
+        )
 
     endpoit = op_args.get("endpoint", None)
     if endpoit is not None:

bigframes/operations/ai_ops.py

@@ -29,7 +29,7 @@ class AIGenerate(base_ops.NaryOp):
     name: ClassVar[str] = "ai_generate"
 
     prompt_context: Tuple[str | None, ...]
-    connection_id: str
+    connection_id: str | None
     endpoint: str | None
     request_type: Literal["dedicated", "shared", "unspecified"]
     model_params: str | None
@@ -57,7 +57,7 @@ class AIGenerateBool(base_ops.NaryOp):
     name: ClassVar[str] = "ai_generate_bool"
 
     prompt_context: Tuple[str | None, ...]
-    connection_id: str
+    connection_id: str | None
     endpoint: str | None
     request_type: Literal["dedicated", "shared", "unspecified"]
     model_params: str | None
@@ -79,7 +79,7 @@ class AIGenerateInt(base_ops.NaryOp):
     name: ClassVar[str] = "ai_generate_int"
 
     prompt_context: Tuple[str | None, ...]
-    connection_id: str
+    connection_id: str | None
     endpoint: str | None
     request_type: Literal["dedicated", "shared", "unspecified"]
     model_params: str | None
@@ -101,7 +101,7 @@ class AIGenerateDouble(base_ops.NaryOp):
     name: ClassVar[str] = "ai_generate_double"
 
     prompt_context: Tuple[str | None, ...]
-    connection_id: str
+    connection_id: str | None
     endpoint: str | None
     request_type: Literal["dedicated", "shared", "unspecified"]
     model_params: str | None

tests/unit/core/compile/sqlglot/expressions/snapshots/test_ai_ops/test_ai_generate/out.sql

@@ -7,7 +7,6 @@ WITH `bfcte_0` AS (
     *,
     AI.GENERATE(
       prompt => (`string_col`, ' is the same as ', `string_col`),
-      connection_id => 'bigframes-dev.us.bigframes-default-connection',
       endpoint => 'gemini-2.5-flash',
       request_type => 'SHARED'
     ) AS `bfcol_1`

tests/unit/core/compile/sqlglot/expressions/snapshots/test_ai_ops/test_ai_generate_bool/out.sql

@@ -7,7 +7,6 @@ WITH `bfcte_0` AS (
     *,
     AI.GENERATE_BOOL(
       prompt => (`string_col`, ' is the same as ', `string_col`),
-      connection_id => 'bigframes-dev.us.bigframes-default-connection',
       endpoint => 'gemini-2.5-flash',
       request_type => 'SHARED'
     ) AS `bfcol_1`

tests/unit/core/compile/sqlglot/expressions/snapshots/test_ai_ops/test_ai_generate_bool_with_connection_id/out.sql

@@ -0,0 +1,18 @@
+WITH `bfcte_0` AS (
+  SELECT
+    `string_col`
+  FROM `bigframes-dev`.`sqlglot_test`.`scalar_types`
+), `bfcte_1` AS (
+  SELECT
+    *,
+    AI.GENERATE_BOOL(
+      prompt => (`string_col`, ' is the same as ', `string_col`),
+      connection_id => 'bigframes-dev.us.bigframes-default-connection',
+      endpoint => 'gemini-2.5-flash',
+      request_type => 'SHARED'
+    ) AS `bfcol_1`
+  FROM `bfcte_0`
+)
+SELECT
+  `bfcol_1` AS `result`
+FROM `bfcte_1`

tests/unit/core/compile/sqlglot/expressions/snapshots/test_ai_ops/test_ai_generate_bool_with_model_param/out.sql

@@ -7,7 +7,6 @@ WITH `bfcte_0` AS (
     *,
     AI.GENERATE_BOOL(
       prompt => (`string_col`, ' is the same as ', `string_col`),
-      connection_id => 'bigframes-dev.us.bigframes-default-connection',
       request_type => 'SHARED',
       model_params => JSON '{}'
     ) AS `bfcol_1`

tests/unit/core/compile/sqlglot/expressions/snapshots/test_ai_ops/test_ai_generate_double/out.sql

@@ -7,7 +7,6 @@ WITH `bfcte_0` AS (
     *,
     AI.GENERATE_DOUBLE(
       prompt => (`string_col`, ' is the same as ', `string_col`),
-      connection_id => 'bigframes-dev.us.bigframes-default-connection',
       endpoint => 'gemini-2.5-flash',
       request_type => 'SHARED'
     ) AS `bfcol_1`

tests/unit/core/compile/sqlglot/expressions/snapshots/test_ai_ops/test_ai_generate_double_with_connection_id/out.sql

@@ -0,0 +1,18 @@
+WITH `bfcte_0` AS (
+  SELECT
+    `string_col`
+  FROM `bigframes-dev`.`sqlglot_test`.`scalar_types`
+), `bfcte_1` AS (
+  SELECT
+    *,
+    AI.GENERATE_DOUBLE(
+      prompt => (`string_col`, ' is the same as ', `string_col`),
+      connection_id => 'bigframes-dev.us.bigframes-default-connection',
+      endpoint => 'gemini-2.5-flash',
+      request_type => 'SHARED'
+    ) AS `bfcol_1`
+  FROM `bfcte_0`
+)
+SELECT
+  `bfcol_1` AS `result`
+FROM `bfcte_1`

tests/unit/core/compile/sqlglot/expressions/snapshots/test_ai_ops/test_ai_generate_double_with_model_param/out.sql

@@ -7,7 +7,6 @@ WITH `bfcte_0` AS (
     *,
     AI.GENERATE_DOUBLE(
       prompt => (`string_col`, ' is the same as ', `string_col`),
-      connection_id => 'bigframes-dev.us.bigframes-default-connection',
       request_type => 'SHARED',
       model_params => JSON '{}'
     ) AS `bfcol_1`