From 1ca6abad454f1308ae16d9c422bfb6c6ec11db45 Mon Sep 17 00:00:00 2001 From: Andrej Simurka Date: Mon, 19 Oct 2026 13:11:19 +0200 Subject: [PATCH] Implemented Azure inference provider --- src/providers/inference/azure/README.md | 3 + src/providers/inference/azure/__init__.py | 4 + src/providers/inference/azure/provider.py | 142 ++++++++++ src/providers/inference/registry.py | 2 + .../providers/inference/azure/__init__.py | 1 + .../inference/azure/test_provider.py | 243 ++++++++++++++++++ .../inference/openai/test_provider.py | 6 +- 7 files changed, 398 insertions(+), 3 deletions(-) create mode 100644 src/providers/inference/azure/provider.py create mode 100644 tests/unit/providers/inference/azure/__init__.py create mode 100644 tests/unit/providers/inference/azure/test_provider.py diff --git a/src/providers/inference/azure/README.md b/src/providers/inference/azure/README.md index 393c6e698..50bf0e822 100644 --- a/src/providers/inference/azure/README.md +++ b/src/providers/inference/azure/README.md @@ -4,3 +4,6 @@ Azure inference provider package. +## [provider.py](provider.py) + +Azure OpenAI runtime inference provider backed by Pydantic AI. diff --git a/src/providers/inference/azure/__init__.py b/src/providers/inference/azure/__init__.py index 2c404ed6f..c732b0aaf 100644 --- a/src/providers/inference/azure/__init__.py +++ b/src/providers/inference/azure/__init__.py @@ -1 +1,5 @@ """Azure inference provider package.""" + +from providers.inference.azure.provider import AzureInferenceProvider + +__all__ = ["AzureInferenceProvider"] diff --git a/src/providers/inference/azure/provider.py b/src/providers/inference/azure/provider.py new file mode 100644 index 000000000..c0325e597 --- /dev/null +++ b/src/providers/inference/azure/provider.py @@ -0,0 +1,142 @@ +"""Azure OpenAI runtime inference provider backed by Pydantic AI.""" + +from __future__ import annotations + +from typing import Any + +from azure.identity.aio import ClientSecretCredential, get_bearer_token_provider +from openai import AsyncAzureOpenAI +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.azure import AzureProvider +from sentence_transformers import CrossEncoder + +from configuration import configuration +from models.config import AzureInferenceProviderEntry +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 AzureInferenceProvider( + InferenceProvider[AzureInferenceProviderEntry, AzureProvider] +): + """Inference provider for Azure OpenAI HTTP APIs.""" + + def _build_provider(self) -> AzureProvider: + """Build the Pydantic AI Azure provider from an LCORE-owned client. + + Returns: + A Pydantic AI Azure provider wrapping an ``AsyncAzureOpenAI`` + client configured with API-key or Entra ID authentication. + + """ + config = self._entry.config + client_kwargs: dict[str, Any] = { + "api_version": config.api_version, + } + if config.base_url is not None: + client_kwargs["azure_endpoint"] = config.base_url + + if (entra_id := config.entra_id or configuration.azure_entra_id) is not None: + credential = ClientSecretCredential( + tenant_id=entra_id.tenant_id.get_secret_value(), + client_id=entra_id.client_id.get_secret_value(), + client_secret=entra_id.client_secret.get_secret_value(), + ) + client_kwargs["azure_ad_token_provider"] = get_bearer_token_provider( + credential, + entra_id.scope, + ) + else: + client_kwargs["api_key"] = ( + config.api_key.get_secret_value() + if config.api_key is not None + else None + ) + + return AzureProvider(openai_client=AsyncAzureOpenAI(**client_kwargs)) + + async def _discover_models(self) -> list[DiscoveredModel]: + """List models from the Azure 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 Azure 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 163022c86..937894bfa 100644 --- a/src/providers/inference/registry.py +++ b/src/providers/inference/registry.py @@ -5,6 +5,7 @@ from typing import Any, Final, Optional from models.config import InferenceConfiguration +from providers.inference.azure.provider import AzureInferenceProvider from providers.inference.base import InferenceProvider from providers.inference.catalog import CatalogModel from providers.inference.errors import ProviderNotFoundError @@ -14,6 +15,7 @@ PROVIDER_CLASSES: Final[dict[ProviderType, type[InferenceProvider[Any, Any]]]] = { "openai": OpenAIInferenceProvider, + "azure": AzureInferenceProvider, } diff --git a/tests/unit/providers/inference/azure/__init__.py b/tests/unit/providers/inference/azure/__init__.py new file mode 100644 index 000000000..9eb8eb3ea --- /dev/null +++ b/tests/unit/providers/inference/azure/__init__.py @@ -0,0 +1 @@ +"""Unit tests for the Azure inference provider.""" diff --git a/tests/unit/providers/inference/azure/test_provider.py b/tests/unit/providers/inference/azure/test_provider.py new file mode 100644 index 000000000..4916f94f3 --- /dev/null +++ b/tests/unit/providers/inference/azure/test_provider.py @@ -0,0 +1,243 @@ +"""Unit tests for the Azure inference provider.""" + +# pylint: disable=redefined-outer-name,protected-access + +from typing import Any + +import pytest +from openai import AsyncAzureOpenAI +from pydantic import SecretStr +from pydantic_ai.embeddings.openai import OpenAIEmbeddingModel +from pydantic_ai.models.openai import OpenAIResponsesModel +from pydantic_ai.providers.azure import AzureProvider + +from models.config import AzureEntraIdConfiguration, AzureInferenceProviderEntry +from providers.inference.azure.provider import AzureInferenceProvider +from providers.inference.catalog import CatalogModel +from providers.inference.errors import ModelTypeNotSupportedError + +AZURE_ENDPOINT = "https://example.openai.azure.com/" +API_VERSION = "2024-12-01-preview" +ENTRA_SCOPE = "https://cognitiveservices.azure.com/.default" + + +def _entry(**config: Any) -> AzureInferenceProviderEntry: + """Build an Azure provider configuration.""" + return AzureInferenceProviderEntry.model_validate( + { + "type": "azure", + "id": "azure", + "config": { + "base_url": AZURE_ENDPOINT, + "api_version": API_VERSION, + "api_key": "test-key", + **config, + }, + } + ) + + +def _entra_id( + *, + tenant_id: str = "tenant", + client_id: str = "client", + client_secret: str = "secret", + scope: str = ENTRA_SCOPE, +) -> AzureEntraIdConfiguration: + """Build Entra ID configuration for nested or legacy auth.""" + return AzureEntraIdConfiguration( + tenant_id=SecretStr(tenant_id), + client_id=SecretStr(client_id), + client_secret=SecretStr(client_secret), + scope=scope, + ) + + +def _patch_entra_dependencies(mocker: Any) -> tuple[Any, Any, Any]: + """Mock Azure credential, token-provider, and SDK client construction.""" + token_provider = mocker.AsyncMock(return_value="token") + credential = mocker.patch( + "providers.inference.azure.provider.ClientSecretCredential" + ) + get_token_provider = mocker.patch( + "providers.inference.azure.provider.get_bearer_token_provider", + return_value=token_provider, + ) + client = mocker.patch( + "providers.inference.azure.provider.AsyncAzureOpenAI", + return_value=mocker.MagicMock(spec=AsyncAzureOpenAI), + ) + return credential, get_token_provider, client + + +@pytest.fixture(autouse=True) +def _default_legacy_entra(mocker: Any) -> None: + """Leave legacy Entra ID unset unless a test overrides it.""" + mocker.patch( + "providers.inference.azure.provider.configuration", + mocker.MagicMock(azure_entra_id=None), + ) + + +@pytest.fixture +def provider() -> AzureInferenceProvider: + """Create an Azure provider configured for API-key authentication.""" + return AzureInferenceProvider(_entry()) + + +def test_build_provider_uses_api_key(provider: AzureInferenceProvider) -> None: + """Builds a Pydantic AI Azure provider with API-key authentication.""" + pai_provider = provider._build_provider() + + assert isinstance(pai_provider, AzureProvider) + assert isinstance(pai_provider.client, AsyncAzureOpenAI) + assert pai_provider.client.api_key == "test-key" + assert pai_provider.client._azure_ad_token_provider is None + + +@pytest.mark.parametrize( + ("nested_entra", "legacy_entra"), + [ + ( + _entra_id( + tenant_id="nested-tenant", + client_id="nested-client", + client_secret="nested-secret", + ), + _entra_id( + tenant_id="legacy-tenant", + client_id="legacy-client", + client_secret="legacy-secret", + ), + ), + ( + None, + _entra_id( + tenant_id="legacy-tenant", + client_id="legacy-client", + client_secret="legacy-secret", + ), + ), + ], + ids=["nested-takes-precedence", "legacy-fallback"], +) +def test_build_provider_uses_expected_entra_credentials( + mocker: Any, + nested_entra: AzureEntraIdConfiguration | None, + legacy_entra: AzureEntraIdConfiguration, +) -> None: + """Nested Entra ID takes precedence; legacy configuration is the fallback.""" + mocker.patch( + "providers.inference.azure.provider.configuration", + mocker.MagicMock(azure_entra_id=legacy_entra), + ) + credential, get_token_provider, client = _patch_entra_dependencies(mocker) + + pai_provider = AzureInferenceProvider( + _entry(entra_id=nested_entra) + )._build_provider() + + expected = nested_entra or legacy_entra + assert isinstance(pai_provider, AzureProvider) + credential.assert_called_once_with( + tenant_id=expected.tenant_id.get_secret_value(), + client_id=expected.client_id.get_secret_value(), + client_secret=expected.client_secret.get_secret_value(), + ) + get_token_provider.assert_called_once_with( + credential.return_value, + expected.scope, + ) + client.assert_called_once_with( + azure_endpoint=AZURE_ENDPOINT, + api_version=API_VERSION, + azure_ad_token_provider=get_token_provider.return_value, + ) + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + ("model_id", "expected_type"), + [ + ("gpt-5.6-luna", "llm"), + ("text-embedding-3-small", "embedding"), + ("custom-model", "llm"), + ], +) +async def test_discover_models_infers_model_types( + provider: AzureInferenceProvider, + mocker: Any, + model_id: str, + expected_type: str, +) -> None: + """Discovers Azure model IDs and applies the model-type inference fallback.""" + response = mocker.MagicMock() + response.data = [mocker.MagicMock(id=model_id)] + + pai_provider = mocker.MagicMock(spec=AzureProvider) + 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: AzureInferenceProvider, +) -> 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-5.6-luna", + model_type="llm", + ) + ) + + assert isinstance(model, OpenAIResponsesModel) + assert model.model_name == "gpt-5.6-luna" + assert model.provider is pai_provider + + +def test_create_embedding_uses_openai_embedding_model( + provider: AzureInferenceProvider, +) -> 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: AzureInferenceProvider, +) -> None: + """Rejects reranking because the Azure 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 == "azure" diff --git a/tests/unit/providers/inference/openai/test_provider.py b/tests/unit/providers/inference/openai/test_provider.py index 05d4f5e60..12224c370 100644 --- a/tests/unit/providers/inference/openai/test_provider.py +++ b/tests/unit/providers/inference/openai/test_provider.py @@ -44,7 +44,7 @@ def test_build_provider_uses_config(provider: OpenAIInferenceProvider) -> None: @pytest.mark.parametrize( ("model_id", "expected_type"), [ - ("gpt-4o-mini", "llm"), + ("gpt-5.6-luna", "llm"), ("text-embedding-3-small", "embedding"), ("custom-model", "llm"), ], @@ -81,13 +81,13 @@ def test_create_llm_uses_openai_responses_model( model = provider._create_llm( CatalogModel( name="chat", - provider_model_id="gpt-4o-mini", + provider_model_id="gpt-5.6-luna", model_type="llm", ) ) assert isinstance(model, OpenAIResponsesModel) - assert model.model_name == "gpt-4o-mini" + assert model.model_name == "gpt-5.6-luna" assert model.provider is pai_provider