diff --git a/tests/memory/long_term_memory_backends/test_in_memory_ltm_backend.py b/tests/memory/long_term_memory_backends/test_in_memory_ltm_backend.py new file mode 100644 index 000000000..0ef2ab3b4 --- /dev/null +++ b/tests/memory/long_term_memory_backends/test_in_memory_ltm_backend.py @@ -0,0 +1,50 @@ +# Copyright (c) 2025 Beijing Volcano Engine Technology Co., Ltd. and/or its affiliates. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import pytest +from llama_index.core.embeddings import MockEmbedding + +from veadk.memory.long_term_memory_backends import in_memory_backend +from veadk.memory.long_term_memory_backends.in_memory_backend import ( + InMemoryLTMBackend, +) + + +@pytest.fixture +def backend(monkeypatch): + monkeypatch.setenv("MODEL_EMBEDDING_API_KEY", "test-key") + monkeypatch.setattr( + in_memory_backend, + "create_embedding_model", + lambda **_: MockEmbedding(embed_dim=8), + ) + return InMemoryLTMBackend(index="test") + + +def test_search_memory_only_returns_memories_of_the_given_user(backend): + backend.save_memory(user_id="alice", event_strings=["alice likes tea"]) + backend.save_memory(user_id="bob", event_strings=["bob likes coffee"]) + + assert backend.search_memory(user_id="alice", query="likes", top_k=5) == [ + "alice likes tea" + ] + assert backend.search_memory(user_id="bob", query="likes", top_k=5) == [ + "bob likes coffee" + ] + + +def test_search_memory_for_user_without_memories_is_empty(backend): + backend.save_memory(user_id="alice", event_strings=["alice likes tea"]) + + assert backend.search_memory(user_id="carol", query="likes", top_k=5) == [] diff --git a/veadk/memory/long_term_memory_backends/in_memory_backend.py b/veadk/memory/long_term_memory_backends/in_memory_backend.py index cf047a11c..1c24a0e95 100644 --- a/veadk/memory/long_term_memory_backends/in_memory_backend.py +++ b/veadk/memory/long_term_memory_backends/in_memory_backend.py @@ -35,7 +35,15 @@ def model_post_init(self, __context: Any) -> None: api_key=self.embedding_config.api_key, api_base=self.embedding_config.api_base, ) - self._vector_index = VectorStoreIndex([], embed_model=self._embed_model) + # One index per user, so memories never leak across users. + self._vector_indexes: dict[str, VectorStoreIndex] = {} + + def _get_vector_index(self, user_id: str) -> VectorStoreIndex: + if user_id not in self._vector_indexes: + self._vector_indexes[user_id] = VectorStoreIndex( + [], embed_model=self._embed_model + ) + return self._vector_indexes[user_id] def precheck_index_naming(self): # no checking @@ -43,17 +51,21 @@ def precheck_index_naming(self): @override def save_memory(self, user_id: str, event_strings: list[str], **kwargs) -> bool: + vector_index = self._get_vector_index(user_id) for event_string in event_strings: document = Document(text=event_string) nodes = self._split_documents([document]) - self._vector_index.insert_nodes(nodes) + vector_index.insert_nodes(nodes) return True @override def search_memory( self, user_id: str, query: str, top_k: int, **kwargs ) -> list[str]: - _retriever = self._vector_index.as_retriever(similarity_top_k=top_k) + vector_index = self._vector_indexes.get(user_id) + if vector_index is None: + return [] + _retriever = vector_index.as_retriever(similarity_top_k=top_k) retrieved_nodes = _retriever.retrieve(query) return [node.text for node in retrieved_nodes]