mirror of
https://github.com/meta-llama/llama-stack.git
synced 2025-10-15 06:37:58 +00:00
28 lines
1.1 KiB
Python
28 lines
1.1 KiB
Python
import pytest
|
|
import pytest_asyncio
|
|
from llama_stack.distribution.store import *
|
|
from llama_stack.apis.memory_banks import GraphMemoryBankDef, VectorMemoryBankDef
|
|
from llama_stack.distribution.utils.config_dirs import DISTRIBS_BASE_DIR
|
|
from llama_stack.providers.utils.kvstore import kvstore_impl, SqliteKVStoreConfig
|
|
from llama_stack.distribution.datatypes import * # noqa: F403
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_registry():
|
|
registry = DiskRegistry(await kvstore_impl(SqliteKVStoreConfig()))
|
|
bank = VectorMemoryBankDef(
|
|
identifier="test_bank",
|
|
embedding_model="all-MiniLM-L6-v2",
|
|
chunk_size_in_tokens=512,
|
|
overlap_size_in_tokens=64,
|
|
provider_id="bar",
|
|
)
|
|
|
|
await registry.register(bank)
|
|
result_bank = await registry.get("test_bank")
|
|
# assert result_bank == bank
|
|
assert result_bank.identifier == bank.identifier
|
|
assert result_bank.embedding_model == bank.embedding_model
|
|
assert result_bank.chunk_size_in_tokens == bank.chunk_size_in_tokens
|
|
assert result_bank.overlap_size_in_tokens == bank.overlap_size_in_tokens
|
|
assert result_bank.provider_id == bank.provider_id
|