refactor(bigframes): reduce code size of AI functions and tests (#17055)

The changes include:
* Make all optional parameters default to None, which reduces the length
of compiled SQLs
* "Upper" the parameters at the API layer if necessary, so that we don't
need to repeat this logic in both compilers later
* Make all optional parameters default to None in the operator
definitions too, saving some lines in the unit tests

---------

Co-authored-by: gemini-code-assist[bot] <176961590+gemini-code-assist[bot]@users.noreply.github.com>

Author: sycai

packages/bigframes/bigframes/bigquery/_operations/ai.py

@@ -48,7 +48,7 @@ def generate(
     *,
     connection_id: str | None = None,
     endpoint: str | None = None,
-    request_type: Literal["dedicated", "shared", "unspecified"] = "unspecified",
+    request_type: Literal["dedicated", "shared", "unspecified"] | None = None,
     model_params: Mapping[Any, Any] | None = None,
     output_schema: Mapping[str, str] | None = None,
 ) -> series.Series:
@@ -129,7 +129,7 @@ def generate(
         prompt_context=tuple(prompt_context),
         connection_id=connection_id,
         endpoint=endpoint,
-        request_type=request_type,
+        request_type=_upper_optional(request_type),
         model_params=json.dumps(model_params) if model_params else None,
         output_schema=output_schema_str,
     )
@@ -143,7 +143,7 @@ def generate_bool(
     *,
     connection_id: str | None = None,
     endpoint: str | None = None,
-    request_type: Literal["dedicated", "shared", "unspecified"] = "unspecified",
+    request_type: Literal["dedicated", "shared", "unspecified"] | None = None,
     model_params: Mapping[Any, Any] | None = None,
 ) -> series.Series:
     """
@@ -207,7 +207,7 @@ def generate_bool(
         prompt_context=tuple(prompt_context),
         connection_id=connection_id,
         endpoint=endpoint,
-        request_type=request_type,
+        request_type=_upper_optional(request_type),
         model_params=json.dumps(model_params) if model_params else None,
     )
 
@@ -220,7 +220,7 @@ def generate_int(
     *,
     connection_id: str | None = None,
     endpoint: str | None = None,
-    request_type: Literal["dedicated", "shared", "unspecified"] = "unspecified",
+    request_type: Literal["dedicated", "shared", "unspecified"] | None = None,
     model_params: Mapping[Any, Any] | None = None,
 ) -> series.Series:
     """
@@ -281,7 +281,7 @@ def generate_int(
         prompt_context=tuple(prompt_context),
         connection_id=connection_id,
         endpoint=endpoint,
-        request_type=request_type,
+        request_type=_upper_optional(request_type),
         model_params=json.dumps(model_params) if model_params else None,
     )
 
@@ -294,7 +294,7 @@ def generate_double(
     *,
     connection_id: str | None = None,
     endpoint: str | None = None,
-    request_type: Literal["dedicated", "shared", "unspecified"] = "unspecified",
+    request_type: Literal["dedicated", "shared", "unspecified"] | None = None,
     model_params: Mapping[Any, Any] | None = None,
 ) -> series.Series:
     """
@@ -355,7 +355,7 @@ def generate_double(
         prompt_context=tuple(prompt_context),
         connection_id=connection_id,
         endpoint=endpoint,
-        request_type=request_type,
+        request_type=_upper_optional(request_type),
         model_params=json.dumps(model_params) if model_params else None,
     )
 
@@ -753,7 +753,7 @@ def embed(
     operator = ai_ops.AIEmbed(
         endpoint=endpoint,
         model=model,
-        task_type=task_type,
+        task_type=_upper_optional(task_type),
         title=title,
         model_params=json.dumps(model_params) if model_params else None,
         connection_id=connection_id,
@@ -775,7 +775,7 @@ def if_(
     *,
     connection_id: str | None = None,
     endpoint: str | None = None,
-    optimization_mode: Literal["minimize_cost", "maximize_quality"] = "minimize_cost",
+    optimization_mode: Literal["minimize_cost", "maximize_quality"] | None = None,
     max_error_ratio: float | None = None,
 ) -> series.Series:
     """
@@ -830,7 +830,7 @@ def if_(
         prompt_context=tuple(prompt_context),
         connection_id=connection_id,
         endpoint=endpoint,
-        optimization_mode=optimization_mode,
+        optimization_mode=_upper_optional(optimization_mode),
         max_error_ratio=max_error_ratio,
     )
 
@@ -904,7 +904,7 @@ def classify(
         examples=example_tuples,
         connection_id=connection_id,
         endpoint=endpoint,
-        optimization_mode=optimization_mode,
+        optimization_mode=_upper_optional(optimization_mode),
         max_error_ratio=max_error_ratio,
     )
 
@@ -1225,3 +1225,9 @@ def _to_dataframe(
         return data
 
     raise ValueError(f"Unsupported data type: {type(data)}")
+
+
+def _upper_optional(value: str | None) -> str | None:
+    if value is None:
+        return None
+    return value.upper()

packages/bigframes/bigframes/core/compile/ibis_compiler/scalar_op_registry.py

@@ -1920,7 +1920,7 @@ def ai_generate(
         _construct_prompt(values, op.prompt_context),  # type: ignore
         op.connection_id,  # type: ignore
         op.endpoint,  # type: ignore
-        op.request_type.upper(),  # type: ignore
+        op.request_type,  # type: ignore
         op.model_params,  # type: ignore
         op.output_schema,  # type: ignore
     ).to_expr()
@@ -1934,7 +1934,7 @@ def ai_generate_bool(
         _construct_prompt(values, op.prompt_context),  # type: ignore
         op.connection_id,  # type: ignore
         op.endpoint,  # type: ignore
-        op.request_type.upper(),  # type: ignore
+        op.request_type,  # type: ignore
         op.model_params,  # type: ignore
     ).to_expr()
 
@@ -1947,7 +1947,7 @@ def ai_generate_int(
         _construct_prompt(values, op.prompt_context),  # type: ignore
         op.connection_id,  # type: ignore
         op.endpoint,  # type: ignore
-        op.request_type.upper(),  # type: ignore
+        op.request_type,  # type: ignore
         op.model_params,  # type: ignore
     ).to_expr()
 
@@ -1960,7 +1960,7 @@ def ai_generate_double(
         _construct_prompt(values, op.prompt_context),  # type: ignore
         op.connection_id,  # type: ignore
         op.endpoint,  # type: ignore
-        op.request_type.upper(),  # type: ignore
+        op.request_type,  # type: ignore
         op.model_params,  # type: ignore
     ).to_expr()
 
@@ -1972,7 +1972,7 @@ def ai_embed(value: ibis_types.Value, op: ops.AIEmbed) -> ibis_types.StructValue
         connection_id=op.connection_id,  # type: ignore
         endpoint=op.endpoint,  # type: ignore
         model=op.model,  # type: ignore
-        task_type=op.task_type.upper() if op.task_type is not None else None,  # type: ignore
+        task_type=op.task_type,  # type: ignore
         title=op.title,  # type: ignore
         model_params=op.model_params,  # type: ignore
     ).to_expr()
@@ -1984,7 +1984,7 @@ def ai_if(*values: ibis_types.Value, op: ops.AIIf) -> ibis_types.StructValue:
         _construct_prompt(values, op.prompt_context),  # type: ignore
         op.connection_id,  # type: ignore
         op.endpoint,  # type: ignore
-        op.optimization_mode.upper() if op.optimization_mode is not None else None,  # type: ignore
+        op.optimization_mode,  # type: ignore
         op.max_error_ratio,  # type: ignore
     ).to_expr()
 
