mirror of
https://github.com/meta-llama/llama-stack.git
synced 2025-07-27 14:38:49 +00:00
d035fe93c6
inference providers each have a static list of supported / known models. some also
have access to a dynamic list of currently available models. this change gives
prodivers using the ModelRegistryHelper the ability to combine their static and
dynamic lists.
for instance, OpenAIInferenceAdapter can implement
```
def query_available_models(self) -> list[str]:
return [entry.model for entry in self.openai_client.models.list()]
```
to augment its static list w/ a current list from openai.
257 lines
9.6 KiB
Python
257 lines
9.6 KiB
Python
# Copyright (c) Meta Platforms, Inc. and affiliates.
|
|
# All rights reserved.
|
|
#
|
|
# This source code is licensed under the terms described in the LICENSE file in
|
|
# the root directory of this source tree.
|
|
|
|
#
|
|
# ModelRegistryHelper provides mixin functionality for registering and
|
|
# unregistering models. It maintains a mapping of model ID / aliases to
|
|
# provider model IDs.
|
|
#
|
|
# Test cases -
|
|
# - Looking up an alias that does not exist should return None.
|
|
# - Registering a model + provider ID should add the model to the registry. If
|
|
# provider ID is known or an alias for a provider ID.
|
|
# - Registering an existing model should return an error. Unless it's a
|
|
# dulicate entry.
|
|
# - Unregistering a model should remove it from the registry.
|
|
# - Unregistering a model that does not exist should return an error.
|
|
# - Supported model ID and their aliases are registered during initialization.
|
|
# Only aliases are added afterwards.
|
|
#
|
|
# Questions -
|
|
# - Should we be allowed to register models w/o provider model IDs? No.
|
|
# According to POST /v1/models, required params are
|
|
# - identifier
|
|
# - provider_resource_id
|
|
# - provider_id
|
|
# - type
|
|
# - metadata
|
|
# - model_type
|
|
#
|
|
# TODO: llama_model functionality
|
|
#
|
|
|
|
import pytest
|
|
|
|
from llama_stack.apis.models import Model
|
|
from llama_stack.providers.utils.inference.model_registry import ModelRegistryHelper, ProviderModelEntry
|
|
|
|
|
|
@pytest.fixture
|
|
def known_model() -> Model:
|
|
return Model(
|
|
provider_id="provider",
|
|
identifier="known-model",
|
|
provider_resource_id="known-provider-id",
|
|
)
|
|
|
|
|
|
@pytest.fixture
|
|
def known_model2() -> Model:
|
|
return Model(
|
|
provider_id="provider",
|
|
identifier="known-model2",
|
|
provider_resource_id="known-provider-id2",
|
|
)
|
|
|
|
|
|
@pytest.fixture
|
|
def known_provider_model(known_model: Model) -> ProviderModelEntry:
|
|
return ProviderModelEntry(
|
|
provider_model_id=known_model.provider_resource_id,
|
|
aliases=[known_model.model_id],
|
|
)
|
|
|
|
|
|
@pytest.fixture
|
|
def known_provider_model2(known_model2: Model) -> ProviderModelEntry:
|
|
return ProviderModelEntry(
|
|
provider_model_id=known_model2.provider_resource_id,
|
|
# aliases=[],
|
|
)
|
|
|
|
|
|
@pytest.fixture
|
|
def unknown_model() -> Model:
|
|
return Model(
|
|
provider_id="provider",
|
|
identifier="unknown-model",
|
|
provider_resource_id="unknown-provider-id",
|
|
)
|
|
|
|
|
|
@pytest.fixture
|
|
def helper(known_provider_model: ProviderModelEntry, known_provider_model2: ProviderModelEntry) -> ModelRegistryHelper:
|
|
return ModelRegistryHelper([known_provider_model, known_provider_model2])
|
|
|
|
|
|
class MockModelRegistryHelperWithDynamicModels(ModelRegistryHelper):
|
|
"""Test helper that simulates a provider with dynamically available models."""
|
|
|
|
def __init__(self, model_entries: list[ProviderModelEntry], available_models: list[str]):
|
|
super().__init__(model_entries)
|
|
self._available_models = available_models
|
|
|
|
async def query_available_models(self) -> list[str]:
|
|
return self._available_models
|
|
|
|
|
|
@pytest.fixture
|
|
def dynamic_model() -> Model:
|
|
"""A model that's not in static config but available dynamically."""
|
|
return Model(
|
|
provider_id="provider",
|
|
identifier="dynamic-model",
|
|
provider_resource_id="dynamic-provider-id",
|
|
)
|
|
|
|
|
|
@pytest.fixture
|
|
def helper_with_dynamic_models(
|
|
known_provider_model: ProviderModelEntry, known_provider_model2: ProviderModelEntry, dynamic_model: Model
|
|
) -> MockModelRegistryHelperWithDynamicModels:
|
|
"""Helper that includes dynamically available models."""
|
|
return MockModelRegistryHelperWithDynamicModels(
|
|
[known_provider_model, known_provider_model2], [dynamic_model.provider_resource_id]
|
|
)
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_lookup_unknown_model(helper: ModelRegistryHelper, unknown_model: Model) -> None:
|
|
assert helper.get_provider_model_id(unknown_model.model_id) is None
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_register_unknown_provider_model(helper: ModelRegistryHelper, unknown_model: Model) -> None:
|
|
with pytest.raises(ValueError):
|
|
await helper.register_model(unknown_model)
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_register_model(helper: ModelRegistryHelper, known_model: Model) -> None:
|
|
model = Model(
|
|
provider_id=known_model.provider_id,
|
|
identifier="new-model",
|
|
provider_resource_id=known_model.provider_resource_id,
|
|
)
|
|
assert helper.get_provider_model_id(model.model_id) is None
|
|
await helper.register_model(model)
|
|
assert helper.get_provider_model_id(model.model_id) == model.provider_resource_id
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_register_model_from_alias(helper: ModelRegistryHelper, known_model: Model) -> None:
|
|
model = Model(
|
|
provider_id=known_model.provider_id,
|
|
identifier="new-model",
|
|
provider_resource_id=known_model.model_id, # use known model's id as an alias for the supported model id
|
|
)
|
|
assert helper.get_provider_model_id(model.model_id) is None
|
|
await helper.register_model(model)
|
|
assert helper.get_provider_model_id(model.model_id) == known_model.provider_resource_id
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_register_model_existing(helper: ModelRegistryHelper, known_model: Model) -> None:
|
|
await helper.register_model(known_model)
|
|
assert helper.get_provider_model_id(known_model.model_id) == known_model.provider_resource_id
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_register_model_existing_different(
|
|
helper: ModelRegistryHelper, known_model: Model, known_model2: Model
|
|
) -> None:
|
|
known_model.provider_resource_id = known_model2.provider_resource_id
|
|
with pytest.raises(ValueError):
|
|
await helper.register_model(known_model)
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_unregister_model(helper: ModelRegistryHelper, known_model: Model) -> None:
|
|
await helper.register_model(known_model) # duplicate entry
|
|
assert helper.get_provider_model_id(known_model.model_id) == known_model.provider_model_id
|
|
await helper.unregister_model(known_model.model_id)
|
|
assert helper.get_provider_model_id(known_model.model_id) is None
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_unregister_unknown_model(helper: ModelRegistryHelper, unknown_model: Model) -> None:
|
|
with pytest.raises(ValueError):
|
|
await helper.unregister_model(unknown_model.model_id)
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_register_model_during_init(helper: ModelRegistryHelper, known_model: Model) -> None:
|
|
assert helper.get_provider_model_id(known_model.provider_resource_id) == known_model.provider_model_id
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_unregister_model_during_init(helper: ModelRegistryHelper, known_model: Model) -> None:
|
|
assert helper.get_provider_model_id(known_model.provider_resource_id) == known_model.provider_model_id
|
|
await helper.unregister_model(known_model.provider_resource_id)
|
|
assert helper.get_provider_model_id(known_model.provider_resource_id) is None
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_register_model_from_query_available_models(
|
|
helper_with_dynamic_models: MockModelRegistryHelperWithDynamicModels, dynamic_model: Model
|
|
) -> None:
|
|
"""Test that models returned by query_available_models can be registered."""
|
|
# Verify the model is not in static config
|
|
assert helper_with_dynamic_models.get_provider_model_id(dynamic_model.provider_resource_id) is None
|
|
|
|
# But it should be available via query_available_models
|
|
available_models = await helper_with_dynamic_models.query_available_models()
|
|
assert dynamic_model.provider_resource_id in available_models
|
|
|
|
# Registration should succeed
|
|
registered_model = await helper_with_dynamic_models.register_model(dynamic_model)
|
|
assert registered_model == dynamic_model
|
|
|
|
# Model should now be registered and accessible
|
|
assert (
|
|
helper_with_dynamic_models.get_provider_model_id(dynamic_model.model_id) == dynamic_model.provider_resource_id
|
|
)
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_register_model_not_in_static_or_dynamic(
|
|
helper_with_dynamic_models: MockModelRegistryHelperWithDynamicModels, unknown_model: Model
|
|
) -> None:
|
|
"""Test that models not in static config or dynamic models are rejected."""
|
|
# Verify the model is not in static config
|
|
assert helper_with_dynamic_models.get_provider_model_id(unknown_model.provider_resource_id) is None
|
|
|
|
# And not in dynamic models
|
|
available_models = await helper_with_dynamic_models.query_available_models()
|
|
assert unknown_model.provider_resource_id not in available_models
|
|
|
|
# Registration should fail with comprehensive error message
|
|
with pytest.raises(Exception) as exc_info: # UnsupportedModelError
|
|
await helper_with_dynamic_models.register_model(unknown_model)
|
|
|
|
# Error should include both static and dynamic models
|
|
error_str = str(exc_info.value)
|
|
assert "dynamic-provider-id" in error_str # dynamic model should be in error
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_register_alias_for_dynamic_model(
|
|
helper_with_dynamic_models: MockModelRegistryHelperWithDynamicModels, dynamic_model: Model
|
|
) -> None:
|
|
"""Test that we can register an alias that maps to a dynamically available model."""
|
|
# Create a model with a different identifier but same provider_resource_id
|
|
alias_model = Model(
|
|
provider_id=dynamic_model.provider_id,
|
|
identifier="dynamic-model-alias",
|
|
provider_resource_id=dynamic_model.provider_resource_id,
|
|
)
|
|
|
|
# Registration should succeed since the provider_resource_id is available dynamically
|
|
registered_model = await helper_with_dynamic_models.register_model(alias_model)
|
|
assert registered_model == alias_model
|
|
|
|
# Both the original provider_resource_id and the new alias should work
|
|
assert helper_with_dynamic_models.get_provider_model_id(alias_model.model_id) == dynamic_model.provider_resource_id
|