feat(genai): Add Provisioned Throughput samples for Gen AI SDK and Vertex AI SDK (#13192)

holtskinner · web-flow · commit 1d4f38a87591 · 2025-02-27T18:20:15.000+01:00
* docs(genai): Add Provisioned Throughput sample for Gen AI SDK

* Add Provisioned Throughput sample for Vertex AI SDK

* Update based on comments
diff --git a/genai/provisioned_throughput/noxfile_config.py b/genai/provisioned_throughput/noxfile_config.py
@@ -0,0 +1,42 @@
+# Copyright 2021 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
+#
+#      http://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.
+
+# Default TEST_CONFIG_OVERRIDE for python repos.
+
+# You can copy this file into your directory, then it will be imported from
+# the noxfile.py.
+
+# The source of truth:
+# https://github.com/GoogleCloudPlatform/python-docs-samples/blob/main/noxfile_config.py
+
+TEST_CONFIG_OVERRIDE = {
+    # You can opt out from the test for specific Python versions.
+    "ignored_versions": ["2.7", "3.7", "3.8", "3.10", "3.11", "3.13"],
+    # Old samples are opted out of enforcing Python type hints
+    # All new samples should feature them
+    "enforce_type_hints": True,
+    # An envvar key for determining the project id to use. Change it
+    # to 'BUILD_SPECIFIC_GCLOUD_PROJECT' if you want to opt in using a
+    # build specific Cloud project. You can also use your own string
+    # to use your own Cloud project.
+    "gcloud_project_env": "GOOGLE_CLOUD_PROJECT",
+    # 'gcloud_project_env': 'BUILD_SPECIFIC_GCLOUD_PROJECT',
+    # If you need to use a specific version of pip,
+    # change pip_version_override to the string representation
+    # of the version number, for example, "20.2.4"
+    "pip_version_override": None,
+    # A dictionary you want to inject into your test. Don't put any
+    # secrets here. These values will override predefined values.
+    "envs": {},
+}
diff --git a/genai/provisioned_throughput/provisionedthroughput_with_txt.py b/genai/provisioned_throughput/provisionedthroughput_with_txt.py
@@ -0,0 +1,50 @@
+# 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() -> str:
+    # [START googlegenaisdk_provisionedthroughput_with_txt]
+
+    from google import genai
+    from google.genai.types import HttpOptions
+
+    client = genai.Client(
+        http_options=HttpOptions(
+            api_version="v1",
+            headers={
+                # Options:
+                # - "dedicated": Use Provisioned Throughput
+                # - "shared": Use pay-as-you-go
+                # https://cloud.google.com/vertex-ai/generative-ai/docs/use-provisioned-throughput
+                "X-Vertex-AI-LLM-Request-Type": "shared"
+            },
+        )
+    )
+    response = client.models.generate_content(
+        model="gemini-2.0-flash-001",
+        contents="How does AI work?",
+    )
+    print(response.text)
+
+    # Example response:
+    # Okay, let's break down how AI works. It's a broad field, so I'll focus on the ...
+    #
+    # Here's a simplified overview:
+    # ...
+    # [END googlegenaisdk_provisionedthroughput_with_txt]
+    return response.text
+
+
+if __name__ == "__main__":
+    generate_content()
diff --git a/genai/provisioned_throughput/requirements-test.txt b/genai/provisioned_throughput/requirements-test.txt
@@ -0,0 +1,2 @@
+google-api-core==2.24.0
+pytest==8.2.0
diff --git a/genai/provisioned_throughput/requirements.txt b/genai/provisioned_throughput/requirements.txt
@@ -0,0 +1 @@
+google-genai==1.3.0
diff --git a/genai/provisioned_throughput/test_provisioned_throughput_examples.py b/genai/provisioned_throughput/test_provisioned_throughput_examples.py
@@ -0,0 +1,31 @@
+# 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.
+
+#
+# Using Google Cloud Vertex AI to test the code samples.
+#
+
+import os
+
+import provisionedthroughput_with_txt
+
+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"
+
+
+def test_provisionedthroughput_with_txt() -> None:
+    response = provisionedthroughput_with_txt.generate_content()
+    assert response
diff --git a/generative_ai/provisioned_throughput/noxfile_config.py b/generative_ai/provisioned_throughput/noxfile_config.py
@@ -0,0 +1,42 @@
+# Copyright 2021 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
+#
+#      http://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.
+
+# Default TEST_CONFIG_OVERRIDE for python repos.
+
+# You can copy this file into your directory, then it will be imported from
+# the noxfile.py.
+
+# The source of truth:
+# https://github.com/GoogleCloudPlatform/python-docs-samples/blob/main/noxfile_config.py
+
+TEST_CONFIG_OVERRIDE = {
+    # You can opt out from the test for specific Python versions.
+    "ignored_versions": ["2.7", "3.7", "3.8", "3.10", "3.11", "3.13"],
+    # Old samples are opted out of enforcing Python type hints
+    # All new samples should feature them
+    "enforce_type_hints": True,
+    # An envvar key for determining the project id to use. Change it
+    # to 'BUILD_SPECIFIC_GCLOUD_PROJECT' if you want to opt in using a
+    # build specific Cloud project. You can also use your own string
+    # to use your own Cloud project.
+    "gcloud_project_env": "GOOGLE_CLOUD_PROJECT",
+    # 'gcloud_project_env': 'BUILD_SPECIFIC_GCLOUD_PROJECT',
+    # If you need to use a specific version of pip,
+    # change pip_version_override to the string representation
+    # of the version number, for example, "20.2.4"
+    "pip_version_override": None,
+    # A dictionary you want to inject into your test. Don't put any
+    # secrets here. These values will override predefined values.
+    "envs": {},
+}
diff --git a/generative_ai/provisioned_throughput/provisioned_throughput_with_txt.py b/generative_ai/provisioned_throughput/provisioned_throughput_with_txt.py
@@ -0,0 +1,55 @@
+# 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.
+import os
+
+PROJECT_ID = os.getenv("GOOGLE_CLOUD_PROJECT")
+
+
+def generate_content() -> str:
+    # [START generativeaionvertexai_provisioned_throughput_with_txt]
+    import vertexai
+    from vertexai.generative_models import GenerativeModel
+
+    # TODO(developer): Update and un-comment below line
+    # PROJECT_ID = "your-project-id"
+    vertexai.init(
+        project=PROJECT_ID,
+        location="us-central1",
+        # Options:
+        # - "dedicated": Use Provisioned Throughput
+        # - "shared": Use pay-as-you-go
+        # https://cloud.google.com/vertex-ai/generative-ai/docs/use-provisioned-throughput
+        request_metadata=[("x-vertex-ai-llm-request-type", "shared")],
+    )
+
+    model = GenerativeModel("gemini-1.5-flash-002")
+
+    response = model.generate_content(
+        "What's a good name for a flower shop that specializes in selling bouquets of dried flowers?"
+    )
+
+    print(response.text)
+    # Example response:
+    # **Emphasizing the Dried Aspect:**
+    # * Everlasting Blooms
+    # * Dried & Delightful
+    # * The Petal Preserve
+    # ...
+
+    # [END generativeaionvertexai_provisioned_throughput_with_txt]
+    return response.text
+
+
+if __name__ == "__main__":
+    generate_content()
diff --git a/generative_ai/provisioned_throughput/requirements-test.txt b/generative_ai/provisioned_throughput/requirements-test.txt
@@ -0,0 +1,4 @@
+backoff==2.2.1
+google-api-core==2.19.0
+pytest==8.2.0
+pytest-asyncio==0.23.6
diff --git a/generative_ai/provisioned_throughput/requirements.txt b/generative_ai/provisioned_throughput/requirements.txt
@@ -0,0 +1,2 @@
+google-cloud-aiplatform==1.82.0
+google-auth==2.38.0
diff --git a/generative_ai/provisioned_throughput/test_provisioned_throughput_examples.py b/generative_ai/provisioned_throughput/test_provisioned_throughput_examples.py
@@ -0,0 +1,21 @@
+# 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.
+
+
+import provisioned_throughput_with_txt
+
+
+def test_provisioned_throughput_with_txt() -> None:
+    response = provisioned_throughput_with_txt.generate_content()
+    assert response

Original file line number	Diff line number	Diff line change
`@@ -0,0 +1,2 @@`
	`1`	`+google-api-core==2.24.0`
	`2`	`+pytest==8.2.0`
Original file line number	Diff line number	Diff line change
`@@ -0,0 +1,2 @@`
	`1`	`+google-cloud-aiplatform==1.82.0`
	`2`	`+google-auth==2.38.0`