feat: GA Context Cache Python SDK

FUTURE_COPYBARA_INTEGRATE_REVIEW=#4861 from googleapis:release-please--branches--main 039f2cb
PiperOrigin-RevId: 718130866

Author: vertex-sdk-bot

google/cloud/aiplatform/compat/__init__.py

@@ -181,8 +181,7 @@
         services.featurestore_online_serving_service_client_v1
     )
     services.featurestore_service_client = services.featurestore_service_client_v1
-    # TODO(b/342585299): Temporary code. Switch to v1 once v1 is available.
-    services.gen_ai_cache_service_client = services.gen_ai_cache_service_client_v1beta1
+    services.gen_ai_cache_service_client = services.gen_ai_cache_service_client_v1
     services.job_service_client = services.job_service_client_v1
     services.model_garden_service_client = services.model_garden_service_client_v1
     services.model_service_client = services.model_service_client_v1
@@ -203,8 +202,7 @@
     types.annotation_spec = types.annotation_spec_v1
     types.artifact = types.artifact_v1
     types.batch_prediction_job = types.batch_prediction_job_v1
-    # TODO(b/342585299): Temporary code. Switch to v1 once v1 is available.
-    types.cached_content = types.cached_content_v1beta1
+    types.cached_content = types.cached_content_v1
     types.completion_stats = types.completion_stats_v1
     types.context = types.context_v1
     types.custom_job = types.custom_job_v1

google/cloud/aiplatform/compat/services/__init__.py

@@ -137,6 +137,9 @@
 from google.cloud.aiplatform_v1.services.featurestore_service import (
     client as featurestore_service_client_v1,
 )
+from google.cloud.aiplatform_v1.services.gen_ai_cache_service import (
+    client as gen_ai_cache_service_client_v1,
+)
 from google.cloud.aiplatform_v1.services.index_service import (
     client as index_service_client_v1,
 )

google/cloud/aiplatform/compat/types/__init__.py

@@ -118,6 +118,7 @@
     annotation_spec as annotation_spec_v1,
     artifact as artifact_v1,
     batch_prediction_job as batch_prediction_job_v1,
+    cached_content as cached_content_v1,
     completion_stats as completion_stats_v1,
     context as context_v1,
     custom_job as custom_job_v1,

google/cloud/aiplatform/utils/__init__.py

@@ -77,6 +77,7 @@
     feature_registry_service_client_v1,
     featurestore_online_serving_service_client_v1,
     featurestore_service_client_v1,
+    gen_ai_cache_service_client_v1,
     index_service_client_v1,
     index_endpoint_service_client_v1,
     job_service_client_v1,
