Stroke out functional approach for CachedContent CURD ops

Author: mayureshagashe2105

google/generativeai/caching.py

@@ -14,63 +14,251 @@
 # limitations under the License.
 from __future__ import annotations
 
-from typing import Optional, Iterable
-
-import google.ai.generativelanguage as glm
+import dataclasses
+import datetime
+from typing import Any, Iterable, Optional
 
+from google.generativeai.types.model_types import idecode_time
 from google.generativeai.types import caching_types
 from google.generativeai.types import content_types
+from google.generativeai.utils import flatten_update_paths
 from google.generativeai.client import get_default_cache_client
 
+from google.protobuf import field_mask_pb2
+import google.ai.generativelanguage as glm
+
+
+@dataclasses.dataclass
+class CachedContent:
+    """Cached content resource."""
+
+    name: str
+    model: str
+    create_time: datetime.datetime
+    update_time: datetime.datetime
+    expire_time: datetime.datetime
+
+    # NOTE: Automatic CachedContent deletion using contextmanager is not P0(P1+).
+    # Adding basic support for now.
+    def __enter__(self):
+        return self
+
+    def __exit__(self, exc_type, exc_value, exc_tb):
+        self.delete()
+
+    def _to_dict(self) -> glm.CachedContent:
+        proto_paths = {
+            "name": self.name,
+            "model": self.model,
+        }
+        return glm.CachedContent(**proto_paths)
+
+    def _apply_update(self, path, value):
+        parts = path.split(".")
+        for part in parts[:-1]:
+            self = getattr(self, part)
+        if parts[-1] == "ttl":
+            value = self.expire_time + datetime.timedelta(seconds=value["seconds"])
+            parts[-1] = "expire_time"
+        setattr(self, parts[-1], value)
+    
+    @classmethod
+    def _decode_cached_content(cls, cached_content: glm.CachedContent) -> CachedContent:
+        # not supposed to get INPUT_ONLY repeated fields, but local gapic lib build
+        # is returning these, hence setting including_default_value_fields to False
+        cached_content = type(cached_content).to_dict(
+            cached_content, including_default_value_fields=False
+        )
+
+        idecode_time(cached_content, "create_time")
+        idecode_time(cached_content, "update_time")
+        # always decode `expire_time` as Timestamp is returned
+        # regardless of what was sent on input
+        idecode_time(cached_content, "expire_time")
+        return cls(**cached_content)
+
+    @staticmethod
+    def _prepare_create_request(
+        name: str,
+        model: str,
+        system_instruction: Optional[content_types.ContentType] = None,
+        contents: Optional[content_types.ContentsType] = None,
+        tools: Optional[content_types.FunctionLibraryType] = None,
+        tool_config: Optional[content_types.ToolConfigType] = None,
+        ttl: Optional[caching_types.ExpirationTypes] = datetime.timedelta(hours=1),
+    ) -> glm.CreateCachedContentRequest:
+        """Prepares a CreateCachedContentRequest."""
+        if "cachedContents/" not in name:
+            name = "cachedContents/" + name
+
+        if "/" not in model:
+            model = "models/" + model
+
+        if system_instruction:
+            system_instruction = content_types.to_content(system_instruction)
+
+        tools_lib = content_types.to_function_library(tools)
+        if tools_lib:
+            tools_lib = tools_lib.to_proto()
+
+        if tool_config:
+            tool_config = content_types.to_tool_config(tool_config)
+
+        if contents:
+            contents = content_types.to_contents(contents)
+
+        if ttl:
+            ttl = caching_types.to_ttl(ttl)
+
+        cached_content = glm.CachedContent(
+            name=name,
+            model=model,
+            system_instruction=system_instruction,
+            contents=contents,
+            tools=tools_lib,
+            tool_config=tool_config,
+            ttl=ttl,
+        )
+
+        return glm.CreateCachedContentRequest(cached_content=cached_content)
+
+    @classmethod
+    def create(
+        cls,
+        name: str,
+        model: str,
+        system_instruction: Optional[content_types.ContentType] = None,
+        contents: Optional[content_types.ContentsType] = None,
+        tools: Optional[content_types.FunctionLibraryType] = None,
+        tool_config: Optional[content_types.ToolConfigType] = None,
+        ttl: Optional[caching_types.ExpirationTypes] = datetime.timedelta(hours=1),
+        client: glm.CacheServiceClient | None = None,
+    ) -> CachedContent:
+        """Creates CachedContent resource.
+
+        Args:
+            name: The resource name referring to the cached content.
+                  Format: cachedContents/{id}.
+            model: The name of the `Model` to use for cached content
+                    Format: models/{model}. Cached content resource can be only
+                    used with model it was created for.
+            system_instruction: Developer set system instruction.
+            contents: Contents to cache.
+            tools: A list of `Tools` the model may use to generate response.
+            tool_config: Config to apply to all tools.
+            ttl: TTL for cached resource (in seconds). Defaults to 1 hour.
+
+        Returns:
+            `CachedContent` resource with specified name.
+        """
+        if client is None:
+            client = get_default_cache_client()
+
+        request = cls._prepare_create_request(
+            name=name,
+            model=model,
+            system_instruction=system_instruction,
+            contents=contents,
+            tools=tools,
+            tool_config=tool_config,
+            ttl=ttl,
+        )
+
+        response = client.create_cached_content(request)
+        return cls._decode_cached_content(response)
+
+    @classmethod
+    def get(cls, name: str, client: glm.CacheServiceClient | None = None) -> CachedContent:
+        """Fetches required `CachedContent` resource.
+
+        Args:
+            name: name: The resource name referring to the cached content.
 
-# alias for `caching_types.CachedContent`.
-CachedContent = caching_types.CachedContent
+        Returns:
+            `CachedContent` resource with specified name.
+        """
+        if client is None:
+            client = get_default_cache_client()
 
