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
14 changes: 11 additions & 3 deletions src/providers/inference/azure/provider.py
Original file line number Diff line number Diff line change
Expand Up @@ -4,8 +4,9 @@

from typing import TYPE_CHECKING, Any

from azure.core.exceptions import ClientAuthenticationError
from azure.identity.aio import ClientSecretCredential, get_bearer_token_provider
from openai import AsyncAzureOpenAI
from openai import APIError, AsyncAzureOpenAI
from pydantic_ai.embeddings import EmbeddingModel
from pydantic_ai.embeddings.openai import OpenAIEmbeddingModel
from pydantic_ai.models import Model
Expand All @@ -17,7 +18,7 @@
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.errors import ModelDiscoveryError, ModelTypeNotSupportedError
from providers.inference.model_type_inference import (
infer_model_type_from_pai_known_names,
)
Expand Down Expand Up @@ -71,8 +72,15 @@ async def _discover_models(self) -> list[DiscoveredModel]:
Returns:
Discovered models with inferred types, before allow-list filtering.

Raises:
ModelDiscoveryError: If the Azure OpenAI models API or Entra ID
authentication fails.

"""
response = await self._get_provider.client.models.list()
try:
response = await self._get_provider.client.models.list()
except (APIError, ClientAuthenticationError) as exc:
Comment thread
asimurka marked this conversation as resolved.
Resolved
raise ModelDiscoveryError(self.provider_id) from exc
return [
DiscoveredModel(
provider_model_id=model.id,
Expand Down
16 changes: 12 additions & 4 deletions src/providers/inference/bedrock/provider.py
Original file line number Diff line number Diff line change
Expand Up @@ -9,6 +9,7 @@
import boto3
from botocore.client import BaseClient
from botocore.config import Config
from botocore.exceptions import BotoCoreError, ClientError
from botocore.session import Session
from botocore.tokens import FrozenAuthToken
from pydantic_ai.embeddings import EmbeddingModel
Expand All @@ -27,7 +28,7 @@
to_discovered_model,
)
from providers.inference.catalog import CatalogModel, DiscoveredModel
from providers.inference.errors import ModelTypeNotSupportedError
from providers.inference.errors import ModelDiscoveryError, ModelTypeNotSupportedError

if TYPE_CHECKING:
from providers.inference.types import ModelType
Expand Down Expand Up @@ -120,10 +121,17 @@ async def _discover_models(self) -> list[DiscoveredModel]:
Returns:
Discovered models before allow-list filtering.

Raises:
ModelDiscoveryError: If Bedrock control-plane listing or AWS
credentials/transport fail.

"""
client = self._create_bedrock_client()
foundations = await asyncio.to_thread(list_foundation_models, client)
profiles = await asyncio.to_thread(list_inference_profiles, client)
try:
client = self._create_bedrock_client()
foundations = await asyncio.to_thread(list_foundation_models, client)
profiles = await asyncio.to_thread(list_inference_profiles, client)
except (ClientError, BotoCoreError) as exc:
raise ModelDiscoveryError(self.provider_id) from exc
profile_ids = first_profile_id_by_foundation_model_id(profiles)

discovered: list[DiscoveredModel] = []
Expand Down
14 changes: 14 additions & 0 deletions src/providers/inference/errors.py
Original file line number Diff line number Diff line change
Expand Up @@ -70,3 +70,17 @@ def __init__(self, model_type: ModelType) -> None:
f"No {model_type!r} model is configured or available in the catalog"
)
self.model_type = model_type


class ModelDiscoveryError(Exception):
"""Raised when a provider cannot discover available models."""

def __init__(self, provider_id: ProviderId) -> None:
"""Initialize with the provider whose model discovery failed.

Args:
provider_id: Provider whose model discovery failed.

"""
super().__init__(f"Failed to discover models for provider {provider_id!r}")
self.provider_id = provider_id
11 changes: 9 additions & 2 deletions src/providers/inference/openai/provider.py
Original file line number Diff line number Diff line change
Expand Up @@ -4,6 +4,7 @@

from typing import TYPE_CHECKING, Any

from openai import APIError
from pydantic_ai.embeddings import EmbeddingModel
from pydantic_ai.embeddings.openai import OpenAIEmbeddingModel
from pydantic_ai.models import Model
Expand All @@ -14,7 +15,7 @@
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.errors import ModelDiscoveryError, ModelTypeNotSupportedError
from providers.inference.model_type_inference import (
infer_model_type_from_pai_known_names,
)
Expand Down Expand Up @@ -48,8 +49,14 @@ async def _discover_models(self) -> list[DiscoveredModel]:
Returns:
Discovered models with inferred types, before allow-list filtering.

Raises:
ModelDiscoveryError: If the OpenAI models API fails.

