pydantic-ai-slim 0.0.25__py3-none-any.whl → 0.0.27__py3-none-any.whl

pydantic_ai/models/fallback.py

@@ -0,0 +1,116 @@
+from __future__ import annotations as _annotations
+
+from collections.abc import AsyncIterator
+from contextlib import AsyncExitStack, asynccontextmanager
+from dataclasses import dataclass, field
+from typing import TYPE_CHECKING, Callable
+
+from ..exceptions import FallbackExceptionGroup, ModelHTTPError
+from . import KnownModelName, Model, ModelRequestParameters, StreamedResponse, infer_model
+
+if TYPE_CHECKING:
+    from ..messages import ModelMessage, ModelResponse
+    from ..settings import ModelSettings
+    from ..usage import Usage
+
+
+@dataclass(init=False)
+class FallbackModel(Model):
+    """A model that uses one or more fallback models upon failure.
+
+    Apart from `__init__`, all methods are private or match those of the base class.
+    """
+
+    models: list[Model]
+
+    _model_name: str = field(repr=False)
+    _fallback_on: Callable[[Exception], bool]
+
+    def __init__(
+        self,
+        default_model: Model | KnownModelName,
+        *fallback_models: Model | KnownModelName,
+        fallback_on: Callable[[Exception], bool] | tuple[type[Exception], ...] = (ModelHTTPError,),
+    ):
+        """Initialize a fallback model instance.
+
+        Args:
+            default_model: The name or instance of the default model to use.
+            fallback_models: The names or instances of the fallback models to use upon failure.
+            fallback_on: A callable or tuple of exceptions that should trigger a fallback.
+        """
+        self.models = [infer_model(default_model), *[infer_model(m) for m in fallback_models]]
+        self._model_name = f'FallBackModel[{", ".join(model.model_name for model in self.models)}]'
+
+        if isinstance(fallback_on, tuple):
+            self._fallback_on = _default_fallback_condition_factory(fallback_on)
+        else:
+            self._fallback_on = fallback_on
+
+    async def request(
+        self,
+        messages: list[ModelMessage],
+        model_settings: ModelSettings | None,
+        model_request_parameters: ModelRequestParameters,
+    ) -> tuple[ModelResponse, Usage]:
+        """Try each model in sequence until one succeeds.
+
+        In case of failure, raise a FallbackExceptionGroup with all exceptions.
+        """
+        exceptions: list[Exception] = []
+
+        for model in self.models:
+            try:
+                return await model.request(messages, model_settings, model_request_parameters)
+            except Exception as exc:
+                if self._fallback_on(exc):
+                    exceptions.append(exc)
+                    continue
+                raise exc
+
+        raise FallbackExceptionGroup('All models from FallbackModel failed', exceptions)
+
+    @asynccontextmanager
+    async def request_stream(
+        self,
+        messages: list[ModelMessage],
+        model_settings: ModelSettings | None,
+        model_request_parameters: ModelRequestParameters,
+    ) -> AsyncIterator[StreamedResponse]:
+        """Try each model in sequence until one succeeds."""
+        exceptions: list[Exception] = []
+
+        for model in self.models:
+            async with AsyncExitStack() as stack:
+                try:
+                    response = await stack.enter_async_context(
+                        model.request_stream(messages, model_settings, model_request_parameters)
+                    )
+                except Exception as exc:
+                    if self._fallback_on(exc):
+                        exceptions.append(exc)
+                        continue
+                    raise exc
+                yield response
+                return
+
+        raise FallbackExceptionGroup('All models from FallbackModel failed', exceptions)
+
+    @property
+    def model_name(self) -> str:
+        """The model name."""
+        return self._model_name
+
+    @property
+    def system(self) -> str | None:
+        """The system / model provider, n/a for fallback models."""
+        return None
+
+
+def _default_fallback_condition_factory(exceptions: tuple[type[Exception], ...]) -> Callable[[Exception], bool]:
+    """Create a default fallback condition for the given exceptions."""
+
+    def fallback_condition(exception: Exception) -> bool:
+        return isinstance(exception, exceptions)
+
+    return fallback_condition

pydantic_ai/models/function.py

@@ -2,7 +2,7 @@ from __future__ import annotations as _annotations
 
 import inspect
 import re
