From 6d2a7e8e8b1d32c686b2c11a3cd9db92dd309bba Mon Sep 17 00:00:00 2001 From: Andrej Simurka Date: Tue, 13 Oct 2026 11:06:02 +0200 Subject: [PATCH] Implemented OpenAI inference provider --- src/providers/inference/base.py | 15 ++ .../inference/model_type_inference.py | 48 +++++++ src/providers/inference/openai/__init__.py | 4 + src/providers/inference/openai/provider.py | 118 ++++++++++++++++ src/providers/inference/registry.py | 5 +- .../providers/inference/openai/__init__.py | 1 + .../inference/openai/test_provider.py | 128 ++++++++++++++++++ .../inference/test_model_type_inference.py | 23 ++++ 8 files changed, 341 insertions(+), 1 deletion(-) create mode 100644 src/providers/inference/model_type_inference.py create mode 100644 src/providers/inference/openai/provider.py create mode 100644 tests/unit/providers/inference/openai/__init__.py create mode 100644 tests/unit/providers/inference/openai/test_provider.py create mode 100644 tests/unit/providers/inference/test_model_type_inference.py diff --git a/src/providers/inference/base.py b/src/providers/inference/base.py index 0e10f85fd..08b3b8a19 100644 --- a/src/providers/inference/base.py +++ b/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. diff --git a/src/providers/inference/model_type_inference.py b/src/providers/inference/model_type_inference.py new file mode 100644 index 000000000..b870f2564 --- /dev/null +++ b/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 diff --git a/src/providers/inference/openai/__init__.py b/src/providers/inference/openai/__init__.py index 58e09e810..2368adec3 100644 --- a/src/providers/inference/openai/__init__.py +++ b/src/providers/inference/openai/__init__.py @@ -1 +1,5 @@ """OpenAI inference provider package.""" + +from providers.inference.openai.provider import OpenAIInferenceProvider + +__all__ = ["OpenAIInferenceProvider"] diff --git a/src/providers/inference/openai/provider.py b/src/providers/inference/openai/provider.py new file mode 100644 index 000000000..a54fc9235 --- /dev/null +++ b/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" diff --git a/src/providers/inference/registry.py b/src/providers/inference/registry.py index 34e8cf6dc..163022c86 100644 --- a/src/providers/inference/registry.py +++ b/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): diff --git a/tests/unit/providers/inference/openai/__init__.py b/tests/unit/providers/inference/openai/__init__.py new file mode 100644 index 000000000..6eb06d7cc --- /dev/null +++ b/tests/unit/providers/inference/openai/__init__.py @@ -0,0 +1 @@ +"""Unit tests for the OpenAI inference provider.""" diff --git a/tests/unit/providers/inference/openai/test_provider.py b/tests/unit/providers/inference/openai/test_provider.py new file mode 100644 index 000000000..05d4f5e60 --- /dev/null +++ b/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" diff --git a/tests/unit/providers/inference/test_model_type_inference.py b/tests/unit/providers/inference/test_model_type_inference.py new file mode 100644 index 000000000..fa1a54059 --- /dev/null +++ b/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