+        if "cachedContents/" not in name:
+            name = "cachedContents/" + name
 
-def get_cached_content(name: str, client: glm.CacheServiceClient | None = None) -> CachedContent:
-    """Fetches required `CachedContent` resource.
+        request = glm.GetCachedContentRequest(name=name)
+        response = client.get_cached_content(request)
+        return cls._decode_cached_content(response)
+    
+    @classmethod
+    def list(
+        cls,
+        page_size: Optional[int] = 1,
+        client: glm.CacheServiceClient | None = None
+    ) -> Iterable[CachedContent]:
+        """Lists `CachedContent` objects associated with the project.
 
-    Args:
-        name: name: The resource name referring to the cached content.
+        Args:
+            page_size: The maximum number of permissions to return (per page). 
+            The service may return fewer `CachedContent` objects.
 
-    Returns:
-        `CachedContent` resource with specified name.
-    """
-    return CachedContent.get(name=name, client=client)
+        Returns:
+            A paginated list of `CachedContent` objects.
+        """
+        if client is None:
+            client = get_default_cache_client()
 
+        request = glm.ListCachedContentsRequest(page_size=page_size)
+        for cached_content in client.list_cached_contents(request):
+            yield cls._decode_cached_content(cached_content)
 
-def delete_cached_content(name: str, client: glm.CacheServiceClient | None = None) -> None:
-    """Deletes `CachedContent` resource.
+    def delete(self, client: glm.CachedServiceClient | None = None) -> None:
+        """Deletes `CachedContent` resource.
 
-    Args:
-        name: The resource name referring to the cached content.
-              Format: cachedContents/{id}.
-    """
-    if client is None:
-        client = get_default_cache_client()
+        Args:
+            name: The resource name referring to the cached content.
+                Format: cachedContents/{id}.
+        """
+        if client is None:
+            client = get_default_cache_client()
 
-    if "cachedContents/" not in name:
-        name = "cachedContents/" + name
+        request = glm.DeleteCachedContentRequest(name=self.name)
+        client.delete_cached_content(request)
+        return
 
-    request = glm.DeleteCachedContentRequest(name=name)
-    client.delete_cached_content(request)
-    return
+    def update(
+        self,
+        updates: dict[str, Any],
+        client: glm.CacheServiceClient | None = None,
+    ) -> CachedContent:
+        """Updates requested `CachedContent` resource.
 
+        Args:
+            updates: The list of fields to update.
+                     Currently only `ttl/expire_time` is supported as an update path.
 
-def list_cached_contents(
-    page_size: Optional[int] = 1, client: glm.CacheServiceClient | None = None
-) -> Iterable[CachedContent]:
-    """Lists `CachedContent` objects associated with the project.
+        Returns:
+            `CachedContent` object with specified updates.
+        """
+        if client is None:
+            client = get_default_cache_client()
 
-    Args:
-        page_size: The maximum number of permissions to return (per page). The service may return fewer `CachedContent` objects.
+        updates = flatten_update_paths(updates)
+        for update_path in updates:
+            if update_path == "ttl":
+                updates = updates.copy()
+                update_path_val = updates.get(update_path)
+                updates[update_path] = caching_types.to_ttl(update_path_val)
+            else:
+                raise ValueError(
+                    f"As of now, only `ttl` can be updated for `CachedContent`. Got: `{update_path}` instead."
+                )
+        field_mask = field_mask_pb2.FieldMask()
 
-    Returns:
-        A paginated list of `CachedContent` objects.
-    """
-    if client is None:
-        client = get_default_cache_client()
+        for path in updates.keys():
+            field_mask.paths.append(path)
+        for path, value in updates.items():
+            self._apply_update(path, value)
 
-    request = glm.ListCachedContentsRequest(page_size=page_size)
-    for cached_content in client.list_cached_contents(request):
-        yield caching_types.decode_cached_content(cached_content)
+        request = glm.UpdateCachedContentRequest(
+            cached_content=self._to_dict(), update_mask=field_mask
+        )
+        client.update_cached_content(request)
+        return self

google/generativeai/generative_models.py

@@ -19,8 +19,6 @@
 from google.generativeai.types import generation_types
 from google.generativeai.types import helper_types
 from google.generativeai.types import safety_types
-from google.generativeai.types import caching_types
-
 
 class GenerativeModel:
     """
@@ -198,15 +196,15 @@ def from_cached_content(
     @classmethod
     def from_cached_content(
         cls,
-        cached_content: caching_types.CachedContent,
+        cached_content: caching.CachedContent,
         generation_config: generation_types.GenerationConfigType | None = None,
         safety_settings: safety_types.SafetySettingOptions | None = None,
     ) -> GenerativeModel: ...
 
     @classmethod
     def from_cached_content(
         cls,
-        cached_content: str | caching_types.CachedContent,
+        cached_content: str | caching.CachedContent,
         generation_config: generation_types.GenerationConfigType | None = None,
         safety_settings: safety_types.SafetySettingOptions | None = None,
     ) -> GenerativeModel:
@@ -219,7 +217,7 @@ def from_cached_content(
             `GenerativeModel` object with `cached_content` as its context.
         """
         if isinstance(cached_content, str):
-            cached_content = caching.get_cached_content(name=cached_content)
+            cached_content = caching.CachedContent.get(name=cached_content)
 
         # call __new__ with the cached_content to set the model's context. This is done to avoid
         # the exposing `cached_content` as a public attribute.