Merge pull request #2960 from asimurka/LCORE-4703-openai-inference-provider

LCORE-4703: Implemented OpenAI inference provider

Author: tisnik

src/providers/inference/base.py

@@ -41,6 +41,21 @@ def __init__(self, entry: EntryT) -> None:
         self._models: list[CatalogModel] = []
         self._entry: EntryT = entry
 
+    @property
+    def _get_provider(self) -> ProviderT:
+        """Return the initialized PAI provider.
+
+        Raises:
+            RuntimeError: If the provider has not been initialized or does not
+                use a PAI provider instance.
+
+        """
+        if self._provider is None:
+            raise RuntimeError(
+                f"Provider {self.provider_id!r} has not been initialized"
+            )
+        return self._provider
+
     async def initialize(self) -> None:
         """Build the PydanticAI provider and load the model catalog.
 

src/providers/inference/model_type_inference.py

@@ -0,0 +1,48 @@
+"""Helpers for inferring catalog model types from known model ids."""
+
+from __future__ import annotations
+
+from typing import Optional, TypeAliasType, get_args
+
+from pydantic_ai.embeddings import KnownEmbeddingModelName
+from pydantic_ai.models import KnownModelName
+
+from providers.inference.types import ModelType
+
+
+def _model_ids_from_known_names(known_names: TypeAliasType) -> frozenset[str]:
+    """Extract unprefixed model ids from a Pydantic AI known-name type alias.
+
+    Args:
+        known_names: A Pydantic AI known-name type alias.
+
+    Returns:
+        The set of model ids with provider prefixes removed.
+
+    """
+    return frozenset(name.split(":", 1)[-1] for name in get_args(known_names.__value__))
+
+
+KNOWN_LLM_IDS: frozenset[str] = _model_ids_from_known_names(KnownModelName)
+KNOWN_EMBEDDING_IDS: frozenset[str] = _model_ids_from_known_names(
+    KnownEmbeddingModelName
+)
+
+
+def infer_model_type_from_pai_known_names(
+    provider_model_id: str,
+) -> Optional[ModelType]:
+    """Infer model type from Pydantic AI's known model names.
+
+    Args:
+        provider_model_id: Model id as reported by the inference provider.
+
+    Returns:
+        The inferred model type, or None if the model id is unknown.
+
+    """
+    if provider_model_id in KNOWN_EMBEDDING_IDS:
+        return "embedding"
+    if provider_model_id in KNOWN_LLM_IDS:
+        return "llm"
+    return None

src/providers/inference/openai/__init__.py

@@ -1 +1,5 @@
 """OpenAI inference provider package."""
+
+from providers.inference.openai.provider import OpenAIInferenceProvider
+
+__all__ = ["OpenAIInferenceProvider"]

src/providers/inference/openai/provider.py

@@ -0,0 +1,118 @@
+"""OpenAI runtime inference provider backed by Pydantic AI."""
+
+from __future__ import annotations
+
+from typing import Any
+
+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 OpenAIResponsesModel
+from pydantic_ai.providers.openai import OpenAIProvider
+from sentence_transformers import CrossEncoder
+
+from models.config import OpenAIInferenceProviderEntry
+from providers.inference.base import InferenceProvider
+from providers.inference.catalog import CatalogModel, DiscoveredModel
+from providers.inference.errors import ModelTypeNotSupportedError
+from providers.inference.model_type_inference import (
+    infer_model_type_from_pai_known_names,
+)
+from providers.inference.types import ModelType
+
+
+class OpenAIInferenceProvider(
+    InferenceProvider[OpenAIInferenceProviderEntry, OpenAIProvider]
+):
+    """Inference provider for OpenAI HTTP APIs."""
+
+    def _build_provider(self) -> OpenAIProvider:
+        """Build the Pydantic AI OpenAI provider from configuration.
+
+        Returns:
+            A Pydantic AI OpenAI provider using the configured API key and
+            base URL.
+
+        """
+        config = self._entry.config
+        api_key = (
+            config.api_key.get_secret_value() if config.api_key is not None else None
+        )
+        return OpenAIProvider(api_key=api_key, base_url=config.base_url)
+
+    async def _discover_models(self) -> list[DiscoveredModel]:
+        """List models from the OpenAI models endpoint.
+
+        Returns:
+            Discovered models with inferred types, before allow-list filtering.
+
+        """
+        response = await self._get_provider.client.models.list()
+        return [
+            DiscoveredModel(
+                provider_model_id=model.id,
+                model_type=self._infer_model_type(model.id),
+            )
+            for model in response.data
+        ]
+
+    def _create_llm(self, model: CatalogModel) -> Model:
+        """Create a Responses API chat model for a catalog entry.
+
+        Args:
+            model: Catalog entry describing the LLM to instantiate.
+
+        Returns:
+            A Pydantic AI OpenAI Responses model bound to this provider.
+
+        """
+        return OpenAIResponsesModel(
+            model.provider_model_id,
+            provider=self._get_provider,
+        )
+
+    def _create_embedding(self, model: CatalogModel) -> EmbeddingModel:
+        """Create an embedding model for a catalog entry.
+
+        Args:
+            model: Catalog entry describing the embedding model to instantiate.
+
+        Returns:
+            A Pydantic AI OpenAI embedding model bound to this provider.
+
+        """
+        return OpenAIEmbeddingModel(
+            model.provider_model_id,
+            provider=self._get_provider,
+        )
+
+    def _create_rerank(self, model: CatalogModel) -> CrossEncoder:
+        """Reject rerank requests; this provider has no rerank API.
+
+        Args:
+            model: Catalog entry that requested rerank support.
+
+        Raises:
+            ModelTypeNotSupportedError: Always, because rerank is unsupported.
+
+        """
+        _ = model
+        raise ModelTypeNotSupportedError("rerank", self.provider_id)
+
+    def _infer_model_type(
+        self,
+        provider_model_id: str,
+        **kwargs: Any,
+    ) -> ModelType:
+        """Infer model type from Pydantic AI known LLM and embedding names.
+
+        Args:
+            provider_model_id: Model id from OpenAI or from an allow-list entry.
+            **kwargs: Unused; accepted for base-class compatibility.
+
+        Returns:
+            embedding when the id is a known embedding model, otherwise llm.
+
+        """
+        _ = kwargs
+        return infer_model_type_from_pai_known_names(provider_model_id) or "llm"

src/providers/inference/registry.py

@@ -8,10 +8,13 @@
 from providers.inference.base import InferenceProvider
 from providers.inference.catalog import CatalogModel
 from providers.inference.errors import ProviderNotFoundError
+from providers.inference.openai.provider import OpenAIInferenceProvider
 from providers.inference.types import ModelType, ProviderId, ProviderType
 from utils.types import Singleton
 
-PROVIDER_CLASSES: Final[dict[ProviderType, type[InferenceProvider[Any, Any]]]] = {}
+PROVIDER_CLASSES: Final[dict[ProviderType, type[InferenceProvider[Any, Any]]]] = {
+    "openai": OpenAIInferenceProvider,
+}
 
 
 class ProviderRegistry(metaclass=Singleton):

tests/unit/providers/inference/openai/__init__.py

@@ -0,0 +1 @@
+"""Unit tests for the OpenAI inference provider."""

tests/unit/providers/inference/openai/test_provider.py

@@ -0,0 +1,128 @@
+"""Unit tests for the OpenAI inference provider."""
+
+# pylint: disable=redefined-outer-name,protected-access
+
+from typing import Any
+
+import pytest
+from pydantic_ai.embeddings.openai import OpenAIEmbeddingModel
+from pydantic_ai.models.openai import OpenAIResponsesModel
+from pydantic_ai.providers.openai import OpenAIProvider
+
+from models.config import OpenAIInferenceProviderEntry
+from providers.inference.catalog import CatalogModel
+from providers.inference.errors import ModelTypeNotSupportedError
+from providers.inference.openai.provider import OpenAIInferenceProvider
+
+
+@pytest.fixture
+def provider() -> OpenAIInferenceProvider:
+    """Create an OpenAI provider."""
+    entry = OpenAIInferenceProviderEntry.model_validate(
+        {
+            "type": "openai",
+            "id": "openai",
+            "config": {
+                "api_key": "test-key",
+                "base_url": "https://example.test/v1",
+            },
+        }
+    )
+    return OpenAIInferenceProvider(entry)
+
+
+def test_build_provider_uses_config(provider: OpenAIInferenceProvider) -> None:
+    """Builds a Pydantic AI provider with configured credentials and URL."""
+    pai_provider = provider._build_provider()
+
+    assert isinstance(pai_provider, OpenAIProvider)
+    assert str(pai_provider.client.base_url).rstrip("/") == "https://example.test/v1"
+    assert pai_provider.client.api_key == "test-key"
+
+
+@pytest.mark.asyncio
+@pytest.mark.parametrize(
+    ("model_id", "expected_type"),
+    [
+        ("gpt-4o-mini", "llm"),
+        ("text-embedding-3-small", "embedding"),
+        ("custom-model", "llm"),
+    ],
+)
+async def test_discover_models_infers_model_types(
+    provider: OpenAIInferenceProvider,
+    mocker: Any,
+    model_id: str,
+    expected_type: str,
+) -> None:
+    """Discovers OpenAI model IDs and applies the model-type inference fallback."""
+    response = mocker.MagicMock()
+    response.data = [mocker.MagicMock(id=model_id)]
+
+    pai_provider = mocker.MagicMock(spec=OpenAIProvider)
+    pai_provider.client.models.list = mocker.AsyncMock(return_value=response)
+    provider._provider = pai_provider
+
+    discovered = await provider._discover_models()
+
+    assert len(discovered) == 1
+    assert discovered[0].provider_model_id == model_id
+    assert discovered[0].model_type == expected_type
+    pai_provider.client.models.list.assert_awaited_once_with()
+
+
+def test_create_llm_uses_openai_responses_model(
+    provider: OpenAIInferenceProvider,
+) -> None:
+    """Creates a Pydantic AI Responses model using the provider model ID."""
+    pai_provider = provider._build_provider()
+    provider._provider = pai_provider
+
+    model = provider._create_llm(
+        CatalogModel(
+            name="chat",
+            provider_model_id="gpt-4o-mini",
+            model_type="llm",
+        )
+    )
+
+    assert isinstance(model, OpenAIResponsesModel)
+    assert model.model_name == "gpt-4o-mini"
+    assert model.provider is pai_provider
+
+
+def test_create_embedding_uses_openai_embedding_model(
+    provider: OpenAIInferenceProvider,
+) -> None:
+    """Creates a Pydantic AI embedding model using the provider model ID."""
+    pai_provider = provider._build_provider()
+    provider._provider = pai_provider
+
+    model = provider._create_embedding(
+        CatalogModel(
+            name="embed",
+            provider_model_id="text-embedding-3-small",
+            model_type="embedding",
+        )
+    )
+
+    assert isinstance(model, OpenAIEmbeddingModel)
+    assert model.model_name == "text-embedding-3-small"
+    assert model._provider is pai_provider
+
+
+def test_create_rerank_is_unsupported(
+    provider: OpenAIInferenceProvider,
+) -> None:
+    """Rejects reranking because the OpenAI provider does not support it."""
+    with pytest.raises(ModelTypeNotSupportedError) as exc_info:
+        provider._create_rerank(
+            CatalogModel(
+                name="rerank",
+                provider_model_id="rerank-model",
+                model_type="rerank",
+            )
+        )
+
+    assert exc_info.value.model_type == "rerank"
+    assert exc_info.value.provider_id == "openai"

tests/unit/providers/inference/test_model_type_inference.py

@@ -0,0 +1,23 @@
+"""Unit tests for model type inference helpers."""
+
+from providers.inference.model_type_inference import (
+    KNOWN_EMBEDDING_IDS,
+    KNOWN_LLM_IDS,
+    infer_model_type_from_pai_known_names,
+)
+
+
+def test_known_ids_include_openai_model_ids() -> None:
+    """PAI known-name tables are exposed as prefix-stripped model ids."""
+    assert "gpt-4o-mini" in KNOWN_LLM_IDS
+    assert "text-embedding-3-small" in KNOWN_EMBEDDING_IDS
+    assert "openai:gpt-4o-mini" not in KNOWN_LLM_IDS
+
+
+def test_infer_model_type_from_pai_known_names() -> None:
+    """Known ids map to embedding or llm; unknown ids return None."""
+    assert (
+        infer_model_type_from_pai_known_names("text-embedding-3-small") == "embedding"
+    )
+    assert infer_model_type_from_pai_known_names("gpt-4o-mini") == "llm"
+    assert infer_model_type_from_pai_known_names("totally-custom-chat") is None