mirror of
https://github.com/meta-llama/llama-stack.git
synced 2025-10-15 06:37:58 +00:00
52 lines
1.9 KiB
Python
52 lines
1.9 KiB
Python
# Copyright (c) Meta Platforms, Inc. and affiliates.
|
|
# All rights reserved.
|
|
#
|
|
# This source code is licensed under the terms described in the LICENSE file in
|
|
# the root directory of this source tree.
|
|
|
|
import os
|
|
|
|
import pytest
|
|
from llama_stack.distribution.store import * # noqa: F403
|
|
from llama_stack.apis.memory_banks import VectorMemoryBankDef
|
|
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():
|
|
config = SqliteKVStoreConfig(db_path="/tmp/test_registry.db")
|
|
# delete the file if it exists
|
|
if os.path.exists(config.db_path):
|
|
os.remove(config.db_path)
|
|
registry = DiskRegistry(await kvstore_impl(config))
|
|
bank = VectorMemoryBankDef(
|
|
identifier="test_bank",
|
|
embedding_model="all-MiniLM-L6-v2",
|
|
chunk_size_in_tokens=512,
|
|
overlap_size_in_tokens=64,
|
|
provider_id="bar",
|
|
)
|
|
model = ModelDefWithProvider(
|
|
identifier="test_model",
|
|
llama_model="Llama3.2-3B-Instruct",
|
|
provider_id="foo",
|
|
)
|
|
|
|
await registry.register(bank)
|
|
await registry.register(model)
|
|
results = await registry.get("test_bank")
|
|
assert len(results) == 1
|
|
result_bank = results[0]
|
|
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
|
|
|
|
results = await registry.get("test_model")
|
|
assert len(results) == 1
|
|
result_model = results[0]
|
|
assert result_model.identifier == model.identifier
|
|
assert result_model.llama_model == model.llama_model
|
|
assert result_model.provider_id == model.provider_id
|