79 lines
2.7 KiB
Python
79 lines
2.7 KiB
Python
from pathlib import Path
|
|
|
|
import src.ai.memory as memory_module
|
|
|
|
|
|
class _FakeJSONStore:
|
|
def __init__(self, storage_path: Path):
|
|
self.storage_path = storage_path
|
|
|
|
|
|
def test_memory_retries_once_on_chroma_conflict(monkeypatch, tmp_path: Path):
|
|
class _FlakyChroma:
|
|
calls = 0
|
|
|
|
def __init__(self, persist_directory: Path):
|
|
type(self).calls += 1
|
|
self.persist_directory = persist_directory
|
|
if type(self).calls == 1:
|
|
raise RuntimeError("table embeddings already exists")
|
|
|
|
monkeypatch.setattr(memory_module, "ChromaVectorStore", _FlakyChroma)
|
|
monkeypatch.setattr(memory_module, "JSONVectorStore", _FakeJSONStore)
|
|
|
|
memory = memory_module.MemorySystem(
|
|
storage_path=tmp_path / "long_term_memory.json",
|
|
use_vector_db=True,
|
|
)
|
|
|
|
assert isinstance(memory.vector_store, _FlakyChroma)
|
|
assert _FlakyChroma.calls == 2
|
|
|
|
|
|
def test_memory_rebuilds_chroma_directory_when_conflict_persists(monkeypatch, tmp_path: Path):
|
|
class _RecoverOnThirdTryChroma:
|
|
calls = 0
|
|
|
|
def __init__(self, persist_directory: Path):
|
|
type(self).calls += 1
|
|
self.persist_directory = persist_directory
|
|
if type(self).calls <= 2:
|
|
raise RuntimeError("table embeddings already exists")
|
|
|
|
monkeypatch.setattr(memory_module, "ChromaVectorStore", _RecoverOnThirdTryChroma)
|
|
monkeypatch.setattr(memory_module, "JSONVectorStore", _FakeJSONStore)
|
|
|
|
storage_root = tmp_path / "ai"
|
|
chroma_dir = storage_root / "chroma_db"
|
|
chroma_dir.mkdir(parents=True, exist_ok=True)
|
|
(chroma_dir / "marker.txt").write_text("legacy-data", encoding="utf-8")
|
|
|
|
memory = memory_module.MemorySystem(
|
|
storage_path=storage_root / "long_term_memory.json",
|
|
use_vector_db=True,
|
|
)
|
|
|
|
assert isinstance(memory.vector_store, _RecoverOnThirdTryChroma)
|
|
assert _RecoverOnThirdTryChroma.calls == 3
|
|
|
|
backups = sorted(storage_root.glob("chroma_db_backup_conflict_*"))
|
|
assert backups, "expected backup directory after repair"
|
|
assert (backups[-1] / "marker.txt").exists()
|
|
assert chroma_dir.exists()
|
|
|
|
|
|
def test_memory_falls_back_to_json_if_chroma_repair_fails(monkeypatch, tmp_path: Path):
|
|
class _AlwaysFailChroma:
|
|
def __init__(self, persist_directory: Path):
|
|
raise RuntimeError("table embeddings already exists")
|
|
|
|
monkeypatch.setattr(memory_module, "ChromaVectorStore", _AlwaysFailChroma)
|
|
monkeypatch.setattr(memory_module, "JSONVectorStore", _FakeJSONStore)
|
|
|
|
memory = memory_module.MemorySystem(
|
|
storage_path=tmp_path / "long_term_memory.json",
|
|
use_vector_db=True,
|
|
)
|
|
|
|
assert isinstance(memory.vector_store, _FakeJSONStore)
|