-from collections.abc import AsyncIterator, Awaitable, Iterable
+from collections.abc import AsyncIterator, Awaitable, Iterable, Sequence
 from contextlib import asynccontextmanager
 from dataclasses import dataclass, field
 from datetime import datetime
@@ -14,6 +14,9 @@ from typing_extensions import TypeAlias, assert_never, overload
 from .. import _utils, usage
 from .._utils import PeekableAsyncStream
 from ..messages import (
+    AudioUrl,
+    BinaryContent,
+    ImageUrl,
     ModelMessage,
     ModelRequest,
     ModelResponse,
@@ -23,6 +26,7 @@ from ..messages import (
     TextPart,
     ToolCallPart,
     ToolReturnPart,
+    UserContent,
     UserPromptPart,
 )
 from ..settings import ModelSettings
@@ -44,15 +48,23 @@ class FunctionModel(Model):
     _system: str | None = field(default=None, repr=False)
 
     @overload
-    def __init__(self, function: FunctionDef) -> None: ...
+    def __init__(self, function: FunctionDef, *, model_name: str | None = None) -> None: ...
 
     @overload
-    def __init__(self, *, stream_function: StreamFunctionDef) -> None: ...
+    def __init__(self, *, stream_function: StreamFunctionDef, model_name: str | None = None) -> None: ...
 
     @overload
-    def __init__(self, function: FunctionDef, *, stream_function: StreamFunctionDef) -> None: ...
+    def __init__(
+        self, function: FunctionDef, *, stream_function: StreamFunctionDef, model_name: str | None = None
+    ) -> None: ...
 
-    def __init__(self, function: FunctionDef | None = None, *, stream_function: StreamFunctionDef | None = None):
+    def __init__(
+        self,
+        function: FunctionDef | None = None,
+        *,
+        stream_function: StreamFunctionDef | None = None,
+        model_name: str | None = None,
+    ):
         """Initialize a `FunctionModel`.
 
         Either `function` or `stream_function` must be provided, providing both is allowed.
@@ -60,6 +72,7 @@ class FunctionModel(Model):
         Args:
             function: The function to call for non-streamed requests.
             stream_function: The function to call for streamed requests.
+            model_name: The name of the model. If not provided, a name is generated from the function names.
         """
         if function is None and stream_function is None:
             raise TypeError('Either `function` or `stream_function` must be provided')
@@ -68,7 +81,7 @@ class FunctionModel(Model):
 
         function_name = self.function.__name__ if self.function is not None else ''
         stream_function_name = self.stream_function.__name__ if self.stream_function is not None else ''
-        self._model_name = f'function:{function_name}:{stream_function_name}'
+        self._model_name = model_name or f'function:{function_name}:{stream_function_name}'
 
     async def request(
         self,
@@ -91,7 +104,7 @@ class FunctionModel(Model):
             response_ = await _utils.run_in_executor(self.function, messages, agent_info)
             assert isinstance(response_, ModelResponse), response_
             response = response_
-        response.model_name = f'function:{self.function.__name__}'
+        response.model_name = self._model_name
         # TODO is `messages` right here? Should it just be new messages?
         return response, _estimate_usage(chain(messages, [response]))
 
@@ -119,7 +132,7 @@ class FunctionModel(Model):
         if isinstance(first, _utils.Unset):
             raise ValueError('Stream function must return at least one item')
 
-        yield FunctionStreamedResponse(_model_name=f'function:{self.stream_function.__name__}', _iter=response_stream)
+        yield FunctionStreamedResponse(_model_name=self._model_name, _iter=response_stream)
 
     @property
     def model_name(self) -> str:
@@ -262,7 +275,12 @@ def _estimate_usage(messages: Iterable[ModelMessage]) -> usage.Usage:
     )
 
 
-def _estimate_string_tokens(content: str) -> int:
+def _estimate_string_tokens(content: str | Sequence[UserContent]) -> int:
     if not content:
         return 0
-    return len(re.split(r'[\s",.:]+', content.strip()))
+    if isinstance(content, str):
+        return len(re.split(r'[\s",.:]+', content.strip()))
+    # TODO(Marcelo): We need to study how we can estimate the tokens for these types of content.
+    else:  # pragma: no cover
+        assert isinstance(content, (AudioUrl, ImageUrl, BinaryContent))
+        return 0

pydantic_ai/models/gemini.py

@@ -1,5 +1,6 @@
 from __future__ import annotations as _annotations
 
+import base64
 import os
 import re
 from collections.abc import AsyncIterator, Sequence
@@ -14,8 +15,11 @@ import pydantic
 from httpx import USE_CLIENT_DEFAULT, AsyncClient as AsyncHTTPClient, Response as HTTPResponse
 from typing_extensions import NotRequired, TypedDict, assert_never
 
-from .. import UnexpectedModelBehavior, _utils, exceptions, usage
+from .. import ModelHTTPError, UnexpectedModelBehavior, UserError, _utils, usage
 from ..messages import (
+    AudioUrl,
+    BinaryContent,
+    ImageUrl,
     ModelMessage,
     ModelRequest,
     ModelResponse,
@@ -108,7 +112,7 @@ class GeminiModel(Model):
             if env_api_key := os.getenv('GEMINI_API_KEY'):
                 api_key = env_api_key
             else:
-                raise exceptions.UserError('API key must be provided or set in the GEMINI_API_KEY environment variable')
+                raise UserError('API key must be provided or set in the GEMINI_API_KEY environment variable')
         self.http_client = http_client or cached_async_http_client()
         self._auth = ApiKeyAuth(api_key)
         self._url = url_template.format(model=model_name)
@@ -185,7 +189,7 @@ class GeminiModel(Model):
     ) -> AsyncIterator[HTTPResponse]:
         tools = self._get_tools(model_request_parameters)
         tool_config = self._get_tool_config(model_request_parameters, tools)
-        sys_prompt_parts, contents = self._message_to_gemini_content(messages)
+        sys_prompt_parts, contents = await self._message_to_gemini_content(messages)
 
         request_data = _GeminiRequest(contents=contents)
         if sys_prompt_parts:
@@ -229,9 +233,11 @@ class GeminiModel(Model):
             headers=headers,
             timeout=model_settings.get('timeout', USE_CLIENT_DEFAULT),
         ) as r:
