From 5eb4874810ec83f7d76e63800143041f1cd24ba7 Mon Sep 17 00:00:00 2001 From: asimurka Date: Tue, 27 Oct 2026 16:04:13 +0100 Subject: [PATCH] Added model resolution logic --- src/providers/inference/README.md | 4 + src/providers/inference/errors.py | 18 +- src/providers/inference/registry.py | 137 ++++++- src/providers/inference/resolve.py | 60 ++++ tests/unit/providers/inference/README.md | 4 + tests/unit/providers/inference/_helpers.py | 4 +- tests/unit/providers/inference/test_errors.py | 9 + .../unit/providers/inference/test_registry.py | 337 +++++++++++++++++- .../unit/providers/inference/test_resolve.py | 55 +++ 9 files changed, 619 insertions(+), 9 deletions(-) create mode 100644 src/providers/inference/resolve.py create mode 100644 tests/unit/providers/inference/test_resolve.py diff --git a/src/providers/inference/README.md b/src/providers/inference/README.md index 4aff05530..7d3e89816 100644 --- a/src/providers/inference/README.md +++ b/src/providers/inference/README.md @@ -24,6 +24,10 @@ Helpers for inferring catalog model types from known model ids. Singleton registry of initialized inference providers. +## [resolve.py](resolve.py) + +Soft-fallback helpers for request-path model resolution. + ## [types.py](types.py) Shared literals and aliases for the inference provider registry. diff --git a/src/providers/inference/errors.py b/src/providers/inference/errors.py index 9204dc0e3..ac839c2ea 100644 --- a/src/providers/inference/errors.py +++ b/src/providers/inference/errors.py @@ -9,7 +9,7 @@ class ModelNotAvailableError(Exception): def __init__(self, name: str, model_type: str, provider_id: str) -> None: """Initialize with the unavailable model identity. - Parameters: + Args: name: Local catalog model name that was requested. model_type: Requested model type (llm, embedding, or rerank). provider_id: Provider that rejected the request. @@ -54,3 +54,19 @@ def __init__(self, provider_id: ProviderId) -> None: """ super().__init__(f"Inference provider {provider_id!r} is not registered") self.provider_id = provider_id + + +class NoModelConfiguredError(Exception): + """Raised when no model can be resolved after soft fallbacks.""" + + def __init__(self, model_type: ModelType) -> None: + """Initialize with the unresolved model capability. + + Args: + model_type: Model type that could not be resolved. + + """ + super().__init__( + f"No {model_type!r} model is configured or available in the catalog" + ) + self.model_type = model_type diff --git a/src/providers/inference/registry.py b/src/providers/inference/registry.py index 7daee4e27..b3bef22a6 100644 --- a/src/providers/inference/registry.py +++ b/src/providers/inference/registry.py @@ -2,12 +2,20 @@ from __future__ import annotations -from typing import TYPE_CHECKING, Any, Final, Optional +from typing import TYPE_CHECKING, Any, Final, Literal, Optional, overload + +from pydantic_ai.embeddings import EmbeddingModel +from pydantic_ai.models import Model +from sentence_transformers import CrossEncoder from providers.inference.azure.provider import AzureInferenceProvider from providers.inference.bedrock.provider import BedrockInferenceProvider -from providers.inference.errors import ProviderNotFoundError +from providers.inference.errors import ( + NoModelConfiguredError, + ProviderNotFoundError, +) from providers.inference.openai.provider import OpenAIInferenceProvider +from providers.inference.resolve import configured_defaults, conversation_last_used from providers.inference.vertexai.provider import VertexAIInferenceProvider from providers.inference.watsonx.provider import WatsonxInferenceProvider from utils.types import Singleton @@ -100,3 +108,128 @@ def list_models( for provider in self._providers.values(): models.extend(provider.list_catalog_models(model_type=model_type)) return models + + @overload + def resolve_model_instance( + self, + model_type: Literal["llm"], + provider_id: Optional[ProviderId] = None, + model_id: Optional[str] = None, + conversation_id: Optional[str] = None, + ) -> Model: ... + + @overload + def resolve_model_instance( + self, + model_type: Literal["embedding"], + provider_id: ProviderId, + model_id: str, + ) -> EmbeddingModel: ... + + @overload + def resolve_model_instance( + self, + model_type: Literal["rerank"], + provider_id: ProviderId, + model_id: str, + ) -> CrossEncoder: ... + + def resolve_model_instance( + self, + model_type: ModelType, + provider_id: Optional[ProviderId] = None, + model_id: Optional[str] = None, + conversation_id: Optional[str] = None, + ) -> Model | EmbeddingModel | CrossEncoder: + """Resolve a typed model instance from the cached provider catalog. + + Embedding and rerank require both ``provider_id`` and ``model_id``. + For llm those ids may be omitted; resolution then falls through + conversation last-used, configured defaults, then the first available llm. + + Args: + model_type: Required model capability to resolve. + provider_id: Local provider id; required for embedding and rerank. + model_id: Local catalog model name; required for embedding and rerank. + conversation_id: Optional conversation used for llm last-used fallback. + + Returns: + A per-request Model, EmbeddingModel, or CrossEncoder instance. + + Raises: + ProviderNotFoundError: Explicit provider id is not registered. + ModelNotAvailableError: Explicit model is missing from the catalog. + NoModelConfiguredError: No llm remains after soft fallbacks. + + """ + if provider_id is not None and model_id is not None: + return self._get_model_instance(provider_id, model_id, model_type) + + # Soft fallbacks (llm only): conversation last-used → defaults → first llm. + if last_used := conversation_last_used(conversation_id): + last_provider_id, last_model_id = last_used + if self._is_llm_available(last_provider_id, last_model_id): + return self._get_model_instance(last_provider_id, last_model_id, "llm") + + if defaults := configured_defaults(): + default_provider_id, default_model_id = defaults + if self._is_llm_available(default_provider_id, default_model_id): + return self._get_model_instance( + default_provider_id, default_model_id, "llm" + ) + + if instance := self._first_available_llm(): + return instance + + raise NoModelConfiguredError("llm") + + def _get_model_instance( + self, + provider_id: ProviderId, + model_id: str, + model_type: ModelType, + ) -> Model | EmbeddingModel | CrossEncoder: + """Get a typed model instance from a registered provider. + + Args: + provider_id: Local provider id. + model_id: Local catalog model name. + model_type: Required model capability. + + Returns: + The instantiated model for the requested type. + + Raises: + ProviderNotFoundError: Provider id is not registered. + ModelNotAvailableError: Model id is not in the catalog. + + """ + return self.get_provider(provider_id).get_model_instance(model_id, model_type) + + def _is_llm_available(self, provider_id: ProviderId, model_id: str) -> bool: + """Return whether an llm catalog pair exists for the provider. + + Args: + provider_id: Candidate provider id. + model_id: Candidate local catalog model name. + + Returns: + True when the provider is registered and exposes the llm. + + """ + if provider := self.find_provider(provider_id): + return provider.get_catalog_model(model_id, "llm") is not None + return False + + def _first_available_llm(self) -> Optional[Model]: + """Return the first llm instance from registration order. + + Returns: + The first available llm instance, or None when the catalog has none. + + """ + for provider in self.list_providers(): + models = provider.list_catalog_models(model_type="llm") + if models: + return provider.get_model_instance(models[0].name, "llm") + return None diff --git a/src/providers/inference/resolve.py b/src/providers/inference/resolve.py new file mode 100644 index 000000000..04ad13ce8 --- /dev/null +++ b/src/providers/inference/resolve.py @@ -0,0 +1,60 @@ +"""Soft-fallback helpers for request-path model resolution.""" + +from __future__ import annotations + +from typing import Optional + +from configuration import configuration +from log import get_logger +from providers.inference.types import ProviderId +from utils.endpoints import retrieve_conversation + +logger = get_logger(__name__) + + +def conversation_last_used( + conversation_id: Optional[str], +) -> Optional[tuple[ProviderId, str]]: + """Load last-used provider/model ids from a conversation when present. + + Args: + conversation_id: Conversation id to look up, or None to skip. + + Returns: + ``(provider_id, model_id)`` from the conversation, or None when the + conversation id is omitted or not found. + + """ + if conversation_id is None: + return None + + conversation = retrieve_conversation(conversation_id) + if conversation is None: + logger.warning( + "Conversation %r not found while resolving llm; soft-skipping " + "last-used fallback", + conversation_id, + ) + return None + + return conversation.last_used_provider, conversation.last_used_model + + +def configured_defaults() -> Optional[tuple[ProviderId, str]]: + """Return configured default provider/model ids when both are set. + + Returns: + ``(provider_id, model_id)`` from inference configuration, or None + when either default is unset. + + """ + inference = configuration.inference + provider_id = inference.default_provider + model_id = inference.default_model + if not provider_id or not model_id: + logger.warning( + "Inference default provider/model unset; soft-skipping configured " + "default fallback" + ) + return None + return provider_id, model_id diff --git a/tests/unit/providers/inference/README.md b/tests/unit/providers/inference/README.md index bc6f7d43d..d55f1848c 100644 --- a/tests/unit/providers/inference/README.md +++ b/tests/unit/providers/inference/README.md @@ -28,3 +28,7 @@ Unit tests for model type inference helpers. Unit tests for the ProviderRegistry singleton. +## [test_resolve.py](test_resolve.py) + +Unit tests for soft-fallback resolution helpers. + diff --git a/tests/unit/providers/inference/_helpers.py b/tests/unit/providers/inference/_helpers.py index 38f21085a..7d9ec9a6b 100644 --- a/tests/unit/providers/inference/_helpers.py +++ b/tests/unit/providers/inference/_helpers.py @@ -28,8 +28,8 @@ async def _discover_models(self) -> list[DiscoveredModel]: return [] def _create_llm(self, model: CatalogModel) -> Model: - """Return a PydanticAI test chat model.""" - return TestModel() + """Return a PydanticAI test chat model tagged with the catalog name.""" + return TestModel(model_name=model.name) def _create_embedding(self, model: CatalogModel) -> EmbeddingModel: """Return a PydanticAI test embedding model.""" diff --git a/tests/unit/providers/inference/test_errors.py b/tests/unit/providers/inference/test_errors.py index 16ef96f2c..d4057e292 100644 --- a/tests/unit/providers/inference/test_errors.py +++ b/tests/unit/providers/inference/test_errors.py @@ -6,6 +6,7 @@ from providers.inference.errors import ( ModelNotAvailableError, ModelTypeNotSupportedError, + NoModelConfiguredError, ProviderNotFoundError, ) from providers.inference.types import ModelType @@ -41,6 +42,14 @@ def test_unknown_provider_type_error_message_and_attributes() -> None: assert str(error) == "Inference provider 'fake-openai' is not registered" +def test_no_model_configured_error_message_and_attributes() -> None: + """Error message and attributes identify the unresolved model type.""" + error = NoModelConfiguredError("llm") + + assert error.model_type == "llm" + assert str(error) == ("No 'llm' model is configured or available in the catalog") + + @pytest.mark.parametrize("model_type", ["llm", "embedding", "rerank"]) def test_get_model_instance_raises_model_not_available_error( model_type: ModelType, diff --git a/tests/unit/providers/inference/test_registry.py b/tests/unit/providers/inference/test_registry.py index 916768f3c..fbce07581 100644 --- a/tests/unit/providers/inference/test_registry.py +++ b/tests/unit/providers/inference/test_registry.py @@ -1,9 +1,12 @@ """Unit tests for the ProviderRegistry singleton.""" from collections.abc import Generator -from typing import Optional +from typing import Any, Optional import pytest +from pydantic_ai.embeddings import TestEmbeddingModel +from pydantic_ai.models.test import TestModel +from sentence_transformers import CrossEncoder from models.config import ( AllowedModel, @@ -11,9 +14,13 @@ OpenAIInferenceProviderEntry, VllmInferenceProviderEntry, ) -from providers.inference.errors import ProviderNotFoundError +from providers.inference.errors import ( + ModelNotAvailableError, + NoModelConfiguredError, + ProviderNotFoundError, +) from providers.inference.registry import PROVIDER_CLASSES, ProviderRegistry -from providers.inference.types import ProviderType +from providers.inference.types import ModelType, ProviderType from tests.unit.providers.inference._helpers import FakeInferenceProvider from utils.types import Singleton @@ -57,7 +64,17 @@ def _openai_entry( name="chat", provider_model_id="gpt-4o-mini", model_type="llm", - ) + ), + AllowedModel( + name="embed", + provider_model_id="text-embedding-3-small", + model_type="embedding", + ), + AllowedModel( + name="rerank", + provider_model_id="cross-encoder", + model_type="rerank", + ), ] ), ) # pyright: ignore[reportCallIssue] @@ -95,6 +112,81 @@ def _inference_config( ) +def _llm_model(name: str, provider_model_id: str = "gpt-4o") -> AllowedModel: + """Return an llm allow-list entry with the given local name.""" + return AllowedModel( + name=name, + provider_model_id=provider_model_id, + model_type="llm", + ) + + +def _embed_only_models() -> list[AllowedModel]: + """Return an allow-list with embedding models only.""" + return [ + AllowedModel( + name="embed-only", + provider_model_id="text-embedding-3-small", + model_type="embedding", + ) + ] + + +def _rerank_only_models() -> list[AllowedModel]: + """Return an allow-list with rerank models only.""" + return [ + AllowedModel( + name="rerank-only", + provider_model_id="cross-encoder", + model_type="rerank", + ) + ] + + +def _patch_inference_defaults( + mocker: Any, + *, + default_model: Optional[str], + default_provider: Optional[str], +) -> None: + """Patch resolve.configuration with fixed default provider/model ids.""" + mocker.patch( + "providers.inference.resolve.configuration", + mocker.Mock( + inference=mocker.Mock( + default_model=default_model, + default_provider=default_provider, + ) + ), + ) + + +def _patch_conversation( + mocker: Any, + *, + last_used_provider: str = "openai", + last_used_model: str = "conv-chat", +) -> None: + """Patch conversation lookup used by soft last-used fallback.""" + conversation = mocker.Mock() + conversation.last_used_provider = last_used_provider + conversation.last_used_model = last_used_model + mocker.patch( + "providers.inference.resolve.retrieve_conversation", + return_value=conversation, + ) + + +async def _ready_registry( + *providers: OpenAIInferenceProviderEntry, +) -> ProviderRegistry: + """Register FakeInferenceProvider and initialize with the given providers.""" + _register_fake("openai") + registry = ProviderRegistry() + await registry.initialize(_inference_config(*providers)) + return registry + + @pytest.mark.asyncio async def test_initialize_registers_configured_providers() -> None: """initialize registers a provider for each configured entry.""" @@ -195,3 +287,240 @@ async def test_list_models_aggregates_providers_and_filters_by_type() -> None: "embed-a", "embed-b", ] + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + ("model_type", "model_id", "expected_type"), + [ + ("llm", "chat", TestModel), + ("embedding", "embed", TestEmbeddingModel), + ("rerank", "rerank", CrossEncoder), + ], +) +async def test_resolve_model_instance_explicit_success( + model_type: ModelType, + model_id: str, + expected_type: type, +) -> None: + """Explicit provider/model ids return the typed model instance.""" + registry = await _ready_registry(_openai_entry()) + + instance = registry.resolve_model_instance( + model_type, + provider_id="openai", + model_id=model_id, + ) + + assert isinstance(instance, expected_type) + + +@pytest.mark.asyncio +async def test_resolve_model_instance_unknown_provider_raises() -> None: + """Explicit unknown provider_id raises ProviderNotFoundError.""" + registry = await _ready_registry(_openai_entry()) + + with pytest.raises(ProviderNotFoundError) as exc_info: + registry.resolve_model_instance( + "llm", + provider_id="missing", + model_id="chat", + ) + + assert exc_info.value.provider_id == "missing" + + +@pytest.mark.asyncio +async def test_resolve_model_instance_unknown_model_raises() -> None: + """Explicit unknown model_id raises ModelNotAvailableError.""" + registry = await _ready_registry(_openai_entry()) + + with pytest.raises(ModelNotAvailableError) as exc_info: + registry.resolve_model_instance( + "llm", + provider_id="openai", + model_id="missing", + ) + + assert exc_info.value.name == "missing" + assert exc_info.value.model_type == "llm" + assert exc_info.value.provider_id == "openai" + + +@pytest.mark.asyncio +async def test_resolve_model_instance_type_mismatch_raises() -> None: + """Explicit model present under another type raises ModelNotAvailableError.""" + registry = await _ready_registry(_openai_entry()) + + with pytest.raises(ModelNotAvailableError) as exc_info: + registry.resolve_model_instance( + "embedding", + provider_id="openai", + model_id="chat", + ) + + assert exc_info.value.name == "chat" + assert exc_info.value.model_type == "embedding" + + +@pytest.mark.asyncio +async def test_resolve_model_instance_prefers_conversation_over_defaults( + mocker: Any, +) -> None: + """Conversation last-used wins when both it and configured defaults are available.""" + registry = await _ready_registry( + _openai_entry( + allowed_models=[ + _llm_model("conv-chat"), + _llm_model("default-chat"), + _llm_model("first-chat"), + ] + ) + ) + _patch_conversation(mocker, last_used_model="conv-chat") + _patch_inference_defaults( + mocker, default_model="default-chat", default_provider="openai" + ) + + instance = registry.resolve_model_instance("llm", conversation_id="conv-1") + + assert isinstance(instance, TestModel) + assert instance.model_name == "conv-chat" + + +@pytest.mark.asyncio +async def test_resolve_model_instance_skips_unavailable_conversation_model( + mocker: Any, +) -> None: + """Missing conversation last-used model falls through to configured defaults.""" + registry = await _ready_registry( + _openai_entry( + allowed_models=[ + _llm_model("default-chat"), + _llm_model("first-chat"), + ] + ) + ) + _patch_conversation(mocker, last_used_model="gone-chat") + _patch_inference_defaults( + mocker, default_model="default-chat", default_provider="openai" + ) + + instance = registry.resolve_model_instance("llm", conversation_id="conv-1") + + assert isinstance(instance, TestModel) + assert instance.model_name == "default-chat" + + +@pytest.mark.asyncio +async def test_resolve_model_instance_skips_unregistered_conversation_provider( + mocker: Any, +) -> None: + """Unregistered conversation last-used provider falls through to defaults.""" + registry = await _ready_registry( + _openai_entry( + allowed_models=[ + _llm_model("default-chat"), + _llm_model("first-chat"), + ] + ) + ) + _patch_conversation( + mocker, last_used_provider="missing", last_used_model="conv-chat" + ) + _patch_inference_defaults( + mocker, default_model="default-chat", default_provider="openai" + ) + + instance = registry.resolve_model_instance("llm", conversation_id="conv-1") + + assert isinstance(instance, TestModel) + assert instance.model_name == "default-chat" + + +@pytest.mark.asyncio +async def test_resolve_model_instance_skips_unavailable_default_model( + mocker: Any, +) -> None: + """Unavailable conversation and default models fall through to first llm.""" + registry = await _ready_registry( + _openai_entry(allowed_models=[_llm_model("first-chat")]) + ) + _patch_conversation(mocker, last_used_model="gone-chat") + _patch_inference_defaults( + mocker, default_model="gone-default", default_provider="openai" + ) + + instance = registry.resolve_model_instance("llm", conversation_id="conv-1") + + assert isinstance(instance, TestModel) + assert instance.model_name == "first-chat" + + +@pytest.mark.asyncio +async def test_resolve_model_instance_skips_unregistered_default_provider( + mocker: Any, +) -> None: + """Unregistered configured default provider falls through to first llm.""" + registry = await _ready_registry( + _openai_entry(allowed_models=[_llm_model("first-chat")]) + ) + _patch_conversation(mocker, last_used_model="gone-chat") + _patch_inference_defaults( + mocker, default_model="default-chat", default_provider="missing" + ) + + instance = registry.resolve_model_instance("llm", conversation_id="conv-1") + + assert isinstance(instance, TestModel) + assert instance.model_name == "first-chat" + + +@pytest.mark.asyncio +async def test_resolve_model_instance_selects_first_available_llm_in_order( + mocker: Any, +) -> None: + """First-available llm uses registration order and skips non-llm providers.""" + registry = await _ready_registry( + _openai_entry("embed-provider", allowed_models=_embed_only_models()), + _openai_entry("rerank-provider", allowed_models=_rerank_only_models()), + _openai_entry( + "llm-provider", + allowed_models=[ + _llm_model("first-chat"), + _llm_model("second-chat"), + ], + ), + _openai_entry( + "later-provider", + allowed_models=[_llm_model("later-chat")], + ), + ) + _patch_inference_defaults(mocker, default_model=None, default_provider=None) + + instance = registry.resolve_model_instance("llm") + + assert isinstance(instance, TestModel) + assert instance.model_name == "first-chat" + + +@pytest.mark.asyncio +async def test_resolve_model_instance_raises_when_nothing_resolvable( + mocker: Any, +) -> None: + """Exhausted soft fallbacks raise NoModelConfiguredError.""" + registry = await _ready_registry( + _openai_entry("embed-provider", allowed_models=_embed_only_models()), + _openai_entry("rerank-provider", allowed_models=_rerank_only_models()), + ) + _patch_conversation( + mocker, last_used_provider="missing", last_used_model="gone-chat" + ) + _patch_inference_defaults( + mocker, default_model="gone-default", default_provider="missing" + ) + + with pytest.raises(NoModelConfiguredError) as exc_info: + registry.resolve_model_instance("llm", conversation_id="conv-1") + + assert exc_info.value.model_type == "llm" diff --git a/tests/unit/providers/inference/test_resolve.py b/tests/unit/providers/inference/test_resolve.py new file mode 100644 index 000000000..b3643b43b --- /dev/null +++ b/tests/unit/providers/inference/test_resolve.py @@ -0,0 +1,55 @@ +"""Unit tests for soft-fallback resolution helpers.""" + +from typing import Any + +from providers.inference.resolve import configured_defaults, conversation_last_used + + +def test_conversation_last_used_returns_none_without_conversation_id() -> None: + """conversation_last_used skips lookup when conversation_id is omitted.""" + assert conversation_last_used(None) is None + + +def test_conversation_last_used_returns_ids(mocker: Any) -> None: + """conversation_last_used returns last-used provider/model from the conversation.""" + conversation = mocker.Mock() + conversation.last_used_provider = "openai" + conversation.last_used_model = "chat" + mocker.patch( + "providers.inference.resolve.retrieve_conversation", + return_value=conversation, + ) + + assert conversation_last_used("conv-1") == ("openai", "chat") + + +def test_conversation_last_used_soft_skips_missing_conversation(mocker: Any) -> None: + """Missing conversations soft-skip and return None.""" + mocker.patch( + "providers.inference.resolve.retrieve_conversation", + return_value=None, + ) + + assert conversation_last_used("missing") is None + + +def test_configured_defaults_returns_ids(mocker: Any) -> None: + """configured_defaults returns both configured default ids.""" + mocker.patch( + "providers.inference.resolve.configuration", + mocker.Mock( + inference=mocker.Mock(default_model="chat", default_provider="openai") + ), + ) + + assert configured_defaults() == ("openai", "chat") + + +def test_configured_defaults_soft_skips_unset(mocker: Any) -> None: + """Unset defaults soft-skip and return None.""" + mocker.patch( + "providers.inference.resolve.configuration", + mocker.Mock(inference=mocker.Mock(default_model=None, default_provider=None)), + ) + + assert configured_defaults() is None