Skip to content

Commit

Permalink
chore: update return type to Optional[str]
Browse files Browse the repository at this point in the history
Updated the return type of the methods `get_provider_model_id` and
`get_llama_model` in the `ModelRegistryHelper` class to `Optional[str]`
to indicate that they may return a string or None when no match is
found. This change improves the clarity of the methods' behavior and
supports better type safety.

Signed-off-by: Sébastien Han <seb@redhat.com>
  • Loading branch information
leseb committed Feb 6, 2025
1 parent 3922999 commit b3fa9ae
Showing 1 changed file with 2 additions and 2 deletions.
4 changes: 2 additions & 2 deletions llama_stack/providers/utils/inference/model_registry.py
Original file line number Diff line number Diff line change
Expand Up @@ -57,13 +57,13 @@ def __init__(self, model_aliases: List[ModelAlias]):
self.alias_to_provider_id_map[alias_obj.llama_model] = alias_obj.provider_model_id
self.provider_id_to_llama_model_map[alias_obj.provider_model_id] = alias_obj.llama_model

def get_provider_model_id(self, identifier: str) -> str:
def get_provider_model_id(self, identifier: str) -> Optional[str]:
if identifier in self.alias_to_provider_id_map:
return self.alias_to_provider_id_map[identifier]
else:
return None

def get_llama_model(self, provider_model_id: str) -> str:
def get_llama_model(self, provider_model_id: str) -> Optional[str]:
if provider_model_id in self.provider_id_to_llama_model_map:
return self.provider_id_to_llama_model_map[provider_model_id]
else:
Expand Down

0 comments on commit b3fa9ae

Please sign in to comment.