Gemini takes deterministic filepaths to support caching

Author: tomasruizt

llmlib/llmlib/base_llm.py

@@ -12,7 +12,7 @@ class Message:
     role: Literal["user", "assistant"]
     msg: str
     img_name: str | None = None
-    img: Image.Image | None = None
+    img: Path | Image.Image | None = None
     video: Path | BytesIO | None = None
 
     @classmethod

llmlib/llmlib/gemini/gemini_code.py

@@ -287,14 +287,19 @@ def dump_files_return_paths(msg: Message) -> list[Path]:
     """Can return 0, 1 or 2 paths"""
     paths: list[Path] = []
     if msg.has_image():
-        temp_file = tempfile.mktemp(suffix=".jpg")
-        msg.img.save(temp_file)
-        paths.append(Path(temp_file))
+        if isinstance(msg.img, Path):
+            paths.append(msg.img)
+        else:
+            raise PathNeededError()
     if msg.has_video():
         if isinstance(msg.video, Path):
             paths.append(msg.video)
         else:
-            temp_file = tempfile.mktemp(suffix=".mp4")
-            msg.video.save(temp_file)
-            paths.append(Path(temp_file))
+            raise PathNeededError()
     return paths
+
+
+def PathNeededError():
+    return ValueError(
+        "To support caching based on filename, please provide a deterministic filepath."
+    )

tests/helpers.py

@@ -76,8 +76,7 @@ def mona_lisa_message() -> Message:
 
 
 def pyramid_message() -> Message:
-    img_name = "pyramid.jpg"
-    img = get_test_img(img_name)
+    img = file_for_test("pyramid.jpg")
     msg = Message(role="user", msg="What is in the image?", img=img, img_name="")
     return msg