Wrap provider model discovery failures in ModelDiscoveryError

Author: asimurka

src/providers/inference/azure/provider.py

@@ -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
@@ -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,
 )
@@ -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:
+            raise ModelDiscoveryError(self.provider_id) from exc
         return [
             DiscoveredModel(
                 provider_model_id=model.id,

src/providers/inference/bedrock/provider.py

@@ -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
@@ -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
@@ -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] = []

src/providers/inference/errors.py

@@ -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

src/providers/inference/openai/provider.py

@@ -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
@@ -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,
 )
@@ -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:
+            raise ModelDiscoveryError(self.provider_id) from exc
         return [
             DiscoveredModel(
                 provider_model_id=model.id,

src/providers/inference/vertexai/provider.py

@@ -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
@@ -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,
 )
@@ -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:

src/providers/inference/vllm/provider.py

@@ -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
@@ -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,
 )
@@ -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:
+            raise ModelDiscoveryError(self.provider_id) from exc
         return [
             DiscoveredModel(
                 provider_model_id=model.id,

src/providers/inference/watsonx/provider.py

@@ -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
@@ -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
 
@@ -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(
+                pai.client.models.list(
+                    extra_query={"project_id": project_id} if project_id else None,
+                ),
+                self._list_foundation_model_functions(),
+            )
+        except APIError as exc:
+            raise ModelDiscoveryError(self.provider_id) from exc
 
         discovered: list[DiscoveredModel] = []
         for model in models_response.data:

tests/unit/providers/inference/azure/test_provider.py

@@ -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
@@ -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"
@@ -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:

tests/unit/providers/inference/bedrock/test_provider.py

@@ -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
@@ -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:
@@ -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,
+    )
+
+    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,