phoenix-oss/llama-stack-mirror
0
0
Fork 1
mirror of https://github.com/meta-llama/llama-stack.git synced 2025-08-01 00:05:18 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions

linter fixes

Browse source
This commit is contained in:
Dinesh Yeduguru 2024-11-04 16:47:46 -08:00
parent 70b33155d3
commit 180d5205e5
3 changed files with 37 additions and 23 deletions
Download patch file Download diff file Expand all files Collapse all files

8
llama_stack/distribution/store/registry.py
View file

@ -5,8 +5,7 @@
# the root directory of this source tree. # the root directory of this source tree.
import json import json
import asyncio from typing import Dict, List, Protocol
from typing import Protocol, Dict, List
import pydantic import pydantic
@ -74,12 +73,15 @@ class DiskDistributionRegistry(DistributionRegistry):
existing_objects.append(obj) existing_objects.append(obj)
objects_json = [obj.model_dump_json() for obj in existing_objects] # Fixed variable name objects_json = [
obj.model_dump_json() for obj in existing_objects
] # Fixed variable name
await self.kvstore.set( await self.kvstore.set(
KEY_FORMAT.format(obj.identifier), json.dumps(objects_json) KEY_FORMAT.format(obj.identifier), json.dumps(objects_json)
) )
return True return True
class CachedDiskDistributionRegistry(DiskDistributionRegistry): class CachedDiskDistributionRegistry(DiskDistributionRegistry):
def __init__(self, kvstore: KVStore): def __init__(self, kvstore: KVStore):
super().__init__(kvstore) super().__init__(kvstore)

26
llama_stack/distribution/store/tests/test_registry.py
View file

@ -5,14 +5,15 @@
# the root directory of this source tree. # the root directory of this source tree.
import os import os
import asyncio
import pytest import pytest
import pytest_asyncio import pytest_asyncio
from llama_stack.distribution.store import * from llama_stack.distribution.store import * # noqa F403
from llama_stack.apis.memory_banks import VectorMemoryBankDef
from llama_stack.apis.inference import ModelDefWithProvider from llama_stack.apis.inference import ModelDefWithProvider
from llama_stack.apis.memory_banks import VectorMemoryBankDef
from llama_stack.providers.utils.kvstore import kvstore_impl, SqliteKVStoreConfig from llama_stack.providers.utils.kvstore import kvstore_impl, SqliteKVStoreConfig
from llama_stack.distribution.datatypes import * from llama_stack.distribution.datatypes import * # noqa F403
@pytest.fixture @pytest.fixture
def config(): def config():
@ -21,18 +22,21 @@ def config():
os.remove(config.db_path) os.remove(config.db_path)
return config return config
@pytest_asyncio.fixture @pytest_asyncio.fixture
async def registry(config): async def registry(config):
registry = DiskDistributionRegistry(await kvstore_impl(config)) registry = DiskDistributionRegistry(await kvstore_impl(config))
await registry.initialize() await registry.initialize()
return registry return registry
@pytest_asyncio.fixture @pytest_asyncio.fixture
async def cached_registry(config): async def cached_registry(config):
registry = CachedDiskDistributionRegistry(await kvstore_impl(config)) registry = CachedDiskDistributionRegistry(await kvstore_impl(config))
await registry.initialize() await registry.initialize()
return registry return registry
@pytest.fixture @pytest.fixture
def sample_bank(): def sample_bank():
return VectorMemoryBankDef( return VectorMemoryBankDef(
@ -40,23 +44,26 @@ def sample_bank():
embedding_model="all-MiniLM-L6-v2", embedding_model="all-MiniLM-L6-v2",
chunk_size_in_tokens=512, chunk_size_in_tokens=512,
overlap_size_in_tokens=64, overlap_size_in_tokens=64,
provider_id="test-provider" provider_id="test-provider",
) )
@pytest.fixture @pytest.fixture
def sample_model(): def sample_model():
return ModelDefWithProvider( return ModelDefWithProvider(
identifier="test_model", identifier="test_model",
llama_model="Llama3.2-3B-Instruct", llama_model="Llama3.2-3B-Instruct",
provider_id="test-provider" provider_id="test-provider",
) )
@pytest.mark.asyncio @pytest.mark.asyncio
async def test_registry_initialization(registry): async def test_registry_initialization(registry):
# Test empty registry # Test empty registry
results = await registry.get("nonexistent") results = await registry.get("nonexistent")
assert len(results) == 0 assert len(results) == 0
@pytest.mark.asyncio @pytest.mark.asyncio
async def test_basic_registration(registry, sample_bank, sample_model): async def test_basic_registration(registry, sample_bank, sample_model):
print(f"Registering {sample_bank}") print(f"Registering {sample_bank}")
@ -80,6 +87,7 @@ async def test_basic_registration(registry, sample_bank, sample_model):
assert result_model.llama_model == sample_model.llama_model assert result_model.llama_model == sample_model.llama_model
assert result_model.provider_id == sample_model.provider_id assert result_model.provider_id == sample_model.provider_id
@pytest.mark.asyncio @pytest.mark.asyncio
async def test_cached_registry_initialization(config, sample_bank, sample_model): async def test_cached_registry_initialization(config, sample_bank, sample_model):
# First populate the disk registry # First populate the disk registry
@ -101,6 +109,7 @@ async def test_cached_registry_initialization(config, sample_bank, sample_model)
assert result_bank.overlap_size_in_tokens == sample_bank.overlap_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 assert result_bank.provider_id == sample_bank.provider_id
@pytest.mark.asyncio @pytest.mark.asyncio
async def test_cached_registry_updates(config): async def test_cached_registry_updates(config):
cached_registry = CachedDiskDistributionRegistry(await kvstore_impl(config)) cached_registry = CachedDiskDistributionRegistry(await kvstore_impl(config))
@ -131,6 +140,7 @@ async def test_cached_registry_updates(config):
assert result_bank.identifier == new_bank.identifier assert result_bank.identifier == new_bank.identifier
assert result_bank.provider_id == new_bank.provider_id assert result_bank.provider_id == new_bank.provider_id
@pytest.mark.asyncio @pytest.mark.asyncio
async def test_duplicate_provider_registration(config): async def test_duplicate_provider_registration(config):
cached_registry = CachedDiskDistributionRegistry(await kvstore_impl(config)) cached_registry = CachedDiskDistributionRegistry(await kvstore_impl(config))
@ -156,4 +166,6 @@ async def test_duplicate_provider_registration(config):
results = await cached_registry.get("test_bank_2") results = await cached_registry.get("test_bank_2")
assert len(results) == 1 # Still only one result assert len(results) == 1 # Still only one result
assert results[0].embedding_model == original_bank.embedding_model # Original values preserved assert (
results[0].embedding_model == original_bank.embedding_model
) # Original values preserved