From 352b19899b37f0637ee1dc7e321b59b15fa98694 Mon Sep 17 00:00:00 2001 From: Douwe Maan Date: Fri, 9 Oct 2026 12:13:15 -0600 Subject: [PATCH] Raise `ModelAPIError` for transport errors that escaped Google, Mistral, Cohere, Hugging Face, Anthropic, Groq and Bedrock (#9336) --- docs/models/overview.md | 5 +- .../pydantic_ai/embeddings/cohere.py | 8 ++ .../pydantic_ai/embeddings/google.py | 8 +- .../pydantic_ai/models/__init__.py | 6 +- .../pydantic_ai/models/_continuation.py | 6 +- .../pydantic_ai/models/_transport_errors.py | 11 +++ .../pydantic_ai/models/anthropic.py | 30 +++---- .../pydantic_ai/models/bedrock.py | 10 ++- pydantic_ai_slim/pydantic_ai/models/cohere.py | 8 ++ pydantic_ai_slim/pydantic_ai/models/google.py | 18 +++-- pydantic_ai_slim/pydantic_ai/models/groq.py | 9 +++ .../pydantic_ai/models/huggingface.py | 10 ++- .../pydantic_ai/models/mistral.py | 5 ++ tests/models/test_anthropic.py | 49 +++++++++++- tests/models/test_bedrock.py | 78 ++++++++++++++++++- tests/models/test_cohere.py | 18 +++++ tests/models/test_continuation_stream.py | 22 +++++- tests/models/test_google.py | 66 ++++++++++++++++ tests/models/test_groq.py | 36 ++++++++- tests/models/test_huggingface.py | 77 +++++++++++++++++- tests/models/test_mistral.py | 68 +++++++++++++++- tests/test_embeddings.py | 35 +++++++++ tests/test_streaming.py | 68 +++++++++------- 23 files changed, 588 insertions(+), 63 deletions(-) create mode 100644 pydantic_ai_slim/pydantic_ai/models/_transport_errors.py diff --git a/docs/models/overview.md b/docs/models/overview.md index 744ec3e7db..11430a342f 100644 --- a/docs/models/overview.md +++ b/docs/models/overview.md @@ -311,8 +311,9 @@ in sequence until one succeeds. Pydantic AI can switch to the next model when th raises an exception (like a 4xx/5xx API error) **or** when the response content indicates a semantic failure (like a truncated response or a failed native tool call). -By default, fallback triggers on [`ModelAPIError`][pydantic_ai.exceptions.ModelAPIError] (4xx/5xx API errors), -so you don't need to configure anything for the most common use case. +By default, fallback triggers on [`ModelAPIError`][pydantic_ai.exceptions.ModelAPIError] (4xx/5xx API errors, and +connection failures or timeouts before the response starts), so you don't need to configure anything for the most +common use case. This behavior is controlled by the `fallback_on` parameter (see [`FallbackModel`][pydantic_ai.models.fallback.FallbackModel]), which accepts exception types, diff --git a/pydantic_ai_slim/pydantic_ai/embeddings/cohere.py b/pydantic_ai_slim/pydantic_ai/embeddings/cohere.py index 76b2b518a4..16048ad154 100644 --- a/pydantic_ai_slim/pydantic_ai/embeddings/cohere.py +++ b/pydantic_ai_slim/pydantic_ai/embeddings/cohere.py @@ -5,6 +5,7 @@ from typing import Any, Literal, cast from pydantic_ai.exceptions import ModelAPIError, ModelHTTPError, UnexpectedModelBehavior from pydantic_ai.models import check_allow_model_requests +from pydantic_ai.models._transport_errors import transport_error_message from pydantic_ai.providers import Provider, infer_provider from pydantic_ai.usage import RequestUsage @@ -27,6 +28,10 @@ except ImportError as _import_error: 'you can use the `cohere` optional group — `pip install "pydantic-ai-slim[cohere]"`' ) from _import_error +# Below the guard on purpose: `cohere` requires `httpx`, so without the extra the error above +# is what users should see, not `ModuleNotFoundError: httpx`. +import httpx + @contextmanager def _map_api_errors(model_name: str) -> Generator[None]: @@ -41,6 +46,9 @@ def _map_api_errors(model_name: str) -> Generator[None]: # Neither `embed()` nor `count_tokens()` exercises this fallback. raise ModelAPIError(model_name=model_name, message=str(e)) from e # pragma: no cover + except httpx.TransportError as e: + # `cohere` doesn't wrap connection errors and timeouts in its own exceptions. + raise ModelAPIError(model_name=model_name, message=transport_error_message(e)) from e LatestCohereEmbeddingModelNames = Literal[ diff --git a/pydantic_ai_slim/pydantic_ai/embeddings/google.py b/pydantic_ai_slim/pydantic_ai/embeddings/google.py index d6106f6f40..cec71a2e95 100644 --- a/pydantic_ai_slim/pydantic_ai/embeddings/google.py +++ b/pydantic_ai_slim/pydantic_ai/embeddings/google.py @@ -4,9 +4,12 @@ from contextlib import contextmanager from dataclasses import dataclass, field from typing import Literal, cast -from pydantic_ai.exceptions import ModelHTTPError, UnexpectedModelBehavior +import httpx2 + +from pydantic_ai.exceptions import ModelAPIError, ModelHTTPError, UnexpectedModelBehavior from pydantic_ai.models import check_allow_model_requests from pydantic_ai.models._decode_errors import map_decode_errors +from pydantic_ai.models._transport_errors import transport_error_message from pydantic_ai.providers import Provider, infer_provider from pydantic_ai.usage import RequestUsage @@ -39,6 +42,9 @@ def _map_api_errors(model_name: str) -> Generator[None]: ) from e raise + except httpx2.TransportError as e: + # `google.genai` doesn't wrap connection errors and timeouts in its own exceptions. + raise ModelAPIError(model_name=model_name, message=transport_error_message(e)) from e LatestGoogleGLAEmbeddingModelNames = Literal['gemini-embedding-001', 'gemini-embedding-2-preview', 'gemini-embedding-2'] diff --git a/pydantic_ai_slim/pydantic_ai/models/__init__.py b/pydantic_ai_slim/pydantic_ai/models/__init__.py index 7ec2f7b056..e3f6fc032b 100644 --- a/pydantic_ai_slim/pydantic_ai/models/__init__.py +++ b/pydantic_ai_slim/pydantic_ai/models/__init__.py @@ -40,7 +40,7 @@ from .._output import StructuredTextOutputSchema from .._parts_manager import ModelResponsePartsManager from .._run_context import RunContext from .._warnings import PydanticAIDeprecationWarning as PydanticAIDeprecationWarning -from ..exceptions import UserError +from ..exceptions import ModelAPIError, UserError from ..messages import ( STANDING_PROMPT_PLANTED_KEY, BaseToolCallPart, @@ -1358,6 +1358,10 @@ class StreamedResponse(ABC): except self.get_stream_cancel_errors(): if not self.cancelled: raise + except ModelAPIError as e: + # Adapters map transport errors to `ModelAPIError`, so one caused by `cancel()` arrives wrapped. + if not (self.cancelled and isinstance(e.__cause__, self.get_stream_cancel_errors())): + raise else: # Only natural `StopAsyncIteration` on a stream that wasn't # cancelled flips `_finished`. Early `break` / `aclose()` (raising diff --git a/pydantic_ai_slim/pydantic_ai/models/_continuation.py b/pydantic_ai_slim/pydantic_ai/models/_continuation.py index daffc62709..be50bd18d3 100644 --- a/pydantic_ai_slim/pydantic_ai/models/_continuation.py +++ b/pydantic_ai_slim/pydantic_ai/models/_continuation.py @@ -33,7 +33,7 @@ from typing import Any, Literal from .. import _utils from .._run_context import RunContext -from ..exceptions import UnexpectedModelBehavior +from ..exceptions import ModelAPIError, UnexpectedModelBehavior from ..messages import ( FinalResultEvent, ModelMessage, @@ -384,6 +384,10 @@ class _ContinuationStreamedResponse(StreamedResponse): except self.get_stream_cancel_errors(): if not self.cancelled: raise + except ModelAPIError as e: + # Adapters map transport errors to `ModelAPIError`, so one caused by `cancel()` arrives wrapped. + if not (self.cancelled and isinstance(e.__cause__, self.get_stream_cancel_errors())): + raise else: if not self._cancelled: self._finished = True diff --git a/pydantic_ai_slim/pydantic_ai/models/_transport_errors.py b/pydantic_ai_slim/pydantic_ai/models/_transport_errors.py new file mode 100644 index 0000000000..cf59142277 --- /dev/null +++ b/pydantic_ai_slim/pydantic_ai/models/_transport_errors.py @@ -0,0 +1,11 @@ +"""Describe a transport failure the provider SDK didn't wrap in its own exception.""" + +from __future__ import annotations as _annotations + + +def transport_error_message(e: Exception) -> str: + """Return the message for the `ModelAPIError` raised in place of a transport failure. + + Some transport errors, like a bare `httpx.ReadTimeout()`, stringify to `''`, so fall back to the class name. + """ + return str(e) or type(e).__name__ diff --git a/pydantic_ai_slim/pydantic_ai/models/anthropic.py b/pydantic_ai_slim/pydantic_ai/models/anthropic.py index e8fbdc91eb..64fb717a72 100644 --- a/pydantic_ai_slim/pydantic_ai/models/anthropic.py +++ b/pydantic_ai_slim/pydantic_ai/models/anthropic.py @@ -102,6 +102,7 @@ from ._anthropic_containers import is_tool_result_only as _is_tool_result_only from ._decode_errors import MapStreamDecodeErrors, map_decode_errors from ._prompt_cache import excess_cache_points, previous_tail_needing_breakpoint, split_cache_setting from ._tool_choice import resolve_tool_choice, support_tool_forcing, tool_forcing_unavailable_reason +from ._transport_errors import transport_error_message _FINISH_REASON_MAP: dict[BetaStopReason, FinishReason | None] = { 'compaction': 'stop', @@ -427,6 +428,10 @@ def _map_api_errors(model_name: str, model_id_namespace: str = 'anthropic') -> G raise ModelAPIError(model_name=model_name, message=e.message) from e # pragma: lax no cover except APIConnectionError as e: raise ModelAPIError(model_name=model_name, message=e.message) from e + except httpx2.TransportError as e: + # `anthropic` wraps transport failures in `APIConnectionError` only until the response starts; one that breaks + # off a stream mid-way surfaces as the raw `httpx2` error. + raise ModelAPIError(model_name=model_name, message=transport_error_message(e)) from e LatestAnthropicModelNames = ModelParam @@ -1118,20 +1123,17 @@ class AnthropicModel(Model[AsyncAnthropicClient]): model_request_parameters, ) model_settings = cast(AnthropicModelSettings, model_settings or {}) - # A non-streaming request's transport errors reach us as the SDK's `APIConnectionError`, but a stream's don't. - try: - response = await self._messages_create(messages, False, model_settings, model_request_parameters) - if isinstance(response, BetaMessage): - return self._process_response(response, model_request_parameters, model_settings) - # The request was streamed behind the scenes, see `_messages_create`. - async with response.source: - streamed_response = await self._process_streamed_response( - response, model_request_parameters, model_settings - ) - async for _ in streamed_response: - pass - except httpx2.TransportError as e: - raise ModelAPIError(model_name=self.model_name, message=str(e)) from e + response = await self._messages_create(messages, False, model_settings, model_request_parameters) + if isinstance(response, BetaMessage): + return self._process_response(response, model_request_parameters, model_settings) + # The request was streamed behind the scenes, see `_messages_create`. `_map_api_errors` maps a transport error + # that breaks the stream off, both while it opens and while it's read. + async with response.source: + streamed_response = await self._process_streamed_response( + response, model_request_parameters, model_settings + ) + async for _ in streamed_response: + pass return streamed_response.get() async def count_tokens( diff --git a/pydantic_ai_slim/pydantic_ai/models/bedrock.py b/pydantic_ai_slim/pydantic_ai/models/bedrock.py index ed9a57efe3..dbc0c155ef 100644 --- a/pydantic_ai_slim/pydantic_ai/models/bedrock.py +++ b/pydantic_ai_slim/pydantic_ai/models/bedrock.py @@ -27,6 +27,7 @@ try: HTTPClientError, ) from botocore.model import StructureShape + from urllib3.exceptions import HTTPError as Urllib3HTTPError except ImportError as _import_error: raise ImportError( 'Please install `boto3` to use the Bedrock model, ' @@ -89,6 +90,7 @@ from pydantic_ai.models._tool_choice import ( support_tool_forcing, tool_forcing_unavailable_reason, ) +from pydantic_ai.models._transport_errors import transport_error_message from pydantic_ai.native_tools import AbstractNativeTool, CodeExecutionTool from pydantic_ai.output import StructuredOutputMode from pydantic_ai.profiles import DEFAULT_THINKING_TAGS, ModelProfile @@ -177,6 +179,10 @@ def _map_api_errors(model_name: str, model_id_namespace: str = 'bedrock') -> Gen except (HTTPClientError, BotocoreConnectionError) as e: # botocore raises transport failures (timeouts, connection errors) as `BotoCoreError`, not `ClientError`. raise ModelAPIError(model_name=model_name, message=str(e)) from e + except Urllib3HTTPError as e: + # botocore reads an event stream straight from the urllib3 response, so a connection that breaks off or + # times out mid-stream raises the raw urllib3 error rather than a botocore one. + raise ModelAPIError(model_name=model_name, message=transport_error_message(e)) from e except EventStreamParserError as e: # botocore's event stream parser raises its own `ParserError`, not a `BotoCoreError`, for a streamed 200 whose # body isn't an event stream, e.g. keep-alive whitespace before an upstream failure. @@ -1838,7 +1844,9 @@ class BedrockStreamedResponse(StreamedResponse): _provider_response_id: str | None = None def get_stream_cancel_errors(self) -> tuple[type[BaseException], ...]: - return (BotoCoreError, ClientError) + # botocore reads the event stream straight from the urllib3 response, so tearing it down mid-read can raise a + # raw urllib3 error too. + return (BotoCoreError, ClientError, Urllib3HTTPError) async def close_stream(self) -> None: await anyio.to_thread.run_sync(self._event_stream.close) diff --git a/pydantic_ai_slim/pydantic_ai/models/cohere.py b/pydantic_ai_slim/pydantic_ai/models/cohere.py index 438ee68ac9..1ad521cf28 100644 --- a/pydantic_ai_slim/pydantic_ai/models/cohere.py +++ b/pydantic_ai_slim/pydantic_ai/models/cohere.py @@ -47,6 +47,7 @@ from . import ( check_allow_model_requests, ) from ._tool_choice import resolve_tool_choice +from ._transport_errors import transport_error_message try: from cohere import ( @@ -75,6 +76,10 @@ except ImportError as _import_error: 'you can use the `cohere` optional group — `pip install "pydantic-ai-slim[cohere]"`' ) from _import_error +# Below the guard on purpose: `cohere` requires `httpx`, so without the extra the error above +# is what users should see, not `ModuleNotFoundError: httpx`. +import httpx + LatestCohereModelNames = Literal[ 'c4ai-aya-expanse-32b', 'c4ai-aya-expanse-8b', @@ -214,6 +219,9 @@ class CohereModel(Model[AsyncClientV2]): status_code=status_code, model_name=self.model_name, body=e.body, headers=e.headers ) from e raise ModelAPIError(model_name=self.model_name, message=str(e)) from e + except httpx.TransportError as e: + # `cohere` doesn't wrap connection errors and timeouts in its own exceptions. + raise ModelAPIError(model_name=self.model_name, message=transport_error_message(e)) from e def _get_tool_choice( self, diff --git a/pydantic_ai_slim/pydantic_ai/models/google.py b/pydantic_ai_slim/pydantic_ai/models/google.py index 6410b45814..7377d018eb 100644 --- a/pydantic_ai_slim/pydantic_ai/models/google.py +++ b/pydantic_ai_slim/pydantic_ai/models/google.py @@ -11,6 +11,8 @@ from functools import cached_property from typing import Any, Literal, assert_never, cast, get_args, overload from uuid import uuid4 +import httpx2 + from .. import UnexpectedModelBehavior, _utils, usage from .._run_context import RunContext from ..exceptions import ModelAPIError, ModelHTTPError, UserError @@ -75,6 +77,7 @@ from . import ( ) from ._decode_errors import MapStreamDecodeErrors, map_decode_errors from ._tool_choice import resolve_tool_choice +from ._transport_errors import transport_error_message try: from google.genai import Client, errors @@ -402,8 +405,13 @@ def _resolve_google_cloud_service_tier(model_settings: GoogleModelSettings) -> G return 'pt_then_on_demand' -def _map_api_error(e: errors.APIError, model_name: str, model_id_namespace: str = 'google') -> ModelAPIError: - """Map a `google.genai` API error to the pydantic-ai exception to raise in its place.""" +def _map_api_error( + e: errors.APIError | httpx2.TransportError, model_name: str, model_id_namespace: str = 'google' +) -> ModelAPIError: + """Map a `google.genai` API error or a transport failure to the pydantic-ai exception to raise in its place.""" + if isinstance(e, httpx2.TransportError): + # `google.genai` doesn't wrap connection errors and timeouts in its own exceptions. + return ModelAPIError(model_name=model_name, message=transport_error_message(e)) if (status_code := e.code) >= 400: headers = dict(e.response.headers) if e.response is not None else None # pyright: ignore[reportUnknownMemberType,reportUnknownArgumentType] details = e.details # pyright: ignore[reportUnknownMemberType,reportUnknownVariableType] @@ -964,7 +972,7 @@ class GoogleModel(Model[Client]): try: with map_decode_errors(self._model_name, errors.UnknownApiResponseError): return await func(model=self._model_name, contents=contents, config=config) # pyright: ignore[reportReturnType] - except errors.APIError as e: + except (errors.APIError, httpx2.TransportError) as e: raise _map_api_error(e, self._model_name, self._provider.model_id_namespace) from e def _translate_thinking( @@ -1194,7 +1202,7 @@ class GoogleModel(Model[Client]): try: with map_decode_errors(self._model_name, errors.UnknownApiResponseError): first_chunk = await peekable_response.peek() - except errors.APIError as e: + except (errors.APIError, httpx2.TransportError) as e: raise _map_api_error(e, self._model_name, self._provider.model_id_namespace) from e if isinstance(first_chunk, _utils.Unset): raise UnexpectedModelBehavior('Streamed response ended without content or tool calls') # pragma: no cover @@ -1714,7 +1722,7 @@ class GeminiStreamedResponse(StreamedResponse): grounding_metadata, self.provider_name, self._file_search_tool_call_ids ): yield self._parts_manager.handle_part(vendor_part_id=uuid4(), part=part) - except errors.APIError as e: + except (errors.APIError, httpx2.TransportError) as e: raise _map_api_error(e, self._model_name, self._model_id_namespace) from e def _reserve_part(self, part: NativeToolReturnPart) -> None: diff --git a/pydantic_ai_slim/pydantic_ai/models/groq.py b/pydantic_ai_slim/pydantic_ai/models/groq.py index 377590004f..ff0b6e4b3f 100644 --- a/pydantic_ai_slim/pydantic_ai/models/groq.py +++ b/pydantic_ai_slim/pydantic_ai/models/groq.py @@ -64,6 +64,7 @@ from . import ( ) from ._decode_errors import MapStreamDecodeErrors, map_decode_errors from ._tool_choice import resolve_tool_choice +from ._transport_errors import transport_error_message try: from groq import NOT_GIVEN, APIConnectionError, APIError, APIStatusError, AsyncGroq, AsyncStream, NotGiven @@ -79,6 +80,10 @@ except ImportError as _import_error: 'you can use the `groq` optional group — `pip install "pydantic-ai-slim[groq]"`' ) from _import_error +# Below the guard on purpose: `groq` requires `httpx`, so without the extra the error above +# is what users should see, not `ModuleNotFoundError: httpx`. +import httpx + @contextmanager def _map_api_errors(model_name: str, model_id_namespace: str = 'groq') -> Generator[None]: @@ -105,6 +110,10 @@ def _map_api_errors(model_name: str, model_id_namespace: str = 'groq') -> Genera # The SDK raises the base `APIError` for an error object inside a stream, after the HTTP 200 has already # been received, so there is no status code to report. raise ModelAPIError(model_name=model_name, message=e.message) from e + except httpx.TransportError as e: + # `groq` wraps transport failures in `APIConnectionError` only until the response starts; one that breaks + # off a stream mid-way surfaces as the raw `httpx` error. + raise ModelAPIError(model_name=model_name, message=transport_error_message(e)) from e ProductionGroqModelNames = Literal[ diff --git a/pydantic_ai_slim/pydantic_ai/models/huggingface.py b/pydantic_ai_slim/pydantic_ai/models/huggingface.py index 3298699d6f..60c29c8fca 100644 --- a/pydantic_ai_slim/pydantic_ai/models/huggingface.py +++ b/pydantic_ai_slim/pydantic_ai/models/huggingface.py @@ -54,6 +54,7 @@ from . import ( ) from ._decode_errors import MapStreamDecodeErrors, map_decode_errors from ._tool_choice import resolve_tool_choice +from ._transport_errors import transport_error_message try: from huggingface_hub import ( @@ -70,7 +71,7 @@ try: ChatCompletionStreamOutput, TextGenerationOutputFinishReason, ) - from huggingface_hub.errors import HfHubHTTPError, TextGenerationError + from huggingface_hub.errors import HfHubHTTPError, InferenceTimeoutError, TextGenerationError except ImportError as _import_error: raise ImportError( @@ -78,6 +79,10 @@ except ImportError as _import_error: 'you can use the `huggingface` optional group — `pip install "pydantic-ai-slim[huggingface]"`' ) from _import_error +# Below the guard on purpose: `huggingface_hub` requires `httpx`, so without the extra the error above +# is what users should see, not `ModuleNotFoundError: httpx`. +import httpx + @contextmanager def _map_api_errors(model_name: str) -> Generator[None]: @@ -94,6 +99,9 @@ def _map_api_errors(model_name: str) -> Generator[None]: # Raised for an error object inside a stream, after the HTTP 200 has already been received, so there is no # status code to report. raise ModelAPIError(model_name=model_name, message=str(e)) from e + except (httpx.TransportError, InferenceTimeoutError) as e: + # `huggingface_hub` doesn't wrap connection errors and read timeouts in its own exceptions. + raise ModelAPIError(model_name=model_name, message=transport_error_message(e)) from e __all__ = ( diff --git a/pydantic_ai_slim/pydantic_ai/models/mistral.py b/pydantic_ai_slim/pydantic_ai/models/mistral.py index 11a8aaf454..0fc54b3eed 100644 --- a/pydantic_ai_slim/pydantic_ai/models/mistral.py +++ b/pydantic_ai_slim/pydantic_ai/models/mistral.py @@ -6,6 +6,7 @@ from dataclasses import dataclass, field from datetime import datetime from typing import Any, Literal, assert_never, cast +import httpx2 import pydantic_core from pydantic import JsonValue @@ -65,6 +66,7 @@ from . import ( get_user_agent, ) from ._tool_choice import resolve_tool_choice +from ._transport_errors import transport_error_message try: from mistralai.client import Mistral @@ -128,6 +130,9 @@ def _map_api_errors(model_name: str) -> Generator[None]: status_code=status_code, model_name=model_name, body=e.body, headers=dict(e.headers) ) from e raise ModelAPIError(model_name=model_name, message=e.message) from e + except httpx2.TransportError as e: + # `mistralai` doesn't wrap connection errors and timeouts in its own exceptions. + raise ModelAPIError(model_name=model_name, message=transport_error_message(e)) from e LatestMistralModelNames = Literal[ diff --git a/tests/models/test_anthropic.py b/tests/models/test_anthropic.py index af78f2a07d..69ad3315f9 100644 --- a/tests/models/test_anthropic.py +++ b/tests/models/test_anthropic.py @@ -3,7 +3,7 @@ from __future__ import annotations as _annotations import json import os import re -from collections.abc import Callable, Sequence +from collections.abc import AsyncIterator, Callable, Sequence from dataclasses import dataclass, field from datetime import UTC, datetime from decimal import Decimal @@ -268,9 +268,10 @@ async def test_anthropic_read_error_is_raised_when_not_cancelled(): _enabled_server_tool_names=frozenset(), ) - with pytest.raises(httpx2.ReadError): + with pytest.raises(ModelAPIError) as exc_info: async for _event in response: pass + assert isinstance(exc_info.value.__cause__, httpx2.ReadError) @dataclass @@ -3229,6 +3230,50 @@ def test_model_connection_error(allow_model_requests: None) -> None: assert 'Connection to https://api.anthropic.com timed out' in str(exc_info.value.message) +async def test_stream_transport_error_mid_stream(allow_model_requests: None) -> None: + """A connection that breaks off mid-stream surfaces as `ModelAPIError`, not the raw `httpx2` error. + + The SDK wraps transport failures in `APIConnectionError` only until the response starts, and a cassette can't + replay a broken-off connection, so a mock transport raises it. + """ + events: list[dict[str, Any]] = [ + { + 'type': 'message_start', + 'message': { + 'id': 'msg_1', + 'type': 'message', + 'role': 'assistant', + 'model': 'claude-sonnet-4-5', + 'content': [], + 'stop_reason': None, + 'stop_sequence': None, + 'usage': {'input_tokens': 1, 'output_tokens': 1}, + }, + }, + {'type': 'content_block_start', 'index': 0, 'content_block': {'type': 'text', 'text': ''}}, + {'type': 'content_block_delta', 'index': 0, 'delta': {'type': 'text_delta', 'text': 'Hello'}}, + ] + + class StreamBreakingOff(httpx2.AsyncByteStream): + async def __aiter__(self) -> AsyncIterator[bytes]: + yield ''.join(f'event: {e["type"]}\ndata: {json.dumps(e)}\n\n' for e in events).encode() + raise httpx2.ReadError('connection reset') + + def handler(request: httpx2.Request) -> httpx2.Response: + return httpx2.Response(200, headers={'content-type': 'text/event-stream'}, stream=StreamBreakingOff()) + + async with httpx2.AsyncClient(transport=httpx2.MockTransport(handler)) as http_client: + client = AsyncAnthropic(api_key='test-key', base_url='http://localhost', http_client=http_client) + agent = Agent(AnthropicModel('claude-sonnet-4-5', provider=AnthropicProvider(anthropic_client=client))) + with pytest.raises(ModelAPIError) as exc_info: + async with agent.run_stream('hello') as result: + await result.get_output() + + assert type(exc_info.value) is ModelAPIError + assert exc_info.value.message == 'connection reset' + assert isinstance(exc_info.value.__cause__, httpx2.ReadError) + + async def test_count_tokens_connection_error(allow_model_requests: None) -> None: mock_client = MockAnthropic.create_mock( APIConnectionError( diff --git a/tests/models/test_bedrock.py b/tests/models/test_bedrock.py index 7c08d1f27c..cea288eb7d 100644 --- a/tests/models/test_bedrock.py +++ b/tests/models/test_bedrock.py @@ -90,7 +90,8 @@ with try_import() as imports_successful: from cassetter import Cassette from mypy_boto3_bedrock_runtime import BedrockRuntimeClient from mypy_boto3_bedrock_runtime.type_defs import MessageUnionTypeDef, SystemContentBlockTypeDef, ToolTypeDef - from urllib3 import HTTPResponse + from urllib3 import HTTPConnectionPool, HTTPResponse + from urllib3.exceptions import ProtocolError as Urllib3ProtocolError, ReadTimeoutError as Urllib3ReadTimeoutError from pydantic_ai.models.bedrock import ( BedrockConverseModel, @@ -4199,6 +4200,81 @@ async def test_bedrock_stream_whitespace_only_leading_delta( assert output == expected_output +@pytest.mark.parametrize('error_kind', ['read-timeout', 'connection-reset']) +async def test_bedrock_stream_transport_error_mid_stream( + allow_model_requests: None, bedrock_provider: BedrockProvider, mocker: MockerFixture, error_kind: str +): + """botocore reads an event stream straight from the urllib3 response, so a connection that breaks off mid-stream + raises the raw urllib3 error; it surfaces as `ModelAPIError`. + + Not a VCR test: a cassette can't replay a broken-off connection. + """ + error = ( + Urllib3ReadTimeoutError(HTTPConnectionPool('localhost'), '', 'Read timed out.') + if error_kind == 'read-timeout' + else Urllib3ProtocolError('Connection broken: reset') + ) + model = BedrockConverseModel('us.amazon.nova-micro-v1:0', provider=bedrock_provider) + + def _stream() -> Iterator[dict[str, Any]]: + yield {'messageStart': {'role': 'assistant'}} + yield {'contentBlockDelta': {'contentBlockIndex': 0, 'delta': {'text': 'Hello'}}} + raise error + + mock_converse_stream = mocker.patch.object(model.client, 'converse_stream') + mock_converse_stream.return_value = {'stream': _stream(), 'ResponseMetadata': {'RequestId': 'stub'}} + + with pytest.raises(ModelAPIError) as exc_info: + async with Agent(model).run_stream('hello') as result: + await result.get_output() + + assert type(exc_info.value) is ModelAPIError + assert exc_info.value.__cause__ is error + + +async def test_bedrock_stream_cancel_suppresses_urllib3_error( + allow_model_requests: None, bedrock_provider: BedrockProvider, mocker: MockerFixture +): + """`cancel()` closing the event stream mid-read can make the next read raise a raw urllib3 error; it cancels cleanly. + + Not a VCR test: a cassette can't replay a connection torn down mid-read. + """ + + class _EventStream: + def __init__(self) -> None: + self.closed = False + self._events = iter( + [ + {'messageStart': {'role': 'assistant'}}, + {'contentBlockDelta': {'contentBlockIndex': 0, 'delta': {'text': 'Hello'}}}, + ] + ) + + def __iter__(self) -> _EventStream: + return self + + def __next__(self) -> dict[str, Any]: + if self.closed: + raise Urllib3ProtocolError('Connection broken: closed') + return next(self._events) + + def close(self) -> None: + self.closed = True + + model = BedrockConverseModel('us.amazon.nova-micro-v1:0', provider=bedrock_provider) + mock_converse_stream = mocker.patch.object(model.client, 'converse_stream') + mock_converse_stream.return_value = {'stream': _EventStream(), 'ResponseMetadata': {'RequestId': 'stub'}} + + async with model.request_stream([ModelRequest.user_text_prompt('hi')], None, ModelRequestParameters()) as stream: + iterator = stream.__aiter__() + await iterator.__anext__() + await stream.cancel() + async for _ in iterator: + pass + + assert stream.get().state == 'interrupted' + + @pytest.mark.vcr() async def test_bedrock_error(allow_model_requests: None, bedrock_provider: BedrockProvider): """Test that errors convert to ModelHTTPError.""" diff --git a/tests/models/test_cohere.py b/tests/models/test_cohere.py index de224ba731..97b88fbbad 100644 --- a/tests/models/test_cohere.py +++ b/tests/models/test_cohere.py @@ -8,6 +8,7 @@ from datetime import UTC, datetime from decimal import Decimal from typing import Any, cast +import httpx import pytest from pydantic_ai import ( @@ -647,6 +648,23 @@ def test_model_non_http_error(allow_model_requests: None) -> None: assert exc_info.value.model_name == 'command-r' +async def test_model_transport_error(allow_model_requests: None) -> None: + """`cohere` doesn't wrap transport failures; a cassette can't replay one, so a mock transport raises it.""" + + def handler(request: httpx.Request) -> httpx.Response: + raise httpx.ReadTimeout('timed out', request=request) + + async with httpx.AsyncClient(transport=httpx.MockTransport(handler)) as http_client: + cohere_client = AsyncClientV2(api_key='test-key', base_url='http://localhost', httpx_client=http_client) + agent = Agent(CohereModel('command-r', provider=CohereProvider(cohere_client=cohere_client))) + with pytest.raises(ModelAPIError) as exc_info: + await agent.run('hello') + + assert type(exc_info.value) is ModelAPIError + assert exc_info.value.message == 'timed out' + assert isinstance(exc_info.value.__cause__, httpx.ReadTimeout) + + @pytest.mark.vcr() async def test_request_simple_success_with_vcr(allow_model_requests: None, co_api_key: str): m = CohereModel('command-r7b-12-2024', provider=CohereProvider(api_key=co_api_key)) diff --git a/tests/models/test_continuation_stream.py b/tests/models/test_continuation_stream.py index 1afc20e1ee..a2c794c80b 100644 --- a/tests/models/test_continuation_stream.py +++ b/tests/models/test_continuation_stream.py @@ -28,7 +28,7 @@ import httpx import pytest from pydantic_ai import Agent -from pydantic_ai.exceptions import UnexpectedModelBehavior, UsageLimitExceeded +from pydantic_ai.exceptions import ModelAPIError, UnexpectedModelBehavior, UsageLimitExceeded from pydantic_ai.messages import ( ModelMessage, ModelResponse, @@ -1206,9 +1206,16 @@ async def test_cancel_during_between_segment_sleep_skips_next_request() -> None: assert stream.get().state == 'interrupted' -async def test_cancel_suppresses_non_httpx_segment_error() -> None: +@pytest.mark.parametrize( + ('wrapped', 'cancel'), + [(False, True), (True, True), (True, False)], + ids=['raw', 'wrapped-in-model-api-error', 'wrapped-not-cancelled'], +) +async def test_cancel_suppresses_non_httpx_segment_error(wrapped: bool, cancel: bool) -> None: """`cancel()` tearing down a non-httpx sub-stream suppresses the sub's own transport error type. + Adapters wrap transport errors in `ModelAPIError`, so the error is suppressed when it arrives as that too. + The composite's cancel-guard defaults to httpx errors only; a sub on a different transport (Bedrock botocore, xAI grpc) reports its cancel errors via its own `get_stream_cancel_errors()`. The composite must consult the in-flight sub's error types, or the error escapes to the user @@ -1226,7 +1233,10 @@ async def test_cancel_suppresses_non_httpx_segment_error() -> None: yield self._events[0] # After the first event, `cancel()` tore the connection down; the next pull raises the # transport's own (non-httpx) error, exactly what a real botocore/grpc stream does. - raise _BotoError('connection torn down') + error = _BotoError('connection torn down') + if wrapped: + raise ModelAPIError('fake', 'connection torn down') from error + raise error class _NonHttpxModel(_FakeModel): @asynccontextmanager @@ -1254,6 +1264,12 @@ async def test_cancel_suppresses_non_httpx_segment_error() -> None: iterator = stream.__aiter__() await iterator.__anext__() # first segment in flight + if not cancel: + # Without a `cancel()`, the same error is a real failure and propagates. + with pytest.raises(ModelAPIError): + async for _ in iterator: + pass + return await stream.cancel() # Draining resumes the closed segment, which raises the non-httpx error; the guard must suppress it. diff --git a/tests/models/test_google.py b/tests/models/test_google.py index f1188d2721..7bead2f252 100644 --- a/tests/models/test_google.py +++ b/tests/models/test_google.py @@ -18,8 +18,11 @@ import pytest from cassetter import Cassette from httpx import Timeout from httpx2 import ( + AsyncByteStream as HTTPX2AsyncByteStream, AsyncClient as HTTPX2AsyncClient, + ConnectError as HTTPX2ConnectError, MockTransport as HTTPX2MockTransport, + ReadError as HTTPX2ReadError, Request as HTTPX2Request, Response as HTTPX2Response, ) @@ -5370,6 +5373,69 @@ async def test_google_stream_api_error_before_first_chunk_is_wrapped(allow_model assert len(requests) == 1 +class _HTTPX2StreamBreakingOff(HTTPX2AsyncByteStream): + """A response body that sends `first` and then breaks off with `error`.""" + + def __init__(self, first: bytes, error: Exception): + self.first = first + self.error = error + + async def __aiter__(self) -> AsyncIterator[bytes]: + yield self.first + raise self.error + + +@pytest.mark.parametrize('stream', [False, True]) +async def test_google_connect_error_is_wrapped(allow_model_requests: None, stream: bool): + """`google.genai` doesn't wrap transport failures; a cassette can't replay one, so a mock transport raises it.""" + + async def handler(request: HTTPX2Request) -> HTTPX2Response: + raise HTTPX2ConnectError('connection refused', request=request) + + async with HTTPX2AsyncClient(transport=HTTPX2MockTransport(handler)) as http_client: + model = GoogleModel( + 'gemini-2.5-flash', + provider=GoogleProvider(api_key='test-key', http_client=http_client, base_url='http://localhost'), + ) + with pytest.raises(ModelAPIError) as exc_info: + if stream: + async with Agent(model).run_stream('test'): + pass # pragma: no cover + else: + await Agent(model).run('test') + + assert type(exc_info.value) is ModelAPIError + assert exc_info.value.message == 'connection refused' + assert isinstance(exc_info.value.__cause__, HTTPX2ConnectError) + + +async def test_google_stream_read_error_mid_stream_is_wrapped(allow_model_requests: None): + """A connection that breaks off mid-stream surfaces as `ModelAPIError`, not the raw `httpx2` error.""" + chunk = {'candidates': [{'content': {'role': 'model', 'parts': [{'text': 'Hello'}]}, 'index': 0}]} + + async def handler(request: HTTPX2Request) -> HTTPX2Response: + return HTTPX2Response( + 200, + headers={'content-type': 'text/event-stream'}, + stream=_HTTPX2StreamBreakingOff( + f'data: {json.dumps(chunk)}\r\n\r\n'.encode(), HTTPX2ReadError('connection reset') + ), + ) + + async with HTTPX2AsyncClient(transport=HTTPX2MockTransport(handler)) as http_client: + model = GoogleModel( + 'gemini-2.5-flash', + provider=GoogleProvider(api_key='test-key', http_client=http_client, base_url='http://localhost'), + ) + with pytest.raises(ModelAPIError) as exc_info: + async with Agent(model).run_stream('test') as result: + await result.get_output() + + assert type(exc_info.value) is ModelAPIError + assert exc_info.value.message == 'connection reset' + assert isinstance(exc_info.value.__cause__, HTTPX2ReadError) + + async def test_google_count_tokens_api_error_is_wrapped(allow_model_requests: None): """An API error from `count_tokens` is mapped like one from the request, not raised as the SDK's own error.""" error_response = {'error': {'code': 429, 'message': 'Resource exhausted', 'status': 'RESOURCE_EXHAUSTED'}} diff --git a/tests/models/test_groq.py b/tests/models/test_groq.py index 5c17deed27..a1846976b7 100644 --- a/tests/models/test_groq.py +++ b/tests/models/test_groq.py @@ -3,7 +3,7 @@ from __future__ import annotations as _annotations import json import os import re -from collections.abc import Sequence +from collections.abc import AsyncIterator, Sequence from dataclasses import dataclass from datetime import UTC, datetime from decimal import Decimal @@ -726,6 +726,40 @@ def test_model_connection_error(allow_model_requests: None) -> None: assert 'Connection to https://api.groq.com timed out' in str(exc_info.value.message) +async def test_stream_transport_error_mid_stream(allow_model_requests: None) -> None: + """A connection that breaks off mid-stream surfaces as `ModelAPIError`, not the raw `httpx` error. + + The SDK wraps transport failures in `APIConnectionError` only until the response starts, and a cassette can't + replay a broken-off connection, so a mock transport raises it. + """ + chunk = { + 'id': '1', + 'object': 'chat.completion.chunk', + 'created': 0, + 'model': 'llama-3.3-70b-versatile', + 'choices': [{'index': 0, 'delta': {'role': 'assistant', 'content': 'Hello'}, 'finish_reason': None}], + } + + class StreamBreakingOff(httpx.AsyncByteStream): + async def __aiter__(self) -> AsyncIterator[bytes]: + yield f'data: {json.dumps(chunk)}\n\n'.encode() + raise httpx.RemoteProtocolError('peer closed connection without sending complete message body') + + def handler(request: httpx.Request) -> httpx.Response: + return httpx.Response(200, headers={'content-type': 'text/event-stream'}, stream=StreamBreakingOff()) + + async with httpx.AsyncClient(transport=httpx.MockTransport(handler)) as http_client: + provider = GroqProvider(api_key='test-key', base_url='http://localhost', http_client=http_client) + agent = Agent(GroqModel('llama-3.3-70b-versatile', provider=provider)) + with pytest.raises(ModelAPIError) as exc_info: + async with agent.run_stream('hello') as result: + await result.get_output() + + assert type(exc_info.value) is ModelAPIError + assert exc_info.value.message == 'peer closed connection without sending complete message body' + assert isinstance(exc_info.value.__cause__, httpx.RemoteProtocolError) + + _STREAM_ERROR_SSE_CHUNK = ( b'data: {"id":"1","object":"chat.completion.chunk","created":0,"model":"llama-3.3-70b-versatile",' b'"choices":[{"index":0,"delta":{"role":"assistant","content":"Hello"},"finish_reason":null}]}\n\n' diff --git a/tests/models/test_huggingface.py b/tests/models/test_huggingface.py index 4d12259c40..3529184faf 100644 --- a/tests/models/test_huggingface.py +++ b/tests/models/test_huggingface.py @@ -1,7 +1,7 @@ from __future__ import annotations as _annotations import json -from collections.abc import Sequence +from collections.abc import AsyncIterator, Callable, Sequence from dataclasses import asdict, dataclass, field from datetime import UTC, datetime from decimal import Decimal @@ -47,6 +47,7 @@ from ..conftest import IsDatetime, IsInstance, IsNow, IsStr, message, raise_if_e from .mock_async_stream import MockAsyncStream with try_import() as imports_successful: + import huggingface_hub.utils._http from huggingface_hub import ( AsyncInferenceClient, ChatCompletionInputMessage, @@ -647,6 +648,80 @@ async def test_image_as_binary_content_input( ) +class _StreamBreakingOff(httpx.AsyncByteStream): + """A response body that sends `first` and then breaks off with `error`.""" + + def __init__(self, first: bytes, error: Exception): + self.first = first + self.error = error + + async def __aiter__(self) -> AsyncIterator[bytes]: + yield self.first + raise self.error + + +_TEXT_CHUNK = { + 'id': '1', + 'object': 'chat.completion.chunk', + 'created': 0, + 'model': 'm', + 'system_fingerprint': 'x', + 'choices': [{'index': 0, 'delta': {'role': 'assistant', 'content': 'Hello'}, 'finish_reason': None}], +} + + +def _connect_error(request: httpx.Request) -> httpx.Response: + raise httpx.ConnectError('connection refused', request=request) + + +def _timeout(request: httpx.Request) -> httpx.Response: + # `huggingface_hub` converts `TimeoutError` to `InferenceTimeoutError`. + raise TimeoutError + + +def _stream_breaking_off(request: httpx.Request) -> httpx.Response: + return httpx.Response( + 200, + headers={'content-type': 'text/event-stream'}, + stream=_StreamBreakingOff(f'data: {json.dumps(_TEXT_CHUNK)}\n\n'.encode(), httpx.ReadError('connection reset')), + ) + + +@pytest.mark.parametrize( + ('handler', 'stream', 'cause_name'), + [ + pytest.param(_connect_error, False, 'ConnectError', id='connect'), + pytest.param(_connect_error, True, 'ConnectError', id='stream-connect'), + pytest.param(_timeout, False, 'InferenceTimeoutError', id='timeout'), + pytest.param(_stream_breaking_off, True, 'ReadError', id='mid-stream'), + ], +) +async def test_model_transport_error( + allow_model_requests: None, + monkeypatch: pytest.MonkeyPatch, + handler: Callable[[httpx.Request], httpx.Response], + stream: bool, + cause_name: str, +) -> None: + """`huggingface_hub` doesn't wrap transport failures; a cassette can't replay one, so a mock transport raises it.""" + monkeypatch.setattr( + huggingface_hub.utils._http, + '_GLOBAL_ASYNC_CLIENT_FACTORY', + lambda: httpx.AsyncClient(transport=httpx.MockTransport(handler)), + ) + hf_client = AsyncInferenceClient(base_url='http://localhost/v1', api_key='test-key') + agent = Agent(HuggingFaceModel('m', provider=HuggingFaceProvider(hf_client=hf_client, api_key='test-key'))) + with pytest.raises(ModelAPIError) as exc_info: + if stream: + async with agent.run_stream('hello') as result: + await result.get_output() + else: + await agent.run('hello') + + assert type(exc_info.value) is ModelAPIError + assert type(exc_info.value.__cause__).__name__ == cause_name + + def test_model_status_error(allow_model_requests: None) -> None: error = HfHubHTTPError( message='test_error', diff --git a/tests/models/test_mistral.py b/tests/models/test_mistral.py index 68bbd3c49b..476fb58dbe 100644 --- a/tests/models/test_mistral.py +++ b/tests/models/test_mistral.py @@ -2,7 +2,7 @@ from __future__ import annotations as _annotations import json import re -from collections.abc import Sequence +from collections.abc import AsyncIterator, Callable, Sequence from dataclasses import dataclass, field from datetime import UTC, datetime from decimal import Decimal @@ -3018,6 +3018,72 @@ async def test_uploaded_file_input(allow_model_requests: None): await agent.run(['hello', UploadedFile(file_id='file-123', provider_name='anthropic')]) +class _StreamBreakingOff(httpx2.AsyncByteStream): + """A response body that sends `first` and then breaks off with `error`.""" + + def __init__(self, first: bytes, error: Exception): + self.first = first + self.error = error + + async def __aiter__(self) -> AsyncIterator[bytes]: + yield self.first + raise self.error + + +_TEXT_CHUNK = { + 'id': '1', + 'object': 'chat.completion.chunk', + 'created': 0, + 'model': 'mistral-small-latest', + 'choices': [{'index': 0, 'delta': {'role': 'assistant', 'content': 'Hello'}, 'finish_reason': None}], +} + + +def _connect_error(request: httpx2.Request) -> httpx2.Response: + raise httpx2.ConnectError('connection refused', request=request) + + +def _stream_breaking_off(request: httpx2.Request) -> httpx2.Response: + return httpx2.Response( + 200, + headers={'content-type': 'text/event-stream'}, + stream=_StreamBreakingOff( + f'data: {json.dumps(_TEXT_CHUNK)}\n\n'.encode(), httpx2.ReadError('connection reset') + ), + ) + + +@pytest.mark.parametrize( + ('handler', 'stream', 'message', 'cause'), + [ + pytest.param(_connect_error, False, 'connection refused', httpx2.ConnectError, id='connect'), + pytest.param(_connect_error, True, 'connection refused', httpx2.ConnectError, id='stream-connect'), + pytest.param(_stream_breaking_off, True, 'connection reset', httpx2.ReadError, id='mid-stream'), + ], +) +async def test_model_transport_error( + allow_model_requests: None, + handler: Callable[[httpx2.Request], httpx2.Response], + stream: bool, + message: str, + cause: type[Exception], +) -> None: + """`mistralai` doesn't wrap transport failures; a cassette can't replay one, so a mock transport raises it.""" + async with httpx2.AsyncClient(transport=httpx2.MockTransport(handler)) as http_client: + provider = MistralProvider(api_key='test-key', http_client=http_client, base_url='http://localhost') + agent = Agent(MistralModel('mistral-small-latest', provider=provider)) + with pytest.raises(ModelAPIError) as exc_info: + if stream: + async with agent.run_stream('hello') as result: + await result.get_output() + else: + await agent.run('hello') + + assert type(exc_info.value) is ModelAPIError + assert exc_info.value.message == message + assert isinstance(exc_info.value.__cause__, cause) + + def test_model_status_error(allow_model_requests: None) -> None: response = httpx.Response(500, headers={'retry-after': '30', 'x-request-id': 'rid-1'}, content=b'test error') mock_client = MockMistralAI.create_mock(SDKError('test error', raw_response=response)) diff --git a/tests/test_embeddings.py b/tests/test_embeddings.py index 8cc230bb2e..54384f8164 100644 --- a/tests/test_embeddings.py +++ b/tests/test_embeddings.py @@ -56,6 +56,7 @@ with try_import() as openai_imports_successful: with try_import() as cohere_imports_successful: import cohere + from cohere import AsyncClientV2 from cohere.types.embed_by_type_response import EmbedByTypeResponse from cohere.types.embed_by_type_response_embeddings import EmbedByTypeResponseEmbeddings @@ -632,6 +633,24 @@ class TestCohere: with pytest.raises(ModelHTTPError, match='not found,'): await embedder.embed_query('Hello, world!') + async def test_embed_connect_error(self): + """`cohere` doesn't wrap transport failures; a cassette can't replay one, so a mock transport raises it.""" + + def handler(request: httpx.Request) -> httpx.Response: + raise httpx.ConnectError('connection refused', request=request) + + async with httpx.AsyncClient(transport=httpx.MockTransport(handler)) as http_client: + cohere_client = AsyncClientV2(api_key='test-key', base_url='http://localhost', httpx_client=http_client) + embedder = Embedder( + CohereEmbeddingModel('embed-v4.0', provider=CohereProvider(cohere_client=cohere_client)) + ) + with pytest.raises(ModelAPIError) as exc_info: + await embedder.embed_query('Hello, world!') + + assert type(exc_info.value) is ModelAPIError + assert exc_info.value.message == 'connection refused' + assert isinstance(exc_info.value.__cause__, httpx.ConnectError) + async def test_query_with_cohere_truncate(self, co_api_key: str): model = CohereEmbeddingModel('embed-v4.0', provider=CohereProvider(api_key=co_api_key)) embedder = Embedder(model) @@ -1835,6 +1854,22 @@ class TestGoogle: with pytest.raises(ModelHTTPError, match='not found'): await embedder.embed_query('Hello, world!') + async def test_embed_connect_error(self): + """`google.genai` doesn't wrap transport failures; a cassette can't replay one, so a mock transport raises it.""" + + def handler(request: httpx2.Request) -> httpx2.Response: + raise httpx2.ConnectError('connection refused', request=request) + + async with httpx2.AsyncClient(transport=httpx2.MockTransport(handler)) as http_client: + provider = GoogleProvider(api_key='test-key', http_client=http_client, base_url='http://localhost') + embedder = Embedder(GoogleEmbeddingModel('gemini-embedding-001', provider=provider)) + with pytest.raises(ModelAPIError) as exc_info: + await embedder.embed_query('Hello, world!') + + assert type(exc_info.value) is ModelAPIError + assert exc_info.value.message == 'connection refused' + assert isinstance(exc_info.value.__cause__, httpx2.ConnectError) + async def test_count_tokens_error(self, gemini_api_key: str): model = GoogleEmbeddingModel('nonexistent-model', provider=GoogleProvider(api_key=gemini_api_key)) embedder = Embedder(model) diff --git a/tests/test_streaming.py b/tests/test_streaming.py index 43fe0770c4..0d62a0848c 100644 --- a/tests/test_streaming.py +++ b/tests/test_streaming.py @@ -6530,35 +6530,39 @@ async def test_run_stream_cancel_guard_suppresses_transport_error(): ) +@dataclass +class _TransportErrorStream(models.StreamedResponse): + """Streams one text delta, then raises `error` as an unwrapped transport error would.""" + + error: Exception + + async def _get_event_iterator(self) -> AsyncIterator[Any]: + for event in self._parts_manager.handle_text_delta(vendor_part_id=0, content='x'): + yield event + raise self.error + + async def close_stream(self) -> None: + pass + + @property + def model_name(self) -> str: + return 'httpx2' + + @property + def provider_name(self) -> str: + return 'httpx2' + + @property + def provider_url(self) -> str | None: + return None + + @property + def timestamp(self) -> _datetime: + return _datetime.now(tz=datetime.UTC) + + async def test_stream_cancel_guard_suppresses_httpx2_transport_error(): - @dataclass - class _HTTPX2Stream(models.StreamedResponse): - async def _get_event_iterator(self) -> AsyncIterator[Any]: - for event in self._parts_manager.handle_text_delta(vendor_part_id=0, content='x'): - yield event - assert self.cancelled - raise httpx2.StreamClosed() - - async def close_stream(self) -> None: - pass - - @property - def model_name(self) -> str: - return 'httpx2' - - @property - def provider_name(self) -> str: - return 'httpx2' - - @property - def provider_url(self) -> str | None: - return None - - @property - def timestamp(self) -> _datetime: - return _datetime.now(tz=datetime.UTC) - - stream = _HTTPX2Stream(models.ModelRequestParameters()) + stream = _TransportErrorStream(models.ModelRequestParameters(), httpx2.StreamClosed()) iterator = stream.__aiter__() await iterator.__anext__() await stream.cancel() @@ -6568,6 +6572,14 @@ async def test_stream_cancel_guard_suppresses_httpx2_transport_error(): assert stream.get().state == 'interrupted' +async def test_stream_cancel_guard_reraises_transport_error_when_not_cancelled(): + """A transport error the adapter didn't wrap still propagates when the stream wasn't cancelled.""" + stream = _TransportErrorStream(models.ModelRequestParameters(), httpx2.ReadError('connection reset')) + with pytest.raises(httpx2.ReadError): + async for _ in stream: + pass + + async def test_run_stream_cancel_after_complete(): agent = Agent(TestModel())