"""
response = await self._get_provider.client.models.list()
try:
response = await self._get_provider.client.models.list()
except APIError as exc:
Comment thread
asimurka marked this conversation as resolved.
Resolved
raise ModelDiscoveryError(self.provider_id) from exc
return [
DiscoveredModel(
provider_model_id=model.id,
Expand Down
34 changes: 21 additions & 13 deletions src/providers/inference/vertexai/provider.py
Original file line number Diff line number Diff line change
Expand Up @@ -4,6 +4,8 @@

from typing import TYPE_CHECKING, Any

from google.auth.exceptions import GoogleAuthError
from google.genai.errors import APIError as GoogleGenAIAPIError
from pydantic_ai.embeddings import EmbeddingModel
from pydantic_ai.embeddings.google import GoogleEmbeddingModel
from pydantic_ai.models import Model
Expand All @@ -14,7 +16,7 @@
from models.config import VertexAIInferenceProviderEntry
from providers.inference.base import InferenceProvider
from providers.inference.catalog import CatalogModel, DiscoveredModel
from providers.inference.errors import ModelTypeNotSupportedError
from providers.inference.errors import ModelDiscoveryError, ModelTypeNotSupportedError
from providers.inference.model_type_inference import (
infer_model_type_from_pai_known_names,
)
Expand Down Expand Up @@ -51,21 +53,27 @@ async def _discover_models(self) -> list[DiscoveredModel]:
Returns:
Discovered models with inferred types, before allow-list filtering.

Raises:
ModelDiscoveryError: If the Vertex AI models API or Google auth fails.

"""
discovered: list[DiscoveredModel] = []
pager = await self._get_provider.client.aio.models.list()
async for model in pager:
if not model.name:
continue
discovered.append(
DiscoveredModel(
provider_model_id=model.name,
model_type=self._infer_model_type(
model.name,
supported_actions=model.supported_actions,
),
try:
pager = await self._get_provider.client.aio.models.list()
async for model in pager:
if not model.name:
continue
discovered.append(
DiscoveredModel(
provider_model_id=model.name,
model_type=self._infer_model_type(
model.name,
supported_actions=model.supported_actions,
),
)
)
)
except (GoogleGenAIAPIError, GoogleAuthError) as exc:
raise ModelDiscoveryError(self.provider_id) from exc
return discovered

def _create_llm(self, model: CatalogModel) -> Model:
Expand Down
11 changes: 9 additions & 2 deletions src/providers/inference/vllm/provider.py
Original file line number Diff line number Diff line change
Expand Up @@ -4,6 +4,7 @@

from typing import Any

from openai import APIError
from pydantic_ai.embeddings import EmbeddingModel
from pydantic_ai.embeddings.openai import OpenAIEmbeddingModel
from pydantic_ai.models import Model
Expand All @@ -14,7 +15,7 @@
from models.config import VllmInferenceProviderEntry
from providers.inference.base import InferenceProvider
from providers.inference.catalog import CatalogModel, DiscoveredModel
from providers.inference.errors import ModelTypeNotSupportedError
from providers.inference.errors import ModelDiscoveryError, ModelTypeNotSupportedError
from providers.inference.model_type_inference import (
infer_model_type_from_pai_known_names,
)
Expand Down Expand Up @@ -46,8 +47,14 @@ async def _discover_models(self) -> list[DiscoveredModel]:
Returns:
Discovered models with inferred types, before allow-list filtering.

Raises:
ModelDiscoveryError: If the vLLM models API fails.

"""
response = await self._get_provider.client.models.list()
try:
response = await self._get_provider.client.models.list()
except APIError as exc:
Comment thread
asimurka marked this conversation as resolved.
Resolved
raise ModelDiscoveryError(self.provider_id) from exc
return [
DiscoveredModel(
provider_model_id=model.id,
Expand Down
21 changes: 14 additions & 7 deletions src/providers/inference/watsonx/provider.py
Original file line number Diff line number Diff line change
Expand Up @@ -5,6 +5,7 @@
import asyncio
from typing import TYPE_CHECKING, Any, Optional

from openai import APIError
from pydantic_ai.embeddings import EmbeddingModel
from pydantic_ai.embeddings.openai import OpenAIEmbeddingModel
from pydantic_ai.models import Model
Expand All @@ -14,7 +15,7 @@
from models.config import WatsonxInferenceProviderEntry
from providers.inference.base import InferenceProvider
from providers.inference.catalog import CatalogModel, DiscoveredModel
from providers.inference.errors import ModelTypeNotSupportedError
from providers.inference.errors import ModelDiscoveryError, ModelTypeNotSupportedError
from providers.inference.watsonx.pai_provider import WatsonxProvider
from providers.inference.watsonx.specs import FoundationModelSpecsResponse

Expand Down Expand Up @@ -49,15 +50,21 @@ async def _discover_models(self) -> list[DiscoveredModel]:
Returns:
Discovered models before allow-list filtering.

Raises:
ModelDiscoveryError: If WatsonX models or foundation-spec APIs fail.

"""
project_id = self._entry.config.project_id
pai = self._get_provider
models_response, functions_by_id = await asyncio.gather(
pai.client.models.list(
extra_query={"project_id": project_id} if project_id else None,
),
self._list_foundation_model_functions(),
)
try:
models_response, functions_by_id = await asyncio.gather(
Comment thread
asimurka marked this conversation as resolved.
Resolved
pai.client.models.list(
extra_query={"project_id": project_id} if project_id else None,
),
self._list_foundation_model_functions(),
)
except APIError as exc:
Comment thread
asimurka marked this conversation as resolved.
Resolved
raise ModelDiscoveryError(self.provider_id) from exc

discovered: list[DiscoveredModel] = []
for model in models_response.data:
Expand Down
38 changes: 36 additions & 2 deletions tests/unit/providers/inference/azure/test_provider.py
Original file line number Diff line number Diff line change
Expand Up @@ -4,8 +4,10 @@

from typing import Any

import httpx
import pytest
from openai import AsyncAzureOpenAI
from azure.core.exceptions import ClientAuthenticationError
from openai import APIConnectionError, AsyncAzureOpenAI
from pydantic import SecretStr
from pydantic_ai.embeddings.openai import OpenAIEmbeddingModel
from pydantic_ai.models.openai import OpenAIResponsesModel
Expand All @@ -14,7 +16,7 @@
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
from providers.inference.errors import ModelDiscoveryError, ModelTypeNotSupportedError

AZURE_ENDPOINT = "https://example.openai.azure.com/"
API_VERSION = "2024-12-01-preview"
Expand Down Expand Up @@ -190,6 +192,38 @@ async def test_discover_models_infers_model_types(
pai_provider.client.models.list.assert_awaited_once_with()


@pytest.mark.asyncio
@pytest.mark.parametrize(
"cause",
[
APIConnectionError(
message="connection failed",
request=httpx.Request( # pyright: ignore[reportArgumentType]
"GET", f"{AZURE_ENDPOINT}openai/models"
),
),
ClientAuthenticationError(message="token acquisition failed"),
],
ids=["connection-error", "entra-auth-error"],
)
async def test_discover_models_wraps_foreign_failures(
provider: AzureInferenceProvider,
mocker: Any,
cause: Exception,
) -> None:
"""OpenAI API subclasses and Entra auth failures become ModelDiscoveryError."""
pai_provider = mocker.MagicMock(spec=AzureProvider)
pai_provider.client.models.list = mocker.AsyncMock(side_effect=cause)
provider._provider = pai_provider

with pytest.raises(ModelDiscoveryError) as exc_info:
await provider._discover_models()

assert exc_info.value.provider_id == "azure"
assert str(exc_info.value) == "Failed to discover models for provider 'azure'"
assert exc_info.value.__cause__ is cause


def test_create_llm_uses_openai_responses_model(
provider: AzureInferenceProvider,
) -> None:
Expand Down
35 changes: 34 additions & 1 deletion tests/unit/providers/inference/bedrock/test_provider.py
Original file line number Diff line number Diff line change
Expand Up @@ -5,6 +5,7 @@
from typing import Any, Optional

import pytest
from botocore.exceptions import ClientError, EndpointConnectionError
from pydantic_ai.embeddings.bedrock import BedrockEmbeddingModel
from pydantic_ai.models.bedrock import BedrockConverseModel
from pydantic_ai.providers.bedrock import BedrockProvider
Expand All @@ -15,7 +16,7 @@
_BearerTokenSession,
)
from providers.inference.catalog import CatalogModel
from providers.inference.errors import ModelTypeNotSupportedError
from providers.inference.errors import ModelDiscoveryError, ModelTypeNotSupportedError


def _entry(**config: Any) -> BedrockInferenceProviderEntry:
Expand Down Expand Up @@ -259,6 +260,38 @@ async def test_discover_models_remaps_profile_only_and_skips_unusable(
]


@pytest.mark.asyncio
@pytest.mark.parametrize(
"cause",
[
ClientError(
{"Error": {"Code": "AccessDeniedException", "Message": "denied"}},
"ListFoundationModels",
),
EndpointConnectionError(endpoint_url="https://bedrock.example.test"),
],
ids=["client-error", "endpoint-connection-error"],
)
async def test_discover_models_wraps_foreign_failures(
provider: BedrockInferenceProvider,
mocker: Any,
cause: Exception,
) -> None:
"""Bedrock service and transport subclasses become ModelDiscoveryError."""
mocker.patch.object(
provider,
"_create_bedrock_client",
side_effect=cause,
)
Comment thread
asimurka marked this conversation as resolved.
Resolved

with pytest.raises(ModelDiscoveryError) as exc_info:
await provider._discover_models()

assert exc_info.value.provider_id == "bedrock"
assert str(exc_info.value) == "Failed to discover models for provider 'bedrock'"
assert exc_info.value.__cause__ is cause


def test_create_llm_uses_bedrock_converse_model(
provider: BedrockInferenceProvider,
mocker: Any,
Expand Down
Loading
Loading