| 1 | from __future__ import annotations |
| 2 | |
| 3 | import sys |
| 4 | import asyncio |
| 5 | from pathlib import Path |
| 6 | |
| 7 | from langchain_core.documents import Document |
| 8 | |
| 9 | PROJECT_ROOT = Path(__file__).resolve().parents[1] |
| 10 | if str(PROJECT_ROOT) not in sys.path: |
| 11 | sys.path.insert(0, str(PROJECT_ROOT)) |
| 12 | |
| 13 | from plugins._memory.helpers.memory import Memory |
| 14 | |
| 15 | |
| 16 | class FakeFaiss: |
| 17 | def __init__(self, docs: list[Document]): |
| 18 | self.docs = {doc.metadata["id"]: doc for doc in docs} |
| 19 | self.deleted: list[str] = [] |
| 20 | |
| 21 | async def asearch(self, *_args, **_kwargs): |
| 22 | return [] |
| 23 | |
| 24 | async def adelete(self, ids): |
| 25 | for doc_id in ids: |
| 26 | self.deleted.append(doc_id) |
| 27 | self.docs.pop(doc_id, None) |
| 28 | |
| 29 | async def aget_by_ids(self, ids): |
| 30 | return [self.docs[doc_id] for doc_id in ids if doc_id in self.docs] |
| 31 | |
| 32 | def get_all_docs(self): |
| 33 | return self.docs |
| 34 | |
| 35 | def get_by_ids(self, ids): |
| 36 | return [self.docs[doc_id] for doc_id in ids if doc_id in self.docs] |
| 37 | |
| 38 | |
| 39 | class FakeEmbeddings: |
| 40 | def __init__(self): |
| 41 | self.queries: list[str] = [] |
| 42 | |
| 43 | async def aembed_query(self, query: str): |
| 44 | self.queries.append(query) |
| 45 | return [0.25, 0.75] |
| 46 | |
| 47 | |
| 48 | class FakeVectorSearch: |
| 49 | def __init__(self): |
| 50 | self.embedding_function = FakeEmbeddings() |
| 51 | self.docs = [ |
| 52 | (Document(page_content="main", metadata={"area": "main"}), 0.8), |
| 53 | (Document(page_content="solution", metadata={"area": "solutions"}), 0.7), |
| 54 | (Document(page_content="weak", metadata={"area": "main"}), 0.2), |
| 55 | ] |
| 56 | self.embeddings: list[list[float]] = [] |
| 57 | |
| 58 | async def asimilarity_search_with_score_by_vector( |
| 59 | self, embedding, *, k, filter |
| 60 | ): |
| 61 | self.embeddings.append(embedding) |
| 62 | return [(doc, score) for doc, score in self.docs if filter(doc.metadata)][:k] |
| 63 | |
| 64 | |
| 65 | def test_memory_forget_removes_exact_matches_and_derived_fragments(): |
| 66 | main = Document( |
| 67 | page_content="User currently prefers memory cleanup token banana-397.", |
| 68 | metadata={"id": "main-1", "area": "main"}, |
| 69 | ) |
| 70 | fragment = Document( |
| 71 | page_content="Derived note from old preference.", |
| 72 | metadata={ |
| 73 | "id": "fragment-1", |
| 74 | "area": "fragments", |
| 75 | "consolidated_from": ["main-1"], |
| 76 | }, |
| 77 | ) |
| 78 | unrelated = Document( |
| 79 | page_content="Unrelated memory about project setup.", |
| 80 | metadata={"id": "other-1", "area": "main"}, |
| 81 | ) |
| 82 | fake_db = FakeFaiss([main, fragment, unrelated]) |
| 83 | memory = Memory(fake_db, memory_subdir="test") |
| 84 | memory._save_db = lambda: None |
| 85 | |
| 86 | removed = asyncio.run( |
| 87 | memory.delete_documents_by_query( |
| 88 | query="banana-397", |
| 89 | threshold=0.99, |
| 90 | include_exact=True, |
| 91 | cascade=True, |
| 92 | ) |
| 93 | ) |
| 94 | |
| 95 | assert {doc.metadata["id"] for doc in removed} == {"main-1", "fragment-1"} |
| 96 | assert fake_db.deleted == ["main-1", "fragment-1"] |
| 97 | assert set(fake_db.docs) == {"other-1"} |
| 98 | |
| 99 | |
| 100 | def test_memory_delete_cascades_even_when_original_id_is_already_missing(): |
| 101 | replacement = Document( |
| 102 | page_content="User currently prefers concise technical answers.", |
| 103 | metadata={ |
| 104 | "id": "replacement-1", |
| 105 | "area": "main", |
| 106 | "updated_from": "old-pref-1", |
| 107 | }, |
| 108 | ) |
| 109 | fake_db = FakeFaiss([replacement]) |
| 110 | memory = Memory(fake_db, memory_subdir="test") |
| 111 | memory._save_db = lambda: None |
| 112 | |
| 113 | removed = asyncio.run( |
| 114 | memory.delete_documents_by_ids(["old-pref-1"], cascade=True) |
| 115 | ) |
| 116 | |
| 117 | assert [doc.metadata["id"] for doc in removed] == ["replacement-1"] |
| 118 | assert fake_db.deleted == ["replacement-1"] |
| 119 | assert fake_db.docs == {} |
| 120 | |
| 121 | |
| 122 | def test_memory_reuses_one_query_embedding_across_filtered_searches(): |
| 123 | fake_db = FakeVectorSearch() |
| 124 | memory = Memory(fake_db, memory_subdir="test") |
| 125 | |
| 126 | async def search(): |
| 127 | embedding = await memory.embed_query("shared recall query") |
| 128 | memories = await memory.search_similarity_threshold( |
| 129 | query="shared recall query", |
| 130 | limit=12, |
| 131 | threshold=0.7, |
| 132 | filter="area == 'main'", |
| 133 | embedding=embedding, |
| 134 | ) |
| 135 | solutions = await memory.search_similarity_threshold( |
| 136 | query="shared recall query", |
| 137 | limit=8, |
| 138 | threshold=0.7, |
| 139 | filter="area == 'solutions'", |
| 140 | embedding=embedding, |
| 141 | ) |
| 142 | return memories, solutions |
| 143 | |
| 144 | memories, solutions = asyncio.run(search()) |
| 145 | |
| 146 | assert fake_db.embedding_function.queries == ["shared recall query"] |
| 147 | assert fake_db.embeddings == [[0.25, 0.75], [0.25, 0.75]] |
| 148 | assert [doc.page_content for doc in memories] == ["main"] |
| 149 | assert [doc.page_content for doc in solutions] == ["solution"] |