refactor(bigframes): update SQLGlot compiler to process AI func params better (googleapis#16757)

Author: sycai

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

@@ -62,14 +62,8 @@ def _(*exprs: TypedExpr, op: ops.AIIf) -> sge.Expression:
 
 @register_nary_op(ops.AIClassify, pass_op=True)
 def _(*exprs: TypedExpr, op: ops.AIClassify) -> sge.Expression:
-    category_literals = [sge.Literal.string(cat) for cat in op.categories]
-    categories_arg = sge.Kwarg(
-        this="categories", expression=sge.array(*category_literals)
-    )
-
     args = [
         _construct_prompt(exprs, op.prompt_context, param_name="input"),
-        categories_arg,
     ] + _construct_named_args(op)
 
     return sge.func("AI.CLASSIFY", *args)
@@ -105,44 +99,33 @@ def _construct_named_args(op: ops.NaryOp) -> list[sge.Kwarg]:
 
     op_args = asdict(op)
 
-    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)
-            )
-        )
-
-    endpoint = op_args.get("endpoint", None)
-    if endpoint is not None:
-        args.append(sge.Kwarg(this="endpoint", expression=sge.Literal.string(endpoint)))
+    for field, value in op_args.items():
+        if value is None or field == "prompt_context":
+            continue
 
-    request_type = op_args.get("request_type", None)
-    if request_type is not None:
-        args.append(
-            sge.Kwarg(
-                this="request_type", expression=sge.Literal.string(request_type.upper())
+        if field == "categories":
+            category_literals = [sge.Literal.string(cat) for cat in value]
+            categories_arg = sge.Kwarg(
+                this="categories", expression=sge.array(*category_literals)
             )
-        )
-
-    model_params = op_args.get("model_params", None)
-    if model_params is not None:
-        args.append(
-            sge.Kwarg(
-                this="model_params",
-                # sge.JSON requires the SQLGlot version to be at least 25.18.0
-                # PARSE_JSON won't work as the function requires a JSON literal.
-                expression=sge.JSON(this=sge.Literal.string(model_params)),
+            args.append(categories_arg)
+        elif field == "model_params":
+            # model_params is a JSON string, so we need to use the JSON function to pass it as a named argument.
+            args.append(
+                sge.Kwarg(
+                    this="model_params",
+                    # sge.JSON requires the SQLGlot version to be at least 25.18.0
+                    # PARSE_JSON won't work as the function requires a JSON literal.
+                    expression=sge.JSON(this=sge.Literal.string(value)),
+                )
             )
-        )
-
-    output_schema = op_args.get("output_schema", None)
-    if output_schema is not None:
-        args.append(
-            sge.Kwarg(
-                this="output_schema",
-                expression=sge.Literal.string(output_schema),
+        elif field == "request_type":
+            args.append(
+                sge.Kwarg(this=field, expression=sge.Literal.string(value.upper()))
+            )
+        else:
+            args.append(
+                sge.Kwarg(this=field, expression=sge.Literal.string(str(value)))
             )
-        )
 
     return args