-            if r.status_code != 200:
+            if (status_code := r.status_code) != 200:
                 await r.aread()
-                raise exceptions.UnexpectedModelBehavior(f'Unexpected response from gemini {r.status_code}', r.text)
+                if status_code >= 400:
+                    raise ModelHTTPError(status_code=status_code, model_name=self.model_name, body=r.text)
+                raise UnexpectedModelBehavior(f'Unexpected response from gemini {status_code}', r.text)
             yield r
 
     def _process_response(self, response: _GeminiResponse) -> ModelResponse:
@@ -269,7 +275,7 @@ class GeminiModel(Model):
         return GeminiStreamedResponse(_model_name=self._model_name, _content=content, _stream=aiter_bytes)
 
     @classmethod
-    def _message_to_gemini_content(
+    async def _message_to_gemini_content(
         cls, messages: list[ModelMessage]
     ) -> tuple[list[_GeminiTextPart], list[_GeminiContent]]:
         sys_prompt_parts: list[_GeminiTextPart] = []
@@ -282,7 +288,7 @@ class GeminiModel(Model):
                     if isinstance(part, SystemPromptPart):
                         sys_prompt_parts.append(_GeminiTextPart(text=part.content))
                     elif isinstance(part, UserPromptPart):
-                        message_parts.append(_GeminiTextPart(text=part.content))
+                        message_parts.extend(await cls._map_user_prompt(part))
                     elif isinstance(part, ToolReturnPart):
                         message_parts.append(_response_part_from_response(part.tool_name, part.model_response_object()))
                     elif isinstance(part, RetryPromptPart):
@@ -303,6 +309,40 @@ class GeminiModel(Model):
 
         return sys_prompt_parts, contents
 
+    @staticmethod
+    async def _map_user_prompt(part: UserPromptPart) -> list[_GeminiPartUnion]:
+        if isinstance(part.content, str):
+            return [{'text': part.content}]
+        else:
+            content: list[_GeminiPartUnion] = []
+            for item in part.content:
+                if isinstance(item, str):
+                    content.append({'text': item})
+                elif isinstance(item, BinaryContent):
+                    base64_encoded = base64.b64encode(item.data).decode('utf-8')
+                    content.append(
+                        _GeminiInlineDataPart(inline_data={'data': base64_encoded, 'mime_type': item.media_type})
+                    )
+                elif isinstance(item, (AudioUrl, ImageUrl)):
+                    try:
+                        content.append(
+                            _GeminiFileDataPart(file_data={'file_uri': item.url, 'mime_type': item.media_type})
+                        )
+                    except ValueError:
+                        # Download the file if can't find the mime type.
+                        client = cached_async_http_client()
+                        response = await client.get(item.url, follow_redirects=True)
+                        response.raise_for_status()
+                        base64_encoded = base64.b64encode(response.content).decode('utf-8')
+                        content.append(
+                            _GeminiInlineDataPart(
+                                inline_data={'data': base64_encoded, 'mime_type': response.headers['Content-Type']}
+                            )
+                        )
+                else:
+                    assert_never(item)
+        return content
+
 
 class AuthProtocol(Protocol):
     """Abstract definition for Gemini authentication."""
@@ -494,6 +534,28 @@ class _GeminiTextPart(TypedDict):
     text: str
 
 
+class _GeminiInlineData(TypedDict):
+    data: str
+    mime_type: Annotated[str, pydantic.Field(alias='mimeType')]
+
+
+class _GeminiInlineDataPart(TypedDict):
+    """See <https://ai.google.dev/api/caching#Blob>."""
+
+    inline_data: Annotated[_GeminiInlineData, pydantic.Field(alias='inlineData')]
+
+
+class _GeminiFileData(TypedDict):
+    """See <https://ai.google.dev/api/caching#FileData>."""
+
+    file_uri: Annotated[str, pydantic.Field(alias='fileUri')]
+    mime_type: Annotated[str, pydantic.Field(alias='mimeType')]
+
+
+class _GeminiFileDataPart(TypedDict):
+    file_data: Annotated[_GeminiFileData, pydantic.Field(alias='fileData')]
+
+
 class _GeminiFunctionCallPart(TypedDict):
     function_call: Annotated[_GeminiFunctionCall, pydantic.Field(alias='functionCall')]
 
@@ -517,7 +579,7 @@ def _process_response_from_parts(
                 )
             )
         elif 'function_response' in part:
