phoenix-oss/llama-stack-mirror
0
0
Fork 1
mirror of https://github.com/meta-llama/llama-stack.git synced 2025-12-09 03:19:20 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions

Introduce model_store, shield_store, memory_bank_store

Browse source
This commit is contained in:
Ashwin Bharambe 2024-10-06 16:29:33 -07:00 • committed by Ashwin Bharambe
parent e45a417543
commit 91e0063593
19 changed files with 172 additions and 297 deletions
Download patch file Download diff file Expand all files Collapse all files

18
llama_stack/providers/utils/inference/model_registry.py
View file

@ -4,7 +4,7 @@
# This source code is licensed under the terms described in the LICENSE file in
# the root directory of this source tree.
from typing import Dict, List
from typing import Dict
from llama_models.sku_list import resolve_model
@ -15,7 +15,6 @@ class ModelRegistryHelper:
def __init__(self, stack_to_provider_models_map: Dict[str, str]):
self.stack_to_provider_models_map = stack_to_provider_models_map
self.registered_models = []
def map_to_provider_model(self, identifier: str) -> str:
model = resolve_model(identifier)
@ -30,22 +29,7 @@ class ModelRegistryHelper:
return self.stack_to_provider_models_map[identifier]
async def register_model(self, model: ModelDef) -> None:
existing = await self.get_model(model.identifier)
if existing is not None:
return
if model.identifier not in self.stack_to_provider_models_map:
raise ValueError(
f"Unsupported model {model.identifier}. Supported models: {self.stack_to_provider_models_map.keys()}"
)
self.registered_models.append(model)
async def list_models(self) -> List[ModelDef]:
return self.registered_models
async def get_model(self, identifier: str) -> Optional[ModelDef]:
for model in self.registered_models:
if model.identifier == identifier:
return model
return None