feat: use EUC for AI IF, CLASSIFY, and SCORE when connection is not provided

bigframes/bigquery/_operations/ai.py

@@ -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:])

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

@@ -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:

bigframes/operations/ai_ops.py

@@ -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

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

@@ -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`

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

@@ -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`

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

@@ -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`

tests/unit/core/compile/sqlglot/expressions/test_ai_ops.py

@@ -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(

third_party/bigframes_vendored/ibis/expr/operations/ai_ops.py

@@ -113,7 +113,7 @@ class AIIf(Value):
     """Generate True/False based on the prompt"""
 
     prompt: Value
-    connection_id: Value[dt.String]
+    connection_id: Optional[Value[dt.String]]
 
     shape = rlz.shape_like("prompt")
 
@@ -128,7 +128,7 @@ class AIClassify(Value):
 
     input: Value
     categories: Value[dt.Array[dt.String]]
-    connection_id: Value[dt.String]
+    connection_id: Optional[Value[dt.String]]
 
     shape = rlz.shape_like("input")
 
@@ -142,7 +142,7 @@ class AIScore(Value):
     """Generate doubles based on the prompt"""
 
     prompt: Value
-    connection_id: Value[dt.String]
+    connection_id: Optional[Value[dt.String]]
 
     shape = rlz.shape_like("prompt")