chore: GenAI - Explicit Caching - Move caching back to public preview

PiperOrigin-RevId: 642706611

tests/system/vertexai/test_generative_models.py

@@ -29,7 +29,7 @@
 from vertexai.preview import (
     generative_models as preview_generative_models,
 )
-from vertexai._caching import _caching as caching
+from vertexai.preview import caching
 
 GEMINI_MODEL_NAME = "gemini-1.0-pro-002"
 GEMINI_VISION_MODEL_NAME = "gemini-1.0-pro-vision"
@@ -100,6 +100,44 @@ def setup_method(self):
             credentials=credentials,
         )
 
+    def test_generate_content_with_cached_content_from_text(self):
+        cached_content = caching.CachedContent.create(
+            model_name=GEMINI_15_0514_MODEL_NAME,
+            system_instruction="Please answer all the questions like a pirate.",
+            contents=[
+                Content.from_dict(
+                    {
+                        "role": "user",
+                        "parts": [
+                            {
+                                "file_data": {
+                                    "mime_type": "application/pdf",
+                                    "file_uri": "gs://ucaip-samples-us-central1/sdk_system_test_resources/megatro-llm.pdf",
+                                }
+                            }
+                            for _ in range(10)
+                        ]
+                        + [
+                            {"text": "Please try to summarize the previous contents."},
+                        ],
+                    }
+                )
+            ],
+        )
+
+        model = generative_models.GenerativeModel.from_cached_content(
+            cached_content=cached_content
+        )
+
+        response = model.generate_content(
+            "Why is sky blue?",
+            generation_config=generative_models.GenerationConfig(temperature=0),
+        )
+        try:
+            assert response.text
+        finally:
+            cached_content.delete()
+
     def test_generate_content_from_text(self):
         model = generative_models.GenerativeModel(GEMINI_MODEL_NAME)
         response = model.generate_content(
@@ -441,41 +479,3 @@ def test_additional_request_metadata(self):
             generation_config=generative_models.GenerationConfig(temperature=0),
         )
         assert response
-
-    def test_generate_content_with_cached_content_from_text(self):
-        cached_content = caching.CachedContent.create(
-            model_name=GEMINI_15_0514_MODEL_NAME,
-            system_instruction="Please answer all the questions like a pirate.",
-            contents=[
-                Content.from_dict(
-                    {
-                        "role": "user",
-                        "parts": [
-                            {
-                                "file_data": {
-                                    "mime_type": "application/pdf",
-                                    "file_uri": "gs://ucaip-samples-us-central1/sdk_system_test_resources/megatro-llm.pdf",
-                                }
-                            }
-                            for _ in range(10)
-                        ]
-                        + [
-                            {"text": "Please try to summarize the previous contents."},
-                        ],
-                    }
-                )
-            ],
-        )
-
-        model = preview_generative_models.GenerativeModel._from_cached_content(
-            cached_content=cached_content
-        )
-
-        response = model.generate_content(
-            "Why is sky blue?",
-            generation_config=generative_models.GenerationConfig(temperature=0),
-        )
-        try:
-            assert response.text
-        finally:
-            cached_content.delete()

tests/unit/vertexai/test_caching.py

@@ -19,7 +19,7 @@
 import datetime
 import pytest
 import mock
-from vertexai._caching import _caching as caching
+from vertexai.preview import caching
 from google.cloud.aiplatform import initializer
 import vertexai
 from google.cloud.aiplatform_v1beta1.types.cached_content import (

tests/unit/vertexai/test_generative_models.py

@@ -41,7 +41,7 @@
     gen_ai_cache_service,
 )
 from vertexai.generative_models import _function_calling_utils
-from vertexai._caching import _caching as caching
+from vertexai.preview import caching
 
 
 _TEST_PROJECT = "test-project"
@@ -459,7 +459,7 @@ def test_generative_model_from_cached_content(
             "cached-content-id-in-from-cached-content-test"
         )
 
-        model = preview_generative_models.GenerativeModel._from_cached_content(
+        model = preview_generative_models.GenerativeModel.from_cached_content(
             cached_content=cached_content
         )
 
@@ -587,7 +587,7 @@ def test_generate_content_with_cached_content(
             "cached-content-id-in-from-cached-content-test"
         )
 
-        model = preview_generative_models.GenerativeModel._from_cached_content(
+        model = preview_generative_models.GenerativeModel.from_cached_content(
             cached_content=cached_content
         )
 

vertexai/generative_models/_generative_models.py

@@ -49,7 +49,7 @@
 import warnings
 
 if TYPE_CHECKING:
-    from vertexai._caching import _caching as caching
+    from vertexai.preview import caching
 
 try:
     from PIL import Image as PIL_Image  # pylint: disable=g-import-not-at-top
@@ -2601,7 +2601,7 @@ def start_chat(
         )
 
     @classmethod
-    def _from_cached_content(
+    def from_cached_content(
         cls,
         cached_content: "caching.CachedContent",
         *,

vertexai/preview/caching.py

@@ -0,0 +1,20 @@
+# Copyright 2024 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.
+#
+
+from vertexai.caching._caching import CachedContent
+
+__all__ = [
+    "CachedContent",
+]