Explicit Caching patch

google/generativeai/caching.py

@@ -22,13 +22,15 @@
 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 import string_utils
 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
 
 
+@string_utils.prettyprint
 @dataclasses.dataclass
 class CachedContent:
     """Cached content resource."""
@@ -39,29 +41,19 @@ class CachedContent:
     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) -> protos.CachedContent:
+    def _get_update_fields(self, **input_only_update_fields) -> protos.CachedContent:
         proto_paths = {
             "name": self.name,
-            "model": self.model,
         }
+        proto_paths.update(input_only_update_fields)
         return protos.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)
+        if path[-1] != "ttl":
+            setattr(self, parts[-1], value)
 
     @classmethod
     def _decode_cached_content(cls, cached_content: protos.CachedContent) -> CachedContent:
@@ -112,7 +104,7 @@ def _prepare_create_request(
             contents = content_types.to_contents(contents)
 
         if ttl:
-            ttl = caching_types.to_ttl(ttl)
+            ttl = caching_types.to_expiration(ttl)
 
         cached_content = protos.CachedContent(
             name=name,
@@ -236,25 +228,35 @@ def update(
         if client is None:
             client = get_default_cache_client()
 
+        if "ttl" in updates and "expire_time" in updates:
+            raise ValueError(
+                "`expiration` is a _oneof field. Please provide either `ttl` or `expire_time`."
+            )
+
+        field_mask = field_mask_pb2.FieldMask()
+
         updates = flatten_update_paths(updates)
         for update_path in updates:
-            if update_path == "ttl":
+            if update_path == "ttl" or update_path == "expire_time":
                 updates = updates.copy()
                 update_path_val = updates.get(update_path)
-                updates[update_path] = caching_types.to_ttl(update_path_val)
+                updates[update_path] = caching_types.to_expiration(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()
 
-        for path in updates.keys():
-            field_mask.paths.append(path)
+            field_mask.paths.append(update_path)
+
         for path, value in updates.items():
             self._apply_update(path, value)
 
         request = protos.UpdateCachedContentRequest(
-            cached_content=self._to_dict(), update_mask=field_mask
+            cached_content=self._get_update_fields(**updates), update_mask=field_mask
         )
-        client.update_cached_content(request)
+        updated_cc = client.update_cached_content(request)
+        updated_cc = self._decode_cached_content(updated_cc)
+        for path, value in dataclasses.asdict(updated_cc).items():
+            self._apply_update(path, value)
+
         return self

google/generativeai/generative_models.py

@@ -96,14 +96,9 @@ def __init__(
         self._client = None
         self._async_client = None
 
-    def __new__(cls, *args, **kwargs) -> GenerativeModel:
-        self = super().__new__(cls)
-
-        if cached_instance := kwargs.pop("cached_content", None):
-            setattr(self, "_cached_content", cached_instance.name)
-            setattr(cls, "cached_content", property(fget=lambda self: self._cached_content))
-
-        return self
+    @property
+    def cached_content(self) -> str:
+        return getattr(self, "_cached_content", None)
 
     @property
     def model_name(self):
@@ -123,7 +118,7 @@ def maybe_text(content):
                 safety_settings={self._safety_settings},
                 tools={self._tools},
                 system_instruction={maybe_text(self._system_instruction)},
-                cached_content={getattr(self, "cached_content", None)}
+                cached_content={self.cached_content}
             )"""
         )
 
@@ -139,13 +134,11 @@ def _prepare_request(
         tool_config: content_types.ToolConfigType | None,
     ) -> protos.GenerateContentRequest:
         """Creates a `protos.GenerateContentRequest` from raw inputs."""
-        if hasattr(self, "cached_content") and any([self._system_instruction, tools, tool_config]):
+        if hasattr(self, "_cached_content") and any([self._system_instruction, tools, tool_config]):
             raise ValueError(
                 "`tools`, `tool_config`, `system_instruction` cannot be set on a model instantinated with `cached_content` as its context."
             )
 
-        cached_content = getattr(self, "cached_content", None)
-
         tools_lib = self._get_tools_lib(tools)
         if tools_lib is not None:
             tools_lib = tools_lib.to_proto()
@@ -174,7 +167,7 @@ def _prepare_request(
             tools=tools_lib,
             tool_config=tool_config,
             system_instruction=self._system_instruction,
-            cached_content=cached_content,
+            cached_content=self.cached_content,
         )
 
     def _get_tools_lib(
@@ -221,17 +214,16 @@ def from_cached_content(
         if isinstance(cached_content, str):
             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.
-        self = cls.__new__(cls, cached_content=cached_content)
-
         # call __init__ to set the model's `generation_config`, `safety_settings`.
         # `model_name` will be the name of the model for which the `cached_content` was created.
-        self.__init__(
+        self = GenerativeModel(
             model_name=cached_content.model,
             generation_config=generation_config,
             safety_settings=safety_settings,
         )
+
+        # set the model's context.
+        setattr(self, "_cached_content", cached_content.name)
         return self
 
     def generate_content(

google/generativeai/types/caching_types.py

@@ -19,7 +19,7 @@
 from typing_extensions import TypedDict
 import re
 
-__all__ = ["TTL"]
+__all__ = ["ExpirationTypes", "ExpireTime", "TTL"]
 
 
 _VALID_CACHED_CONTENT_NAME = r"([a-z0-9-\.]+)$"
@@ -33,19 +33,38 @@ def valid_cached_content_name(name: str) -> bool:
 
 
 class TTL(TypedDict):
+    # Represents datetime.datetime.now() + desired ttl
     seconds: int
+    nanos: int
 
 
-ExpirationTypes = Union[TTL, int, datetime.timedelta]
+class ExpireTime(TypedDict):
+    # Represents seconds of UTC time since Unix epoch
+    seconds: int
+    nanos: int
+
+
+ExpirationTypes = Union[TTL, ExpireTime, int, datetime.timedelta, datetime.datetime]
 
 
-def to_ttl(expiration: Optional[ExpirationTypes]) -> TTL:
-    if isinstance(expiration, datetime.timedelta):
-        return {"seconds": int(expiration.total_seconds())}
+def to_expiration(expiration: Optional[ExpirationTypes]) -> TTL | ExpireTime:
+    if isinstance(expiration, datetime.timedelta):  # consider `ttl`
+        return {
+            "seconds": int(expiration.total_seconds()),
+            "nanos": int(expiration.microseconds * 1000),
+        }
+    elif isinstance(expiration, datetime.datetime):  # consider `expire_time`
+        timestamp = expiration.timestamp()
+        seconds = int(timestamp)
+        nanos = int((seconds % 1) * 1000)
+        return {
+            "seconds": seconds,
+            "nanos": nanos,
+        }
     elif isinstance(expiration, dict):
         return expiration
-    elif isinstance(expiration, int):
-        return {"seconds": expiration}
+    elif isinstance(expiration, int):  # consider `ttl`
+        return {"seconds": expiration, "nanos": 0}
     else:
         raise TypeError(
             f"Could not convert input to `expire_time` \n'" f"  type: {type(expiration)}\n",

tests/test_caching.py

@@ -210,10 +210,16 @@ def test_update_cached_content_invalid_update_paths(self):
         with self.assertRaises(ValueError):
             cc.update(updates=update_masks)
 
-    def test_update_cached_content_valid_update_paths(self):
-        update_masks = dict(
-            ttl=datetime.timedelta(hours=2),
-        )
+    @parameterized.named_parameters(
+        [
+            dict(testcase_name="ttl", update_masks=dict(ttl=datetime.timedelta(hours=2))),
+            dict(
+                testcase_name="expire_time",
+                update_masks=dict(expire_time=datetime.datetime(2024, 6, 5, 12, 12, 12, 23)),
+            ),
+        ]
+    )
+    def test_update_cached_content_valid_update_paths(self, update_masks):
 
         cc = caching.CachedContent.get(name="cachedContents/test-cached-content")
         cc = cc.update(updates=update_masks)
@@ -229,18 +235,6 @@ def test_delete_cached_content(self):
         cc.delete()
         self.assertIsInstance(self.observed_requests[-1], protos.DeleteCachedContentRequest)
 
-    def test_auto_delete_cached_content_with_context_manager(self):
-        with caching.CachedContent.create(
-            name="test-cached-content",
-            model="models/gemini-1.0-pro-001",
-            contents=["Add 5 and 6"],
-            system_instruction="Always add 10 to the result.",
-            ttl=datetime.timedelta(minutes=30),
-        ) as cc:
-            ...  # some logic
-
-        self.assertIsInstance(self.observed_requests[-1], protos.DeleteCachedContentRequest)
-
 
 if __name__ == "__main__":
     absltest.main()