phoenix-oss/llama-stack-mirror
0
0
Fork 1
mirror of https://github.com/meta-llama/llama-stack.git synced 2025-12-17 14:39:47 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions

Allow models to be registered as long as llama model is provided

Browse source
This commit is contained in:
Dinesh Yeduguru 2024-11-18 11:58:32 -08:00
parent f1b9578f8d
commit ccb5445d2a
2 changed files with 30 additions and 21 deletions
Download patch file Download diff file Expand all files Collapse all files

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

@ -51,7 +51,7 @@ class ModelRegistryHelper(ModelsProtocolPrivate):
if identifier in self.alias_to_provider_id_map:
return self.alias_to_provider_id_map[identifier]
else:
raise ValueError(f"Unknown model: `{identifier}`")
return None
def get_llama_model(self, provider_model_id: str) -> str:
if provider_model_id in self.provider_id_to_llama_model_map:
@ -60,8 +60,18 @@ class ModelRegistryHelper(ModelsProtocolPrivate):
return None
async def register_model(self, model: Model) -> Model:
model.provider_resource_id = self.get_provider_model_id(
model.provider_resource_id
)
provider_resource_id = self.get_provider_model_id(model.provider_resource_id)
if provider_resource_id:
model.provider_resource_id = provider_resource_id
else:
if model.metadata.get("llama_model") is None:
raise ValueError(
f"Model '{model.provider_resource_id}' is not available and no llama_model was specified in metadata. "
"Please specify a llama_model in metadata or use a supported model identifier"
)
# Register the mapping from provider model id to llama model for future lookups
self.provider_id_to_llama_model_map[model.provider_resource_id] = (
model.metadata["llama_model"]
)
return model