Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
15 changes: 15 additions & 0 deletions src/providers/inference/base.py
Original file line number Diff line number Diff line change
Expand Up @@ -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.

Expand Down
48 changes: 48 additions & 0 deletions src/providers/inference/model_type_inference.py
Original file line number Diff line number Diff line change
@@ -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
4 changes: 4 additions & 0 deletions src/providers/inference/openai/__init__.py
Original file line number Diff line number Diff line change
@@ -1 +1,5 @@
"""OpenAI inference provider package."""

from providers.inference.openai.provider import OpenAIInferenceProvider

__all__ = ["OpenAIInferenceProvider"]
118 changes: 118 additions & 0 deletions src/providers/inference/openai/provider.py
Original file line number Diff line number Diff line change
@@ -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),
Comment thread
asimurka marked this conversation as resolved.
)
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"
5 changes: 4 additions & 1 deletion src/providers/inference/registry.py
Original file line number Diff line number Diff line change
Expand Up @@ -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):
Expand Down
1 change: 1 addition & 0 deletions tests/unit/providers/inference/openai/__init__.py
Original file line number Diff line number Diff line change
@@ -0,0 +1 @@
"""Unit tests for the OpenAI inference provider."""
128 changes: 128 additions & 0 deletions tests/unit/providers/inference/openai/test_provider.py
Original file line number Diff line number Diff line change
@@ -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:
Comment thread
asimurka marked this conversation as resolved.
"""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
Comment thread
asimurka marked this conversation as resolved.


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"
23 changes: 23 additions & 0 deletions tests/unit/providers/inference/test_model_type_inference.py
Original file line number Diff line number Diff line change
@@ -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
Loading