pickle stuff in llm caching path bc not everything is jsonifiable

Signed-off-by: Henry Lindeman <[email protected]>
aryn-ai · Jan 30, 2025 · 1789409 · 1789409
1 parent 98ce6a0
commit 1789409
Showing 1 changed file with 10 additions and 7 deletions.
diff --git a/lib/sycamore/sycamore/llms/llms.py b/lib/sycamore/sycamore/llms/llms.py
@@ -129,6 +129,7 @@ def _llm_cache_get(self, prompt: RenderedPrompt, llm_kwargs: Optional[dict]) ->
         key = self._llm_cache_key(prompt, llm_kwargs)
         hit = self._cache.get(key)
         if hit:
+            hit = pickle.loads(hit)
             assert (
                 len(hit) == 5
                 and hit.get("prompt") == RenderedPrompt(messages=prompt.messages)
@@ -156,13 +157,15 @@ def _llm_cache_set(self, prompt: RenderedPrompt, llm_kwargs: Optional[dict], res
         key = self._llm_cache_key(prompt, llm_kwargs)
         self._cache.set(
             key,
-            {
-                "prompt": RenderedPrompt(messages=prompt.messages),
-                "prompt.response_format": self._pickleable_response_format(prompt),
-                "llm_kwargs": llm_kwargs,
-                "model_name": self._model_name,
-                "result": result,
-            },
+            pickle.dumps(
+                {
+                    "prompt": RenderedPrompt(messages=prompt.messages),
+                    "prompt.response_format": self._pickleable_response_format(prompt),
+                    "llm_kwargs": llm_kwargs,
+                    "model_name": self._model_name,
+                    "result": result,
+                }
+            ),
         )
 
     def get_metadata(self, kwargs, response_text, wall_latency, in_tokens, out_tokens) -> dict: