From 28d93ddbc6bcac6e4b2ca3261f894e075bc455b5 Mon Sep 17 00:00:00 2001 From: asimurka Date: Tue, 3 Nov 2026 08:38:46 +0100 Subject: [PATCH 1/4] Implement registry initialization Co-authored-by: Cursor --- src/app/main.py | 3 + src/providers/inference/registry.py | 10 ++ tests/unit/app/README.md | 4 + tests/unit/app/test_main_startup.py | 170 ++++++++++++++++++++++++++++ 4 files changed, 187 insertions(+) create mode 100644 tests/unit/app/test_main_startup.py diff --git a/src/app/main.py b/src/app/main.py index 6b8a13632..5ed8bd4df 100644 --- a/src/app/main.py +++ b/src/app/main.py @@ -27,6 +27,7 @@ from models.api.responses.error import InternalServerErrorResponse from observability.profiling import initialize_pyroscope from observability.sentry import initialize_sentry +from providers.inference.registry import ProviderRegistry from utils.degraded_mode import DegradedModeTracker from utils.ogx_version import check_ogx_version @@ -97,6 +98,8 @@ async def lifespan(_app: FastAPI) -> AsyncIterator[None]: initialize_sentry() initialize_pyroscope() + await ProviderRegistry().initialize(configuration.inference) + ogx_config = configuration.configuration.ogx await AsyncOgxClientHolder().load(ogx_config) client: AsyncOgxClient = AsyncOgxClientHolder().get_client() diff --git a/src/providers/inference/registry.py b/src/providers/inference/registry.py index 9aead1ace..dbad636b4 100644 --- a/src/providers/inference/registry.py +++ b/src/providers/inference/registry.py @@ -2,6 +2,7 @@ from __future__ import annotations +import logging from typing import TYPE_CHECKING, Any, Final, Literal, Optional, overload from pydantic_ai.embeddings import EmbeddingModel @@ -40,6 +41,8 @@ "sentence_transformers": SentenceTransformersInferenceProvider, } +logger = logging.getLogger(__name__) + class ProviderRegistry(metaclass=Singleton): """Registry of initialized inference providers.""" @@ -61,6 +64,13 @@ async def initialize(self, inference: InferenceConfiguration) -> None: await provider.initialize() self._providers[provider.provider_id] = provider + provider_ids = [provider.provider_id for provider in self.list_providers()] + logger.info( + "Inference provider registry initialized with %d provider(s): %s", + len(provider_ids), + provider_ids, + ) + def find_provider( self, provider_id: ProviderId ) -> Optional[InferenceProvider[Any, Any]]: diff --git a/tests/unit/app/README.md b/tests/unit/app/README.md index 5fcadf112..786a9bd31 100644 --- a/tests/unit/app/README.md +++ b/tests/unit/app/README.md @@ -12,6 +12,10 @@ Unit tests for app.database module. Unit tests for the pure ASGI middlewares in main.py. +## [test_main_startup.py](test_main_startup.py) + +Unit tests for FastAPI lifespan wiring of the inference provider registry. + ## [test_routers.py](test_routers.py) Unit tests for routers.py. diff --git a/tests/unit/app/test_main_startup.py b/tests/unit/app/test_main_startup.py new file mode 100644 index 000000000..a4af411ec --- /dev/null +++ b/tests/unit/app/test_main_startup.py @@ -0,0 +1,170 @@ +"""Unit tests for FastAPI lifespan wiring of the inference provider registry.""" + +from collections.abc import Generator +from pathlib import Path +from typing import Any + +import httpx +import pytest +import yaml +from fastapi import FastAPI +from openai import APIConnectionError +from pytest_mock import MockerFixture + +from app.main import lifespan +from providers.inference.errors import ModelDiscoveryError +from providers.inference.registry import ProviderRegistry +from utils.types import Singleton + + +@pytest.fixture(autouse=True) +def _reset_provider_registry() -> Generator[None, None, None]: + """Reset the singleton registry between startup unit tests.""" + # pylint: disable=protected-access + Singleton._instances.pop(ProviderRegistry, None) + + yield + + Singleton._instances.pop(ProviderRegistry, None) + # pylint: enable=protected-access + + +def _write_config(tmp_path: Path, *, inference: dict[str, Any]) -> Path: + """Write a minimal LCS config with the given inference section.""" + config_path = tmp_path / "lightspeed-stack.yaml" + payload: dict[str, Any] = { + "name": "lifespan-test", + "service": { + "host": "localhost", + "port": 8080, + "auth_enabled": False, + "workers": 1, + "color_log": True, + "access_log": True, + }, + "ogx": { + "use_as_library_client": False, + "url": "http://localhost:8321", + "api_key": "xyzzy", + "allow_degraded_mode": True, + }, + "user_data_collection": {"feedback_enabled": False}, + "mcp_servers": [], + "inference": inference, + } + config_path.write_text(yaml.dump(payload), encoding="utf-8") + return config_path + + +def _openai_provider_entry(*, skip_model_discovery: bool) -> dict[str, Any]: + """Build an OpenAI inference.providers[] entry for lifespan configs.""" + return { + "type": "openai", + "id": "openai", + "skip_model_discovery": skip_model_discovery, + "allowed_models": [ + { + "name": "chat", + "provider_model_id": "gpt-4o-mini", + "model_type": "llm", + } + ], + "config": { + "api_key": "test-key", + "base_url": "https://example.test/v1", + }, + } + + +def _patch_lifespan_foreign_dependencies(mocker: MockerFixture) -> None: + """Mock foreign startup dependencies outside the inference registry.""" + mocker.patch("app.main.initialize_sentry") + mocker.patch("app.main.initialize_pyroscope") + mock_holder = mocker.patch("app.main.AsyncOgxClientHolder") + mock_holder_instance = mock_holder.return_value + mock_holder_instance.load = mocker.AsyncMock() + mock_holder_instance.get_client.return_value = mocker.AsyncMock() + mocker.patch( + "app.main.check_ogx_version", + new=mocker.AsyncMock(return_value="0.2.22"), + ) + mocker.patch("app.main.setup_model_metrics", new=mocker.AsyncMock()) + mocker.patch("app.main.initialize_database") + mocker.patch("app.main.create_tables") + mocker.patch( + "app.main.shutdown_background_topic_summary_tasks", + new=mocker.AsyncMock(), + ) + mocker.patch("app.main.A2AStorageFactory.cleanup", new=mocker.AsyncMock()) + mocker.patch("app.main.sentry_sdk.flush") + + +@pytest.mark.asyncio +async def test_lifespan_initializes_inference_registry( + tmp_path: Path, + monkeypatch: pytest.MonkeyPatch, + mocker: MockerFixture, +) -> None: + """Successful startup initializes the registry from inference.providers[].""" + monkeypatch.setenv( + "LIGHTSPEED_STACK_CONFIG_PATH", + str( + _write_config( + tmp_path, + inference={ + "providers": [ + _openai_provider_entry(skip_model_discovery=True), + ] + }, + ) + ), + ) + _patch_lifespan_foreign_dependencies(mocker) + + async with lifespan(FastAPI()): + registry = ProviderRegistry() + assert [provider.provider_id for provider in registry.list_providers()] == [ + "openai" + ] + assert [model.name for model in registry.list_models()] == ["chat"] + + +@pytest.mark.asyncio +async def test_lifespan_propagates_model_discovery_error( + tmp_path: Path, + monkeypatch: pytest.MonkeyPatch, + mocker: MockerFixture, +) -> None: + """OpenAI models.list failure becomes ModelDiscoveryError and aborts startup.""" + monkeypatch.setenv( + "LIGHTSPEED_STACK_CONFIG_PATH", + str( + _write_config( + tmp_path, + inference={ + "providers": [ + _openai_provider_entry(skip_model_discovery=False), + ] + }, + ) + ), + ) + _patch_lifespan_foreign_dependencies(mocker) + mocker.patch( + "openai.resources.models.AsyncModels.list", + new=mocker.AsyncMock( + side_effect=APIConnectionError( + message="connection failed", + request=httpx.Request( # pyright: ignore[reportArgumentType] + "GET", "https://example.test/v1/models" + ), + ) + ), + ) + + with pytest.raises(ModelDiscoveryError) as exc_info: + async with lifespan(FastAPI()): + pass + + assert exc_info.value.provider_id == "openai" + assert ProviderRegistry().list_providers() == [] From 85c7922660013ec927645960919b8551e2338fa5 Mon Sep 17 00:00:00 2001 From: asimurka Date: Tue, 3 Nov 2026 09:43:42 +0100 Subject: [PATCH 2/4] Make sentence-transformers a main dependency --- pyproject.toml | 4 +- src/providers/inference/azure/provider.py | 3 +- src/providers/inference/base.py | 3 +- src/providers/inference/bedrock/provider.py | 3 +- src/providers/inference/openai/provider.py | 3 +- src/providers/inference/registry.py | 3 +- src/providers/inference/vertexai/provider.py | 3 +- src/providers/inference/vllm/provider.py | 6 ++- src/providers/inference/watsonx/provider.py | 3 +- uv.lock | 42 ++++++++++---------- 10 files changed, 41 insertions(+), 32 deletions(-) diff --git a/pyproject.toml b/pyproject.toml index efc9ca9bd..9530078aa 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -86,6 +86,8 @@ dependencies = [ # Used for Pydantic AI "pydantic-ai==2.38.0", "pydantic-ai-skills>=0.11.0,<2.0.0", + # Used by the sentence-transformers inference provider (embedding / rerank) + "sentence-transformers>=5.6.0", # Used for OpenTelemetry instrumentation "opentelemetry-distro>=0.49b0", "opentelemetry-exporter-otlp>=1.34.1", @@ -183,8 +185,6 @@ ogxlibdev = [ "langdetect>=1.0.9", "emoji>=2.1.0", "nltk>=3.8.1", - # API inference: inline::sentence-transformers - "sentence-transformers>=5.6.0", # API vector_io: inline::faiss, inline::sqlite-vec, remote::pgvector "faiss-cpu>=1.11.0", "chardet>=5.2.0", diff --git a/src/providers/inference/azure/provider.py b/src/providers/inference/azure/provider.py index 32c0b348f..0f12c1459 100644 --- a/src/providers/inference/azure/provider.py +++ b/src/providers/inference/azure/provider.py @@ -12,7 +12,6 @@ from pydantic_ai.models import Model from pydantic_ai.models.openai import OpenAIResponsesModel from pydantic_ai.providers.azure import AzureProvider -from sentence_transformers import CrossEncoder from configuration import configuration from models.config import AzureInferenceProviderEntry @@ -24,6 +23,8 @@ ) if TYPE_CHECKING: + from sentence_transformers import CrossEncoder + from providers.inference.types import ModelType diff --git a/src/providers/inference/base.py b/src/providers/inference/base.py index 18face7bb..5ab199770 100644 --- a/src/providers/inference/base.py +++ b/src/providers/inference/base.py @@ -8,7 +8,6 @@ from pydantic_ai.embeddings import EmbeddingModel from pydantic_ai.models import Model from pydantic_ai.providers import Provider -from sentence_transformers import CrossEncoder from log import get_logger from models.config import ( @@ -19,6 +18,8 @@ from providers.inference.errors import ModelNotAvailableError if TYPE_CHECKING: + from sentence_transformers import CrossEncoder + from providers.inference.types import ModelType, ProviderId logger = get_logger(__name__) diff --git a/src/providers/inference/bedrock/provider.py b/src/providers/inference/bedrock/provider.py index 599495176..2824b9323 100644 --- a/src/providers/inference/bedrock/provider.py +++ b/src/providers/inference/bedrock/provider.py @@ -17,7 +17,6 @@ from pydantic_ai.models import Model from pydantic_ai.models.bedrock import BedrockConverseModel from pydantic_ai.providers.bedrock import BedrockProvider -from sentence_transformers import CrossEncoder from models.config import BedrockInferenceProviderEntry from providers.inference.base import InferenceProvider @@ -31,6 +30,8 @@ from providers.inference.errors import ModelDiscoveryError, ModelTypeNotSupportedError if TYPE_CHECKING: + from sentence_transformers import CrossEncoder + from providers.inference.types import ModelType diff --git a/src/providers/inference/openai/provider.py b/src/providers/inference/openai/provider.py index 7add3d1c4..ad7b3dfdc 100644 --- a/src/providers/inference/openai/provider.py +++ b/src/providers/inference/openai/provider.py @@ -10,7 +10,6 @@ from pydantic_ai.models import Model from pydantic_ai.models.openai import OpenAIResponsesModel from pydantic_ai.providers.openai import OpenAIProvider -from sentence_transformers import CrossEncoder from models.config import OpenAIInferenceProviderEntry from providers.inference.base import InferenceProvider @@ -21,6 +20,8 @@ ) if TYPE_CHECKING: + from sentence_transformers import CrossEncoder + from providers.inference.types import ModelType diff --git a/src/providers/inference/registry.py b/src/providers/inference/registry.py index dbad636b4..07842fdf0 100644 --- a/src/providers/inference/registry.py +++ b/src/providers/inference/registry.py @@ -7,7 +7,6 @@ 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 @@ -26,6 +25,8 @@ from utils.types import Singleton if TYPE_CHECKING: + from sentence_transformers import CrossEncoder + from models.config import InferenceConfiguration from providers.inference.base import InferenceProvider from providers.inference.catalog import CatalogModel diff --git a/src/providers/inference/vertexai/provider.py b/src/providers/inference/vertexai/provider.py index 7bf5d03c5..1fdafaf15 100644 --- a/src/providers/inference/vertexai/provider.py +++ b/src/providers/inference/vertexai/provider.py @@ -11,7 +11,6 @@ from pydantic_ai.models import Model from pydantic_ai.models.google import GoogleModel from pydantic_ai.providers.google_cloud import GoogleCloudProvider -from sentence_transformers import CrossEncoder from models.config import VertexAIInferenceProviderEntry from providers.inference.base import InferenceProvider @@ -22,6 +21,8 @@ ) if TYPE_CHECKING: + from sentence_transformers import CrossEncoder + from providers.inference.types import ModelType diff --git a/src/providers/inference/vllm/provider.py b/src/providers/inference/vllm/provider.py index e2fd14558..0995634d3 100644 --- a/src/providers/inference/vllm/provider.py +++ b/src/providers/inference/vllm/provider.py @@ -2,7 +2,7 @@ from __future__ import annotations -from typing import Any +from typing import TYPE_CHECKING, Any from openai import APIError from pydantic_ai.embeddings import EmbeddingModel @@ -10,7 +10,6 @@ from pydantic_ai.models import Model from pydantic_ai.models.openai import OpenAIChatModel from pydantic_ai.providers.vllm import VLLMProvider -from sentence_transformers import CrossEncoder from models.config import VllmInferenceProviderEntry from providers.inference.base import InferenceProvider @@ -21,6 +20,9 @@ ) from providers.inference.types import ModelType +if TYPE_CHECKING: + from sentence_transformers import CrossEncoder + class VllmInferenceProvider( InferenceProvider[VllmInferenceProviderEntry, VLLMProvider] diff --git a/src/providers/inference/watsonx/provider.py b/src/providers/inference/watsonx/provider.py index 5ba4c6a9d..bfd646cce 100644 --- a/src/providers/inference/watsonx/provider.py +++ b/src/providers/inference/watsonx/provider.py @@ -10,7 +10,6 @@ from pydantic_ai.embeddings.openai import OpenAIEmbeddingModel from pydantic_ai.models import Model from pydantic_ai.models.openai import OpenAIChatModel -from sentence_transformers import CrossEncoder from models.config import WatsonxInferenceProviderEntry from providers.inference.base import InferenceProvider @@ -20,6 +19,8 @@ from providers.inference.watsonx.specs import FoundationModelSpecsResponse if TYPE_CHECKING: + from sentence_transformers import CrossEncoder + from providers.inference.types import ModelType diff --git a/uv.lock b/uv.lock index 73c810c19..3790bf4ab 100644 --- a/uv.lock +++ b/uv.lock @@ -1,18 +1,18 @@ version = 1 -revision = 3 +revision = 5 requires-python = ">=3.12, <3.15" resolution-markers = [ "python_full_version >= '3.14' and sys_platform == 'win32'", "python_full_version >= '3.14' and sys_platform == 'emscripten'", "python_full_version >= '3.14' and sys_platform != 'darwin' and sys_platform != 'emscripten' and sys_platform != 'win32'", + "python_full_version >= '3.14' and sys_platform == 'darwin'", "python_full_version == '3.13.*' and sys_platform == 'win32'", "python_full_version == '3.13.*' and sys_platform == 'emscripten'", "python_full_version == '3.13.*' and sys_platform != 'darwin' and sys_platform != 'emscripten' and sys_platform != 'win32'", + "python_full_version == '3.13.*' and sys_platform == 'darwin'", "python_full_version < '3.13' and sys_platform == 'win32'", "python_full_version < '3.13' and sys_platform == 'emscripten'", "python_full_version < '3.13' and sys_platform != 'darwin' and sys_platform != 'emscripten' and sys_platform != 'win32'", - "python_full_version >= '3.14' and sys_platform == 'darwin'", - "python_full_version == '3.13.*' and sys_platform == 'darwin'", "python_full_version < '3.13' and sys_platform == 'darwin'", ] @@ -2443,6 +2443,7 @@ dependencies = [ { name = "requests" }, { name = "rich" }, { name = "semver" }, + { name = "sentence-transformers" }, { name = "sentry-sdk", extra = ["fastapi"] }, { name = "sqlalchemy", extra = ["asyncio"] }, { name = "starlette" }, @@ -2517,7 +2518,6 @@ ogxlibdev = [ { name = "pypdf" }, { name = "pythainlp" }, { name = "requests" }, - { name = "sentence-transformers" }, { name = "torch", version = "2.11.0", source = { registry = "https://download.pytorch.org/whl/cpu" }, marker = "sys_platform == 'darwin'" }, { name = "torch", version = "2.11.0+cpu", source = { registry = "https://download.pytorch.org/whl/cpu" }, marker = "sys_platform != 'darwin'" }, { name = "transformers" }, @@ -2567,6 +2567,7 @@ requires-dist = [ { name = "requests", specifier = ">=2.33.0" }, { name = "rich", specifier = ">=14.0.0" }, { name = "semver", specifier = "<4.0.0" }, + { name = "sentence-transformers", specifier = ">=5.6.0" }, { name = "sentry-sdk", extras = ["fastapi"], specifier = ">=2.58.0" }, { name = "sqlalchemy", extras = ["asyncio"], specifier = ">=2.0.42" }, { name = "starlette", specifier = ">=0.47.1" }, @@ -2641,7 +2642,6 @@ ogxlibdev = [ { name = "pypdf", specifier = ">=6.19.0" }, { name = "pythainlp", specifier = ">=3.0.10" }, { name = "requests", specifier = ">=2.33.0" }, - { name = "sentence-transformers", specifier = ">=5.6.0" }, { name = "torch", specifier = "==2.11.0", index = "https://download.pytorch.org/whl/cpu", conflict = { package = "lightspeed-stack", group = "ogxlibdev" } }, { name = "transformers", specifier = ">=4.34.0" }, { name = "tree-sitter", specifier = ">=0.24.0" }, @@ -5382,8 +5382,8 @@ name = "secretstorage" version = "3.5.0" source = { registry = "https://pypi.org/simple" } dependencies = [ - { name = "cryptography", marker = "sys_platform != 'darwin' and sys_platform != 'emscripten' and sys_platform != 'win32'" }, - { name = "jeepney", marker = "sys_platform != 'darwin' and sys_platform != 'emscripten' and sys_platform != 'win32'" }, + { name = "cryptography" }, + { name = "jeepney" }, ] sdist = { url = "https://files.pythonhosted.org/packages/1c/03/e834bcd866f2f8a49a85eaff47340affa3bfa391ee9912a952a1faa68c7b/secretstorage-3.5.0.tar.gz", hash = "sha256:f04b8e4689cbce351744d5537bf6b1329c6fc68f91fa666f60a380edddcd11be", size = 19884, upload-time = "2025-11-23T19:02:53.191Z" } wheels = [ @@ -5944,13 +5944,13 @@ resolution-markers = [ "python_full_version < '3.13' and sys_platform == 'darwin'", ] dependencies = [ - { name = "filelock", marker = "sys_platform == 'darwin'" }, - { name = "fsspec", marker = "sys_platform == 'darwin'" }, - { name = "jinja2", marker = "sys_platform == 'darwin'" }, - { name = "networkx", marker = "sys_platform == 'darwin'" }, - { name = "setuptools", marker = "sys_platform == 'darwin'" }, - { name = "sympy", marker = "sys_platform == 'darwin'" }, - { name = "typing-extensions", marker = "sys_platform == 'darwin'" }, + { name = "filelock" }, + { name = "fsspec" }, + { name = "jinja2" }, + { name = "networkx" }, + { name = "setuptools" }, + { name = "sympy" }, + { name = "typing-extensions" }, ] wheels = [ { url = "https://download-r2.pytorch.org/whl/cpu/torch-2.11.0-cp312-cp312-macosx_11_0_arm64.whl", hash = "sha256:43b35116802c85fb88d99f4a396b8bd4472bfca1dd82e69499e5a4f9b8b4e252", upload-time = "2026-03-23T15:16:58Z" }, @@ -5976,13 +5976,13 @@ resolution-markers = [ "python_full_version < '3.13' and sys_platform != 'darwin' and sys_platform != 'emscripten' and sys_platform != 'win32'", ] dependencies = [ - { name = "filelock", marker = "sys_platform != 'darwin'" }, - { name = "fsspec", marker = "sys_platform != 'darwin'" }, - { name = "jinja2", marker = "sys_platform != 'darwin'" }, - { name = "networkx", marker = "sys_platform != 'darwin'" }, - { name = "setuptools", marker = "sys_platform != 'darwin'" }, - { name = "sympy", marker = "sys_platform != 'darwin'" }, - { name = "typing-extensions", marker = "sys_platform != 'darwin'" }, + { name = "filelock" }, + { name = "fsspec" }, + { name = "jinja2" }, + { name = "networkx" }, + { name = "setuptools" }, + { name = "sympy" }, + { name = "typing-extensions" }, ] wheels = [ { url = "https://download-r2.pytorch.org/whl/cpu/torch-2.11.0%2Bcpu-cp312-cp312-linux_s390x.whl", hash = "sha256:2db3ae5404e32cb42b5fcbd94f13607761eaec0cf1687fde95095289d1e26cfb", upload-time = "2026-04-28T00:06:06Z" }, From 11f7c0210accb0a022e5e47aa5c2a2b009ea3170 Mon Sep 17 00:00:00 2001 From: asimurka Date: Tue, 3 Nov 2026 09:46:14 +0100 Subject: [PATCH 3/4] Fixed logger import --- src/providers/inference/registry.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/src/providers/inference/registry.py b/src/providers/inference/registry.py index 07842fdf0..239ea6bb2 100644 --- a/src/providers/inference/registry.py +++ b/src/providers/inference/registry.py @@ -2,12 +2,12 @@ from __future__ import annotations -import logging from typing import TYPE_CHECKING, Any, Final, Literal, Optional, overload from pydantic_ai.embeddings import EmbeddingModel from pydantic_ai.models import Model +from log import get_logger from providers.inference.azure.provider import AzureInferenceProvider from providers.inference.bedrock.provider import BedrockInferenceProvider from providers.inference.errors import ( @@ -42,7 +42,7 @@ "sentence_transformers": SentenceTransformersInferenceProvider, } -logger = logging.getLogger(__name__) +logger = get_logger(__name__) class ProviderRegistry(metaclass=Singleton): From cdc5617bc75f20b15964b58d3ba42d60220f2b22 Mon Sep 17 00:00:00 2001 From: asimurka Date: Tue, 3 Nov 2026 09:54:26 +0100 Subject: [PATCH 4/4] Fix mypy issue --- src/providers/inference/sentence_transformers/model_cache.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/src/providers/inference/sentence_transformers/model_cache.py b/src/providers/inference/sentence_transformers/model_cache.py index 32e407904..315e4ad64 100644 --- a/src/providers/inference/sentence_transformers/model_cache.py +++ b/src/providers/inference/sentence_transformers/model_cache.py @@ -2,7 +2,7 @@ from __future__ import annotations -from typing import Literal, Optional, overload +from typing import Literal, Optional, cast, overload from pydantic_ai.embeddings import EmbeddingModel from sentence_transformers import CrossEncoder @@ -69,4 +69,4 @@ def put( if isinstance(model, CrossEncoder): self._rerankers[provider_model_id] = model else: - self._embeddings[provider_model_id] = model + self._embeddings[provider_model_id] = cast(EmbeddingModel, model)