mirror of
https://github.com/agent0ai/agent-zero.git
synced 2026-05-16 19:50:43 +00:00
93 lines
2.7 KiB
Python
93 lines
2.7 KiB
Python
from __future__ import annotations
|
|
|
|
import sys
|
|
import asyncio
|
|
from pathlib import Path
|
|
|
|
from langchain_core.documents import Document
|
|
|
|
PROJECT_ROOT = Path(__file__).resolve().parents[1]
|
|
if str(PROJECT_ROOT) not in sys.path:
|
|
sys.path.insert(0, str(PROJECT_ROOT))
|
|
|
|
from plugins._memory.helpers.memory import Memory
|
|
|
|
|
|
class FakeFaiss:
|
|
def __init__(self, docs: list[Document]):
|
|
self.docs = {doc.metadata["id"]: doc for doc in docs}
|
|
self.deleted: list[str] = []
|
|
|
|
async def asearch(self, *_args, **_kwargs):
|
|
return []
|
|
|
|
async def adelete(self, ids):
|
|
for doc_id in ids:
|
|
self.deleted.append(doc_id)
|
|
self.docs.pop(doc_id, None)
|
|
|
|
async def aget_by_ids(self, ids):
|
|
return [self.docs[doc_id] for doc_id in ids if doc_id in self.docs]
|
|
|
|
def get_all_docs(self):
|
|
return self.docs
|
|
|
|
def get_by_ids(self, ids):
|
|
return [self.docs[doc_id] for doc_id in ids if doc_id in self.docs]
|
|
|
|
|
|
def test_memory_forget_removes_exact_matches_and_derived_fragments():
|
|
main = Document(
|
|
page_content="User currently prefers memory cleanup token banana-397.",
|
|
metadata={"id": "main-1", "area": "main"},
|
|
)
|
|
fragment = Document(
|
|
page_content="Derived note from old preference.",
|
|
metadata={
|
|
"id": "fragment-1",
|
|
"area": "fragments",
|
|
"consolidated_from": ["main-1"],
|
|
},
|
|
)
|
|
unrelated = Document(
|
|
page_content="Unrelated memory about project setup.",
|
|
metadata={"id": "other-1", "area": "main"},
|
|
)
|
|
fake_db = FakeFaiss([main, fragment, unrelated])
|
|
memory = Memory(fake_db, memory_subdir="test")
|
|
memory._save_db = lambda: None
|
|
|
|
removed = asyncio.run(
|
|
memory.delete_documents_by_query(
|
|
query="banana-397",
|
|
threshold=0.99,
|
|
include_exact=True,
|
|
cascade=True,
|
|
)
|
|
)
|
|
|
|
assert {doc.metadata["id"] for doc in removed} == {"main-1", "fragment-1"}
|
|
assert fake_db.deleted == ["main-1", "fragment-1"]
|
|
assert set(fake_db.docs) == {"other-1"}
|
|
|
|
|
|
def test_memory_delete_cascades_even_when_original_id_is_already_missing():
|
|
replacement = Document(
|
|
page_content="User currently prefers concise technical answers.",
|
|
metadata={
|
|
"id": "replacement-1",
|
|
"area": "main",
|
|
"updated_from": "old-pref-1",
|
|
},
|
|
)
|
|
fake_db = FakeFaiss([replacement])
|
|
memory = Memory(fake_db, memory_subdir="test")
|
|
memory._save_db = lambda: None
|
|
|
|
removed = asyncio.run(
|
|
memory.delete_documents_by_ids(["old-pref-1"], cascade=True)
|
|
)
|
|
|
|
assert [doc.metadata["id"] for doc in removed] == ["replacement-1"]
|
|
assert fake_db.deleted == ["replacement-1"]
|
|
assert fake_db.docs == {}
|