linter fixes

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

View file

@ -152,7 +152,7 @@ class CommonRoutingTableImpl(RoutingTable):
async def register_object(self, obj: RoutableObjectWithProvider): async def register_object(self, obj: RoutableObjectWithProvider):
# Get existing objects from registry # Get existing objects from registry
existing_objects = await self.dist_registry.get(obj.identifier) existing_objects = await self.dist_registry.get(obj.identifier)
# Check for existing registration # Check for existing registration
for existing_obj in existing_objects: for existing_obj in existing_objects:
if existing_obj.provider_id == obj.provider_id or not obj.provider_id: if existing_obj.provider_id == obj.provider_id or not obj.provider_id:

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)
@ -88,9 +90,9 @@ class CachedDiskDistributionRegistry(DiskDistributionRegistry):
async def initialize(self) -> None: async def initialize(self) -> None:
start_key = KEY_FORMAT.format("") start_key = KEY_FORMAT.format("")
end_key = KEY_FORMAT.format("\xff") end_key = KEY_FORMAT.format("\xff")
keys = await self.kvstore.range(start_key, end_key) keys = await self.kvstore.range(start_key, end_key)
for key in keys: for key in keys:
identifier = key.split(":")[-1] identifier = key.split(":")[-1]
objects = await super().get(identifier) objects = await super().get(identifier)
@ -106,28 +108,28 @@ class CachedDiskDistributionRegistry(DiskDistributionRegistry):
async def get(self, identifier: str) -> List[RoutableObjectWithProvider]: async def get(self, identifier: str) -> List[RoutableObjectWithProvider]:
if identifier in self.cache: if identifier in self.cache:
return self.cache[identifier] return self.cache[identifier]
objects = await super().get(identifier) objects = await super().get(identifier)
if objects: if objects:
self.cache[identifier] = objects self.cache[identifier] = objects
return objects return objects
async def register(self, obj: RoutableObjectWithProvider) -> bool: async def register(self, obj: RoutableObjectWithProvider) -> bool:
# First update disk # First update disk
success = await super().register(obj) success = await super().register(obj)
if success: if success:
# Then update cache # Then update cache
if obj.identifier not in self.cache: if obj.identifier not in self.cache:
self.cache[obj.identifier] = [] self.cache[obj.identifier] = []
# Check if provider already exists in cache # Check if provider already exists in cache
for cached_obj in self.cache[obj.identifier]: for cached_obj in self.cache[obj.identifier]:
if cached_obj.provider_id == obj.provider_id: if cached_obj.provider_id == obj.provider_id:
return success return success
# If not, update cache # If not, update cache
self.cache[obj.identifier].append(obj) self.cache[obj.identifier].append(obj)
return success return success

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,45 +22,51 @@ 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(
identifier="test_bank", identifier="test_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}")
await registry.register(sample_bank) await registry.register(sample_bank)
print(f"Registering {sample_model}") print(f"Registering {sample_model}")
await registry.register(sample_model) await registry.register(sample_model)
@ -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))
@ -153,7 +163,9 @@ async def test_duplicate_provider_registration(config):
provider_id="baz", # Same provider_id provider_id="baz", # Same provider_id
) )
await cached_registry.register(duplicate_bank) await cached_registry.register(duplicate_bank)
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