fix stuffs

Author: sycai

bigframes/bigquery/_operations/ai.py

@@ -23,7 +23,7 @@
 
 
 def ai_generate_bool(
-    prompt: series.Series | List[str | series.Series] | Tuple[str | series.Series],
+    prompt: series.Series | List[str | series.Series] | Tuple[str | series.Series, ...],
     *,
     connection_id: str | None = None,
     endpoint: str | None = None,
@@ -71,7 +71,7 @@ def ai_generate_bool(
         Name: result, dtype: boolean
 
     Args:
-        prompt (series.Series | List[str|series.Series] | Tuple[str|series.Series]):
+        prompt (series.Series | List[str|series.Series] | Tuple[str|series.Series, ...]):
             A mixture of Series and string literals that specifies the prompt to send to the model.
         connection_id (str, optional):
             Specifies the connection to use to communicate with the model. For example, `myproject.us.myconnection`.
@@ -87,15 +87,15 @@ def ai_generate_bool(
             * "dedicated": function only uses Provisioned Throughput quota. The AI.GENERATE function returns the error Provisioned throughput is not purchased or is not active if Provisioned Throughput quota isn't available.
             * "shared": the function only uses dynamic shared quota (DSQ), even if you have purchased Provisioned Throughput quota.
             * "unspecified":
-                * If you haven't purchased Provisioned Throughput quota, the function uses DSQ quota.
-                * If you have purchased Provisioned Throughput quota, the function uses the Provisioned Throughput quota first. If requests exceed the Provisioned Throughput quota, the overflow traffic uses DSQ quota.
+                If you haven't purchased Provisioned Throughput quota, the function uses DSQ quota.
+                If you have purchased Provisioned Throughput quota, the function uses the Provisioned Throughput quota first. If requests exceed the Provisioned Throughput quota, the overflow traffic uses DSQ quota.
         model_params (Mapping[Any, Any]):
             Provides additional parameters to the model. The MODEL_PARAMS value must conform to the generateContent request body format.
 
     Returns:
         bigframes.series.Series: A new struct Series with the result data. The struct contains these fields:
         * "result": a BOOL value containing the model's response to the prompt. The result is None if the request fails or is filtered by responsible AI.
-        * "full_resposne": a STRING value containing the JSON response from the projects.locations.endpoints.generateContent call to the model. The generated text is in the text element.
+        * "full_response": a STRING value containing the JSON response from the projects.locations.endpoints.generateContent call to the model. The generated text is in the text element.
         * "status": a STRING value that contains the API response status for the corresponding row. This value is empty if the operation was successful.
     """
 

bigframes/core/compile/ibis_compiler/scalar_op_registry.py

@@ -1968,26 +1968,24 @@ def struct_op_impl(
 @scalar_op_compiler.register_nary_op(ops.AIGenerateBool, pass_op=True)
 def ai_generate_bool(
     *values: ibis_types.Value, op: ops.AIGenerateBool
-) -> ibis_dtypes.StructValue:
+) -> ibis_types.StructValue:
 
-    prompt = {}
+    prompt: dict[str, ibis_types.Value | str] = {}
     column_ref_idx = 0
 
     for idx, elem in enumerate(op.prompt_context):
         if elem is None:
-            value = values[column_ref_idx]
+            prompt[f"_field_{idx + 1}"] = values[column_ref_idx]
             column_ref_idx += 1
         else:
-            value = elem
-
-        prompt[f"_field_{idx + 1}"] = value
+            prompt[f"_field_{idx + 1}"] = elem
 
     return ai_ops.AIGenerateBool(
-        ibis.struct(prompt),
-        op.connection_id,
-        op.endpoint,
-        op.request_type.upper(),
-        op.model_params,
+        ibis.struct(prompt),  # type: ignore  
+        op.connection_id,# type: ignore
+        op.endpoint,# type: ignore
+        op.request_type.upper(),# type: ignore
+        op.model_params,# type: ignore
     ).to_expr()
 
 

bigframes/operations/ai_ops.py

@@ -29,7 +29,7 @@ class AIGenerateBool(base_ops.NaryOp):
     name: ClassVar[str] = "ai_generate_bool"
 
     # None are the placeholders for column references.
-    prompt_context: Tuple[str | None]
+    prompt_context: Tuple[str | None, ...]
     connection_id: str
     endpoint: str | None
     request_type: Literal["dedicated", "shared", "unspecified"]

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

@@ -4,6 +4,8 @@
 
 from __future__ import annotations
 
+from typing import Optional
+
 from bigframes_vendored.ibis.common.annotations import attribute
 import bigframes_vendored.ibis.expr.datatypes as dt
 from bigframes_vendored.ibis.expr.operations.core import Value
@@ -17,9 +19,9 @@ class AIGenerateBool(Value):
 
     prompt: Value
     connection_id: Value[dt.String]
-    endpoint: Value[dt.String] | None
+    endpoint: Optional[Value[dt.String]]
     request_type: Value[dt.String]
-    model_params: Value[dt.String] | None
+    model_params: Optional[Value[dt.String]]
 
     shape = rlz.shape_like("prompt")