diff --git a/tests/memory/long_term_memory_backends/test_redis_ltm_backend.py b/tests/memory/long_term_memory_backends/test_redis_ltm_backend.py new file mode 100644 index 000000000..bbe46f45b --- /dev/null +++ b/tests/memory/long_term_memory_backends/test_redis_ltm_backend.py @@ -0,0 +1,118 @@ +# 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 redis.exceptions import ResponseError +from redisvl.index.storage import BaseStorage + +from veadk.memory.long_term_memory_backends import redis_backend +from veadk.memory.long_term_memory_backends.redis_backend import RedisLTMBackend + + +class _FakeSearch: + def __init__(self, indexes, name): + self._indexes, self._name = indexes, name + + def info(self): + if self._name not in self._indexes: + raise ResponseError("Unknown index name") + return {"index_definition": ["prefixes", [self._indexes[self._name]]]} + + def dropindex(self, delete_documents=False): + assert delete_documents is False + del self._indexes[self._name] + + +class _FakeRedis: + indexes: dict = {} + + def __init__(self, **_): + pass + + def ft(self, name): + return _FakeSearch(self.indexes, name) + + +class _FakeRedisVectorStore: + def __init__(self, schema, redis_client): + self.schema = schema + self.index_name = schema.index.name + + +@pytest.fixture +def schemas(monkeypatch): + monkeypatch.setenv("MODEL_EMBEDDING_API_KEY", "test-key") + monkeypatch.setattr( + redis_backend, + "create_embedding_model", + lambda **_: MockEmbedding(embed_dim=8), + ) + captured = {} + + def fake_from_vector_store(vector_store, embed_model): + captured[vector_store.schema.index.name] = vector_store.schema + return None + + _FakeRedis.indexes = {} + monkeypatch.setattr(redis_backend, "Redis", _FakeRedis) + monkeypatch.setattr(redis_backend, "RedisVectorStore", _FakeRedisVectorStore) + monkeypatch.setattr( + redis_backend.VectorStoreIndex, + "from_vector_store", + staticmethod(fake_from_vector_store), + ) + backend = RedisLTMBackend(index="app") + backend._create_vector_index("veadk-ltm/app/u1") + backend._create_vector_index("veadk-ltm/app/u10") + return captured + + +def _key(schema, node_id: str) -> str: + return BaseStorage._key(node_id, schema.index.prefix, schema.index.key_separator) + + +def test_user_index_prefix_does_not_match_keys_of_other_users(schemas): + u1 = schemas["veadk-ltm/app/u1"] + u10 = schemas["veadk-ltm/app/u10"] + + # RediSearch indexes every hash whose key starts with the index prefix. + assert not _key(u10, "node").startswith(u1.index.prefix) + assert _key(u1, "node").startswith(u1.index.prefix) + + +def test_user_keys_keep_existing_format(schemas): + u1 = schemas["veadk-ltm/app/u1"] + + assert _key(u1, "node") == "veadk-ltm/app/u1_node" + + +def test_index_created_with_old_prefix_is_redefined_without_deleting_data(schemas): + _FakeRedis.indexes = {"veadk-ltm/app/u1": "veadk-ltm/app/u1"} + backend = RedisLTMBackend(index="app") + + backend._create_vector_index("veadk-ltm/app/u1") + + # Only the stale index definition was dropped (documents kept), so the + # vector store recreates it with the new prefix. + assert "veadk-ltm/app/u1" not in _FakeRedis.indexes + + +def test_index_with_current_prefix_is_kept(schemas): + _FakeRedis.indexes = {"veadk-ltm/app/u1": "veadk-ltm/app/u1_"} + backend = RedisLTMBackend(index="app") + + backend._create_vector_index("veadk-ltm/app/u1") + + assert _FakeRedis.indexes == {"veadk-ltm/app/u1": "veadk-ltm/app/u1_"} diff --git a/veadk/memory/long_term_memory_backends/redis_backend.py b/veadk/memory/long_term_memory_backends/redis_backend.py index c199f3c6c..1fd09756e 100644 --- a/veadk/memory/long_term_memory_backends/redis_backend.py +++ b/veadk/memory/long_term_memory_backends/redis_backend.py @@ -30,6 +30,7 @@ try: from llama_index.vector_stores.redis import RedisVectorStore from redis import Redis + from redis.exceptions import ResponseError from redisvl.schema import IndexSchema except ImportError: @@ -76,11 +77,18 @@ def _create_vector_index(self, index: str) -> VectorStoreIndex: password=self.redis_config.password, ) + prefix = f"{index}_" + self._drop_index_with_stale_prefix(redis_client, index, prefix) + # Create an index for each user # Should be Optimized in the future schema = IndexSchema.from_dict( { - "index": {"name": index, "prefix": index, "key_separator": "_"}, + # RediSearch indexes every key that starts with the prefix, so + # the prefix must end with the separator. Otherwise the index of + # user `u1` also picks up the keys of user `u10`. Keys keep the + # `{index}_{id}` format, so existing memories are still found. + "index": {"name": index, "prefix": prefix, "key_separator": ""}, "fields": [ {"name": "id", "type": "tag", "attrs": {"sortable": False}}, {"name": "doc_id", "type": "tag", "attrs": {"sortable": False}}, @@ -107,6 +115,34 @@ def _create_vector_index(self, index: str) -> VectorStoreIndex: vector_store=vector_store, embed_model=self._embed_model ) + @staticmethod + def _drop_index_with_stale_prefix( + redis_client: Redis, index: str, prefix: str + ) -> None: + """Drop the definition (not the data) of an index created with an old prefix. + + Indexes created before the prefix fix still match other users' keys. + Dropping only the definition keeps the stored memories, and the index + is recreated with the correct prefix right after. + """ + try: + info = redis_client.ft(index).info() + except ResponseError: + return # index does not exist yet + + definition = info.get("index_definition") or [] + if isinstance(definition, list): + definition = dict(zip(definition[::2], definition[1::2])) + prefixes = [ + p.decode() if isinstance(p, bytes) else p + for p in definition.get("prefixes", definition.get(b"prefixes", [])) + ] + if prefixes and prefixes != [prefix]: + logger.info( + f"Recreate Redis index {index} with prefix {prefix} (was {prefixes})" + ) + redis_client.ft(index).dropindex(delete_documents=False) + @override def save_memory(self, user_id: str, event_strings: list[str], **kwargs) -> bool: index = f"veadk-ltm/{self.index}/{user_id}"