From f7d96cf296e52cd229d80e66cb2e7318e71dfe0a Mon Sep 17 00:00:00 2001 From: asimurka Date: Mon, 2 Nov 2026 10:41:39 +0100 Subject: [PATCH] Wrap provider model discovery failures in ModelDiscoveryError --- src/providers/inference/azure/provider.py | 14 ++++-- src/providers/inference/bedrock/provider.py | 16 ++++-- src/providers/inference/errors.py | 14 ++++++ src/providers/inference/openai/provider.py | 11 ++++- src/providers/inference/vertexai/provider.py | 34 ++++++++----- src/providers/inference/vllm/provider.py | 11 ++++- src/providers/inference/watsonx/provider.py | 21 +++++--- .../inference/azure/test_provider.py | 38 +++++++++++++- .../inference/bedrock/test_provider.py | 35 ++++++++++++- .../inference/openai/test_provider.py | 43 +++++++++++++++- tests/unit/providers/inference/test_errors.py | 13 ++++- .../inference/vertexai/test_provider.py | 31 +++++++++++- .../providers/inference/vllm/test_provider.py | 43 +++++++++++++++- .../inference/watsonx/test_provider.py | 49 ++++++++++++++++++- 14 files changed, 333 insertions(+), 40 deletions(-) diff --git a/src/providers/inference/azure/provider.py b/src/providers/inference/azure/provider.py index 053a4ce29..32c0b348f 100644 --- a/src/providers/inference/azure/provider.py +++ b/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, diff --git a/src/providers/inference/bedrock/provider.py b/src/providers/inference/bedrock/provider.py index c7029c7b2..599495176 100644 --- a/src/providers/inference/bedrock/provider.py +++ b/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] = [] diff --git a/src/providers/inference/errors.py b/src/providers/inference/errors.py index ac839c2ea..73140d9a6 100644 --- a/src/providers/inference/errors.py +++ b/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 diff --git a/src/providers/inference/openai/provider.py b/src/providers/inference/openai/provider.py index d65a20511..7add3d1c4 100644 --- a/src/providers/inference/openai/provider.py +++ b/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, diff --git a/src/providers/inference/vertexai/provider.py b/src/providers/inference/vertexai/provider.py index 99f8458e1..7bf5d03c5 100644 --- a/src/providers/inference/vertexai/provider.py +++ b/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: diff --git a/src/providers/inference/vllm/provider.py b/src/providers/inference/vllm/provider.py index 947a2b0ff..e2fd14558 100644 --- a/src/providers/inference/vllm/provider.py +++ b/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, diff --git a/src/providers/inference/watsonx/provider.py b/src/providers/inference/watsonx/provider.py index 9b09026e9..5ba4c6a9d 100644 --- a/src/providers/inference/watsonx/provider.py +++ b/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: diff --git a/tests/unit/providers/inference/azure/test_provider.py b/tests/unit/providers/inference/azure/test_provider.py index 3dc83ab3a..573ce04fe 100644 --- a/tests/unit/providers/inference/azure/test_provider.py +++ b/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: diff --git a/tests/unit/providers/inference/bedrock/test_provider.py b/tests/unit/providers/inference/bedrock/test_provider.py index 497bd29b0..071f88a66 100644 --- a/tests/unit/providers/inference/bedrock/test_provider.py +++ b/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, diff --git a/tests/unit/providers/inference/openai/test_provider.py b/tests/unit/providers/inference/openai/test_provider.py index 12224c370..4e22c31fb 100644 --- a/tests/unit/providers/inference/openai/test_provider.py +++ b/tests/unit/providers/inference/openai/test_provider.py @@ -4,14 +4,16 @@ from typing import Any +import httpx import pytest +from openai import APIConnectionError, AuthenticationError 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.errors import ModelDiscoveryError, ModelTypeNotSupportedError from providers.inference.openai.provider import OpenAIInferenceProvider @@ -71,6 +73,45 @@ 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", "https://example.test/v1/models" + ), + ), + AuthenticationError( + message="unauthorized", + response=httpx.Response( # pyright: ignore[reportArgumentType] + 401, + request=httpx.Request("GET", "https://example.test/v1/models"), + ), + body=None, + ), + ], + ids=["connection-error", "authentication-error"], +) +async def test_discover_models_wraps_api_error( + provider: OpenAIInferenceProvider, + mocker: Any, + cause: Exception, +) -> None: + """OpenAI models.list API subclasses become ModelDiscoveryError.""" + pai_provider = mocker.MagicMock(spec=OpenAIProvider) + 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 == "openai" + assert str(exc_info.value) == "Failed to discover models for provider 'openai'" + assert exc_info.value.__cause__ is cause + + def test_create_llm_uses_openai_responses_model( provider: OpenAIInferenceProvider, ) -> None: diff --git a/tests/unit/providers/inference/test_errors.py b/tests/unit/providers/inference/test_errors.py index d4057e292..3ceac9e89 100644 --- a/tests/unit/providers/inference/test_errors.py +++ b/tests/unit/providers/inference/test_errors.py @@ -4,6 +4,7 @@ from models.config import InferenceProviderEntryBase from providers.inference.errors import ( + ModelDiscoveryError, ModelNotAvailableError, ModelTypeNotSupportedError, NoModelConfiguredError, @@ -34,8 +35,8 @@ def test_model_type_not_supported_error_message_and_attributes() -> None: assert str(error) == ("Model type 'rerank' is not supported by provider 'openai'") -def test_unknown_provider_type_error_message_and_attributes() -> None: - """Error message and attributes identify the unknown provider type.""" +def test_provider_not_found_error_message_and_attributes() -> None: + """Error message and attributes identify the unknown provider id.""" error = ProviderNotFoundError("fake-openai") assert error.provider_id == "fake-openai" @@ -50,6 +51,14 @@ def test_no_model_configured_error_message_and_attributes() -> None: assert str(error) == ("No 'llm' model is configured or available in the catalog") +def test_model_discovery_error_message_and_attributes() -> None: + """Error message and attributes identify the provider without foreign details.""" + error = ModelDiscoveryError("azure") + + assert error.provider_id == "azure" + assert str(error) == "Failed to discover models for provider 'azure'" + + @pytest.mark.parametrize("model_type", ["llm", "embedding", "rerank"]) def test_get_model_instance_raises_model_not_available_error( model_type: ModelType, diff --git a/tests/unit/providers/inference/vertexai/test_provider.py b/tests/unit/providers/inference/vertexai/test_provider.py index 808f1202b..d7a7e9b36 100644 --- a/tests/unit/providers/inference/vertexai/test_provider.py +++ b/tests/unit/providers/inference/vertexai/test_provider.py @@ -5,13 +5,15 @@ from typing import Any import pytest +from google.auth.exceptions import DefaultCredentialsError +from google.genai.errors import ClientError as GoogleGenAIClientError from pydantic_ai.embeddings.google import GoogleEmbeddingModel from pydantic_ai.models.google import GoogleModel from pydantic_ai.providers.google_cloud import GoogleCloudProvider from models.config import VertexAIInferenceProviderEntry from providers.inference.catalog import CatalogModel -from providers.inference.errors import ModelTypeNotSupportedError +from providers.inference.errors import ModelDiscoveryError, ModelTypeNotSupportedError from providers.inference.vertexai.provider import VertexAIInferenceProvider @@ -107,6 +109,33 @@ async def test_discover_models_keeps_catalog_name( pai_provider.client.aio.models.list.assert_awaited_once_with() +@pytest.mark.asyncio +@pytest.mark.parametrize( + "cause", + [ + GoogleGenAIClientError(401, {"error": {"message": "unauthorized"}}), + DefaultCredentialsError("credentials unavailable"), + ], + ids=["genai-client-error", "default-credentials-error"], +) +async def test_discover_models_wraps_foreign_failures( + provider: VertexAIInferenceProvider, + mocker: Any, + cause: Exception, +) -> None: + """Vertex API/auth subclasses become ModelDiscoveryError.""" + pai_provider = mocker.MagicMock(spec=GoogleCloudProvider) + pai_provider.client.aio.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 == "vertexai" + assert str(exc_info.value) == "Failed to discover models for provider 'vertexai'" + assert exc_info.value.__cause__ is cause + + def test_create_llm_uses_google_model( provider: VertexAIInferenceProvider, ) -> None: diff --git a/tests/unit/providers/inference/vllm/test_provider.py b/tests/unit/providers/inference/vllm/test_provider.py index c247eb705..87ef7249b 100644 --- a/tests/unit/providers/inference/vllm/test_provider.py +++ b/tests/unit/providers/inference/vllm/test_provider.py @@ -4,14 +4,16 @@ from typing import Any +import httpx import pytest +from openai import APIConnectionError, AuthenticationError from pydantic_ai.embeddings.openai import OpenAIEmbeddingModel from pydantic_ai.models.openai import OpenAIChatModel from pydantic_ai.providers.vllm import VLLMProvider from models.config import VllmInferenceProviderEntry from providers.inference.catalog import CatalogModel -from providers.inference.errors import ModelTypeNotSupportedError +from providers.inference.errors import ModelDiscoveryError, ModelTypeNotSupportedError from providers.inference.vllm.provider import VllmInferenceProvider BASE_URL = "https://example.test/v1" @@ -88,6 +90,45 @@ 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"{BASE_URL}/models" + ), + ), + AuthenticationError( + message="unauthorized", + response=httpx.Response( # pyright: ignore[reportArgumentType] + 401, + request=httpx.Request("GET", f"{BASE_URL}/models"), + ), + body=None, + ), + ], + ids=["connection-error", "authentication-error"], +) +async def test_discover_models_wraps_api_error( + provider: VllmInferenceProvider, + mocker: Any, + cause: Exception, +) -> None: + """vLLM models.list API subclasses become ModelDiscoveryError.""" + pai_provider = mocker.MagicMock(spec=VLLMProvider) + 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 == "vllm" + assert str(exc_info.value) == "Failed to discover models for provider 'vllm'" + assert exc_info.value.__cause__ is cause + + def test_create_llm_uses_openai_chat_model( provider: VllmInferenceProvider, ) -> None: diff --git a/tests/unit/providers/inference/watsonx/test_provider.py b/tests/unit/providers/inference/watsonx/test_provider.py index a15fb8d7c..71367bc70 100644 --- a/tests/unit/providers/inference/watsonx/test_provider.py +++ b/tests/unit/providers/inference/watsonx/test_provider.py @@ -4,13 +4,15 @@ from typing import Any, Optional +import httpx import pytest +from openai import APIConnectionError, AuthenticationError from pydantic_ai.embeddings.openai import OpenAIEmbeddingModel from pydantic_ai.models.openai import OpenAIChatModel from models.config import WatsonxInferenceProviderEntry from providers.inference.catalog import CatalogModel -from providers.inference.errors import ModelTypeNotSupportedError +from providers.inference.errors import ModelDiscoveryError, ModelTypeNotSupportedError from providers.inference.watsonx.auth import WatsonxIAMAuth from providers.inference.watsonx.pai_provider import WatsonxProvider from providers.inference.watsonx.provider import WatsonxInferenceProvider @@ -106,6 +108,51 @@ async def test_discover_models_keeps_chat_and_embedding_only( ) +@pytest.mark.asyncio +@pytest.mark.parametrize( + "cause", + [ + APIConnectionError( + message="connection failed", + request=httpx.Request( # pyright: ignore[reportArgumentType] + "GET", "https://us-south.ml.cloud.ibm.com/ml/v1/models" + ), + ), + AuthenticationError( + message="unauthorized", + response=httpx.Response( # pyright: ignore[reportArgumentType] + 401, + request=httpx.Request( + "GET", "https://us-south.ml.cloud.ibm.com/ml/v1/models" + ), + ), + body=None, + ), + ], + ids=["connection-error", "authentication-error"], +) +async def test_discover_models_wraps_api_error( + provider: WatsonxInferenceProvider, + mocker: Any, + cause: Exception, +) -> None: + """WatsonX models.list API subclasses become ModelDiscoveryError.""" + pai_provider = _attach_pai_provider(provider, mocker) + pai_provider.client.models.list = mocker.AsyncMock(side_effect=cause) + mocker.patch.object( + provider, + "_list_foundation_model_functions", + new=mocker.AsyncMock(return_value={}), + ) + + with pytest.raises(ModelDiscoveryError) as exc_info: + await provider._discover_models() + + assert exc_info.value.provider_id == "watsonx" + assert str(exc_info.value) == "Failed to discover models for provider 'watsonx'" + assert exc_info.value.__cause__ is cause + + @pytest.mark.asyncio async def test_list_foundation_model_functions( provider: WatsonxInferenceProvider,