fix: deflake ai_gen_bool multimodel test (#2085)

* fix: deflake ai_gen_bool multimodel test

* fix lint

* fix doctest too

* consolidates tests under system/small

* fix doctest

Author: sycai

bigframes/bigquery/_operations/ai.py

@@ -47,30 +47,13 @@ def generate_bool(
         ...     "col_1": ["apple", "bear", "pear"],
         ...     "col_2": ["fruit", "animal", "animal"]
         ... })
-        >>> bbq.ai_generate_bool((df["col_1"], " is a ", df["col_2"]))
+        >>> bbq.ai.generate_bool((df["col_1"], " is a ", df["col_2"]))
         0    {'result': True, 'full_response': '{"candidate...
         1    {'result': True, 'full_response': '{"candidate...
         2    {'result': False, 'full_response': '{"candidat...
         dtype: struct<result: bool, full_response: string, status: string>[pyarrow]
 
-        >>> bbq.ai_generate_bool((df["col_1"], " is a ", df["col_2"])).struct.field("result")
-        0     True
-        1     True
-        2    False
-        Name: result, dtype: boolean
-
-        >>> model_params = {
-        ...     "generation_config": {
-        ...         "thinking_config": {
-        ...             "thinking_budget": 0
-        ...         }
-        ...     }
-        ... }
-        >>> bbq.ai_generate_bool(
-        ...     (df["col_1"], " is a ", df["col_2"]),
-        ...     endpoint="gemini-2.5-pro",
-        ...     model_params=model_params,
-        ... ).struct.field("result")
+        >>> bbq.ai.generate_bool((df["col_1"], " is a ", df["col_2"])).struct.field("result")
         0     True
         1     True
         2    False

tests/system/large/bigquery/__init__.py

@@ -15,9 +15,10 @@
 import sys
 
 import pandas as pd
-import pandas.testing
+import pyarrow as pa
 import pytest
 
+from bigframes import series
 import bigframes.bigquery as bbq
 import bigframes.pandas as bpd
 
@@ -27,15 +28,17 @@ def test_ai_generate_bool(session):
     s2 = bpd.Series(["fruit", "tree"], session=session)
     prompt = (s1, " is a ", s2)
 
-    result = bbq.ai.generate_bool(prompt, endpoint="gemini-2.5-flash").struct.field(
-        "result"
-    )
+    result = bbq.ai.generate_bool(prompt, endpoint="gemini-2.5-flash")
 
-    pandas.testing.assert_series_equal(
-        result.to_pandas(),
-        pd.Series([True, False], name="result"),
-        check_dtype=False,
-        check_index=False,
+    assert _contains_no_nulls(result)
+    assert result.dtype == pd.ArrowDtype(
+        pa.struct(
+            (
+                pa.field("result", pa.bool_()),
+                pa.field("full_response", pa.string()),
+                pa.field("status", pa.string()),
+            )
+        )
     )
 
 
@@ -52,11 +55,38 @@ def test_ai_generate_bool_with_model_params(session):
 
     result = bbq.ai.generate_bool(
         prompt, endpoint="gemini-2.5-flash", model_params=model_params
-    ).struct.field("result")
+    )
+
+    assert _contains_no_nulls(result)
+    assert result.dtype == pd.ArrowDtype(
+        pa.struct(
+            (
+                pa.field("result", pa.bool_()),
+                pa.field("full_response", pa.string()),
+                pa.field("status", pa.string()),
+            )
+        )
+    )
+
 
-    pandas.testing.assert_series_equal(
-        result.to_pandas(),
-        pd.Series([True, False], name="result"),
-        check_dtype=False,
-        check_index=False,
+def test_ai_generate_bool_multi_model(session):
+    df = session.from_glob_path(
+        "gs://bigframes-dev-testing/a_multimodel/images/*", name="image"
     )
+
+    result = bbq.ai.generate_bool((df["image"], " contains an animal"))
+
+    assert _contains_no_nulls(result)
+    assert result.dtype == pd.ArrowDtype(
+        pa.struct(
+            (
+                pa.field("result", pa.bool_()),
+                pa.field("full_response", pa.string()),
+                pa.field("status", pa.string()),
+            )
+        )
+    )
+
+
+def _contains_no_nulls(s: series.Series) -> bool:
+    return len(s) == s.count()