feat: batch predict with BQ

Author: Valeriy-Burlaka

genai/batch_predict/requirements-test.txt …i/batch_prediction/requirements-test.txtgenai/batch_predict/requirements-test.txt renamed to genai/batch_prediction/requirements-test.txt genai/batch_predict/requirements-test.txt renamed to genai/batch_prediction/requirements-test.txt

@@ -1,3 +1,4 @@
 google-api-core==2.24.0
+google-cloud-bigquery==3.29.0
 google-cloud-storage==2.19.0
 pytest==8.2.0

genai/batch_predict/requirements.txt genai/batch_prediction/requirements.txtgenai/batch_predict/requirements.txt renamed to genai/batch_prediction/requirements.txt

@@ -0,0 +1,64 @@
+# Copyright 2025 Google LLC
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#    https://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+
+def generate_content(output_uri: str) -> str:
+    # [START googlegenaisdk_batchpredict_with_bq]
+    import time
+
+    from google import genai
+
+    client = genai.Client()
+    # TODO(developer): Update and un-comment below line
+    # output_uri = f"bq://your-project.your_dataset.your_table"
+
+    job = client.batches.create(
+        model="gemini-1.5-pro-002",
+        src="bq://storage-samples.generative_ai.batch_requests_for_multimodal_input",
+        config={
+            "dest": output_uri
+        }
+    )
+    print(f"Job name: {job.name}")
+    print(f"Job state: {job.state}")
+    # Example response:
+    # Job name: projects/%PROJECT_ID%/locations/us-central1/batchPredictionJobs/9876453210000000000
+    # Job state: JOB_STATE_PENDING
+
+    # See the documentation: https://googleapis.github.io/python-genai/genai.html#genai.types.BatchJob
+    completed_states = [
+        "JOB_STATE_SUCCEEDED",
+        "JOB_STATE_FAILED",
+        "JOB_STATE_CANCELLED",
+        "JOB_STATE_PAUSED",
+    ]
+    while job.state not in completed_states:
+        time.sleep(30)
+        job = client.batches.get(name=job.name)
+        print(f"Job state: {job.state}")
+    # Example response:
+    # Job state: JOB_STATE_PENDING
+    # Job state: JOB_STATE_RUNNING
+    # Job state: JOB_STATE_RUNNING
+    # ...
+    # Job state: JOB_STATE_SUCCEEDED
+
+    # [END googlegenaisdk_batchpredict_with_bq]
+    return job.state
+
+
+if __name__ == "__main__":
+    generate_content(
+        output_uri="bq://your-project.your_dataset.your_table"
+    )

genai/batch_prediction/submit_with_bq.py

@@ -12,11 +12,9 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-from google.genai.types import BatchJob
 
-
-def create_job(output_uri: str) -> BatchJob:
-    # [START googlegenaisdk_batch_prediction_with_gcs]
+def generate_content(output_uri: str) -> str:
+    # [START googlegenaisdk_batchpredict_with_gcs]
     import time
 
     from google import genai
@@ -26,7 +24,7 @@ def create_job(output_uri: str) -> BatchJob:
     # output_uri = "gs://your-bucket/your-prefix/..."
 
     job = client.batches.create(
-        model="gemini-2.0-flash-001",
+        model="gemini-1.5-pro-002",
         src="gs://cloud-samples-data/batch/prompt_for_batch_gemini_predict.jsonl",
         config={
             "dest": output_uri
@@ -56,11 +54,11 @@ def create_job(output_uri: str) -> BatchJob:
     # ...
     # Job state: JOB_STATE_SUCCEEDED
 
-    # [END googlegenaisdk_batch_prediction_with_gcs]
-    return job
+    # [END googlegenaisdk_batchpredict_with_gcs]
+    return job.state
 
 
 if __name__ == "__main__":
-    create_job(
+    generate_content(
         output_uri="gs://your-bucket/your-prefix/..."
     )

…tch_predict/batch_prediction_with_gcs.py genai/batch_prediction/submit_with_gcs.pygenai/batch_predict/batch_prediction_with_gcs.py renamed to genai/batch_prediction/submit_with_gcs.py genai/batch_predict/batch_prediction_with_gcs.py renamed to genai/batch_prediction/submit_with_gcs.py

@@ -12,27 +12,40 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-from datetime import datetime as dt, UTC
+from datetime import datetime as dt
 import os
 
-from google.cloud import storage
+from google.cloud import bigquery, storage
 from google.genai.types import JobState
+
 import pytest
 
-import batch_prediction_with_gcs
+import submit_with_bq
+import submit_with_gcs
 
 
 os.environ["GOOGLE_GENAI_USE_VERTEXAI"] = "True"
 os.environ["GOOGLE_CLOUD_LOCATION"] = "us-central1"
 # The project name is included in the CICD pipeline
 # os.environ['GOOGLE_CLOUD_PROJECT'] = "add-your-project-name"
+BQ_OUTPUT_DATASET = f"{os.environ['GOOGLE_CLOUD_PROJECT']}.gen_ai_batch_prediction"
 GCS_OUTPUT_BUCKET = "python-docs-samples-tests"
-GCS_OUTPUT_BUCKET = "gemini-batch-prediction-results"
+
+
+@pytest.fixture(scope="session")
+def bq_output_uri():
+    table_name = f"text_output_{dt.now().strftime('%Y_%m_%d_T%H_%M_%S')}"
+    table_uri = f"{BQ_OUTPUT_DATASET}.{table_name}"
+
+    yield f"bq://{table_uri}"
+
+    bq_client = bigquery.Client()
+    bq_client.delete_table(table_uri, not_found_ok=True)
 
 
 @pytest.fixture(scope="session")
 def gcs_output_uri():
-    prefix = f"text_output/{dt.now(UTC)}"
+    prefix = f"text_output/{dt.now()}"
 
     yield f"gs://{GCS_OUTPUT_BUCKET}/{prefix}"
 
@@ -43,9 +56,11 @@ def gcs_output_uri():
         blob.delete()
 
 
+def test_batch_prediction_with_bq(bq_output_uri) -> None:
+    response = submit_with_bq.generate_content(output_uri=bq_output_uri)
+    assert response == JobState.JOB_STATE_SUCCEEDED
+
+
 def test_batch_prediction_with_gcs(gcs_output_uri) -> None:
-    job = batch_prediction_with_gcs.create_job(output_uri=gcs_output_uri)
-    assert job
-    assert job.state == "JOB_STATE_SUCCEEDED"
-    assert job.dest.gcs_uri == gcs_output_uri
-    assert job.state == JobState.JOB_STATE_SUCCEEDED
+    response = submit_with_gcs.generate_content(output_uri=gcs_output_uri)
+    assert response == JobState.JOB_STATE_SUCCEEDED

genai/batch_predict/test_batch_predict.py …i/batch_prediction/test_batch_predict.pygenai/batch_predict/test_batch_predict.py renamed to genai/batch_prediction/test_batch_predict.py genai/batch_predict/test_batch_predict.py renamed to genai/batch_prediction/test_batch_predict.py

@@ -16,7 +16,7 @@
 def generate_content() -> str:
     # [START googlegenaisdk_ctrlgen_with_class_schema]
     from google import genai
-    from pydantic import BaseModel, TypeAdapter
+    from pydantic import BaseModel
 
     class Recipe(BaseModel):
         recipe_name: str