LCORE-4710: Implement registry initialization

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",

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()

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
 
 

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__)

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
 
 

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
 
 

src/providers/inference/registry.py

@@ -6,8 +6,8 @@
 
 from pydantic_ai.embeddings import EmbeddingModel
 from pydantic_ai.models import Model
-from sentence_transformers import CrossEncoder
 
+from log import get_logger
 from providers.inference.azure.provider import AzureInferenceProvider
 from providers.inference.bedrock.provider import BedrockInferenceProvider
 from providers.inference.errors import (
@@ -25,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
@@ -40,6 +42,8 @@
     "sentence_transformers": SentenceTransformersInferenceProvider,
 }
 
+logger = get_logger(__name__)
+
 
 class ProviderRegistry(metaclass=Singleton):
     """Registry of initialized inference providers."""
@@ -61,6 +65,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]]:

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)

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
 
 

src/providers/inference/vllm/provider.py

@@ -2,15 +2,14 @@
 
 from __future__ import annotations
 
-from typing import Any
+from typing import TYPE_CHECKING, Any
 
 from openai import APIError
 from pydantic_ai.embeddings import EmbeddingModel
 from pydantic_ai.embeddings.openai import OpenAIEmbeddingModel
 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]

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
 
 

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.