feat: support string literal inputs for AI functions

Author: sycai

bigframes/bigquery/_operations/ai.py

@@ -24,10 +24,11 @@
 import pandas as pd
 
 from bigframes import clients, dtypes, series, session
-from bigframes.core import convert, log_adapter
+from bigframes.core import convert, global_session, log_adapter
 from bigframes.operations import ai_ops, output_schemas
 
 PROMPT_TYPE = Union[
+    str,
     series.Series,
     pd.Series,
     List[Union[str, series.Series, pd.Series]],
@@ -514,9 +515,14 @@ def _separate_context_and_series(
     Input: ("str1", series1, "str2", "str3", series2)
     Output: ["str1", None, "str2", "str3", None], [series1, series2]
     """
-    if not isinstance(prompt, (list, tuple, series.Series)):
+    if not isinstance(prompt, (str, list, tuple, series.Series)):
         raise ValueError(f"Unsupported prompt type: {type(prompt)}")
 
+    if isinstance(prompt, str):
+        return [None], [
+            series.Series([prompt], session=global_session.get_global_session())
+        ]
+
     if isinstance(prompt, series.Series):
         if prompt.dtype == dtypes.OBJ_REF_DTYPE:
             # Multi-model support

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 unittest import mock
+
 from packaging import version
 import pandas as pd
 import pyarrow as pa
@@ -42,6 +44,27 @@ def test_ai_function_pandas_input(session):
     )
 
 
+def test_ai_function_string_input(session):
+    with mock.patch(
+        "bigframes.core.global_session.get_global_session"
+    ) as mock_get_session:
+        mock_get_session.return_value = session
+        prompt = "Is apple a fruit?"
+
+        result = bbq.ai.generate_bool(prompt, endpoint="gemini-2.5-flash")
+
+        assert _contains_no_nulls(result)
+        assert result.dtype == pd.ArrowDtype(
+            pa.struct(
+                (
+                    pa.field("result", pa.bool_()),
+                    pa.field("full_response", dtypes.JSON_ARROW_TYPE),
+                    pa.field("status", pa.string()),
+                )
+            )
+        )
+
+
 def test_ai_function_compile_model_params(session):
     if version.Version(sqlglot.__version__) < version.Version("25.18.0"):
         pytest.skip(