fix update method

Change-Id: I433c25b2d80cdf6e483b59f61ff29bb8d2dc6595

Author: mayureshagashe2105

google/generativeai/caching.py

@@ -47,21 +47,19 @@ def __enter__(self):
     def __exit__(self, exc_type, exc_value, exc_tb):
         self.delete()
 
-    def _to_dict(self) -> protos.CachedContent:
+    def _to_dict(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 +110,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 +234,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._to_dict(**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/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,18 +33,36 @@ def valid_cached_content_name(name: str) -> bool:
 
 
 class TTL(TypedDict):
+    # Represents datetime.datetime.now() + desired ttl
     seconds: int
+    nanos: int = 0
+
+class ExpireTime(TypedDict):
+    # Represents seconds of UTC time since Unix epoch
+    seconds: int
+    nanos: int = 0
 
 
-ExpirationTypes = Union[TTL, int, datetime.timedelta]
+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:
+    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):
+    elif isinstance(expiration, int):  # consider `ttl`
         return {"seconds": expiration}
     else:
         raise TypeError(

tests/test_caching.py

@@ -210,10 +210,19 @@ 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)