google
diff --git a/‎bigframes/bigquery/_operations/ai.py‎
Lines changed: 6 additions & 6 deletions b/‎bigframes/bigquery/_operations/ai.py‎
Lines changed: 6 additions & 6 deletions
diff --git a/‎bigframes/core/compile/sqlglot/expressions/ai_ops.py‎
Lines changed: 3 additions & 3 deletions b/‎bigframes/core/compile/sqlglot/expressions/ai_ops.py‎
Lines changed: 3 additions & 3 deletions
diff --git a/‎bigframes/operations/ai_ops.py‎
Lines changed: 3 additions & 3 deletions b/‎bigframes/operations/ai_ops.py‎
Lines changed: 3 additions & 3 deletions
diff --git a/‎tests/unit/core/compile/sqlglot/expressions/snapshots/test_ai_ops/test_ai_classify/None/out.sql‎
Lines changed: 3 additions & 0 deletions b/‎tests/unit/core/compile/sqlglot/expressions/snapshots/test_ai_ops/test_ai_classify/None/out.sql‎
Lines changed: 3 additions & 0 deletions
diff --git a/‎…ots/test_ai_ops/test_ai_classify/out.sql‎ ‎….us.bigframes-default-connection/out.sql‎tests/unit/core/compile/sqlglot/expressions/snapshots/test_ai_ops/test_ai_classify/out.sql renamed to tests/unit/core/compile/sqlglot/expressions/snapshots/test_ai_ops/test_ai_classify/bigframes-dev.us.bigframes-default-connection/out.sql b/‎…ots/test_ai_ops/test_ai_classify/out.sql‎ ‎….us.bigframes-default-connection/out.sql‎tests/unit/core/compile/sqlglot/expressions/snapshots/test_ai_ops/test_ai_classify/out.sql renamed to tests/unit/core/compile/sqlglot/expressions/snapshots/test_ai_ops/test_ai_classify/bigframes-dev.us.bigframes-default-connection/out.sql
diff --git a/‎tests/unit/core/compile/sqlglot/expressions/snapshots/test_ai_ops/test_ai_if/None/out.sql‎
Lines changed: 3 additions & 0 deletions b/‎tests/unit/core/compile/sqlglot/expressions/snapshots/test_ai_ops/test_ai_if/None/out.sql‎
Lines changed: 3 additions & 0 deletions
diff --git a/‎…snapshots/test_ai_ops/test_ai_if/out.sql‎ ‎….us.bigframes-default-connection/out.sql‎tests/unit/core/compile/sqlglot/expressions/snapshots/test_ai_ops/test_ai_if/out.sql renamed to tests/unit/core/compile/sqlglot/expressions/snapshots/test_ai_ops/test_ai_if/bigframes-dev.us.bigframes-default-connection/out.sql b/‎…snapshots/test_ai_ops/test_ai_if/out.sql‎ ‎….us.bigframes-default-connection/out.sql‎tests/unit/core/compile/sqlglot/expressions/snapshots/test_ai_ops/test_ai_if/out.sql renamed to tests/unit/core/compile/sqlglot/expressions/snapshots/test_ai_ops/test_ai_if/bigframes-dev.us.bigframes-default-connection/out.sql
diff --git a/‎tests/unit/core/compile/sqlglot/expressions/snapshots/test_ai_ops/test_ai_score/None/out.sql‎
Lines changed: 3 additions & 0 deletions b/‎tests/unit/core/compile/sqlglot/expressions/snapshots/test_ai_ops/test_ai_score/None/out.sql‎
Lines changed: 3 additions & 0 deletions
diff --git a/‎…pshots/test_ai_ops/test_ai_score/out.sql‎ ‎….us.bigframes-default-connection/out.sql‎tests/unit/core/compile/sqlglot/expressions/snapshots/test_ai_ops/test_ai_score/out.sql renamed to tests/unit/core/compile/sqlglot/expressions/snapshots/test_ai_ops/test_ai_score/bigframes-dev.us.bigframes-default-connection/out.sql b/‎…pshots/test_ai_ops/test_ai_score/out.sql‎ ‎….us.bigframes-default-connection/out.sql‎tests/unit/core/compile/sqlglot/expressions/snapshots/test_ai_ops/test_ai_score/out.sql renamed to tests/unit/core/compile/sqlglot/expressions/snapshots/test_ai_ops/test_ai_score/bigframes-dev.us.bigframes-default-connection/out.sql
diff --git a/‎tests/unit/core/compile/sqlglot/expressions/test_ai_ops.py‎
Lines changed: 9 additions & 6 deletions b/‎tests/unit/core/compile/sqlglot/expressions/test_ai_ops.py‎
Lines changed: 9 additions & 6 deletions
@@ -745,7 +745,7 @@ def if_(
             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.
 
     Returns:
         bigframes.series.Series: A new series of bools.
@@ -756,7 +756,7 @@ def if_(
 
     operator = ai_ops.AIIf(
         prompt_context=tuple(prompt_context),
-        connection_id=_resolve_connection_id(series_list[0], connection_id),
+        connection_id=connection_id,
     )
 
     return series_list[0]._apply_nary_op(operator, series_list[1:])
@@ -800,7 +800,7 @@ def classify(
             Categories to classify the input into.
         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.
 
     Returns:
         bigframes.series.Series: A new series of strings.
@@ -812,7 +812,7 @@ def classify(
     operator = ai_ops.AIClassify(
         prompt_context=tuple(prompt_context),
         categories=tuple(categories),
-        connection_id=_resolve_connection_id(series_list[0], connection_id),
+        connection_id=connection_id,
     )
 
     return series_list[0]._apply_nary_op(operator, series_list[1:])
@@ -853,7 +853,7 @@ def score(
             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.
 
     Returns:
         bigframes.series.Series: A new series of double (float) values.
@@ -864,7 +864,7 @@ def score(
 
     operator = ai_ops.AIScore(
         prompt_context=tuple(prompt_context),
-        connection_id=_resolve_connection_id(series_list[0], connection_id),
+        connection_id=connection_id,
     )
 
     return series_list[0]._apply_nary_op(operator, series_list[1:])
 
@@ -113,9 +113,9 @@ def _construct_named_args(op: ops.NaryOp) -> list[sge.Kwarg]:
             )
         )
 
-    endpoit = op_args.get("endpoint", None)
-    if endpoit is not None:
-        args.append(sge.Kwarg(this="endpoint", expression=sge.Literal.string(endpoit)))
+    endpoint = op_args.get("endpoint", None)
+    if endpoint is not None:
+        args.append(sge.Kwarg(this="endpoint", expression=sge.Literal.string(endpoint)))
 
     request_type = op_args.get("request_type", None)
     if request_type is not None:
 
@@ -123,7 +123,7 @@ class AIIf(base_ops.NaryOp):
     name: ClassVar[str] = "ai_if"
 
     prompt_context: Tuple[str | None, ...]
-    connection_id: str
+    connection_id: str | None
 
     def output_type(self, *input_types: dtypes.ExpressionType) -> dtypes.ExpressionType:
         return dtypes.BOOL_DTYPE
@@ -135,7 +135,7 @@ class AIClassify(base_ops.NaryOp):
 
     prompt_context: Tuple[str | None, ...]
     categories: tuple[str, ...]
-    connection_id: str
+    connection_id: str | None
 
     def output_type(self, *input_types: dtypes.ExpressionType) -> dtypes.ExpressionType:
         return dtypes.STRING_DTYPE
@@ -146,7 +146,7 @@ class AIScore(base_ops.NaryOp):
     name: ClassVar[str] = "ai_score"
 
     prompt_context: Tuple[str | None, ...]
-    connection_id: str
+    connection_id: str | None
 
     def output_type(self, *input_types: dtypes.ExpressionType) -> dtypes.ExpressionType:
         return dtypes.FLOAT_DTYPE
@@ -0,0 +1,3 @@
+SELECT
+  AI.CLASSIFY(input => (`string_col`), categories => ['greeting', 'rejection']) AS `result`
+FROM `bigframes-dev`.`sqlglot_test`.`scalar_types` AS `bft_0`
@@ -0,0 +1,3 @@
+SELECT
+  AI.IF(prompt => (`string_col`, ' is the same as ', `string_col`)) AS `result`
+FROM `bigframes-dev`.`sqlglot_test`.`scalar_types` AS `bft_0`
@@ -0,0 +1,3 @@
+SELECT
+  AI.SCORE(prompt => (`string_col`, ' is the same as ', `string_col`)) AS `result`
+FROM `bigframes-dev`.`sqlglot_test`.`scalar_types` AS `bft_0`
@@ -281,12 +281,13 @@ def test_ai_generate_double_with_model_param(
     snapshot.assert_match(sql, "out.sql")
 
 
-def test_ai_if(scalar_types_df: dataframe.DataFrame, snapshot):
+@pytest.mark.parametrize("connection_id", [None, CONNECTION_ID])
+def test_ai_if(scalar_types_df: dataframe.DataFrame, snapshot, connection_id):
     col_name = "string_col"
 
     op = ops.AIIf(
         prompt_context=(None, " is the same as ", None),
-        connection_id=CONNECTION_ID,
+        connection_id=connection_id,
     )
 
     sql = utils._apply_ops_to_sql(
@@ -296,26 +297,28 @@ def test_ai_if(scalar_types_df: dataframe.DataFrame, snapshot):
     snapshot.assert_match(sql, "out.sql")
 
 
-def test_ai_classify(scalar_types_df: dataframe.DataFrame, snapshot):
+@pytest.mark.parametrize("connection_id", [None, CONNECTION_ID])
+def test_ai_classify(scalar_types_df: dataframe.DataFrame, snapshot, connection_id):
     col_name = "string_col"
 
     op = ops.AIClassify(
         prompt_context=(None,),
         categories=("greeting", "rejection"),
-        connection_id=CONNECTION_ID,
+        connection_id=connection_id,
     )
 
     sql = utils._apply_ops_to_sql(scalar_types_df, [op.as_expr(col_name)], ["result"])
 
     snapshot.assert_match(sql, "out.sql")
 
 
-def test_ai_score(scalar_types_df: dataframe.DataFrame, snapshot):
+@pytest.mark.parametrize("connection_id", [None, CONNECTION_ID])
+def test_ai_score(scalar_types_df: dataframe.DataFrame, snapshot, connection_id):
     col_name = "string_col"
 
     op = ops.AIScore(
         prompt_context=(None, " is the same as ", None),
-        connection_id=CONNECTION_ID,
+        connection_id=connection_id,
     )
 
     sql = utils._apply_ops_to_sql(
Original file line number	Diff line number	Diff line change
`@@ -113,9 +113,9 @@ def _construct_named_args(op: ops.NaryOp) -> list[sge.Kwarg]:`
`113`	`113`	`)`
`114`	`114`	`)`
`115`	`115`
`116`		`- endpoit = op_args.get("endpoint", None)`
`117`		`- if endpoit is not None:`
`118`		`- args.append(sge.Kwarg(this="endpoint", expression=sge.Literal.string(endpoit)))`
	`116`	`+ endpoint = op_args.get("endpoint", None)`
	`117`	`+ if endpoint is not None:`
	`118`	`+ args.append(sge.Kwarg(this="endpoint", expression=sge.Literal.string(endpoint)))`
`119`	`119`
`120`	`120`	`request_type = op_args.get("request_type", None)`
`121`	`121`	`if request_type is not None:`
Original file line number	Diff line number	Diff line change
`@@ -0,0 +1,3 @@`
	`1`	`+SELECT`
	`2`	+ AI.CLASSIFY(input => (`string_col`), categories => ['greeting', 'rejection']) AS `result`
	`3`	+FROM `bigframes-dev`.`sqlglot_test`.`scalar_types` AS `bft_0`
Original file line number	Diff line number	Diff line change
`@@ -0,0 +1,3 @@`
	`1`	`+SELECT`
	`2`	+ AI.IF(prompt => (`string_col`, ' is the same as ', `string_col`)) AS `result`
	`3`	+FROM `bigframes-dev`.`sqlglot_test`.`scalar_types` AS `bft_0`
Original file line number	Diff line number	Diff line change
`@@ -0,0 +1,3 @@`
	`1`	`+SELECT`
	`2`	+ AI.SCORE(prompt => (`string_col`, ' is the same as ', `string_col`)) AS `result`
	`3`	+FROM `bigframes-dev`.`sqlglot_test`.`scalar_types` AS `bft_0`