Repository navigation
LCORE-4703: Implemented OpenAI inference provider #2960
New issue
Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.
By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.
Already on GitHub? Sign in to your account
Merged
tisnik
merged 1 commit into
lightspeed-core:main
from
asimurka:LCORE-4703-openai-inference-provider
Oct 19, 2026
+341
−1
Merged
Changes from all commits
Commits
File filter
Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
There are no files selected for viewing
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| 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 |
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| 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"] |
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| 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), | ||
| ) | ||
| 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" | ||
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1 @@ | ||
| """Unit tests for the OpenAI inference provider.""" |
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| 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: | ||
|
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 | ||
|
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
23
tests/unit/providers/inference/test_model_type_inference.py
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| 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 |
Oops, something went wrong.
Add this suggestion to a batch that can be applied as a single commit.
This suggestion is invalid because no changes were made to the code.
Suggestions cannot be applied while the pull request is closed.
Suggestions cannot be applied while viewing a subset of changes.
Only one suggestion per line can be applied in a batch.
Add this suggestion to a batch that can be applied as a single commit.
Applying suggestions on deleted lines is not supported.
You must change the existing code in this line in order to create a valid suggestion.
Outdated suggestions cannot be applied.
This suggestion has been applied or marked resolved.
Suggestions cannot be applied from pending reviews.
Suggestions cannot be applied on multi-line comments.
Suggestions cannot be applied while the pull request is queued to merge.
Suggestion cannot be applied right now. Please check back later.
Uh oh!
There was an error while loading. Please reload this page.