Skip to content

Commit dad234f

Browse files
committed
Fix #2331: purge embeddings when delete_node_by_prams removes memories
`Neo4jCommunityGraphDB.delete_node_by_prams` deleted the Neo4j nodes but never removed their embeddings from `vec_db`, so a deleted memory stayed retrievable through vector search while the delete API still answered "success". `clear()` and `delete_node_by_mem_cube_id()` already purged `vec_db`; `delete_node_by_prams` simply omitted that step. Collect the matching node IDs before the graph delete and purge those vectors afterwards via `_delete_vectors_for_deleted_nodes`. A vector store failure is logged rather than raised, matching the existing `clear()` and `delete_node_by_mem_cube_id()` behaviour, so a committed graph deletion is never reported as failed. The `neo4j` backend has no `vec_db` (it searches Neo4j's native vector index), so the guard is a no-op there. Regression coverage lives in `tests/graph_dbs/test_delete_vec_cleanup.py`; it fails on the parent commit and passes with this change. Closes #2331
1 parent a7367d0 commit dad234f

3 files changed

Lines changed: 261 additions & 12 deletions

File tree

‎src/memos/graph_dbs/neo4j.py‎

Lines changed: 26 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1893,6 +1893,24 @@ def _parse_node(self, node_data: dict[str, Any]) -> dict[str, Any]:
18931893
node["sources"][idx] = json.loads(node["sources"][idx])
18941894
return {"id": node.pop("id"), "memory": node.pop("memory", ""), "metadata": node}
18951895

1896+
def _delete_vectors_for_deleted_nodes(self, node_ids: list[str]) -> None:
1897+
"""Purge the embeddings of just-deleted graph nodes from the vector store.
1898+
1899+
Deleting a graph node alone leaves its vector behind, and vector search is
1900+
what retrieval reads first, so the deleted memory stays retrievable. The
1901+
vector store is an external dependency, so a failure here is logged rather
1902+
than raised: the graph deletion is already committed and must not be
1903+
reported as failed.
1904+
"""
1905+
vec_db = getattr(self, "vec_db", None)
1906+
if not node_ids or vec_db is None:
1907+
return
1908+
try:
1909+
vec_db.delete(node_ids)
1910+
logger.info("[delete_node_by_prams] Deleted %d vectors from VecDB", len(node_ids))
1911+
except Exception as e:
1912+
logger.warning("[delete_node_by_prams] Failed to delete vectors from VecDB: %s", e)
1913+
18961914
def delete_node_by_prams(
18971915
self,
18981916
writable_cube_ids: list[str] | None = None,
@@ -1998,9 +2016,16 @@ def delete_node_by_prams(
19982016
delete_query = f"MATCH (n:Memory) WHERE {final_where} DETACH DELETE n"
19992017
logger.info(f"[delete_node_by_prams] delete_query: {delete_query}")
20002018

2019+
# Collect the IDs in the same transaction, before the nodes are gone, so their
2020+
# embedding vectors can be purged from the vector store afterwards.
2021+
id_collect_query = f"MATCH (n:Memory) WHERE {final_where} RETURN n.id AS id"
2022+
20012023
deleted_count = 0
2024+
deleted_ids: list[str] = []
20022025
try:
20032026
with self.driver.session(database=self.db_name) as session:
2027+
deleted_ids = [record["id"] for record in session.run(id_collect_query, **params)]
2028+
20042029
# Execute delete query
20052030
result = session.run(delete_query, **params)
20062031
# Consume the result to ensure deletion completes and get the summary
@@ -2016,6 +2041,7 @@ def delete_node_by_prams(
20162041
logger.error(f"[delete_node_by_prams] Failed to delete nodes: {e}", exc_info=True)
20172042
raise
20182043

2044+
self._delete_vectors_for_deleted_nodes(deleted_ids)
20192045
logger.info(f"[delete_node_by_prams] Successfully deleted {deleted_count} nodes")
20202046
return deleted_count
20212047

‎src/memos/graph_dbs/neo4j_community.py‎

Lines changed: 11 additions & 12 deletions
Original file line numberDiff line numberDiff line change
@@ -1038,10 +1038,11 @@ def delete_node_by_prams(
10381038
f"[delete_node_by_prams] Deleting nodes - memory_ids: {memory_ids}, file_ids: {file_ids}, filter: {filter}"
10391039
)
10401040

1041-
# First count matching nodes to get accurate count
1042-
count_query = f"MATCH (n:Memory) WHERE {ids_where} RETURN count(n) AS node_count"
1043-
logger.info(f"[delete_node_by_prams] count_query: {count_query}")
1044-
print(f"[delete_node_by_prams] count_query: {count_query}")
1041+
# First collect the matching node IDs so their embedding vectors can be
1042+
# purged from the vector store after the graph deletion.
1043+
id_collect_query = f"MATCH (n:Memory) WHERE {ids_where} RETURN n.id AS id"
1044+
logger.info(f"[delete_node_by_prams] id_collect_query: {id_collect_query}")
1045+
print(f"[delete_node_by_prams] id_collect_query: {id_collect_query}")
10451046

10461047
# Then delete nodes
10471048
delete_query = f"MATCH (n:Memory) WHERE {ids_where} DETACH DELETE n"
@@ -1050,24 +1051,22 @@ def delete_node_by_prams(
10501051
print(f"[delete_node_by_prams] params: {params}")
10511052

10521053
deleted_count = 0
1054+
deleted_ids: list[str] = []
10531055
try:
10541056
with self.driver.session(database=self.db_name) as session:
1055-
# Count nodes before deletion
1056-
count_result = session.run(count_query, **params)
1057-
count_record = count_result.single()
1058-
expected_count = 0
1059-
if count_record:
1060-
expected_count = count_record["node_count"] or 0
1057+
# Collect the IDs before deletion
1058+
deleted_ids = [record["id"] for record in session.run(id_collect_query, **params)]
1059+
# Use the count from before deletion as the actual deleted count
1060+
deleted_count = len(deleted_ids)
10611061

10621062
# Delete nodes
10631063
session.run(delete_query, **params)
1064-
# Use the count from before deletion as the actual deleted count
1065-
deleted_count = expected_count
10661064

10671065
except Exception as e:
10681066
logger.error(f"[delete_node_by_prams] Failed to delete nodes: {e}", exc_info=True)
10691067
raise
10701068

1069+
self._delete_vectors_for_deleted_nodes(deleted_ids)
10711070
logger.info(f"[delete_node_by_prams] Successfully deleted {deleted_count} nodes")
10721071
return deleted_count
10731072

Lines changed: 224 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,224 @@
1+
"""
2+
Regression test for issue #2331: deleting memories must also purge their vectors.
3+
4+
`Neo4jCommunityGraphDB.delete_node_by_prams` removed the Neo4j nodes but never
5+
removed the matching embeddings from `vec_db`, so a deleted memory stayed
6+
retrievable through vector search while the delete API still answered "success".
7+
`clear()` and `delete_node_by_mem_cube_id()` already purged `vec_db`;
8+
`delete_node_by_prams` simply omitted that step.
9+
10+
The Neo4j driver and the vector store are both faked, so these tests need no
11+
Neo4j, no Qdrant, and no network.
12+
"""
13+
14+
from types import SimpleNamespace
15+
from unittest.mock import MagicMock, patch
16+
17+
import pytest
18+
19+
from memos.configs.graph_db import Neo4jCommunityGraphDBConfig, Neo4jGraphDBConfig
20+
21+
22+
class _FakeNeo4jSession:
23+
"""Minimal stand-in for a neo4j session that tracks a set of live node ids."""
24+
25+
def __init__(self, live_ids: set[str], issued_queries: list[str]):
26+
self.live_ids = live_ids
27+
self.issued_queries = issued_queries
28+
29+
def run(self, query, params=None, **kwargs):
30+
self.issued_queries.append(query)
31+
params = params or kwargs
32+
ids = params.get("memory_ids") or params.get("filter_ids")
33+
matched = sorted(i for i in self.live_ids if ids is None or i in ids)
34+
if "count(n) AS node_count" in query:
35+
return _FakeResult([{"node_count": len(matched)}])
36+
if "DETACH DELETE" in query:
37+
self.live_ids.difference_update(matched)
38+
return _FakeResult([], deleted=len(matched))
39+
if "RETURN n.id AS id" in query:
40+
return _FakeResult([{"id": i} for i in matched])
41+
return _FakeResult([])
42+
43+
def __enter__(self):
44+
return self
45+
46+
def __exit__(self, *exc):
47+
return False
48+
49+
50+
class _FakeResult:
51+
"""Minimal stand-in for a neo4j Result (iterable, plus .single()/.consume())."""
52+
53+
def __init__(self, records, deleted: int = 0):
54+
self._records = records
55+
self._deleted = deleted
56+
57+
def __iter__(self):
58+
return iter(self._records)
59+
60+
def single(self):
61+
return self._records[0] if self._records else None
62+
63+
def consume(self):
64+
return SimpleNamespace(counters=SimpleNamespace(nodes_deleted=self._deleted))
65+
66+
67+
def _base_config_kwargs() -> dict:
68+
return {
69+
"uri": "bolt://localhost:7687",
70+
"user": "neo4j",
71+
"password": "test",
72+
"db_name": "test_memory_db",
73+
"auto_create": False,
74+
"use_multi_db": False,
75+
"user_name": "alice",
76+
"embedding_dimension": 3,
77+
}
78+
79+
80+
@pytest.fixture
81+
def community_config():
82+
"""`neo4j-community` is the backend that owns a `vec_db` mirror."""
83+
return Neo4jCommunityGraphDBConfig(
84+
**_base_config_kwargs(),
85+
vec_config={
86+
"backend": "qdrant",
87+
"config": {
88+
"collection_name": "test_vectors",
89+
"path": "/tmp/memos_test_vectors",
90+
},
91+
},
92+
)
93+
94+
95+
@pytest.fixture
96+
def enterprise_config():
97+
"""`neo4j` has no `vec_db`; it searches the native vector index instead."""
98+
return Neo4jGraphDBConfig(**_base_config_kwargs())
99+
100+
101+
@pytest.fixture
102+
def community_db(community_config):
103+
"""Community backend with a faked driver and a faked vector store."""
104+
from memos.graph_dbs.neo4j_community import Neo4jCommunityGraphDB
105+
106+
live_ids = {"mem-1", "mem-2", "mem-3"}
107+
issued_queries: list[str] = []
108+
109+
with patch("neo4j.GraphDatabase") as mock_gd:
110+
mock_gd.driver.return_value = MagicMock()
111+
with patch("memos.graph_dbs.neo4j_community.VecDBFactory.from_config") as mock_factory:
112+
mock_factory.return_value = MagicMock()
113+
db = Neo4jCommunityGraphDB(community_config)
114+
115+
db.driver = MagicMock()
116+
db.driver.session.side_effect = lambda **kwargs: _FakeNeo4jSession(live_ids, issued_queries)
117+
118+
vector_ids = set(live_ids)
119+
db.vec_db = MagicMock()
120+
db.vec_db.delete.side_effect = lambda ids: [vector_ids.discard(i) for i in ids]
121+
122+
db.live_ids = live_ids
123+
db.vector_ids = vector_ids
124+
db.issued_queries = issued_queries
125+
return db
126+
127+
128+
class TestCommunityDeletePurgesVectors:
129+
"""Regression: the community backend deletes vectors together with nodes."""
130+
131+
def test_delete_by_memory_ids_purges_vectors(self, community_db):
132+
deleted = community_db.delete_node_by_prams(memory_ids=["mem-1", "mem-2"])
133+
134+
assert deleted == 2
135+
assert community_db.live_ids == {"mem-3"}
136+
assert community_db.vector_ids == {"mem-3"}, "vectors of deleted nodes must not survive"
137+
assert community_db.vec_db.delete.call_count == 1
138+
assert sorted(community_db.vec_db.delete.call_args[0][0]) == ["mem-1", "mem-2"]
139+
140+
def test_graph_deletion_and_return_value_unchanged(self, community_db):
141+
"""The fix must not change which nodes the graph store deletes."""
142+
deleted = community_db.delete_node_by_prams(memory_ids=["mem-3"])
143+
144+
assert deleted == 1
145+
assert community_db.live_ids == {"mem-1", "mem-2"}
146+
assert community_db.vector_ids == {"mem-1", "mem-2"}
147+
assert any("DETACH DELETE" in q for q in community_db.issued_queries)
148+
149+
def test_vec_db_failure_does_not_mask_successful_delete(self, community_db):
150+
"""A vector store outage must not turn a committed delete into a failure."""
151+
community_db.vec_db.delete.side_effect = RuntimeError("qdrant unreachable")
152+
153+
deleted = community_db.delete_node_by_prams(memory_ids=["mem-1"])
154+
155+
assert deleted == 1
156+
assert community_db.live_ids == {"mem-2", "mem-3"}
157+
158+
def test_delete_by_filter_purges_vectors(self, community_db):
159+
"""Deleting by metadata filter must purge vectors too."""
160+
community_db.get_by_metadata = MagicMock(return_value=["mem-1", "mem-3"])
161+
162+
deleted = community_db.delete_node_by_prams(filter={"user_id": "u_1"})
163+
164+
assert deleted == 2
165+
assert community_db.live_ids == {"mem-2"}
166+
assert community_db.vector_ids == {"mem-2"}
167+
assert sorted(community_db.vec_db.delete.call_args[0][0]) == ["mem-1", "mem-3"]
168+
169+
def test_no_match_does_not_touch_vec_db(self, community_db):
170+
"""An empty match set must not issue a vector delete call."""
171+
deleted = community_db.delete_node_by_prams(memory_ids=["does-not-exist"])
172+
173+
assert deleted == 0
174+
assert community_db.live_ids == {"mem-1", "mem-2", "mem-3"}
175+
assert community_db.vector_ids == {"mem-1", "mem-2", "mem-3"}
176+
community_db.vec_db.delete.assert_not_called()
177+
178+
def test_clear_still_purges_all_vectors(self, community_db):
179+
"""The neighbouring `clear()` path keeps working after the refactor."""
180+
community_db.vec_db.get_by_filter.return_value = [
181+
SimpleNamespace(id="mem-1"),
182+
SimpleNamespace(id="mem-2"),
183+
]
184+
185+
community_db.clear(user_name="alice")
186+
187+
community_db.vec_db.get_by_filter.assert_called_once_with({"user_name": "alice"})
188+
community_db.vec_db.delete.assert_called_once_with(["mem-1", "mem-2"])
189+
assert community_db.live_ids == set()
190+
assert any("DETACH DELETE" in q for q in community_db.issued_queries)
191+
192+
193+
class TestBackendWithoutVecDB:
194+
"""`neo4j` sets up no `vec_db`, so the vector purge must degrade to a no-op."""
195+
196+
@pytest.fixture
197+
def enterprise_db(self, enterprise_config):
198+
from memos.graph_dbs.neo4j import Neo4jGraphDB
199+
200+
live_ids = {"mem-1", "mem-2", "mem-3"}
201+
issued_queries: list[str] = []
202+
203+
with patch("neo4j.GraphDatabase") as mock_gd:
204+
mock_gd.driver.return_value = MagicMock()
205+
db = Neo4jGraphDB(enterprise_config)
206+
207+
db.driver = MagicMock()
208+
db.driver.session.side_effect = lambda **kwargs: _FakeNeo4jSession(live_ids, issued_queries)
209+
db.live_ids = live_ids
210+
db.issued_queries = issued_queries
211+
return db
212+
213+
def test_delete_without_vec_db_succeeds(self, enterprise_db):
214+
deleted = enterprise_db.delete_node_by_prams(memory_ids=["mem-1", "mem-2"])
215+
216+
assert deleted == 2
217+
assert enterprise_db.live_ids == {"mem-3"}
218+
assert any("DETACH DELETE" in q for q in enterprise_db.issued_queries)
219+
220+
def test_no_match_without_vec_db(self, enterprise_db):
221+
deleted = enterprise_db.delete_node_by_prams(memory_ids=["does-not-exist"])
222+
223+
assert deleted == 0
224+
assert enterprise_db.live_ids == {"mem-1", "mem-2", "mem-3"}

0 commit comments

Comments
 (0)