@@ -805,8 +806,7 @@ class GenAiCacheServiceClientWithOverride(ClientWithOverride):
     _version_map = (
         (
             compat.V1,
-            # TODO(b/342585299): Temporary code. Switch to v1 once v1 is available.
-            gen_ai_cache_service_client_v1beta1.GenAiCacheServiceClient,
+            gen_ai_cache_service_client_v1.GenAiCacheServiceClient,
         ),
         (
             compat.V1BETA1,

tests/unit/vertexai/test_caching.py

@@ -22,22 +22,29 @@
 import json
 import mock
 import pytest
-from vertexai.preview import caching
+from vertexai.caching import _caching
 from google.cloud.aiplatform import initializer
 import vertexai
-from google.cloud.aiplatform_v1beta1.types.cached_content import (
+from google.cloud.aiplatform_v1.types.cached_content import (
     CachedContent as GapicCachedContent,
 )
-from google.cloud.aiplatform_v1beta1.types.content import (
+from google.cloud.aiplatform_v1.types.content import (
     Content as GapicContent,
     Part as GapicPart,
 )
-from google.cloud.aiplatform_v1beta1.types.tool import (
+from google.cloud.aiplatform_v1.types.tool import (
     ToolConfig as GapicToolConfig,
 )
-from google.cloud.aiplatform_v1beta1.services import (
+from google.cloud.aiplatform_v1.services import (
     gen_ai_cache_service,
 )
+from vertexai.generative_models._generative_models import (
+    Content,
+    PartsType,
+    Tool,
+    ToolConfig,
+    ContentsType,
+)
 
 
 _TEST_PROJECT = "test-project"
@@ -141,7 +148,7 @@ def list_cached_contents(self, request):
 
 @pytest.mark.usefixtures("google_auth_mock")
 class TestCaching:
-    """Unit tests for caching.CachedContent."""
+    """Unit tests for _caching.CachedContent."""
 
     def setup_method(self):
         vertexai.init(
@@ -156,7 +163,7 @@ def test_constructor_with_full_resource_name(self, mock_get_cached_content):
         full_resource_name = (
             "projects/123/locations/europe-west1/cachedContents/contents-id"
         )
-        cache = caching.CachedContent(
+        cache = _caching.CachedContent(
             cached_content_name=full_resource_name,
         )
 
@@ -166,7 +173,7 @@ def test_constructor_with_full_resource_name(self, mock_get_cached_content):
     def test_constructor_with_only_content_id(self, mock_get_cached_content):
         partial_resource_name = "contents-id"
 
-        cache = caching.CachedContent(
+        cache = _caching.CachedContent(
             cached_content_name=partial_resource_name,
         )
 
@@ -179,7 +186,7 @@ def test_constructor_with_only_content_id(self, mock_get_cached_content):
     def test_get_with_content_id(self, mock_get_cached_content):
         partial_resource_name = "contents-id"
 
-        cache = caching.CachedContent.get(
+        cache = _caching.CachedContent.get(
             cached_content_name=partial_resource_name,
         )
 
@@ -192,14 +199,14 @@ def test_get_with_content_id(self, mock_get_cached_content):
     def test_create_with_real_payload(
         self, mock_create_cached_content, mock_get_cached_content
     ):
-        cache = caching.CachedContent.create(
+        cache = _caching.CachedContent.create(
             model_name="model-name",
-            system_instruction=GapicContent(
-                role="system", parts=[GapicPart(text="system instruction")]
-            ),
+            system_instruction=Content._from_gapic(
+            GapicContent(
+                role="system", parts=[GapicPart(text="system instruction")])),
             tools=[],
             tool_config=GapicToolConfig(),
-            contents=[GapicContent(role="user")],
+            contents=[Content._from_gapic(GapicContent(role="user"))],
             ttl=datetime.timedelta(days=1),
             display_name=_TEST_DISPLAY_NAME,
         )
@@ -219,7 +226,7 @@ def test_create_with_real_payload(
     def test_create_with_real_payload_and_wrapped_type(
         self, mock_create_cached_content, mock_get_cached_content
     ):
-        cache = caching.CachedContent.create(
+        cache = _caching.CachedContent.create(
             model_name="model-name",
             system_instruction="Please answer my questions with cool",
             tools=[],
@@ -239,15 +246,15 @@ def test_create_with_real_payload_and_wrapped_type(
         assert cache.display_name == _TEST_DISPLAY_NAME
 
     def test_list(self, mock_list_cached_contents):
-        cached_contents = caching.CachedContent.list()
+        cached_contents = _caching.CachedContent.list()
         for i, cached_content in enumerate(cached_contents):
             assert cached_content.name == f"cached_content{i + 1}_from_list_request"
             assert cached_content.model_name == f"model-name{i + 1}"
 
     def test_print_a_cached_content(
         self, mock_create_cached_content, mock_get_cached_content
     ):
-        cached_content = caching.CachedContent.create(
+        cached_content = _caching.CachedContent.create(
             model_name="model-name",
             system_instruction="Please answer my questions with cool",
             tools=[],

tests/unit/vertexai/test_generative_models.py

@@ -39,14 +39,14 @@
     gapic_content_types,
     gapic_tool_types,
 )
-from google.cloud.aiplatform_v1beta1.types.cached_content import (
+from google.cloud.aiplatform_v1.types.cached_content import (
     CachedContent as GapicCachedContent,
 )
-from google.cloud.aiplatform_v1beta1.services import (
+from google.cloud.aiplatform_v1.services import (
     gen_ai_cache_service,
 )
 from vertexai.generative_models import _function_calling_utils
-from vertexai.preview import caching
+from vertexai.caching import _caching
 
 
 _TEST_PROJECT = "test-project"
@@ -655,11 +655,11 @@ def test_generative_model_from_cached_content(
         project_location_prefix = (
             f"projects/{_TEST_PROJECT}/locations/{_TEST_LOCATION}/"
         )
-        cached_content = caching.CachedContent(
+        cached_content = _caching.CachedContent(
             "cached-content-id-in-from-cached-content-test"
         )
 
-        model = preview_generative_models.GenerativeModel.from_cached_content(
+        model = generative_models.GenerativeModel.from_cached_content(
             cached_content=cached_content
         )
 
@@ -690,7 +690,7 @@ def test_generative_model_from_cached_content_with_resource_name(
             f"projects/{_TEST_PROJECT}/locations/{_TEST_LOCATION}/"
         )
 
-        model = preview_generative_models.GenerativeModel.from_cached_content(
+        model = generative_models.GenerativeModel.from_cached_content(
             cached_content="cached-content-id-in-from-cached-content-test"
         )
 
@@ -848,7 +848,7 @@ def test_generate_content(
         assert response5.text
 
     @mock.patch.object(
-        target=prediction_service.PredictionServiceClient,
+        target=prediction_service_v1.PredictionServiceClient,
         attribute="generate_content",
         new=lambda self, request: gapic_prediction_service_types.GenerateContentResponse(
             candidates=[
@@ -870,11 +870,11 @@ def test_generate_content_with_cached_content(
         self,
         mock_get_cached_content_fixture,
     ):
-        cached_content = caching.CachedContent(
+        cached_content = _caching.CachedContent(
             "cached-content-id-in-from-cached-content-test"
         )
 
-        model = preview_generative_models.GenerativeModel.from_cached_content(
+        model = generative_models.GenerativeModel.from_cached_content(
             cached_content=cached_content
         )
 

vertexai/caching/__init__.py

@@ -0,0 +1,25 @@
+# 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
+#
+#     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.
+#
+"""Classes for working with the Gemini models."""
+
+# We just want to re-export certain classes
+# pylint: disable=g-multiple-import,g-importing-member
+from vertexai.caching._caching import (
+    CachedContent,
+)
+
+__all__ = [
+    "CachedContent",
+]

vertexai/caching/_caching.py

@@ -22,16 +22,16 @@
 from google.cloud.aiplatform import initializer as aiplatform_initializer
 from google.cloud.aiplatform import utils as aiplatform_utils
 from google.cloud.aiplatform.compat.types import (
-    cached_content_v1beta1 as gca_cached_content,
+    cached_content_v1 as gca_cached_content,
 )
-from google.cloud.aiplatform_v1beta1.services import gen_ai_cache_service
-from google.cloud.aiplatform_v1beta1.types.cached_content import (
+from google.cloud.aiplatform_v1.services import gen_ai_cache_service
+from google.cloud.aiplatform_v1.types.cached_content import (
     CachedContent as GapicCachedContent,
 )
-from google.cloud.aiplatform_v1beta1.types import (
+from google.cloud.aiplatform_v1.types import (
     content as gapic_content_types,
 )
-from google.cloud.aiplatform_v1beta1.types.gen_ai_cache_service import (
+from google.cloud.aiplatform_v1.types.gen_ai_cache_service import (
     CreateCachedContentRequest,
     GetCachedContentRequest,
     UpdateCachedContentRequest,