Skip to content

Commit 56b25f2

Browse files
chore(genai): update Gen AI SDK v1.42.0 & code comments (#13595)
* chore(genai): update code comments for batch prediction samples. * chore(genai): update genai SDK version * Update genai/batch_prediction/get_batch_job.py Co-authored-by: gemini-code-assist[bot] <176961590+gemini-code-assist[bot]@users.noreply.github.com> * chore(genai): rollback changes to model_garden folder * chore(genai): rollback changes to model_garden folder * chore(genai): Update google-genai skd to v1.41.0 * chore(genai): Add missing input file * chore(genai): Rollback changes to Video Gen. - Video samples need model update - SDK update has minor breaking changes. * chore(genai): Update Live Samples testcases For Live Conversation, it is not necessary to have text responses. * chore(genai): Update Batch Prediction samples Tests to use Mocking. Using mocking for faster CICD responses. * chore(genai): lint error - missing line brake from test file. --------- Co-authored-by: gemini-code-assist[bot] <176961590+gemini-code-assist[bot]@users.noreply.github.com>
1 parent a9f2193 commit 56b25f2

Some content is hidden

Large Commits have some content hidden by default. Use the searchbox below for content that may be hidden.

43 files changed

+139
-114
lines changed

genai/batch_prediction/batchpredict_embeddings_with_gcs.py

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -34,7 +34,7 @@ def generate_content(output_uri: str) -> str:
3434
print(f"Job name: {job.name}")
3535
print(f"Job state: {job.state}")
3636
# Example response:
37-
# Job name: projects/%PROJECT_ID%/locations/us-central1/batchPredictionJobs/9876453210000000000
37+
# Job name: projects/.../locations/.../batchPredictionJobs/9876453210000000000
3838
# Job state: JOB_STATE_PENDING
3939

4040
# See the documentation: https://googleapis.github.io/python-genai/genai.html#genai.types.BatchJob

genai/batch_prediction/batchpredict_with_bq.py

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -35,7 +35,7 @@ def generate_content(output_uri: str) -> str:
3535
print(f"Job name: {job.name}")
3636
print(f"Job state: {job.state}")
3737
# Example response:
38-
# Job name: projects/%PROJECT_ID%/locations/us-central1/batchPredictionJobs/9876453210000000000
38+
# Job name: projects/.../locations/.../batchPredictionJobs/9876453210000000000
3939
# Job state: JOB_STATE_PENDING
4040

4141
# See the documentation: https://googleapis.github.io/python-genai/genai.html#genai.types.BatchJob

genai/batch_prediction/batchpredict_with_gcs.py

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -36,7 +36,7 @@ def generate_content(output_uri: str) -> str:
3636
print(f"Job name: {job.name}")
3737
print(f"Job state: {job.state}")
3838
# Example response:
39-
# Job name: projects/%PROJECT_ID%/locations/us-central1/batchPredictionJobs/9876453210000000000
39+
# Job name: projects/.../locations/.../batchPredictionJobs/9876453210000000000
4040
# Job state: JOB_STATE_PENDING
4141

4242
# See the documentation: https://googleapis.github.io/python-genai/genai.html#genai.types.BatchJob

genai/batch_prediction/get_batch_job.py

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -23,7 +23,7 @@ def get_batch_job(batch_job_name: str) -> types.BatchJob:
2323
client = genai.Client(http_options=HttpOptions(api_version="v1"))
2424

2525
# Get the batch job
26-
# Eg. batch_job_name = "projects/123456789012/locations/us-central1/batchPredictionJobs/1234567890123456789"
26+
# Eg. batch_job_name = "projects/123456789012/locations/.../batchPredictionJobs/1234567890123456789"
2727
batch_job = client.batches.get(name=batch_job_name)
2828

2929
print(f"Job state: {batch_job.state}")
Lines changed: 0 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -1,4 +1,2 @@
11
google-api-core==2.24.0
2-
google-cloud-bigquery==3.29.0
3-
google-cloud-storage==2.19.0
42
pytest==8.2.0
Lines changed: 3 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1 +1,3 @@
1-
google-genai==1.27.0
1+
google-cloud-bigquery==3.29.0
2+
google-cloud-storage==2.19.0
3+
google-genai==1.42.0

genai/batch_prediction/test_batch_prediction_examples.py

Lines changed: 83 additions & 46 deletions
Original file line numberDiff line numberDiff line change
@@ -11,87 +11,124 @@
1111
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
1212
# See the License for the specific language governing permissions and
1313
# limitations under the License.
14-
15-
#
16-
# Using Google Cloud Vertex AI to test the code samples.
17-
#
18-
from datetime import datetime as dt
19-
import os
20-
2114
from unittest.mock import MagicMock, patch
2215

23-
from google.cloud import bigquery, storage
2416
from google.genai import types
2517
from google.genai.types import JobState
26-
import pytest
2718

2819
import batchpredict_embeddings_with_gcs
2920
import batchpredict_with_bq
3021
import batchpredict_with_gcs
3122
import get_batch_job
3223

3324

34-
os.environ["GOOGLE_GENAI_USE_VERTEXAI"] = "True"
35-
os.environ["GOOGLE_CLOUD_LOCATION"] = "us-central1"
36-
# The project name is included in the CICD pipeline
37-
# os.environ['GOOGLE_CLOUD_PROJECT'] = "add-your-project-name"
38-
BQ_OUTPUT_DATASET = f"{os.environ['GOOGLE_CLOUD_PROJECT']}.gen_ai_batch_prediction"
39-
GCS_OUTPUT_BUCKET = "python-docs-samples-tests"
40-
41-
42-
@pytest.fixture(scope="session")
43-
def bq_output_uri() -> str:
44-
table_name = f"text_output_{dt.now().strftime('%Y_%m_%d_T%H_%M_%S')}"
45-
table_uri = f"{BQ_OUTPUT_DATASET}.{table_name}"
25+
@patch("google.genai.Client")
26+
@patch("time.sleep", return_value=None)
27+
def test_batch_prediction_embeddings_with_gcs(
28+
mock_sleep: MagicMock, mock_genai_client: MagicMock
29+
) -> None:
30+
# Mock the API response
31+
mock_batch_job_running = types.BatchJob(
32+
name="test-batch-job", state="JOB_STATE_RUNNING"
33+
)
34+
mock_batch_job_succeeded = types.BatchJob(
35+
name="test-batch-job", state="JOB_STATE_SUCCEEDED"
36+
)
4637

47-
yield f"bq://{table_uri}"
38+
mock_genai_client.return_value.batches.create.return_value = (
39+
mock_batch_job_running
40+
)
41+
mock_genai_client.return_value.batches.get.return_value = (
42+
mock_batch_job_succeeded
43+
)
4844

49-
bq_client = bigquery.Client()
50-
bq_client.delete_table(table_uri, not_found_ok=True)
45+
response = batchpredict_embeddings_with_gcs.generate_content(
46+
output_uri="gs://test-bucket/test-prefix"
47+
)
5148

49+
mock_genai_client.assert_called_once_with(
50+
http_options=types.HttpOptions(api_version="v1")
51+
)
52+
mock_genai_client.return_value.batches.create.assert_called_once()
53+
mock_genai_client.return_value.batches.get.assert_called_once()
54+
assert response == JobState.JOB_STATE_SUCCEEDED
5255

53-
@pytest.fixture(scope="session")
54-
def gcs_output_uri() -> str:
55-
prefix = f"text_output/{dt.now()}"
5656

57-
yield f"gs://{GCS_OUTPUT_BUCKET}/{prefix}"
57+
@patch("google.genai.Client")
58+
@patch("time.sleep", return_value=None)
59+
def test_batch_prediction_with_bq(
60+
mock_sleep: MagicMock, mock_genai_client: MagicMock
61+
) -> None:
62+
# Mock the API response
63+
mock_batch_job_running = types.BatchJob(
64+
name="test-batch-job", state="JOB_STATE_RUNNING"
65+
)
66+
mock_batch_job_succeeded = types.BatchJob(
67+
name="test-batch-job", state="JOB_STATE_SUCCEEDED"
68+
)
5869

59-
storage_client = storage.Client()
60-
bucket = storage_client.get_bucket(GCS_OUTPUT_BUCKET)
61-
blobs = bucket.list_blobs(prefix=prefix)
62-
for blob in blobs:
63-
blob.delete()
70+
mock_genai_client.return_value.batches.create.return_value = (
71+
mock_batch_job_running
72+
)
73+
mock_genai_client.return_value.batches.get.return_value = (
74+
mock_batch_job_succeeded
75+
)
6476

77+
response = batchpredict_with_bq.generate_content(
78+
output_uri="bq://test-project.test_dataset.test_table"
79+
)
6580

66-
def test_batch_prediction_embeddings_with_gcs(gcs_output_uri: str) -> None:
67-
response = batchpredict_embeddings_with_gcs.generate_content(
68-
output_uri=gcs_output_uri
81+
mock_genai_client.assert_called_once_with(
82+
http_options=types.HttpOptions(api_version="v1")
6983
)
84+
mock_genai_client.return_value.batches.create.assert_called_once()
85+
mock_genai_client.return_value.batches.get.assert_called_once()
7086
assert response == JobState.JOB_STATE_SUCCEEDED
7187

7288

73-
def test_batch_prediction_with_bq(bq_output_uri: str) -> None:
74-
response = batchpredict_with_bq.generate_content(output_uri=bq_output_uri)
75-
assert response == JobState.JOB_STATE_SUCCEEDED
89+
@patch("google.genai.Client")
90+
@patch("time.sleep", return_value=None)
91+
def test_batch_prediction_with_gcs(
92+
mock_sleep: MagicMock, mock_genai_client: MagicMock
93+
) -> None:
94+
# Mock the API response
95+
mock_batch_job_running = types.BatchJob(
96+
name="test-batch-job", state="JOB_STATE_RUNNING"
97+
)
98+
mock_batch_job_succeeded = types.BatchJob(
99+
name="test-batch-job", state="JOB_STATE_SUCCEEDED"
100+
)
76101

102+
mock_genai_client.return_value.batches.create.return_value = (
103+
mock_batch_job_running
104+
)
105+
mock_genai_client.return_value.batches.get.return_value = (
106+
mock_batch_job_succeeded
107+
)
108+
109+
response = batchpredict_with_gcs.generate_content(
110+
output_uri="gs://test-bucket/test-prefix"
111+
)
77112

78-
def test_batch_prediction_with_gcs(gcs_output_uri: str) -> None:
79-
response = batchpredict_with_gcs.generate_content(output_uri=gcs_output_uri)
113+
mock_genai_client.assert_called_once_with(
114+
http_options=types.HttpOptions(api_version="v1")
115+
)
116+
mock_genai_client.return_value.batches.create.assert_called_once()
117+
mock_genai_client.return_value.batches.get.assert_called_once()
80118
assert response == JobState.JOB_STATE_SUCCEEDED
81119

82120

83121
@patch("google.genai.Client")
84122
def test_get_batch_job(mock_genai_client: MagicMock) -> None:
85123
# Mock the API response
86-
mock_batch_job = types.BatchJob(
87-
name="test-batch-job",
88-
state="JOB_STATE_PENDING"
89-
)
124+
mock_batch_job = types.BatchJob(name="test-batch-job", state="JOB_STATE_PENDING")
90125

91126
mock_genai_client.return_value.batches.get.return_value = mock_batch_job
92127

93128
response = get_batch_job.get_batch_job("test-batch-job")
94129

95-
mock_genai_client.assert_called_once_with(http_options=types.HttpOptions(api_version="v1"))
130+
mock_genai_client.assert_called_once_with(
131+
http_options=types.HttpOptions(api_version="v1")
132+
)
96133
mock_genai_client.return_value.batches.get.assert_called_once()
97134
assert response == mock_batch_job
Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1,2 +1,2 @@
1-
google-genai==1.27.0
1+
google-genai==1.42.0
22
pillow==11.1.0

genai/content_cache/contentcache_create_with_txt_gcs_pdf.py

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -47,7 +47,7 @@ def create_content_cache() -> str:
4747
contents=contents,
4848
system_instruction=system_instruction,
4949
# (Optional) For enhanced security, the content cache can be encrypted using a Cloud KMS key
50-
# kms_key_name = "projects/.../locations/us-central1/keyRings/.../cryptoKeys/..."
50+
# kms_key_name = "projects/.../locations/.../keyRings/.../cryptoKeys/..."
5151
display_name="example-cache",
5252
ttl="86400s",
5353
),
@@ -56,7 +56,7 @@ def create_content_cache() -> str:
5656
print(content_cache.name)
5757
print(content_cache.usage_metadata)
5858
# Example response:
59-
# projects/111111111111/locations/us-central1/cachedContents/1111111111111111111
59+
# projects/111111111111/locations/.../cachedContents/1111111111111111111
6060
# CachedContentUsageMetadata(audio_duration_seconds=None, image_count=167,
6161
# text_count=153, total_token_count=43130, video_duration_seconds=None)
6262
# [END googlegenaisdk_contentcache_create_with_txt_gcs_pdf]

genai/content_cache/contentcache_delete.py

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -19,11 +19,11 @@ def delete_context_caches(cache_name: str) -> str:
1919

2020
client = genai.Client()
2121
# Delete content cache using name
22-
# E.g cache_name = 'projects/111111111111/locations/us-central1/cachedContents/1111111111111111111'
22+
# E.g cache_name = 'projects/111111111111/locations/.../cachedContents/1111111111111111111'
2323
client.caches.delete(name=cache_name)
2424
print("Deleted Cache", cache_name)
2525
# Example response
26-
# Deleted Cache projects/111111111111/locations/us-central1/cachedContents/1111111111111111111
26+
# Deleted Cache projects/111111111111/locations/.../cachedContents/1111111111111111111
2727
# [END googlegenaisdk_contentcache_delete]
2828
return cache_name
2929

0 commit comments

Comments
 (0)