-            raise exceptions.UnexpectedModelBehavior(
+            raise UnexpectedModelBehavior(
                 f'Unsupported response from Gemini, expected all parts to be function calls or text, got: {part!r}'
             )
     return ModelResponse(parts=items, model_name=model_name, timestamp=timestamp or _utils.now_utc())
@@ -549,6 +611,10 @@ def _part_discriminator(v: Any) -> str:
     if isinstance(v, dict):
         if 'text' in v:
             return 'text'
+        elif 'inlineData' in v:
+            return 'inline_data'
+        elif 'fileData' in v:
+            return 'file_data'
         elif 'functionCall' in v or 'function_call' in v:
             return 'function_call'
         elif 'functionResponse' in v or 'function_response' in v:
@@ -564,6 +630,8 @@ _GeminiPartUnion = Annotated[
         Annotated[_GeminiTextPart, pydantic.Tag('text')],
         Annotated[_GeminiFunctionCallPart, pydantic.Tag('function_call')],
         Annotated[_GeminiFunctionResponsePart, pydantic.Tag('function_response')],
+        Annotated[_GeminiInlineDataPart, pydantic.Tag('inline_data')],
+        Annotated[_GeminiFileDataPart, pydantic.Tag('file_data')],
     ],
     pydantic.Discriminator(_part_discriminator),
 ]
@@ -726,7 +794,7 @@ class _GeminiJsonSchema:
             # noinspection PyTypeChecker
             key = re.sub(r'^#/\$defs/', '', ref)
             if key in refs_stack:
-                raise exceptions.UserError('Recursive `$ref`s in JSON Schema are not supported by Gemini')
+                raise UserError('Recursive `$ref`s in JSON Schema are not supported by Gemini')
             refs_stack += (key,)
             schema_def = self.defs[key]
             self._simplify(schema_def, refs_stack)
@@ -760,7 +828,7 @@ class _GeminiJsonSchema:
     def _object(self, schema: dict[str, Any], refs_stack: tuple[str, ...]) -> None:
         ad_props = schema.pop('additionalProperties', None)
         if ad_props:
-            raise exceptions.UserError('Additional properties in JSON Schema are not supported by Gemini')
+            raise UserError('Additional properties in JSON Schema are not supported by Gemini')
 
         if properties := schema.get('properties'):  # pragma: no branch
             for value in properties.values():

pydantic_ai/models/groq.py

@@ -1,5 +1,6 @@
 from __future__ import annotations as _annotations
 
