linter fixes

llama_stack/distribution/store/registry.py

@@ -5,8 +5,7 @@
 # the root directory of this source tree.
 
 import json
-import asyncio
-from typing import Protocol, Dict, List
+from typing import Dict, List, Protocol
 
 import pydantic
 
@@ -74,12 +73,15 @@ class DiskDistributionRegistry(DistributionRegistry):
 
         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(
             KEY_FORMAT.format(obj.identifier), json.dumps(objects_json)
         )
         return True
 
+
 class CachedDiskDistributionRegistry(DiskDistributionRegistry):
     def __init__(self, kvstore: KVStore):
         super().__init__(kvstore)

llama_stack/distribution/store/tests/test_registry.py

@@ -5,14 +5,15 @@
 # the root directory of this source tree.
 
 import os
-import asyncio
+
 import pytest
 import pytest_asyncio
-from llama_stack.distribution.store import *
-from llama_stack.apis.memory_banks import VectorMemoryBankDef
+from llama_stack.distribution.store import *  # noqa F403
 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.distribution.datatypes import *
+from llama_stack.distribution.datatypes import *  # noqa F403
+
 
 @pytest.fixture
 def config():
@@ -21,18 +22,21 @@ def config():
         os.remove(config.db_path)
     return config
 
+
 @pytest_asyncio.fixture
 async def registry(config):
     registry = DiskDistributionRegistry(await kvstore_impl(config))
     await registry.initialize()
     return registry
 
+
 @pytest_asyncio.fixture
 async def cached_registry(config):
     registry = CachedDiskDistributionRegistry(await kvstore_impl(config))
     await registry.initialize()
     return registry
 
+
 @pytest.fixture
 def sample_bank():
     return VectorMemoryBankDef(
@@ -40,23 +44,26 @@ def sample_bank():
         embedding_model="all-MiniLM-L6-v2",
         chunk_size_in_tokens=512,
         overlap_size_in_tokens=64,
-        provider_id="test-provider"
+        provider_id="test-provider",
     )
 
+
 @pytest.fixture
 def sample_model():
     return ModelDefWithProvider(
         identifier="test_model",
         llama_model="Llama3.2-3B-Instruct",
-        provider_id="test-provider"
+        provider_id="test-provider",
     )
 
+
 @pytest.mark.asyncio
 async def test_registry_initialization(registry):
     # Test empty registry
     results = await registry.get("nonexistent")
     assert len(results) == 0
 
+
 @pytest.mark.asyncio
 async def test_basic_registration(registry, sample_bank, sample_model):
     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.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
@@ -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.provider_id == sample_bank.provider_id
 
+
 @pytest.mark.asyncio
 async def test_cached_registry_updates(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.provider_id == new_bank.provider_id
 
+
 @pytest.mark.asyncio
 async def test_duplicate_provider_registration(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")
     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