expose ai module and rename the function

Author: sycai

bigframes/bigquery/__init__.py

@@ -18,7 +18,7 @@
 
 import sys
 
-from bigframes.bigquery._operations.ai import ai_generate_bool
+from bigframes.bigquery._operations import ai
 from bigframes.bigquery._operations.approx_agg import approx_top_count
 from bigframes.bigquery._operations.array import (
     array_agg,
@@ -58,8 +58,6 @@
 from bigframes.core import log_adapter
 
 _functions = [
-    # ai ops
-    ai_generate_bool,
     # approximate aggregate ops
     approx_top_count,
     # array ops
@@ -99,7 +97,7 @@
     struct,
 ]
 
-__all__ = [f.__name__ for f in _functions]
+__all__ = [f.__name__ for f in _functions] + ["ai"]
 
 _module = sys.modules[__name__]
 for f in _functions:

bigframes/bigquery/_operations/ai.py

@@ -18,10 +18,12 @@
 from typing import Any, List, Literal, Mapping, Tuple
 
 from bigframes import clients, dtypes, series
+from bigframes.core import log_adapter
 from bigframes.operations import ai_ops
 
 
-def ai_generate_bool(
+@log_adapter.method_logger(custom_base_name="bigquery_ai")
+def generate_bool(
     prompt: series.Series | List[str | series.Series] | Tuple[str | series.Series, ...],
     *,
     connection_id: str | None = None,

tests/system/large/bigquery/test_ai.py

@@ -23,7 +23,7 @@ def test_ai_generate_bool_multi_model(session):
         "gs://bigframes-dev-testing/a_multimodel/images/*", name="image"
     )
 
-    result = bbq.ai_generate_bool((df["image"], " contains an animal")).struct.field(
+    result = bbq.ai.generate_bool((df["image"], " contains an animal")).struct.field(
         "result"
     )
 

tests/system/small/bigquery/test_ai.py

@@ -25,7 +25,7 @@ def test_ai_generate_bool(session):
     prompt = (s1, " is a ", s2)
     model_params = {"generation_config": {"thinking_config": {"thinking_budget": 0}}}
 
-    result = bbq.ai_generate_bool(
+    result = bbq.ai.generate_bool(
         prompt, endpoint="gemini-2.5-flash", model_params=model_params
     ).struct.field("result")