diff --git a/llama_stack/core/store/registry.py b/llama_stack/core/store/registry.py index 6a81d1eae..69d72d69a 100644 --- a/llama_stack/core/store/registry.py +++ b/llama_stack/core/store/registry.py @@ -14,10 +14,7 @@ from llama_stack.core.datatypes import RoutableObjectWithProvider from llama_stack.core.utils.config_dirs import DISTRIBS_BASE_DIR from llama_stack.log import get_logger from llama_stack.providers.utils.kvstore import KVStore, kvstore_impl -from llama_stack.providers.utils.kvstore.config import ( - KVStoreConfig, - SqliteKVStoreConfig, -) +from llama_stack.providers.utils.kvstore.config import KVStoreConfig, SqliteKVStoreConfig logger = get_logger(__name__, category="core::registry") @@ -31,7 +28,9 @@ class DistributionRegistry(Protocol): def get_cached(self, identifier: str) -> RoutableObjectWithProvider | None: ... - async def update(self, obj: RoutableObjectWithProvider) -> RoutableObjectWithProvider: ... + async def update( + self, obj: RoutableObjectWithProvider + ) -> RoutableObjectWithProvider: ... async def register(self, obj: RoutableObjectWithProvider) -> bool: ... @@ -57,7 +56,9 @@ def _parse_registry_values(values: list[str]) -> list[RoutableObjectWithProvider obj = pydantic.TypeAdapter(RoutableObjectWithProvider).validate_json(value) all_objects.append(obj) except pydantic.ValidationError as e: - logger.error(f"Error parsing registry value, raw value: {value}. Error: {e}") + logger.error( + f"Error parsing registry value, raw value: {value}. Error: {e}" + ) continue return all_objects @@ -70,7 +71,9 @@ class DiskDistributionRegistry(DistributionRegistry): async def initialize(self) -> None: pass - def get_cached(self, type: str, identifier: str) -> RoutableObjectWithProvider | None: + def get_cached( + self, type: str, identifier: str + ) -> RoutableObjectWithProvider | None: # Disk registry does not have a cache raise NotImplementedError("Disk registry does not have a cache") @@ -79,15 +82,23 @@ class DiskDistributionRegistry(DistributionRegistry): values = await self.kvstore.values_in_range(start_key, end_key) return _parse_registry_values(values) - async def get(self, type: str, identifier: str) -> RoutableObjectWithProvider | None: - json_str = await self.kvstore.get(KEY_FORMAT.format(type=type, identifier=identifier)) + async def get( + self, type: str, identifier: str + ) -> RoutableObjectWithProvider | None: + json_str = await self.kvstore.get( + KEY_FORMAT.format(type=type, identifier=identifier) + ) if not json_str: return None try: - return pydantic.TypeAdapter(RoutableObjectWithProvider).validate_json(json_str) + return pydantic.TypeAdapter(RoutableObjectWithProvider).validate_json( + json_str + ) except pydantic.ValidationError as e: - logger.error(f"Error parsing registry value for {type}:{identifier}, raw value: {json_str}. Error: {e}") + logger.error( + f"Error parsing registry value for {type}:{identifier}, raw value: {json_str}. Error: {e}" + ) return None async def update(self, obj: RoutableObjectWithProvider) -> None: @@ -152,7 +163,9 @@ class CachedDiskDistributionRegistry(DiskDistributionRegistry): async def initialize(self) -> None: await self._ensure_initialized() - def get_cached(self, type: str, identifier: str) -> RoutableObjectWithProvider | None: + def get_cached( + self, type: str, identifier: str + ) -> RoutableObjectWithProvider | None: return self.cache.get((type, identifier), None) async def get_all(self) -> list[RoutableObjectWithProvider]: @@ -160,7 +173,9 @@ class CachedDiskDistributionRegistry(DiskDistributionRegistry): async with self._locked_cache() as cache: return list(cache.values()) - async def get(self, type: str, identifier: str) -> RoutableObjectWithProvider | None: + async def get( + self, type: str, identifier: str + ) -> RoutableObjectWithProvider | None: await self._ensure_initialized() cache_key = (type, identifier) @@ -202,7 +217,9 @@ async def create_dist_registry( dist_kvstore = await kvstore_impl(metadata_store) else: dist_kvstore = await kvstore_impl( - SqliteKVStoreConfig(db_path=(DISTRIBS_BASE_DIR / image_name / "kvstore.db").as_posix()) + SqliteKVStoreConfig( + db_path=(DISTRIBS_BASE_DIR / image_name / "kvstore.db").as_posix() + ) ) dist_registry = CachedDiskDistributionRegistry(dist_kvstore) await dist_registry.initialize()