feat(genai): Add new samples for batch predict (#13120)

* feat: batch predict with GCS

* feat: batch predict with BQ

* update comments

* bump the sdk version

* Update genai/batch_prediction/submit_with_gcs.py

Co-authored-by: code-review-assist[bot] <182814678+code-review-assist[bot]@users.noreply.github.com>

* fix suggestion

---------

Co-authored-by: code-review-assist[bot] <182814678+code-review-assist[bot]@users.noreply.github.com>

Author: Valeriy-Burlaka

genai/batch_prediction/requirements-test.txt

@@ -0,0 +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_prediction/requirements.txt

@@ -0,0 +1 @@
+google-genai==0.8.0

genai/batch_prediction/submit_with_bq.py

@@ -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_gcs.py

@@ -0,0 +1,66 @@
+# 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_gcs]
+    import time
+
+    from google import genai
+
+    client = genai.Client()
+    # TODO(developer): Update and un-comment below line
+    # output_uri = "gs://your-bucket/your-prefix"
+
+    # See the documentation: https://googleapis.github.io/python-genai/genai.html#genai.batches.Batches.create
+    job = client.batches.create(
+        model="gemini-1.5-pro-002",
+        # Source link: https://storage.cloud.google.com/cloud-samples-data/batch/prompt_for_batch_gemini_predict.jsonl
+        src="gs://cloud-samples-data/batch/prompt_for_batch_gemini_predict.jsonl",
+        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_gcs]
+    return job.state
+
+
+if __name__ == "__main__":
+    generate_content(
+        output_uri="gs://your-bucket/your-prefix"
+    )

genai/batch_prediction/test_batch_predict.py

@@ -0,0 +1,66 @@
+# 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.
+
+from datetime import datetime as dt
+import os
+
+from google.cloud import bigquery, storage
+from google.genai.types import JobState
+
+import pytest
+
+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"
+
+
+@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()}"
+
+    yield f"gs://{GCS_OUTPUT_BUCKET}/{prefix}"
+
+    storage_client = storage.Client()
+    bucket = storage_client.get_bucket(GCS_OUTPUT_BUCKET)
+    blobs = bucket.list_blobs(prefix=prefix)
+    for blob in blobs:
+        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:
+    response = submit_with_gcs.generate_content(output_uri=gcs_output_uri)
+    assert response == JobState.JOB_STATE_SUCCEEDED

genai/controlled_generation/ctrlgen_with_class_schema.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