+import base64
 from collections.abc import AsyncIterable, AsyncIterator, Iterable
 from contextlib import asynccontextmanager
 from dataclasses import dataclass, field
@@ -10,9 +11,11 @@ from typing import Literal, Union, cast, overload
 from httpx import AsyncClient as AsyncHTTPClient
 from typing_extensions import assert_never
 
-from .. import UnexpectedModelBehavior, _utils, usage
+from .. import ModelHTTPError, UnexpectedModelBehavior, _utils, usage
 from .._utils import guard_tool_call_id as _guard_tool_call_id
 from ..messages import (
+    BinaryContent,
+    ImageUrl,
     ModelMessage,
     ModelRequest,
     ModelResponse,
@@ -36,9 +39,9 @@ from . import (
 )
 
 try:
-    from groq import NOT_GIVEN, AsyncGroq, AsyncStream
+    from groq import NOT_GIVEN, APIStatusError, AsyncGroq, AsyncStream
     from groq.types import chat
-    from groq.types.chat import ChatCompletion, ChatCompletionChunk
+    from groq.types.chat.chat_completion_content_part_image_param import ImageURL
 except ImportError as _import_error:
     raise ImportError(
         'Please install `groq` to use the Groq model, '
@@ -163,7 +166,7 @@ class GroqModel(Model):
         stream: Literal[True],
         model_settings: GroqModelSettings,
         model_request_parameters: ModelRequestParameters,
-    ) -> AsyncStream[ChatCompletionChunk]:
+    ) -> AsyncStream[chat.ChatCompletionChunk]:
         pass
 
     @overload
@@ -182,7 +185,7 @@ class GroqModel(Model):
         stream: bool,
         model_settings: GroqModelSettings,
         model_request_parameters: ModelRequestParameters,
-    ) -> chat.ChatCompletion | AsyncStream[ChatCompletionChunk]:
+    ) -> chat.ChatCompletion | AsyncStream[chat.ChatCompletionChunk]:
         tools = self._get_tools(model_request_parameters)
         # standalone function to make it easier to override
         if not tools:
@@ -194,23 +197,28 @@ class GroqModel(Model):
 
         groq_messages = list(chain(*(self._map_message(m) for m in messages)))
 
-        return await self.client.chat.completions.create(
-            model=str(self._model_name),
-            messages=groq_messages,
-            n=1,
-            parallel_tool_calls=model_settings.get('parallel_tool_calls', NOT_GIVEN),
-            tools=tools or NOT_GIVEN,
-            tool_choice=tool_choice or NOT_GIVEN,
-            stream=stream,
-            max_tokens=model_settings.get('max_tokens', NOT_GIVEN),
-            temperature=model_settings.get('temperature', NOT_GIVEN),
-            top_p=model_settings.get('top_p', NOT_GIVEN),
-            timeout=model_settings.get('timeout', NOT_GIVEN),
-            seed=model_settings.get('seed', NOT_GIVEN),
-            presence_penalty=model_settings.get('presence_penalty', NOT_GIVEN),
-            frequency_penalty=model_settings.get('frequency_penalty', NOT_GIVEN),
-            logit_bias=model_settings.get('logit_bias', NOT_GIVEN),
-        )
+        try:
+            return await self.client.chat.completions.create(
+                model=str(self._model_name),
+                messages=groq_messages,
+                n=1,
+                parallel_tool_calls=model_settings.get('parallel_tool_calls', NOT_GIVEN),
+                tools=tools or NOT_GIVEN,
+                tool_choice=tool_choice or NOT_GIVEN,
+                stream=stream,
+                max_tokens=model_settings.get('max_tokens', NOT_GIVEN),
+                temperature=model_settings.get('temperature', NOT_GIVEN),
+                top_p=model_settings.get('top_p', NOT_GIVEN),
+                timeout=model_settings.get('timeout', NOT_GIVEN),
+                seed=model_settings.get('seed', NOT_GIVEN),
+                presence_penalty=model_settings.get('presence_penalty', NOT_GIVEN),
+                frequency_penalty=model_settings.get('frequency_penalty', NOT_GIVEN),
+                logit_bias=model_settings.get('logit_bias', NOT_GIVEN),
+            )
+        except APIStatusError as e:
+            if (status_code := e.status_code) >= 400:
+                raise ModelHTTPError(status_code=status_code, model_name=self.model_name, body=e.body) from e
+            raise
 
     def _process_response(self, response: chat.ChatCompletion) -> ModelResponse:
         """Process a non-streamed response, and prepare a message to return."""
