diff --git a/memory.py b/memory.py index 08dd5d7e..9732028e 100644 --- a/memory.py +++ b/memory.py @@ -84,7 +84,8 @@ def _index_memory(self, memory_id: str, text: str, *, kind: str, status: str = " try: self._collection.upsert( ids=[memory_id], documents=[text], - metadatas=[{"kind": kind, "status": status, "workspace_id": self.workspace_id, + metadatas=[{"kind": kind, "status": status, "user_id": self.user_id, + "workspace_id": self.workspace_id, "session_id": self.session_id or "", "memory_id": memory_id}], ) except Exception as exc: @@ -137,7 +138,7 @@ def add_file(self, rel_path: str, content: str) -> str: self._files_collection.upsert( ids=[file_id], documents=[text], metadatas=[{"rel_path": rel_path, "content_hash": stable_hash(content), - "workspace_id": self.workspace_id}], + "user_id": self.user_id, "workspace_id": self.workspace_id}], ) except Exception as exc: self._last_error = f"file index write failed: {exc}" @@ -149,7 +150,9 @@ def remove_stale_files(self, valid_paths: set[str]) -> int: try: current = self._files_collection.get(include=["metadatas"]) stale = [doc_id for doc_id, meta in zip(current.get("ids", []), current.get("metadatas", [])) - if meta.get("workspace_id") == self.workspace_id and meta.get("rel_path") not in valid_paths] + if meta.get("user_id") == self.user_id + and meta.get("workspace_id") == self.workspace_id + and meta.get("rel_path") not in valid_paths] if stale: self._files_collection.delete(ids=stale) except Exception as exc: @@ -168,18 +171,33 @@ def recall(self, query: str, n_results: int = 2, *, include_scoped: bool = False limit = max(1, min(int(n_results), 100)) if self._collection is not None: try: + where = {"$and": [ + {"user_id": self.user_id}, + {"workspace_id": self.workspace_id}, + {"status": "active"}, + {"session_id": self.session_id or ""} if self.session_id is None else { + "$or": [{"session_id": ""}, {"session_id": self.session_id}], + }, + ]} result = self._collection.query( query_texts=[query], n_results=min(limit * 4, max(1, self._collection.count())), - where={"$and": [{"workspace_id": self.workspace_id}, {"status": "active"}]}, + where=where, ) docs = result.get("documents", [[]]) - if docs and docs[0]: + ids = result.get("ids", [[]]) + if docs and docs[0] and ids and ids[0]: rows = self.store.list_memories(status="active", limit=10000, workspace_id=self.workspace_id, session_id=self.session_id, user_id=self.user_id) - metadata = {row["content"]: row.get("metadata", {}) for row in rows} - return [doc for doc in docs[0] if not doc.startswith("FILE:") and ( - include_scoped or metadata.get(doc, {}).get("visibility", "public") == "public")][:limit] + by_id = {row["id"]: row for row in rows} + scoped_documents = [] + for memory_id, doc in zip(ids[0], docs[0]): + row = by_id.get(memory_id) + if row is None or doc.startswith("FILE:"): + continue + if include_scoped or row.get("metadata", {}).get("visibility", "public") == "public": + scoped_documents.append(doc) + return scoped_documents[:limit] except Exception as exc: self._last_error = f"memory index query failed: {exc}" rows = self.store.list_memories(status="active", limit=10000, workspace_id=self.workspace_id, diff --git a/tests/test_multi_party_memory.py b/tests/test_multi_party_memory.py index 3b126190..98e3e94b 100644 --- a/tests/test_multi_party_memory.py +++ b/tests/test_multi_party_memory.py @@ -1,9 +1,13 @@ +import os import tempfile import unittest from pathlib import Path +from unittest.mock import patch from learning_engine import LearningEngine from memory import MemoryBank +from fastapi.testclient import TestClient +import server class MultiPartyMemoryTests(unittest.TestCase): @@ -91,6 +95,61 @@ def test_memory_bank_reads_are_scoped_to_its_user(self): self.assertEqual(bob.get_all()[0], []) self.assertEqual(bob.store.list_events(workspace_id="project", user_id="bob"), []) + def test_chroma_recall_filters_user_and_session_before_returning_documents(self): + with patch.dict(os.environ, {"KYROZEN_VECTOR_PATH": str(Path(self.directory.name) / "chroma")}): + alice = MemoryBank(self.memory.db_path, user_id="alice", workspace_id="shared", + session_id="alice-session") + bob = MemoryBank(self.memory.db_path, user_id="bob", workspace_id="shared", + session_id="bob-session") + self.assertIsNotNone(alice._collection) + alice.add_log( + "PRIVATE_MARKER_ALICE", + kind="fact", + metadata={"visibility": "private", "speaker": "alice"}, + ) + alice.add_log("SESSION_MARKER_ALICE", kind="fact") + MemoryBank(self.memory.db_path, user_id="alice", workspace_id="shared").add_log( + "GLOBAL_MARKER_ALICE", kind="fact" + ) + + self.assertEqual(bob.recall("PRIVATE_MARKER_ALICE", n_results=5), []) + self.assertEqual(bob.recall_records("PRIVATE_MARKER_ALICE", n_results=5), []) + other_session = MemoryBank(self.memory.db_path, user_id="alice", workspace_id="shared", + session_id="other-session") + self.assertNotIn("SESSION_MARKER_ALICE", other_session.recall("SESSION_MARKER_ALICE", n_results=5)) + self.assertIn("GLOBAL_MARKER_ALICE", alice.recall("GLOBAL_MARKER_ALICE", n_results=5)) + self.assertIn( + "PRIVATE_MARKER_ALICE", + [item["content"] for item in alice.recall_records( + "PRIVATE_MARKER_ALICE", n_results=5, + speaker="alice", authorized_speakers={"alice"})], + ) + + def test_api_memory_recall_does_not_return_another_user_vector_document(self): + with tempfile.TemporaryDirectory(prefix="openkyrozen-memory-api-") as directory: + root = Path(directory) + with patch.dict(os.environ, {"KYROZEN_VECTOR_PATH": str(root / "chroma")}): + owner = MemoryBank(root / "state.sqlite3", user_id=server._SERVER_ACTOR_ID, + workspace_id="api-project", session_id="api-session") + other = MemoryBank(root / "state.sqlite3", user_id="other-user", + workspace_id="api-project", session_id="api-session") + other.add_log( + "API_PRIVATE_MARKER_OTHER", + kind="fact", + metadata={"visibility": "private", "speaker": "other-user"}, + ) + original = server._agent.memory_bank + server._agent.memory_bank = owner + try: + response = TestClient(server.app).get( + "/api/v2/memory", + params={"q": "API_PRIVATE_MARKER_OTHER", "session_id": "api-session"}, + ) + finally: + server._agent.memory_bank = original + self.assertEqual(response.status_code, 200, response.text) + self.assertEqual(response.json()["results"], []) + if __name__ == "__main__": unittest.main()