mirror of
https://github.com/meta-llama/llama-stack.git
synced 2025-08-02 08:44:44 +00:00
Add Redis KVStore integration test
Related to #506 Add integration tests for Redis KVStore to ensure functional parity with other KVStores. * **`llama_stack/distribution/store/tests/test_registry.py`**: - Add fixture for `RedisKVStoreConfig`. - Add tests for `RedisKVStoreConfig` covering initialization, connection, CRUD operations, namespaced keys, expiration handling, error handling, and concurrency. - Add tests for `redis_registry` and `cached_redis_registry` to verify registration, retrieval, and updates of memory banks and models. * **`llama_stack/providers/tests/memory/test_memory.py`**: - Add fixture for `RedisKVStoreConfig`. - Add tests for `redis_kvstore` covering initialization, CRUD operations, namespaced keys, expiration handling, error handling, and concurrency.
This commit is contained in:
Vishwanath Martur
parent
815f4af6cf
commit
45c6453a8b
2 changed files with 269 additions and 1 deletions
|
|
@ -11,7 +11,7 @@ import pytest_asyncio
|
|||
from llama_stack.distribution.store import * # noqa F403
|
||||
from llama_stack.apis.inference import Model
|
||||
from llama_stack.apis.memory_banks import VectorMemoryBank
|
||||
from llama_stack.providers.utils.kvstore import kvstore_impl, SqliteKVStoreConfig
|
||||
from llama_stack.providers.utils.kvstore import kvstore_impl, SqliteKVStoreConfig, RedisKVStoreConfig
|
||||
from llama_stack.distribution.datatypes import * # noqa F403
|
||||
|
||||
|
||||
|
|
@ -23,6 +23,11 @@ def config():
|
|||
return config
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def redis_config():
|
||||
return RedisKVStoreConfig(host="localhost", port=6379)
|
||||
|
||||
|
||||
@pytest_asyncio.fixture
|
||||
async def registry(config):
|
||||
registry = DiskDistributionRegistry(await kvstore_impl(config))
|
||||
|
|
@ -30,6 +35,13 @@ async def registry(config):
|
|||
return registry
|
||||
|
||||
|
||||
@pytest_asyncio.fixture
|
||||
async def redis_registry(redis_config):
|
||||
registry = DiskDistributionRegistry(await kvstore_impl(redis_config))
|
||||
await registry.initialize()
|
||||
return registry
|
||||
|
||||
|
||||
@pytest_asyncio.fixture
|
||||
async def cached_registry(config):
|
||||
registry = CachedDiskDistributionRegistry(await kvstore_impl(config))
|
||||
|
|
@ -37,6 +49,13 @@ async def cached_registry(config):
|
|||
return registry
|
||||
|
||||
|
||||
@pytest_asyncio.fixture
|
||||
async def cached_redis_registry(redis_config):
|
||||
registry = CachedDiskDistributionRegistry(await kvstore_impl(redis_config))
|
||||
await registry.initialize()
|
||||
return registry
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def sample_bank():
|
||||
return VectorMemoryBank(
|
||||
|
|
@ -65,6 +84,13 @@ async def test_registry_initialization(registry):
|
|||
assert len(results) == 0
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_redis_registry_initialization(redis_registry):
|
||||
# Test empty registry
|
||||
results = await redis_registry.get("nonexistent", "nonexistent")
|
||||
assert len(results) == 0
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_basic_registration(registry, sample_bank, sample_model):
|
||||
print(f"Registering {sample_bank}")
|
||||
|
|
@ -88,6 +114,29 @@ async def test_basic_registration(registry, sample_bank, sample_model):
|
|||
assert result_model.provider_id == sample_model.provider_id
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_redis_basic_registration(redis_registry, sample_bank, sample_model):
|
||||
print(f"Registering {sample_bank}")
|
||||
await redis_registry.register(sample_bank)
|
||||
print(f"Registering {sample_model}")
|
||||
await redis_registry.register(sample_model)
|
||||
print("Getting bank")
|
||||
results = await redis_registry.get("memory_bank", "test_bank")
|
||||
assert len(results) == 1
|
||||
result_bank = results[0]
|
||||
assert result_bank.identifier == sample_bank.identifier
|
||||
assert result_bank.embedding_model == sample_bank.embedding_model
|
||||
assert result_bank.chunk_size_in_tokens == sample_bank.chunk_size_in_tokens
|
||||
assert result_bank.overlap_size_in_tokens == sample_bank.overlap_size_in_tokens
|
||||
assert result_bank.provider_id == sample_bank.provider_id
|
||||
|
||||
results = await redis_registry.get("model", "test_model")
|
||||
assert len(results) == 1
|
||||
result_model = results[0]
|
||||
assert result_model.identifier == sample_model.identifier
|
||||
assert result_model.provider_id == sample_model.provider_id
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_cached_registry_initialization(config, sample_bank, sample_model):
|
||||
# First populate the disk registry
|
||||
|
|
@ -110,6 +159,28 @@ async def test_cached_registry_initialization(config, sample_bank, sample_model)
|
|||
assert result_bank.provider_id == sample_bank.provider_id
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_cached_redis_registry_initialization(redis_config, sample_bank, sample_model):
|
||||
# First populate the redis registry
|
||||
redis_registry = DiskDistributionRegistry(await kvstore_impl(redis_config))
|
||||
await redis_registry.initialize()
|
||||
await redis_registry.register(sample_bank)
|
||||
await redis_registry.register(sample_model)
|
||||
|
||||
# Test cached version loads from redis
|
||||
cached_redis_registry = CachedDiskDistributionRegistry(await kvstore_impl(redis_config))
|
||||
await cached_redis_registry.initialize()
|
||||
|
||||
results = await cached_redis_registry.get("memory_bank", "test_bank")
|
||||
assert len(results) == 1
|
||||
result_bank = results[0]
|
||||
assert result_bank.identifier == sample_bank.identifier
|
||||
assert result_bank.embedding_model == sample_bank.embedding_model
|
||||
assert result_bank.chunk_size_in_tokens == sample_bank.chunk_size_in_tokens
|
||||
assert result_bank.overlap_size_in_tokens == sample_bank.overlap_size_in_tokens
|
||||
assert result_bank.provider_id == sample_bank.provider_id
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_cached_registry_updates(config):
|
||||
cached_registry = CachedDiskDistributionRegistry(await kvstore_impl(config))
|
||||
|
|
@ -142,6 +213,38 @@ async def test_cached_registry_updates(config):
|
|||
assert result_bank.provider_id == new_bank.provider_id
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_cached_redis_registry_updates(redis_config):
|
||||
cached_redis_registry = CachedDiskDistributionRegistry(await kvstore_impl(redis_config))
|
||||
await cached_redis_registry.initialize()
|
||||
|
||||
new_bank = VectorMemoryBank(
|
||||
identifier="test_bank_2",
|
||||
embedding_model="all-MiniLM-L6-v2",
|
||||
chunk_size_in_tokens=256,
|
||||
overlap_size_in_tokens=32,
|
||||
provider_resource_id="test_bank_2",
|
||||
provider_id="baz",
|
||||
)
|
||||
await cached_redis_registry.register(new_bank)
|
||||
|
||||
# Verify in cache
|
||||
results = await cached_redis_registry.get("memory_bank", "test_bank_2")
|
||||
assert len(results) == 1
|
||||
result_bank = results[0]
|
||||
assert result_bank.identifier == new_bank.identifier
|
||||
assert result_bank.provider_id == new_bank.provider_id
|
||||
|
||||
# Verify persisted to redis
|
||||
new_registry = DiskDistributionRegistry(await kvstore_impl(redis_config))
|
||||
await new_registry.initialize()
|
||||
results = await new_registry.get("memory_bank", "test_bank_2")
|
||||
assert len(results) == 1
|
||||
result_bank = results[0]
|
||||
assert result_bank.identifier == new_bank.identifier
|
||||
assert result_bank.provider_id == new_bank.provider_id
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_duplicate_provider_registration(config):
|
||||
cached_registry = CachedDiskDistributionRegistry(await kvstore_impl(config))
|
||||
|
|
@ -174,6 +277,38 @@ async def test_duplicate_provider_registration(config):
|
|||
) # Original values preserved
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_redis_duplicate_provider_registration(redis_config):
|
||||
cached_redis_registry = CachedDiskDistributionRegistry(await kvstore_impl(redis_config))
|
||||
await cached_redis_registry.initialize()
|
||||
|
||||
original_bank = VectorMemoryBank(
|
||||
identifier="test_bank_2",
|
||||
embedding_model="all-MiniLM-L6-v2",
|
||||
chunk_size_in_tokens=256,
|
||||
overlap_size_in_tokens=32,
|
||||
provider_resource_id="test_bank_2",
|
||||
provider_id="baz",
|
||||
)
|
||||
await cached_redis_registry.register(original_bank)
|
||||
|
||||
duplicate_bank = VectorMemoryBank(
|
||||
identifier="test_bank_2",
|
||||
embedding_model="different-model",
|
||||
chunk_size_in_tokens=128,
|
||||
overlap_size_in_tokens=16,
|
||||
provider_resource_id="test_bank_2",
|
||||
provider_id="baz", # Same provider_id
|
||||
)
|
||||
await cached_redis_registry.register(duplicate_bank)
|
||||
|
||||
results = await cached_redis_registry.get("memory_bank", "test_bank_2")
|
||||
assert len(results) == 1 # Still only one result
|
||||
assert (
|
||||
results[0].embedding_model == original_bank.embedding_model
|
||||
) # Original values preserved
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_get_all_objects(config):
|
||||
cached_registry = CachedDiskDistributionRegistry(await kvstore_impl(config))
|
||||
|
|
@ -213,3 +348,44 @@ async def test_get_all_objects(config):
|
|||
assert (
|
||||
stored_bank.overlap_size_in_tokens == original_bank.overlap_size_in_tokens
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_redis_get_all_objects(redis_config):
|
||||
cached_redis_registry = CachedDiskDistributionRegistry(await kvstore_impl(redis_config))
|
||||
await cached_redis_registry.initialize()
|
||||
|
||||
# Create multiple test banks
|
||||
test_banks = [
|
||||
VectorMemoryBank(
|
||||
identifier=f"test_bank_{i}",
|
||||
embedding_model="all-MiniLM-L6-v2",
|
||||
chunk_size_in_tokens=256,
|
||||
overlap_size_in_tokens=32,
|
||||
provider_resource_id=f"test_bank_{i}",
|
||||
provider_id=f"provider_{i}",
|
||||
)
|
||||
for i in range(3)
|
||||
]
|
||||
|
||||
# Register all banks
|
||||
for bank in test_banks:
|
||||
await cached_redis_registry.register(bank)
|
||||
|
||||
# Test get_all retrieval
|
||||
all_results = await cached_redis_registry.get_all()
|
||||
assert len(all_results) == 3
|
||||
|
||||
# Verify each bank was stored correctly
|
||||
for original_bank in test_banks:
|
||||
matching_banks = [
|
||||
b for b in all_results if b.identifier == original_bank.identifier
|
||||
]
|
||||
assert len(matching_banks) == 1
|
||||
stored_bank = matching_banks[0]
|
||||
assert stored_bank.embedding_model == original_bank.embedding_model
|
||||
assert stored_bank.provider_id == original_bank.provider_id
|
||||
assert stored_bank.chunk_size_in_tokens == original_bank.chunk_size_in_tokens
|
||||
assert (
|
||||
stored_bank.overlap_size_in_tokens == original_bank.overlap_size_in_tokens
|
||||
)
|
||||
|
|
|
|||
|
|
@ -5,12 +5,14 @@
|
|||
# the root directory of this source tree.
|
||||
|
||||
import uuid
|
||||
from datetime import datetime, timedelta
|
||||
|
||||
import pytest
|
||||
|
||||
from llama_stack.apis.memory import * # noqa: F403
|
||||
from llama_stack.distribution.datatypes import * # noqa: F403
|
||||
from llama_stack.apis.memory_banks.memory_banks import VectorMemoryBankParams
|
||||
from llama_stack.providers.utils.kvstore import RedisKVStoreConfig, kvstore_impl
|
||||
|
||||
# How to run this test:
|
||||
#
|
||||
|
|
@ -45,6 +47,18 @@ def sample_documents():
|
|||
]
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def redis_config():
|
||||
return RedisKVStoreConfig(host="localhost", port=6379)
|
||||
|
||||
|
||||
@pytest_asyncio.fixture
|
||||
async def redis_kvstore(redis_config):
|
||||
kvstore = await kvstore_impl(redis_config)
|
||||
await kvstore.initialize()
|
||||
return kvstore
|
||||
|
||||
|
||||
async def register_memory_bank(
|
||||
banks_impl: MemoryBanks, inference_model: str
|
||||
) -> MemoryBank:
|
||||
|
|
@ -177,6 +191,84 @@ class TestMemory:
|
|||
print("The scores are:", response5.scores)
|
||||
assert all(score >= 0.01 for score in response5.scores)
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_redis_kvstore_initialization(self, redis_kvstore):
|
||||
# Test empty Redis KVStore
|
||||
result = await redis_kvstore.get("nonexistent")
|
||||
assert result is None
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_redis_kvstore_crud_operations(self, redis_kvstore):
|
||||
# Test set and get operations
|
||||
await redis_kvstore.set("key1", "value1")
|
||||
result = await redis_kvstore.get("key1")
|
||||
assert result == "value1"
|
||||
|
||||
# Test delete operation
|
||||
await redis_kvstore.delete("key1")
|
||||
result = await redis_kvstore.get("key1")
|
||||
assert result is None
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_redis_kvstore_namespaced_keys(self, redis_config):
|
||||
redis_config.namespace = "test_namespace"
|
||||
kvstore = await kvstore_impl(redis_config)
|
||||
await kvstore.initialize()
|
||||
|
||||
# Test namespaced set and get operations
|
||||
await kvstore.set("key1", "value1")
|
||||
result = await kvstore.get("key1")
|
||||
assert result == "value1"
|
||||
|
||||
# Test delete operation with namespaced key
|
||||
await kvstore.delete("key1")
|
||||
result = await kvstore.get("key1")
|
||||
assert result is None
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_redis_kvstore_expiration_handling(self, redis_kvstore):
|
||||
# Test set operation with expiration
|
||||
expiration = datetime.now() + timedelta(seconds=1)
|
||||
await redis_kvstore.set("key1", "value1", expiration=expiration)
|
||||
|
||||
# Verify key is retrievable before expiration
|
||||
result = await redis_kvstore.get("key1")
|
||||
assert result == "value1"
|
||||
|
||||
# Wait for expiration and verify key is no longer retrievable
|
||||
await asyncio.sleep(2)
|
||||
result = await redis_kvstore.get("key1")
|
||||
assert result is None
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_redis_kvstore_error_handling(self, redis_config):
|
||||
# Test connection to invalid Redis server
|
||||
invalid_config = RedisKVStoreConfig(host="invalid_host", port=6379)
|
||||
with pytest.raises(RuntimeError):
|
||||
await kvstore_impl(invalid_config)
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_redis_kvstore_concurrency(self, redis_kvstore):
|
||||
# Test concurrent access to Redis KVStore
|
||||
async def set_value(key, value):
|
||||
await redis_kvstore.set(key, value)
|
||||
|
||||
async def get_value(key):
|
||||
return await redis_kvstore.get(key)
|
||||
|
||||
tasks = []
|
||||
for i in range(10):
|
||||
tasks.append(set_value(f"key{i}", f"value{i}"))
|
||||
await asyncio.gather(*tasks)
|
||||
|
||||
tasks = []
|
||||
for i in range(10):
|
||||
tasks.append(get_value(f"key{i}"))
|
||||
results = await asyncio.gather(*tasks)
|
||||
|
||||
for i in range(10):
|
||||
assert results[i] == f"value{i}"
|
||||
|
||||
|
||||
def assert_valid_response(response: QueryDocumentsResponse):
|
||||
assert isinstance(response, QueryDocumentsResponse)
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue