feat: GenAI - ContextCaching - allow from_cached_content to take the cached_content resource name

Zhenyi Qi · copybara-github · commit 8f53902b933f · 2024-07-16T11:18:43.000-07:00
PiperOrigin-RevId: 652913296
diff --git a/tests/unit/vertexai/test_generative_models.py b/tests/unit/vertexai/test_generative_models.py
@@ -483,6 +483,37 @@ def test_generative_model_from_cached_content(
             == "cached-content-id-in-from-cached-content-test"
         )
 
+    def test_generative_model_from_cached_content_with_resource_name(
+        self, mock_get_cached_content_fixture
+    ):
+        project_location_prefix = (
+            f"projects/{_TEST_PROJECT}/locations/{_TEST_LOCATION}/"
+        )
+
+        model = preview_generative_models.GenerativeModel.from_cached_content(
+            cached_content="cached-content-id-in-from-cached-content-test"
+        )
+
+        assert (
+            model._prediction_resource_name
+            == project_location_prefix
+            + "publishers/google/models/"
+            + "gemini-pro-from-mock-get-cached-content"
+        )
+        assert (
+            model._cached_content.model_name
+            == "gemini-pro-from-mock-get-cached-content"
+        )
+        assert (
+            model._cached_content.resource_name
+            == f"projects/{_TEST_PROJECT}/locations/{_TEST_LOCATION}/"
+            "cachedContents/cached-content-id-in-from-cached-content-test"
+        )
+        assert (
+            model._cached_content.name
+            == "cached-content-id-in-from-cached-content-test"
+        )
+
     @mock.patch.object(
         target=prediction_service.PredictionServiceClient,
         attribute="generate_content",
diff --git a/vertexai/generative_models/_generative_models.py b/vertexai/generative_models/_generative_models.py
@@ -2687,11 +2687,29 @@ def start_chat(
     @classmethod
     def from_cached_content(
         cls,
-        cached_content: "caching.CachedContent",
+        cached_content: Union[str, "caching.CachedContent"],
         *,
         generation_config: Optional[GenerationConfigType] = None,
         safety_settings: Optional[SafetySettingsType] = None,
     ) -> "_GenerativeModel":
+        """Creates a model from cached content.
+
+        Creates a model instance with an existing cached content. The cached
+        content becomes the prefix of the requesting contents.
+
+        Args:
+            cached_content: The cached content resource name or object.
+            generation_config: The generation config to use for this model.
+            safety_settings: The safety settings to use for this model.
+
+        Returns:
+            A model instance with the cached content wtih cached content as
+            prefix of all its requests.
+        """
+        if isinstance(cached_content, str):
+            from vertexai.preview import caching
+
+            cached_content = caching.CachedContent.get(cached_content)
         model_name = cached_content.model_name
         model = cls(
             model_name=model_name,
