diff --git a/packages/bigframes/tests/system/small/bigquery/test_ai.py b/packages/bigframes/tests/system/small/bigquery/test_ai.py index 05ebea141440..0b6738dec80a 100644 --- a/packages/bigframes/tests/system/small/bigquery/test_ai.py +++ b/packages/bigframes/tests/system/small/bigquery/test_ai.py @@ -12,6 +12,8 @@ # See the License for the specific language governing permissions and # limitations under the License. +from __future__ import annotations + import uuid from unittest import mock @@ -26,6 +28,16 @@ from bigframes.testing import utils as test_utils +@pytest.fixture +def use_ibis_compiler(): + original_setting = bpd.options.experiments.sql_compiler + bpd.options.experiments.sql_compiler = "legacy" + try: + yield + finally: + bpd.options.experiments.sql_compiler = original_setting + + def _create_mock_obj_ref_df(session, uris, name="image", connection=None): df = bpd.DataFrame({name: uris}, session=session) # Convert string URIs to ObjectRef structs @@ -146,6 +158,19 @@ def test_ai_generate(session): ) +def test_ai_generate_access_full_response_with_ibis(session, use_ibis_compiler): + country = bpd.Series(["Japan", "Canada"], session=session) + prompt = ("What's the capital city of ", country, "? one word only") + + result = ( + bbq.ai.generate(prompt, endpoint="gemini-2.5-flash") + .struct.field("full_response") + .to_pandas() + ) + + assert _contains_no_nulls(result) + + def test_ai_generate_with_output_schema(session): country = bpd.Series(["Japan", "Canada"], session=session) prompt = ("Describe ", country) @@ -200,6 +225,20 @@ def test_ai_generate_bool(session): ) +def test_ai_generate_bool_access_full_response_with_ibis(session, use_ibis_compiler): + s1 = bpd.Series(["apple", "bear"], session=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("full_response") + .to_pandas() + ) + + assert _contains_no_nulls(result) + + def test_ai_generate_bool_multi_model(session, bq_connection): df = _create_mock_obj_ref_df( session, @@ -241,6 +280,19 @@ def test_ai_generate_int(session): ) +def test_ai_generate_int_access_full_response_with_ibis(session, use_ibis_compiler): + s = bpd.Series(["Cat"], session=session) + prompt = ("How many legs does a ", s, " have?") + + result = ( + bbq.ai.generate_int(prompt, endpoint="gemini-2.5-flash") + .struct.field("full_response") + .to_pandas() + ) + + assert _contains_no_nulls(result) + + def test_ai_generate_int_multi_model(session, bq_connection): df = _create_mock_obj_ref_df( session, @@ -284,6 +336,19 @@ def test_ai_generate_double(session): ) +def test_ai_generate_double_access_full_response_with_ibis(session, use_ibis_compiler): + s = bpd.Series(["Cat"], session=session) + prompt = ("How many legs does a ", s, " have?") + + result = ( + bbq.ai.generate_double(prompt, endpoint="gemini-2.5-flash") + .struct.field("full_response") + .to_pandas() + ) + + assert _contains_no_nulls(result) + + def test_ai_generate_double_multi_model(session, bq_connection): df = _create_mock_obj_ref_df( session, @@ -521,5 +586,5 @@ def test_ai_similarity_both_contents_are_string_literals(session): assert result.dtype == dtypes.FLOAT_DTYPE -def _contains_no_nulls(s: series.Series) -> bool: +def _contains_no_nulls(s: series.Series | pd.Series) -> bool: return len(s) == s.count() diff --git a/packages/bigframes/third_party/bigframes_vendored/ibis/expr/operations/ai_ops.py b/packages/bigframes/third_party/bigframes_vendored/ibis/expr/operations/ai_ops.py index fcd97c6f61b2..9fa043d0bab5 100644 --- a/packages/bigframes/third_party/bigframes_vendored/ibis/expr/operations/ai_ops.py +++ b/packages/bigframes/third_party/bigframes_vendored/ibis/expr/operations/ai_ops.py @@ -39,7 +39,7 @@ def dtype(self) -> dt.Struct: pyarrow_output_type = pa.struct( ( *output_pa_fields, - pa.field("full_resposne", pa.string()), + pa.field("full_response", pa.string()), pa.field("status", pa.string()), ) ) @@ -62,7 +62,7 @@ class AIGenerateBool(Value): @attribute def dtype(self) -> dt.Struct: return dt.Struct.from_tuples( - (("result", dt.bool), ("full_resposne", dt.string), ("status", dt.string)) + (("result", dt.bool), ("full_response", dt.string), ("status", dt.string)) ) @@ -81,7 +81,7 @@ class AIGenerateInt(Value): @attribute def dtype(self) -> dt.Struct: return dt.Struct.from_tuples( - (("result", dt.int64), ("full_resposne", dt.string), ("status", dt.string)) + (("result", dt.int64), ("full_response", dt.string), ("status", dt.string)) ) @@ -102,7 +102,7 @@ def dtype(self) -> dt.Struct: return dt.Struct.from_tuples( ( ("result", dt.float64), - ("full_resposne", dt.string), + ("full_response", dt.string), ("status", dt.string), ) )