generate/response -> create/job

GoogleCloudPlatform · Jan 31, 2025 · 4769391 · 4769391
1 parent bc39bf4
commit 4769391
Show file tree

Hide file tree

Showing 2 changed files with 7 additions and 7 deletions.
diff --git a/genai/batch_predict/batch_prediction_with_gcs.py b/genai/batch_predict/batch_prediction_with_gcs.py
@@ -15,7 +15,7 @@
 from google.genai.types import BatchJob
 
 
-def generate_content(output_uri: str) -> BatchJob:
+def create_job(output_uri: str) -> BatchJob:
     # [START googlegenaisdk_batch_prediction_with_gcs]
     import time
 
@@ -61,6 +61,6 @@ def generate_content(output_uri: str) -> BatchJob:
 
 
 if __name__ == "__main__":
-    generate_content(
+    create_job(
         output_uri="gs://gemini-batch-prediction-results/..."
     )
diff --git a/genai/batch_predict/test_batch_predict.py b/genai/batch_predict/test_batch_predict.py
@@ -44,8 +44,8 @@ def gcs_output_uri():
 
 
 def test_batch_prediction_with_gcs(gcs_output_uri) -> None:
-    response = batch_prediction_with_gcs.generate_content(output_uri=gcs_output_uri)
-    assert response
-    assert response.state == "JOB_STATE_SUCCEEDED"
-    assert response.dest.gcs_uri == gcs_output_uri
-    assert response.state == JobState.JOB_STATE_SUCCEEDED
+    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