@@ -1999,7 +1999,7 @@ def ai_classify(
         _construct_examples(op.examples),  # type: ignore
         op.connection_id,  # type: ignore
         op.endpoint,  # type: ignore
-        op.optimization_mode.upper() if op.optimization_mode is not None else None,  # type: ignore
+        op.optimization_mode,  # type: ignore
         op.max_error_ratio,  # type: ignore
     ).to_expr()
 

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

@@ -139,16 +139,6 @@ def _construct_named_args(op: ops.ScalarOp) -> list[sge.Kwarg]:
                     expression=sge.JSON(this=sge.Literal.string(value)),
                 )
             )
-        elif field == "optimization_mode":
-            args.append(
-                sge.Kwarg(this=field, expression=sge.Literal.string(value.upper()))
-            )
-        elif field == "max_error_ratio":
-            args.append(sge.Kwarg(this=field, expression=sge.Literal.number(value)))
-        elif field == "request_type":
-            args.append(
-                sge.Kwarg(this=field, expression=sge.Literal.string(value.upper()))
-            )
         elif field == "examples":
             example_expressions = [
                 sge.Tuple(
@@ -160,8 +150,6 @@ def _construct_named_args(op: ops.ScalarOp) -> list[sge.Kwarg]:
                 sge.Kwarg(this=field, expression=sge.array(*example_expressions))
             )
         else:
-            args.append(
-                sge.Kwarg(this=field, expression=sge.Literal.string(str(value)))
-            )
+            args.append(sge.Kwarg(this=field, expression=sge.convert(value)))
 
     return args

packages/bigframes/bigframes/operations/ai_ops.py

@@ -15,7 +15,7 @@
 from __future__ import annotations
 
 import dataclasses
-from typing import ClassVar, Literal, Tuple
+from typing import ClassVar, Tuple
 
 import pandas as pd
 import pyarrow as pa
@@ -29,11 +29,11 @@ class AIGenerate(base_ops.NaryOp):
     name: ClassVar[str] = "ai_generate"
 
     prompt_context: Tuple[str | None, ...]
-    connection_id: str | None
-    endpoint: str | None
-    request_type: Literal["dedicated", "shared", "unspecified"]
-    model_params: str | None
-    output_schema: str | None
+    connection_id: str | None = None
+    endpoint: str | None = None
+    request_type: str | None = None
+    model_params: str | None = None
+    output_schema: str | None = None
 
     def output_type(self, *input_types: dtypes.ExpressionType) -> dtypes.ExpressionType:
         if self.output_schema is None:
@@ -57,10 +57,10 @@ class AIGenerateBool(base_ops.NaryOp):
     name: ClassVar[str] = "ai_generate_bool"
 
     prompt_context: Tuple[str | None, ...]
-    connection_id: str | None
-    endpoint: str | None
-    request_type: Literal["dedicated", "shared", "unspecified"]
-    model_params: str | None
+    connection_id: str | None = None
+    endpoint: str | None = None
+    request_type: str | None = None
+    model_params: str | None = None
 
     def output_type(self, *input_types: dtypes.ExpressionType) -> dtypes.ExpressionType:
         return pd.ArrowDtype(
@@ -79,10 +79,10 @@ class AIGenerateInt(base_ops.NaryOp):
     name: ClassVar[str] = "ai_generate_int"
 
     prompt_context: Tuple[str | None, ...]
-    connection_id: str | None
-    endpoint: str | None
-    request_type: Literal["dedicated", "shared", "unspecified"]
-    model_params: str | None
+    connection_id: str | None = None
+    endpoint: str | None = None
+    request_type: str | None = None
+    model_params: str | None = None
 
     def output_type(self, *input_types: dtypes.ExpressionType) -> dtypes.ExpressionType:
         return pd.ArrowDtype(
@@ -101,10 +101,10 @@ class AIGenerateDouble(base_ops.NaryOp):
     name: ClassVar[str] = "ai_generate_double"
 
     prompt_context: Tuple[str | None, ...]
-    connection_id: str | None
-    endpoint: str | None
-    request_type: Literal["dedicated", "shared", "unspecified"]
-    model_params: str | None
+    connection_id: str | None = None
+    endpoint: str | None = None
+    request_type: str | None = None
+    model_params: str | None = None
 
     def output_type(self, *input_types: dtypes.ExpressionType) -> dtypes.ExpressionType:
         return pd.ArrowDtype(
@@ -122,12 +122,12 @@ def output_type(self, *input_types: dtypes.ExpressionType) -> dtypes.ExpressionT
 class AIEmbed(base_ops.UnaryOp):
     name: ClassVar[str] = "ai_embed"
 
-    endpoint: str | None
-    model: str | None
-    task_type: str | None
-    title: str | None
-    model_params: str | None
-    connection_id: str | None
+    endpoint: str | None = None
+    model: str | None = None
+    task_type: str | None = None
+    title: str | None = None
+    model_params: str | None = None
+    connection_id: str | None = None
 
     def output_type(self, *input_types: dtypes.ExpressionType) -> dtypes.ExpressionType:
         return pd.ArrowDtype(
@@ -145,7 +145,7 @@ class AIIf(base_ops.NaryOp):
     name: ClassVar[str] = "ai_if"
 
     prompt_context: Tuple[str | None, ...]
-    connection_id: str | None
+    connection_id: str | None = None
     endpoint: str | None = None
     optimization_mode: str | None = None
     max_error_ratio: float | None = None
@@ -160,11 +160,11 @@ class AIClassify(base_ops.NaryOp):
 
     prompt_context: Tuple[str | None, ...]
     categories: tuple[str, ...]
-    examples: tuple[tuple[str, str], ...] | None
-    connection_id: str | None
-    endpoint: str | None
-    optimization_mode: str | None
-    max_error_ratio: float | None
+    examples: tuple[tuple[str, str], ...] | None = None
+    connection_id: str | None = None
+    endpoint: str | None = None
+    optimization_mode: str | None = None
+    max_error_ratio: float | None = None
 
     def output_type(self, *input_types: dtypes.ExpressionType) -> dtypes.ExpressionType:
         return dtypes.STRING_DTYPE
@@ -175,9 +175,9 @@ class AIScore(base_ops.NaryOp):
     name: ClassVar[str] = "ai_score"
 
     prompt_context: Tuple[str | None, ...]
-    connection_id: str | None
-    endpoint: str | None
-    max_error_ratio: float | None
+    connection_id: str | None = None
+    endpoint: str | None = None
+    max_error_ratio: float | None = None
 
     def output_type(self, *input_types: dtypes.ExpressionType) -> dtypes.ExpressionType:
         return dtypes.FLOAT_DTYPE
@@ -187,10 +187,10 @@ def output_type(self, *input_types: dtypes.ExpressionType) -> dtypes.ExpressionT
 class AISimilarity(base_ops.BinaryOp):
     name: ClassVar[str] = "ai_similarity"
 
-    endpoint: str | None
-    model: str | None
-    model_params: str | None
-    connection_id: str | None
+    endpoint: str | None = None
+    model: str | None = None
+    model_params: str | None = None
+    connection_id: str | None = None
 
     def output_type(self, *input_types: dtypes.ExpressionType) -> dtypes.ExpressionType:
         return dtypes.FLOAT_DTYPE

packages/bigframes/tests/unit/core/compile/sqlglot/expressions/snapshots/test_ai_ops/test_ai_embed_with_task_type_and_title/out.sql

@@ -2,7 +2,7 @@ SELECT
   AI.EMBED(
     `string_col`,
     endpoint => 'text-embedding-005',
-    task_type => 'retrieval_document',
+    task_type => 'RETRIEVAL_DOCUMENT',
     title => 'My Document',
     model_params => JSON '{"outputDimensionality": 256}'
   ) AS `result`

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

@@ -1,7 +1,6 @@
 SELECT
   AI.GENERATE_BOOL(
     prompt => (`string_col`, ' is the same as ', `string_col`),
-    endpoint => 'gemini-2.5-flash',
-    request_type => 'SHARED'
+    endpoint => 'gemini-2.5-flash'
   ) AS `result`
 FROM `bigframes-dev`.`sqlglot_test`.`scalar_types` AS `bft_0`

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

@@ -2,7 +2,6 @@ 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'
+    endpoint => 'gemini-2.5-flash'
   ) AS `result`
 FROM `bigframes-dev`.`sqlglot_test`.`scalar_types` AS `bft_0`

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

@@ -1,7 +1,6 @@
 SELECT
   AI.GENERATE_BOOL(
     prompt => (`string_col`, ' is the same as ', `string_col`),
-    request_type => 'SHARED',
     model_params => JSON '{}'
   ) AS `result`
 FROM `bigframes-dev`.`sqlglot_test`.`scalar_types` AS `bft_0`

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

@@ -1,7 +1,6 @@
 SELECT
   AI.GENERATE_DOUBLE(
     prompt => (`string_col`, ' is the same as ', `string_col`),
-    endpoint => 'gemini-2.5-flash',
-    request_type => 'SHARED'
+    endpoint => 'gemini-2.5-flash'
   ) AS `result`
 FROM `bigframes-dev`.`sqlglot_test`.`scalar_types` AS `bft_0`