@@ -224,7 +232,7 @@ class GroqModel(Model):
                 items.append(ToolCallPart(tool_name=c.function.name, args=c.function.arguments, tool_call_id=c.id))
         return ModelResponse(items, model_name=response.model, timestamp=timestamp)
 
-    async def _process_streamed_response(self, response: AsyncStream[ChatCompletionChunk]) -> GroqStreamedResponse:
+    async def _process_streamed_response(self, response: AsyncStream[chat.ChatCompletionChunk]) -> GroqStreamedResponse:
         """Process a streamed response, and prepare a streaming response to return."""
         peekable_response = _utils.PeekableAsyncStream(response)
         first_chunk = await peekable_response.peek()
@@ -293,7 +301,7 @@ class GroqModel(Model):
             if isinstance(part, SystemPromptPart):
                 yield chat.ChatCompletionSystemMessageParam(role='system', content=part.content)
             elif isinstance(part, UserPromptPart):
-                yield chat.ChatCompletionUserMessageParam(role='user', content=part.content)
+                yield cls._map_user_prompt(part)
             elif isinstance(part, ToolReturnPart):
                 yield chat.ChatCompletionToolMessageParam(
                     role='tool',
@@ -310,13 +318,37 @@ class GroqModel(Model):
                         content=part.model_response(),
                     )
 
+    @staticmethod
+    def _map_user_prompt(part: UserPromptPart) -> chat.ChatCompletionUserMessageParam:
+        content: str | list[chat.ChatCompletionContentPartParam]
+        if isinstance(part.content, str):
+            content = part.content
+        else:
+            content = []
+            for item in part.content:
+                if isinstance(item, str):
+                    content.append(chat.ChatCompletionContentPartTextParam(text=item, type='text'))
+                elif isinstance(item, ImageUrl):
+                    image_url = ImageURL(url=item.url)
+                    content.append(chat.ChatCompletionContentPartImageParam(image_url=image_url, type='image_url'))
+                elif isinstance(item, BinaryContent):
+                    base64_encoded = base64.b64encode(item.data).decode('utf-8')
+                    if item.is_image:
+                        image_url = ImageURL(url=f'data:{item.media_type};base64,{base64_encoded}')
+                        content.append(chat.ChatCompletionContentPartImageParam(image_url=image_url, type='image_url'))
+                    else:
+                        raise RuntimeError('Only images are supported for binary content in Groq.')
+                else:  # pragma: no cover
+                    raise RuntimeError(f'Unsupported content type: {type(item)}')
+        return chat.ChatCompletionUserMessageParam(role='user', content=content)
+
 
 @dataclass
 class GroqStreamedResponse(StreamedResponse):
     """Implementation of `StreamedResponse` for Groq models."""
 
     _model_name: GroqModelName
-    _response: AsyncIterable[ChatCompletionChunk]
+    _response: AsyncIterable[chat.ChatCompletionChunk]
     _timestamp: datetime
 
     async def _get_event_iterator(self) -> AsyncIterator[ModelResponseStreamEvent]:
@@ -355,9 +387,9 @@ class GroqStreamedResponse(StreamedResponse):
         return self._timestamp
 
 
-def _map_usage(completion: ChatCompletionChunk | ChatCompletion) -> usage.Usage:
+def _map_usage(completion: chat.ChatCompletionChunk | chat.ChatCompletion) -> usage.Usage:
     response_usage = None
-    if isinstance(completion, ChatCompletion):
+    if isinstance(completion, chat.ChatCompletion):
         response_usage = completion.usage
     elif completion.x_groq is not None:
         response_usage = completion.x_groq.usage

pydantic_ai/models/mistral.py

@@ -1,5 +1,6 @@
 from __future__ import annotations as _annotations
 
+import base64
 import os
 from collections.abc import AsyncIterable, AsyncIterator, Iterable
 from contextlib import asynccontextmanager
@@ -12,9 +13,11 @@ import pydantic_core
 from httpx import AsyncClient as AsyncHTTPClient, Timeout
 from typing_extensions import assert_never
 
