Add param to load_memory_cache to stop pkl files without explicit loading (#9055)

Author: isaacbmiller

dspy/clients/cache.py

@@ -170,7 +170,11 @@ def save_memory_cache(self, filepath: str) -> None:
             with open(filepath, "wb") as f:
                 cloudpickle.dump(self.memory_cache, f)
 
-    def load_memory_cache(self, filepath: str) -> None:
+    def load_memory_cache(self, filepath: str, allow_pickle: bool = False) -> None:
+        if not allow_pickle:
+            raise ValueError("Loading untrusted .pkl files can run arbitrary code, which may be dangerous. \
+            Set `allow_pickle=True` to load if you are running in a trusted environment and the file is from a trusted source.")
+
         if not self.enable_memory_cache:
             return
 

tests/clients/test_cache.py

@@ -189,8 +189,12 @@ def test_save_and_load_memory_cache(cache, tmp_path):
         memory_max_entries=100,
     )
 
-    # Load the memory cache
-    new_cache.load_memory_cache(str(temp_cache_file))
+    # Load the memory cache without allowing pickle (default)
+    with pytest.raises(ValueError):
+        new_cache.load_memory_cache(str(temp_cache_file))
+
+    # Load the memory cache with allow_pickle=True
+    new_cache.load_memory_cache(str(temp_cache_file), allow_pickle=True)
 
     # Verify items are in the new memory cache
     for req in requests: