diff --git a/py/src/braintrust/prompt_cache/disk_cache.py b/py/src/braintrust/prompt_cache/disk_cache.py index b74080249..cddc46edb 100644 --- a/py/src/braintrust/prompt_cache/disk_cache.py +++ b/py/src/braintrust/prompt_cache/disk_cache.py @@ -128,10 +128,15 @@ def set(self, key: str, value: T) -> None: os.makedirs(self._dir, exist_ok=True) file_path = self._get_entry_path(key) + # Serialize before opening the file: opening it for writing truncates any + # existing entry, so a serialization error would otherwise leave an empty + # file behind and destroy the previously cached value. + if self._serializer is not None: + value = self._serializer(value) + payload = json.dumps(value).encode("utf-8") + with gzip.open(file_path, "wb") as f: - if self._serializer is not None: - value = self._serializer(value) - f.write(json.dumps(value).encode("utf-8")) + f.write(payload) self._evict_if_full() except Exception as e: diff --git a/py/src/braintrust/prompt_cache/test_disk_cache.py b/py/src/braintrust/prompt_cache/test_disk_cache.py index 7f49986ed..618b9b2ab 100644 --- a/py/src/braintrust/prompt_cache/test_disk_cache.py +++ b/py/src/braintrust/prompt_cache/test_disk_cache.py @@ -24,6 +24,21 @@ def tearDown(self): except Exception: pass + def test_failed_serialization_keeps_existing_entry(self): + self.cache.set("key", {"a": 1}) + + # A value that json cannot serialize must not clobber the stored entry. + self.cache.set("key", {"a": object()}) + + assert self.cache.get("key") == {"a": 1} + + def test_failed_serialization_does_not_create_entry(self): + self.cache.set("key", {"a": object()}) + + with self.assertRaises(KeyError): + self.cache.get("key") + assert os.listdir(self.cache_dir) == [] + def test_keys_with_invalid_paths(self): data = {"1": "2"} weird_keys = [