-from .. import UnexpectedModelBehavior, _utils
+from .. import ModelHTTPError, UnexpectedModelBehavior, _utils
 from .._utils import now_utc as _now_utc
 from ..messages import (
+    BinaryContent,
+    ImageUrl,
     ModelMessage,
     ModelRequest,
     ModelResponse,
@@ -45,6 +48,8 @@ try:
         Content as MistralContent,
         ContentChunk as MistralContentChunk,
         FunctionCall as MistralFunctionCall,
+        ImageURL as MistralImageURL,
+        ImageURLChunk as MistralImageURLChunk,
         Mistral,
         OptionalNullable as MistralOptionalNullable,
         TextChunk as MistralTextChunk,
@@ -54,6 +59,7 @@ try:
         ChatCompletionResponse as MistralChatCompletionResponse,
         CompletionEvent as MistralCompletionEvent,
         Messages as MistralMessages,
+        SDKError,
         Tool as MistralTool,
         ToolCall as MistralToolCall,
     )
@@ -179,19 +185,25 @@ class MistralModel(Model):
         model_request_parameters: ModelRequestParameters,
     ) -> MistralChatCompletionResponse:
         """Make a non-streaming request to the model."""
-        response = await self.client.chat.complete_async(
-            model=str(self._model_name),
-            messages=list(chain(*(self._map_message(m) for m in messages))),
-            n=1,
-            tools=self._map_function_and_result_tools_definition(model_request_parameters) or UNSET,
-            tool_choice=self._get_tool_choice(model_request_parameters),
-            stream=False,
-            max_tokens=model_settings.get('max_tokens', UNSET),
-            temperature=model_settings.get('temperature', UNSET),
-            top_p=model_settings.get('top_p', 1),
-            timeout_ms=self._get_timeout_ms(model_settings.get('timeout')),
-            random_seed=model_settings.get('seed', UNSET),
-        )
+        try:
+            response = await self.client.chat.complete_async(
+                model=str(self._model_name),
+                messages=list(chain(*(self._map_message(m) for m in messages))),
+                n=1,
+                tools=self._map_function_and_result_tools_definition(model_request_parameters) or UNSET,
+                tool_choice=self._get_tool_choice(model_request_parameters),
+                stream=False,
+                max_tokens=model_settings.get('max_tokens', UNSET),
+                temperature=model_settings.get('temperature', UNSET),
+                top_p=model_settings.get('top_p', 1),
+                timeout_ms=self._get_timeout_ms(model_settings.get('timeout')),
+                random_seed=model_settings.get('seed', UNSET),
+            )
+        except SDKError as e:
+            if (status_code := e.status_code) >= 400:
+                raise ModelHTTPError(status_code=status_code, model_name=self.model_name, body=e.body) from e
+            raise
+
         assert response, 'A unexpected empty response from Mistral.'
         return response
 
@@ -423,7 +435,7 @@ class MistralModel(Model):
             if isinstance(part, SystemPromptPart):
                 yield MistralSystemMessage(content=part.content)
             elif isinstance(part, UserPromptPart):
-                yield MistralUserMessage(content=part.content)
+                yield cls._map_user_prompt(part)
             elif isinstance(part, ToolReturnPart):
                 yield MistralToolMessage(
                     tool_call_id=part.tool_call_id,
@@ -460,6 +472,29 @@ class MistralModel(Model):
         else:
             assert_never(message)
 
+    @staticmethod
+    def _map_user_prompt(part: UserPromptPart) -> MistralUserMessage:
+        content: str | list[MistralContentChunk]
+        if isinstance(part.content, str):
+            content = part.content
+        else:
+            content = []
+            for item in part.content:
+                if isinstance(item, str):
+                    content.append(MistralTextChunk(text=item))
+                elif isinstance(item, ImageUrl):
+                    content.append(MistralImageURLChunk(image_url=MistralImageURL(url=item.url)))
+                elif isinstance(item, BinaryContent):
+                    base64_encoded = base64.b64encode(item.data).decode('utf-8')
+                    if item.is_image:
+                        image_url = MistralImageURL(url=f'data:{item.media_type};base64,{base64_encoded}')
+                        content.append(MistralImageURLChunk(image_url=image_url, type='image_url'))
+                    else:
+                        raise RuntimeError('Only image binary content is supported for Mistral.')
+                else:  # pragma: no cover
+                    raise RuntimeError(f'Unsupported content type: {type(item)}')
+        return MistralUserMessage(content=content)
+
 
 MistralToolCallId = Union[str, None]