From d6dcfade1b6f7abab86a17d5232158ab2a0e05d2 Mon Sep 17 00:00:00 2001
From: eavanvalkenburg <github@vanvalkenburg.eu>
Date: Mon, 31 Aug 2026 11:32:08 +0200
Subject: [PATCH 1/4] Python: migrate Mistral integration to SDK

Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com>

Copilot-Session: 3cc4a2d5-640a-4c05-8fc9-abe08f8134ee
---
 python/packages/mistral/AGENTS.md             |   4 +-
 .../agent_framework_mistral/_chat_client.py   | 333 +++++++++---------
 .../_embedding_client.py                      | 210 +++++------
 python/packages/mistral/pyproject.toml        |   4 +-
 .../tests/mistral/test_mistral_chat_client.py | 143 +++++---
 .../mistral/test_mistral_embedding_client.py  |  58 +--
 python/uv.lock                                |  41 ++-
 7 files changed, 436 insertions(+), 357 deletions(-)

diff --git a/python/packages/mistral/AGENTS.md b/python/packages/mistral/AGENTS.md
index 92e837ba7aa..1304b6267e0 100644
--- a/python/packages/mistral/AGENTS.md
+++ b/python/packages/mistral/AGENTS.md
@@ -4,8 +4,8 @@ Integration with Mistral AI for chat completions and embedding generation.
 
 ## Implementation Notes
 
-- Talks to the Mistral REST API directly over `httpx`; the official `mistralai` SDK is not used
-  because its pinned OpenTelemetry requirements conflict with the rest of the framework.
+- Uses the official `mistralai` SDK for chat completion and embedding requests.
+- Framework message, option, response, and exception translation stays in this package.
 
 ## Main Classes
 
diff --git a/python/packages/mistral/agent_framework_mistral/_chat_client.py b/python/packages/mistral/agent_framework_mistral/_chat_client.py
index 67c13cdbc83..2e52b39cbf5 100644
--- a/python/packages/mistral/agent_framework_mistral/_chat_client.py
+++ b/python/packages/mistral/agent_framework_mistral/_chat_client.py
@@ -7,9 +7,10 @@
 import logging
 import re
 import sys
+import warnings
 from collections.abc import AsyncIterable, Awaitable, Mapping, Sequence
 from datetime import datetime, timezone
-from typing import Any, ClassVar, Generic, Literal, cast
+from typing import Any, ClassVar, Generic, Literal, NoReturn, cast
 
 import httpx
 from agent_framework import (
@@ -40,6 +41,18 @@
     ChatClientInvalidResponseException,
 )
 from agent_framework.observability import ChatTelemetryLayer
+from mistralai.client import Mistral
+from mistralai.client.errors import MistralError
+from mistralai.client.models import (
+    AssistantMessage,
+    ChatCompletionResponse,
+    CompletionChunk,
+    DeltaMessage,
+    TextChunk,
+    ThinkChunk,
+    ToolCall,
+    UsageInfo,
+)
 from pydantic import BaseModel
 
 from ._feature_usage import FeatureIndex
@@ -114,6 +127,9 @@ class MistralChatOptions(ChatOptions[ResponseModelT], Generic[ResponseModelT], t
     prompt_cache_key: str
     """Cache key shared by requests with the same prompt prefix."""
 
+    service_tier: Literal["auto", "standard_only"]
+    """Capacity tier used to serve the request."""
+
     reasoning_effort: Literal["none", "minimal", "low", "medium", "high", "xhigh"]
     """Effort level for models that support reasoning."""
 
@@ -151,10 +167,6 @@ class MistralSettings(TypedDict, total=False):
 # endregion
 
 _MISTRAL_API_BASE_URL = "https://api.mistral.ai"
-_CHAT_COMPLETIONS_PATH = "/v1/chat/completions"
-_DEFAULT_TIMEOUT_SECONDS = 60.0
-_SSE_DATA_PREFIX = "data:"
-_SSE_DONE = "[DONE]"
 
 # Keys mapping to a different Mistral chat-completion parameter name
 _OPTION_TRANSLATIONS: dict[str, str] = {
@@ -213,26 +225,25 @@ def _sanitize_tool_call_id(call_id: str) -> str:
     return hashlib.sha256(call_id.encode("utf-8")).hexdigest()[:9]
 
 
-def _tool_call_id_of(tool_call: Mapping[str, Any]) -> str:
+def _tool_call_id_of(tool_call: ToolCall) -> str:
     """Return the wire tool call ID, treating null/"null" placeholders as missing."""
-    call_id = tool_call.get("id")
+    call_id = tool_call.id
     if isinstance(call_id, str) and call_id and call_id != "null":
         return call_id
     return ""
 
 
-def _function_call_content(tool_call: Mapping[str, Any]) -> Content:
-    function: Mapping[str, Any] = tool_call.get("function") or {}
-    arguments = function.get("arguments")
+def _function_call_content(tool_call: ToolCall) -> Content:
+    arguments = tool_call.function.arguments
     if isinstance(arguments, str):
         normalized_arguments: str | dict[str, Any] = arguments
     elif isinstance(arguments, dict):
-        normalized_arguments = cast("dict[str, Any]", arguments)
+        normalized_arguments = arguments
     else:
         normalized_arguments = str(cast(object, arguments))
     return Content.from_function_call(
         call_id=_tool_call_id_of(tool_call),
-        name=function.get("name") or "",
+        name=tool_call.function.name,
         arguments=normalized_arguments,
         raw_representation=tool_call,
     )
@@ -247,10 +258,10 @@ class _StreamedToolCalls:
     """
 
     def __init__(self) -> None:
-        self._pending: dict[tuple[int, int | str], dict[str, Any]] = {}
+        self._pending: dict[tuple[int, int | str], ToolCall] = {}
         self._auto_key_count = 0
 
-    def add(self, choice_index: int, fragment: Mapping[str, Any]) -> list[Content]:
+    def add(self, choice_index: int, fragment: ToolCall) -> list[Content]:
         """Fold a fragment into its pending call; returns calls completed by an index reuse."""
         flushed: list[Content] = []
         key = self._key_for(choice_index, fragment)
@@ -261,7 +272,7 @@ def add(self, choice_index: int, fragment: Mapping[str, Any]) -> list[Content]:
                 flushed.append(_function_call_content(self._pending.pop(key)))
                 pending = None
         if pending is None:
-            self._pending[key] = {**fragment, "function": dict(fragment.get("function") or {})}
+            self._pending[key] = fragment.model_copy(deep=True)
         else:
             self._merge(pending, fragment)
         return flushed
@@ -275,8 +286,8 @@ def flush_all(self) -> list[Content]:
         self._pending.clear()
         return contents
 
-    def _key_for(self, choice_index: int, fragment: Mapping[str, Any]) -> tuple[int, int | str]:
-        index = fragment.get("index")
+    def _key_for(self, choice_index: int, fragment: ToolCall) -> tuple[int, int | str]:
+        index = fragment.index
         if isinstance(index, int):
             return (choice_index, index)
         if fragment_id := _tool_call_id_of(fragment):
@@ -291,23 +302,19 @@ def _key_for(self, choice_index: int, fragment: Mapping[str, Any]) -> tuple[int,
         return (choice_index, f"auto-{self._auto_key_count}")
 
     @staticmethod
-    def _merge(pending: dict[str, Any], fragment: Mapping[str, Any]) -> None:
+    def _merge(pending: ToolCall, fragment: ToolCall) -> None:
         if fragment_id := _tool_call_id_of(fragment):
-            pending["id"] = fragment_id
-        function: Mapping[str, Any] = fragment.get("function") or {}
-        pending_function: dict[str, Any] = pending["function"]
-        if (name := function.get("name")) and not pending_function.get("name"):
-            pending_function["name"] = name
-        new_arguments = function.get("arguments")
-        old_arguments = pending_function.get("arguments")
-        if new_arguments is None:
-            return
+            pending.id = fragment_id
+        if fragment.function.name and not pending.function.name:
+            pending.function.name = fragment.function.name
+        new_arguments = fragment.function.arguments
+        old_arguments = pending.function.arguments
         if isinstance(old_arguments, str) and isinstance(new_arguments, str):
-            pending_function["arguments"] = old_arguments + new_arguments
+            pending.function.arguments = old_arguments + new_arguments
         elif isinstance(old_arguments, dict) and isinstance(new_arguments, dict):
-            cast("dict[str, Any]", old_arguments).update(cast("dict[str, Any]", new_arguments))
+            old_arguments.update(new_arguments)
         else:
-            pending_function["arguments"] = new_arguments
+            pending.function.arguments = new_arguments
 
 
 class RawMistralChatClient(
@@ -316,7 +323,7 @@ class RawMistralChatClient(
 ):
     """A raw Mistral AI chat client.
 
-    Talks to the Mistral REST API directly over HTTP; the ``mistralai`` SDK is not required.
+    Uses the official ``mistralai`` SDK without the framework's batteries-included layers.
 
     Use this when you want full control over the request pipeline. For instance, to opt out of
     telemetry, use custom middleware, or compose your own layers. If you want the full-featured
@@ -325,7 +332,7 @@ class RawMistralChatClient(
 
     OTEL_PROVIDER_NAME: ClassVar[str] = "mistralai"
 
-    INJECTABLE: ClassVar[set[str]] = {"client"}
+    INJECTABLE: ClassVar[set[str]] = {"client", "http_client"}
 
     def __init__(
         self,
@@ -333,7 +340,8 @@ def __init__(
         model: str | None = None,
         api_key: str | SecretString | None = None,
         server_url: str | None = None,
-        client: httpx.AsyncClient | None = None,
+        client: Mistral | None = None,
+        http_client: httpx.AsyncClient | None = None,
         additional_properties: dict[str, Any] | None = None,
         env_file_path: str | None = None,
         env_file_encoding: str | None = None,
@@ -346,17 +354,35 @@ def __init__(
             api_key: Mistral API key. Defaults to ``MISTRAL_API_KEY`` environment variable.
             server_url: Optional server URL override. Defaults to ``MISTRAL_SERVER_URL``
                 environment variable, or the Mistral default.
-            client: Optional pre-configured ``httpx.AsyncClient``. When provided, api_key is
-                not required and the client is expected to carry its own auth headers and
-                base URL.
+            client: Optional pre-configured ``mistralai.client.Mistral``.
+            http_client: Optional pre-configured ``httpx.AsyncClient``. When provided, api_key
+                is not required and the client is expected to carry its own auth headers.
+                Passing an HTTP client via ``client`` remains supported but is deprecated.
             additional_properties: Additional properties stored on the client instance.
             env_file_path: Path to ``.env`` file for settings.
             env_file_encoding: Encoding for ``.env`` file.
         """
+        if isinstance(client, httpx.AsyncClient):
+            warnings.warn(
+                "Passing an httpx.AsyncClient via 'client' is deprecated; pass it via 'http_client' instead.",
+                DeprecationWarning,
+                stacklevel=2,
+            )
+            if http_client is not None:
+                raise ValueError("Provide either 'client' or 'http_client', not both.")
+            http_client = client
+            client = None
+        if client is not None and not isinstance(client, Mistral):
+            raise TypeError(
+                f"The 'client' parameter accepts a mistralai.client.Mistral instance; got {type(client).__name__}."
+            )
+        if client is not None and http_client is not None:
+            raise ValueError("Provide either 'client' or 'http_client', not both.")
+
         mistral_settings = load_settings(
             MistralSettings,
             env_prefix="MISTRAL_",
-            required_fields=[] if client is not None else ["api_key"],
+            required_fields=[] if client is not None or http_client is not None else ["api_key"],
             api_key=api_key,
             chat_model=model,
             server_url=server_url,
@@ -366,31 +392,32 @@ def __init__(
 
         self.model = mistral_settings.get("chat_model")
         self.server_url = mistral_settings.get("server_url")
-        self._owns_client = client is None
+        self._owns_client = not isinstance(client, Mistral)
 
-        if client is not None:
+        if isinstance(client, Mistral):
             self.client = client
             if self.server_url is None:
-                client_base_url = str(client.base_url).rstrip("/")
-                self.server_url = client_base_url or None
+                self.server_url = client.sdk_configuration.get_server_details()[0]
         else:
-            resolved_api_key: SecretString = mistral_settings["api_key"]  # type: ignore[assignment]
-            self.client = httpx.AsyncClient(
-                base_url=self.server_url or _MISTRAL_API_BASE_URL,
-                headers={
-                    "Authorization": f"Bearer {resolved_api_key.get_secret_value()}",
-                    "User-Agent": get_user_agent(),
-                    "Accept": "application/json",
-                },
-                timeout=_DEFAULT_TIMEOUT_SECONDS,
-            )
+            client_kwargs: dict[str, Any] = {}
+            if resolved_api_key := mistral_settings.get("api_key"):
+                client_kwargs["api_key"] = resolved_api_key.get_secret_value()
+            if http_client is not None:
+                client_kwargs["async_client"] = http_client
+                if self.server_url is None:
+                    client_base_url = str(http_client.base_url).rstrip("/")
+                    self.server_url = client_base_url or None
+            if self.server_url:
+                client_kwargs["server_url"] = self.server_url
+            self.client = Mistral(**client_kwargs)
 
         super().__init__(additional_properties=additional_properties)
 
     async def close(self) -> None:
-        """Close the internally created HTTP client."""
+        """Close the internally created Mistral SDK client."""
         if self._owns_client:
-            await self.client.aclose()
+            await self.client.__aexit__(None, None, None)  # type: ignore[no-untyped-call]
+            self.client.__exit__(None, None, None)  # type: ignore[no-untyped-call]
 
     @override
     def service_url(self) -> str:
@@ -411,18 +438,17 @@ def _inner_get_response(
             async def _stream() -> AsyncIterable[ChatResponseUpdate]:
                 validated = await self._validate_options(options)
                 request = self._prepare_request(messages, validated, **kwargs)
-                request["stream"] = True
                 mark_feature_used(FeatureIndex.MISTRAL)
+                request.setdefault("http_headers", {"User-Agent": get_user_agent()})
                 tool_calls = _StreamedToolCalls()
                 try:
-                    async with self.client.stream("POST", _CHAT_COMPLETIONS_PATH, json=request) as response:
-                        await self._raise_for_status(response)
-                        async for line in response.aiter_lines():
-                            chunk = self._parse_sse_line(line)
-                            if chunk is not None:
-                                yield self._parse_chunk(chunk, tool_calls)
+                    response = await self.client.chat.stream_async(**request)
+                    async for event in response:
+                        yield self._parse_chunk(event.data, tool_calls)
                     if remaining := tool_calls.flush_all():
                         yield ChatResponseUpdate(contents=remaining, role="assistant")
+                except MistralError as ex:
+                    self._raise_sdk_error(ex, streaming=True)
                 except ChatClientException:
                     raise
                 except Exception as ex:
@@ -437,19 +463,17 @@ async def _get_response() -> ChatResponse:
             validated = await self._validate_options(options)
             request = self._prepare_request(messages, validated, **kwargs)
             mark_feature_used(FeatureIndex.MISTRAL)
+            request.setdefault("http_headers", {"User-Agent": get_user_agent()})
             try:
-                response = await self.client.post(_CHAT_COMPLETIONS_PATH, json=request)
-                await self._raise_for_status(response)
+                response = await self.client.chat.complete_async(**request)
+            except MistralError as ex:
+                self._raise_sdk_error(ex)
             except ChatClientException:
                 raise
             except Exception as ex:
                 raise ChatClientException(f"Mistral chat request failed: {ex}", inner_exception=ex) from ex
             try:
-                raw_payload = response.json()
-                if not isinstance(raw_payload, Mapping):
-                    raise ChatClientInvalidResponseException("Mistral chat response must be a JSON object.")
-                payload = cast("Mapping[str, Any]", raw_payload)
-                return self._parse_response(payload, response_format=validated.get("response_format"))
+                return self._parse_response(response, response_format=validated.get("response_format"))
             except ChatClientException:
                 raise
             except Exception as ex:
@@ -461,36 +485,22 @@ async def _get_response() -> ChatResponse:
         return _get_response()
 
     @staticmethod
-    async def _raise_for_status(response: httpx.Response) -> None:
-        if response.status_code < 400:
-            return
-        body = (await response.aread()).decode("utf-8", errors="replace")
-        message = f"Mistral chat request failed with status {response.status_code}: {body[:2000]}"
-        if response.status_code in (401, 403):
-            raise ChatClientInvalidAuthException(message)
-        if response.status_code < 500:
-            raise ChatClientInvalidRequestException(message)
-        raise ChatClientException(message)
-
-    @staticmethod
-    def _parse_sse_line(line: str) -> dict[str, Any] | None:
-        """Parse one server-sent-events line into a completion chunk, or None to skip."""
-        line = line.strip()
-        if not line.startswith(_SSE_DATA_PREFIX):
-            return None
-        data = line[len(_SSE_DATA_PREFIX) :].strip()
-        if not data or data == _SSE_DONE:
-            return None
-        try:
-            parsed = json.loads(data)
-        except json.JSONDecodeError as ex:
+    def _raise_sdk_error(ex: MistralError, *, streaming: bool = False) -> NoReturn:
+        status_code = ex.raw_response.status_code
+        if status_code < 400:
             raise ChatClientInvalidResponseException(
-                "Mistral streaming chat response contained malformed SSE data.",
+                f"Mistral chat response was invalid: {ex}",
                 inner_exception=ex,
             ) from ex
-        if not isinstance(parsed, dict):
-            raise ChatClientInvalidResponseException("Mistral streaming chat SSE data must be a JSON object.")
-        return cast("dict[str, Any]", parsed)
+
+        request_kind = "streaming chat" if streaming else "chat"
+        body = ex.body or str(ex)
+        message = f"Mistral {request_kind} request failed with status {status_code}: {body[:2000]}"
+        if status_code in (401, 403):
+            raise ChatClientInvalidAuthException(message)
+        if status_code < 500:
+            raise ChatClientInvalidRequestException(message)
+        raise ChatClientException(message, inner_exception=ex)
 
     # region Request preparation
 
@@ -760,28 +770,27 @@ def _prepare_response_format(self, response_format: Any) -> dict[str, Any] | Non
 
     def _parse_response(
         self,
-        response: Mapping[str, Any],
+        response: ChatCompletionResponse,
         *,
         response_format: Any | None = None,
     ) -> ChatResponse:
-        """Convert a Mistral chat-completion response payload to a framework ChatResponse."""
-        choices = cast("Sequence[Mapping[str, Any]]", response.get("choices") or ())
-        choice: Mapping[str, Any] = choices[0] if choices else {}
-        message: Mapping[str, Any] = choice.get("message") or {}
-        contents = self._parse_message_contents(message)
-        finish_reason = _map_finish_reason(choice.get("finish_reason"))
+        """Convert a Mistral SDK response to a framework ChatResponse."""
+        choice = response.choices[0] if response.choices else None
+        contents = self._parse_message_contents(choice.message) if choice and choice.message else []
+        finish_reason = _map_finish_reason(choice.finish_reason) if choice else None
         return ChatResponse(
-            response_id=response.get("id"),
-            messages=[Message(role="assistant", contents=contents, raw_representation=choice or None)],
-            usage_details=self._parse_usage(response.get("usage")),
-            model=response.get("model") or self.model,
-            created_at=self._format_created_at(response.get("created")),
+            response_id=response.id,
+            messages=[Message(role="assistant", contents=contents, raw_representation=choice)],
+            usage_details=self._parse_usage(response.usage),
+            model=response.model or self.model,
+            created_at=self._format_created_at(response.created),
             finish_reason=finish_reason,
             response_format=response_format,
+            additional_properties=self._parse_response_metadata(response.usage),
             raw_representation=response,
         )
 
-    def _parse_chunk(self, chunk: Mapping[str, Any], tool_calls: _StreamedToolCalls) -> ChatResponseUpdate:
+    def _parse_chunk(self, chunk: CompletionChunk, tool_calls: _StreamedToolCalls) -> ChatResponseUpdate:
         """Convert a Mistral streaming completion chunk to a framework ChatResponseUpdate.
 
         Tool-call fragments are folded into ``tool_calls`` keyed by (choice, index) and
@@ -789,52 +798,49 @@ def _parse_chunk(self, chunk: Mapping[str, Any], tool_calls: _StreamedToolCalls)
         """
         contents: list[Content] = []
         finish_reason: FinishReason | None = None
-        choices = cast("Sequence[Mapping[str, Any]]", chunk.get("choices") or ())
-        for choice in choices:
-            choice_index = index if isinstance(index := choice.get("index"), int) else 0
-            delta: Mapping[str, Any] = choice.get("delta") or {}
-            contents.extend(self._parse_content_chunks(delta))
-            for fragment in cast("Sequence[Mapping[str, Any]]", delta.get("tool_calls") or ()):
-                contents.extend(tool_calls.add(choice_index, fragment))
-            if reason := choice.get("finish_reason"):
-                contents.extend(tool_calls.flush_choice(choice_index))
-                if finish_reason is None:
-                    finish_reason = _map_finish_reason(reason)
-        if usage := self._parse_usage(chunk.get("usage")):
+        if chunk.choices:
+            choice = chunk.choices[0]
+            contents.extend(self._parse_content_chunks(choice.delta))
+            for fragment in choice.delta.tool_calls if isinstance(choice.delta.tool_calls, list) else ():
+                contents.extend(tool_calls.add(choice.index, fragment))
+            if reason := choice.finish_reason:
+                contents.extend(tool_calls.flush_choice(choice.index))
+                finish_reason = _map_finish_reason(reason)
+        if usage := self._parse_usage(chunk.usage):
             contents.append(Content.from_usage(usage_details=usage, raw_representation=chunk))
         return ChatResponseUpdate(
             contents=contents,
             role="assistant",
-            response_id=chunk.get("id"),
-            model=chunk.get("model"),
-            created_at=self._format_created_at(chunk.get("created")),
+            response_id=chunk.id,
+            model=chunk.model,
+            created_at=self._format_created_at(chunk.created),
             finish_reason=finish_reason,
+            additional_properties=self._parse_response_metadata(chunk.usage),
             raw_representation=chunk,
         )
 
-    def _parse_message_contents(self, message: Mapping[str, Any]) -> list[Content]:
+    def _parse_message_contents(self, message: AssistantMessage) -> list[Content]:
         contents = self._parse_content_chunks(message)
-        tool_calls = cast("Sequence[Mapping[str, Any]]", message.get("tool_calls") or ())
-        contents.extend(_function_call_content(tool_call) for tool_call in tool_calls)
+        if isinstance(message.tool_calls, list):
+            contents.extend(_function_call_content(tool_call) for tool_call in message.tool_calls)
         return contents
 
-    def _parse_content_chunks(self, message: Mapping[str, Any]) -> list[Content]:
+    def _parse_content_chunks(self, message: AssistantMessage | DeltaMessage) -> list[Content]:
         contents: list[Content] = []
-        content = message.get("content")
+        content = message.content
         if isinstance(content, str):
             if content:
                 contents.append(Content.from_text(text=content))
-        elif content:
-            for chunk in cast("Sequence[Mapping[str, Any]]", content):
-                chunk_type = chunk.get("type")
-                if chunk_type == "text":
-                    if text := chunk.get("text"):
-                        contents.append(Content.from_text(text=text, raw_representation=chunk))
-                elif chunk_type == "thinking":
+        elif isinstance(content, list):
+            for chunk in content:
+                if isinstance(chunk, TextChunk):
+                    if chunk.text:
+                        contents.append(Content.from_text(text=chunk.text, raw_representation=chunk))
+                elif isinstance(chunk, ThinkChunk):
                     if reasoning := self._thinking_to_text(chunk):
                         contents.append(Content.from_text_reasoning(text=reasoning, raw_representation=chunk))
                 else:
-                    logger.debug("Skipping unsupported response chunk from Mistral: %s", chunk_type)
+                    logger.debug("Skipping unsupported response chunk from Mistral: %s", type(chunk).__name__)
         return contents
 
     @staticmethod
@@ -844,36 +850,40 @@ def _format_created_at(created: Any) -> str | None:
         return datetime.fromtimestamp(created, tz=timezone.utc).strftime("%Y-%m-%dT%H:%M:%S.%fZ")
 
     @staticmethod
-    def _thinking_to_text(chunk: Mapping[str, Any]) -> str:
-        thinking = chunk.get("thinking")
-        if isinstance(thinking, str):
-            return thinking
-        if isinstance(thinking, Sequence):
-            return "".join(
-                part.get("text") or ""
-                for part in cast("Sequence[Mapping[str, Any]]", thinking)
-                if isinstance(part, Mapping)
-            )
-        return ""
+    def _thinking_to_text(chunk: ThinkChunk) -> str:
+        return "".join(part.text for part in chunk.thinking if isinstance(part, TextChunk))
 
-    def _parse_usage(self, usage: Mapping[str, Any] | None) -> UsageDetails | None:
+    def _parse_usage(self, usage: UsageInfo | None) -> UsageDetails | None:
         if not usage:
             return None
         details: UsageDetails = {}
-        if (value := usage.get("prompt_tokens")) is not None:
-            details["input_token_count"] = value
-        if (value := usage.get("completion_tokens")) is not None:
-            details["output_token_count"] = value
-        if (value := usage.get("total_tokens")) is not None:
-            details["total_token_count"] = value
-        prompt_tokens_details = usage.get("prompt_tokens_details")
+        fields_set = usage.model_fields_set
+        if "prompt_tokens" in fields_set and usage.prompt_tokens is not None:
+            details["input_token_count"] = usage.prompt_tokens
+        if "completion_tokens" in fields_set and usage.completion_tokens is not None:
+            details["output_token_count"] = usage.completion_tokens
+        if "total_tokens" in fields_set and usage.total_tokens is not None:
+            details["total_token_count"] = usage.total_tokens
+        if isinstance(usage.prompt_audio_seconds, int) and not isinstance(usage.prompt_audio_seconds, bool):
+            cast("dict[str, Any]", details)["prompt_audio_seconds"] = usage.prompt_audio_seconds
+        prompt_tokens_details = getattr(usage, "prompt_tokens_details", None)
         if isinstance(prompt_tokens_details, Mapping):
-            cached_tokens = cast("Mapping[str, Any]", prompt_tokens_details).get("cached_tokens")
+            prompt_details = cast("Mapping[str, Any]", prompt_tokens_details)
+            cached_tokens = prompt_details.get("cached_tokens")
             if isinstance(cached_tokens, int) and not isinstance(cached_tokens, bool):
-                details["prompt/cached_tokens"] = cached_tokens
+                cast("dict[str, Any]", details)["prompt/cached_tokens"] = cached_tokens
                 details["cache_read_input_token_count"] = cached_tokens
+            audio_tokens = prompt_details.get("audio_tokens")
+            if isinstance(audio_tokens, int) and not isinstance(audio_tokens, bool):
+                cast("dict[str, Any]", details)["prompt/audio_tokens"] = audio_tokens
         return details or None
 
+    @staticmethod
+    def _parse_response_metadata(usage: UsageInfo | None) -> dict[str, Any] | None:
+        if usage and isinstance(usage.service_tier, str):
+            return {"service_tier": usage.service_tier}
+        return None
+
     # endregion
 
 
@@ -924,7 +934,8 @@ def __init__(
         model: str | None = None,
         api_key: str | SecretString | None = None,
         server_url: str | None = None,
-        client: httpx.AsyncClient | None = None,
+        client: Mistral | None = None,
+        http_client: httpx.AsyncClient | None = None,
         additional_properties: dict[str, Any] | None = None,
         middleware: Sequence[ChatAndFunctionMiddlewareTypes] | None = None,
         function_invocation_configuration: FunctionInvocationConfiguration | None = None,
@@ -939,9 +950,10 @@ def __init__(
             api_key: Mistral API key. Defaults to ``MISTRAL_API_KEY`` environment variable.
             server_url: Optional server URL override. Defaults to ``MISTRAL_SERVER_URL``
                 environment variable, or the Mistral default.
-            client: Optional pre-configured ``httpx.AsyncClient``. When provided, api_key is
-                not required and the client is expected to carry its own auth headers and
-                base URL.
+            client: Optional pre-configured ``mistralai.client.Mistral``.
+            http_client: Optional pre-configured ``httpx.AsyncClient``. When provided, api_key
+                is not required and the client is expected to carry its own auth headers.
+                Passing an HTTP client via ``client`` remains supported but is deprecated.
             additional_properties: Additional properties stored on the client instance.
             middleware: Optional middleware chain applied to every call.
             function_invocation_configuration: Optional configuration for the function invocation loop.
@@ -953,6 +965,7 @@ def __init__(
             api_key=api_key,
             server_url=server_url,
             client=client,
+            http_client=http_client,
             additional_properties=additional_properties,
             middleware=middleware,
             function_invocation_configuration=function_invocation_configuration,
diff --git a/python/packages/mistral/agent_framework_mistral/_embedding_client.py b/python/packages/mistral/agent_framework_mistral/_embedding_client.py
index b347f870f49..d335a6af132 100644
--- a/python/packages/mistral/agent_framework_mistral/_embedding_client.py
+++ b/python/packages/mistral/agent_framework_mistral/_embedding_client.py
@@ -5,8 +5,8 @@
 import logging
 import sys
 import warnings
-from collections.abc import Mapping, Sequence
-from typing import Any, ClassVar, Generic, TypedDict, cast
+from collections.abc import Sequence
+from typing import Any, ClassVar, Generic, NoReturn, TypedDict
 
 import httpx
 from agent_framework import (
@@ -26,6 +26,8 @@
     IntegrationInvalidResponseException,
 )
 from agent_framework.observability import EmbeddingTelemetryLayer
+from mistralai.client import Mistral
+from mistralai.client.errors import MistralError
 
 from ._feature_usage import FeatureIndex
 
@@ -38,37 +40,6 @@
 logger = logging.getLogger("agent_framework.mistral")
 
 _MISTRAL_API_BASE_URL = "https://api.mistral.ai"
-_EMBEDDINGS_PATH = "/v1/embeddings"
-_DEFAULT_TIMEOUT_SECONDS = 60.0
-
-
-def _resolve_injected_clients(
-    http_client: httpx.AsyncClient | None,
-    client: Any | None,
-) -> tuple[httpx.AsyncClient | None, Any | None]:
-    """Split the deprecated ``client`` parameter into REST and legacy-SDK forms.
-
-    Returns ``(http_client, sdk_client)``; at most one is set. The SDK form is
-    duck-typed on ``.embeddings`` so the ``mistralai`` dependency stays optional.
-    """
-    if client is None:
-        return http_client, None
-    warnings.warn(
-        "The 'client' parameter is deprecated; pass an httpx.AsyncClient as 'http_client' instead. "
-        "Support for injected mistralai.Mistral clients will be removed in the next major release.",
-        DeprecationWarning,
-        stacklevel=3,
-    )
-    if http_client is not None:
-        raise ValueError("Provide either 'http_client' or the deprecated 'client' parameter, not both.")
-    if isinstance(client, httpx.AsyncClient):
-        return client, None
-    if hasattr(client, "embeddings"):
-        return None, client
-    raise TypeError(
-        "The 'client' parameter accepts an httpx.AsyncClient or a mistralai.Mistral instance; "
-        f"got {type(client).__name__}."
-    )
 
 
 class MistralEmbeddingOptions(EmbeddingGenerationOptions, total=False):
@@ -105,7 +76,7 @@ class MistralEmbeddingSettings(TypedDict, total=False):
         server_url: Optional server URL override. Resolved from ``MISTRAL_SERVER_URL``.
     """
 
-    api_key: str | None
+    api_key: SecretString | None
     embedding_model: str | None
     server_url: str | None
 
@@ -116,7 +87,7 @@ class RawMistralEmbeddingClient(
 ):
     """Raw Mistral AI embedding client without telemetry.
 
-    Talks to the Mistral REST API directly over HTTP; the ``mistralai`` SDK is not required.
+    Uses the official ``mistralai`` SDK without the framework's telemetry layer.
 
     Keyword Args:
         model: The Mistral embedding model (e.g. "mistral-embed").
@@ -126,9 +97,8 @@ class RawMistralEmbeddingClient(
             environment variable, or the Mistral default.
         http_client: Optional pre-configured ``httpx.AsyncClient``. When provided, api_key is
             not required and the client is expected to carry its own auth headers and base URL.
-        client: Deprecated. Accepts an ``httpx.AsyncClient`` (treated as ``http_client``) or a
-            ``mistralai.Mistral`` instance, which keeps working through the legacy SDK path
-            until the next major release.
+        client: Optional pre-configured ``mistralai.client.Mistral``. Passing an HTTP client via
+            this parameter remains supported but is deprecated; use ``http_client`` instead.
         additional_properties: Additional properties stored on the client instance.
         env_file_path: Path to ``.env`` file for settings.
         env_file_encoding: Encoding for ``.env`` file.
@@ -143,20 +113,36 @@ def __init__(
         api_key: str | SecretString | None = None,
         server_url: str | None = None,
         http_client: httpx.AsyncClient | None = None,
-        client: Any | None = None,
+        client: Mistral | None = None,
         additional_properties: dict[str, Any] | None = None,
         env_file_path: str | None = None,
         env_file_encoding: str | None = None,
     ) -> None:
         """Initialize a raw Mistral AI embedding client."""
-        http_client, sdk_client = _resolve_injected_clients(http_client, client)
-        injected = http_client is not None or sdk_client is not None
+        if isinstance(client, httpx.AsyncClient):
+            warnings.warn(
+                "Passing an httpx.AsyncClient via 'client' is deprecated; pass it via 'http_client' instead.",
+                DeprecationWarning,
+                stacklevel=2,
+            )
+            if http_client is not None:
+                raise ValueError("Provide either 'client' or 'http_client', not both.")
+            http_client = client
+            client = None
+        if client is not None and not isinstance(client, Mistral):
+            raise TypeError(
+                f"The 'client' parameter accepts a mistralai.client.Mistral instance; got {type(client).__name__}."
+            )
+        if client is not None and http_client is not None:
+            raise ValueError("Provide either 'client' or 'http_client', not both.")
+
+        injected = client is not None or http_client is not None
         required_fields = ["embedding_model"] if injected else ["embedding_model", "api_key"]
         mistral_settings = load_settings(
             MistralEmbeddingSettings,
             env_prefix="MISTRAL_",
             required_fields=required_fields,
-            api_key=str(api_key) if isinstance(api_key, SecretString) else api_key,
+            api_key=api_key,
             embedding_model=model,
             server_url=server_url,
             env_file_path=env_file_path,
@@ -165,40 +151,54 @@ def __init__(
 
         self.model: str = mistral_settings["embedding_model"]  # type: ignore[assignment]
         self.server_url = mistral_settings.get("server_url")
-        self._owns_client = not injected
-        self._sdk_client = sdk_client
-        self.client: Any
-
-        if sdk_client is not None:
-            self.client = sdk_client
-        elif http_client is not None:
-            self.client = http_client
+        self._owns_client = not isinstance(client, Mistral)
+
+        if isinstance(client, Mistral):
+            self.client = client
             if self.server_url is None:
-                client_base_url = str(http_client.base_url).rstrip("/")
-                self.server_url = client_base_url or None
+                self.server_url = client.sdk_configuration.get_server_details()[0]
         else:
-            resolved_api_key: str = mistral_settings["api_key"]  # type: ignore[assignment]
-            self.client = httpx.AsyncClient(
-                base_url=self.server_url or _MISTRAL_API_BASE_URL,
-                headers={
-                    "Authorization": f"Bearer {resolved_api_key}",
-                    "User-Agent": get_user_agent(),
-                    "Accept": "application/json",
-                },
-                timeout=_DEFAULT_TIMEOUT_SECONDS,
-            )
+            client_kwargs: dict[str, Any] = {}
+            if resolved_api_key := mistral_settings.get("api_key"):
+                client_kwargs["api_key"] = resolved_api_key.get_secret_value()
+            if http_client is not None:
+                client_kwargs["async_client"] = http_client
+                if self.server_url is None:
+                    client_base_url = str(http_client.base_url).rstrip("/")
+                    self.server_url = client_base_url or None
+            if self.server_url:
+                client_kwargs["server_url"] = self.server_url
+            self.client = Mistral(**client_kwargs)
 
         super().__init__(additional_properties=additional_properties)
 
     async def close(self) -> None:
-        """Close the internally created HTTP client."""
+        """Close the internally created Mistral SDK client."""
         if self._owns_client:
-            await self.client.aclose()
+            await self.client.__aexit__(None, None, None)  # type: ignore[no-untyped-call]
+            self.client.__exit__(None, None, None)  # type: ignore[no-untyped-call]
 
     def service_url(self) -> str:
         """Get the URL of the service."""
         return self.server_url or _MISTRAL_API_BASE_URL
 
+    @staticmethod
+    def _raise_sdk_error(ex: MistralError) -> NoReturn:
+        status_code = ex.raw_response.status_code
+        if status_code < 400:
+            raise IntegrationInvalidResponseException(
+                f"Mistral embeddings response was invalid: {ex}",
+                inner_exception=ex,
+            ) from ex
+
+        body = ex.body or str(ex)
+        message = f"Mistral embeddings request failed with status {status_code}: {body[:2000]}"
+        if status_code in (401, 403):
+            raise IntegrationInvalidAuthException(message)
+        if status_code < 500:
+            raise IntegrationInvalidRequestException(message)
+        raise IntegrationException(message, inner_exception=ex)
+
     async def get_embeddings(
         self,
         values: Sequence[str],
@@ -230,53 +230,43 @@ async def get_embeddings(
             raise ValueError("model is required")
 
         mark_feature_used(FeatureIndex.MISTRAL)
-        if self._sdk_client is not None:
-            return await self._get_embeddings_sdk(self._sdk_client, model, values, opts, options)
-
-        request: dict[str, Any] = {"model": model, "input": list(values)}
+        request: dict[str, Any] = {
+            "model": model,
+            "inputs": list(values),
+            "http_headers": {"User-Agent": get_user_agent()},
+        }
         if "dimensions" in opts:
             request["output_dimension"] = opts["dimensions"]
 
         try:
-            response = await self.client.post(_EMBEDDINGS_PATH, json=request)
-            if response.status_code >= 400:
-                message = (
-                    f"Mistral embeddings request failed with status {response.status_code}: {response.text[:2000]}"
-                )
-                if response.status_code in (401, 403):
-                    raise IntegrationInvalidAuthException(message)
-                if response.status_code < 500:
-                    raise IntegrationInvalidRequestException(message)
-                raise IntegrationException(message)
+            response = await self.client.embeddings.create_async(**request)
+        except MistralError as ex:
+            self._raise_sdk_error(ex)
         except IntegrationException:
             raise
         except Exception as ex:
             raise IntegrationException(f"Mistral embeddings request failed: {ex}", inner_exception=ex) from ex
 
         try:
-            raw_payload = response.json()
-            if not isinstance(raw_payload, Mapping):
-                raise IntegrationInvalidResponseException("Mistral embeddings response must be a JSON object.")
-            payload = cast("Mapping[str, Any]", raw_payload)
             embeddings: list[Embedding[list[float]]] = []
-            data = cast("Sequence[Mapping[str, Any]]", payload.get("data") or ())
-            items = sorted(data, key=lambda item: item.get("index") or 0)
+            items = sorted(response.data or (), key=lambda item: item.index or 0)
             for item in items:
-                vector = [float(v) for v in cast("Sequence[float]", item.get("embedding") or ())]
+                vector = [float(value) for value in item.embedding or ()]
                 embeddings.append(
                     Embedding(
                         vector=vector,
                         dimensions=len(vector),
-                        model=payload.get("model") or model,
+                        model=response.model or model,
                     )
                 )
 
             usage_dict: UsageDetails | None = None
-            if usage := payload.get("usage"):
+            if usage := response.usage:
                 usage_dict = {}
-                if (value := usage.get("prompt_tokens")) is not None:
+                fields_set = getattr(usage, "model_fields_set", None)
+                if (fields_set is None or "prompt_tokens" in fields_set) and (value := usage.prompt_tokens) is not None:
                     usage_dict["input_token_count"] = value
-                if (value := usage.get("total_tokens")) is not None:
+                if (fields_set is None or "total_tokens" in fields_set) and (value := usage.total_tokens) is not None:
                     usage_dict["total_token_count"] = value
 
             return GeneratedEmbeddings(embeddings, options=options, usage=usage_dict or None)
@@ -288,43 +278,6 @@ async def get_embeddings(
                 inner_exception=ex,
             ) from ex
 
-    async def _get_embeddings_sdk(
-        self,
-        sdk_client: Any,
-        model: str,
-        values: Sequence[str],
-        opts: Mapping[str, Any],
-        options: MistralEmbeddingOptionsT | None,
-    ) -> GeneratedEmbeddings[list[float], MistralEmbeddingOptionsT]:
-        """Legacy path for injected mistralai.Mistral clients; removed in the next major release."""
-        kwargs: dict[str, Any] = {"model": model, "inputs": list(values)}
-        if "dimensions" in opts:
-            kwargs["output_dimension"] = opts["dimensions"]
-
-        response = await sdk_client.embeddings.create_async(**kwargs)
-
-        embeddings: list[Embedding[list[float]]] = []
-        if response and response.data:
-            items = sorted(response.data, key=lambda d: d.index if d.index is not None else 0)
-            for item in items:
-                vector = list(item.embedding) if item.embedding else []
-                embeddings.append(
-                    Embedding(
-                        vector=vector,
-                        dimensions=len(vector),
-                        model=response.model or model,
-                    )
-                )
-
-        usage_dict: UsageDetails | None = None
-        if response and response.usage:
-            usage_dict = {
-                "input_token_count": response.usage.prompt_tokens,
-                "total_token_count": response.usage.total_tokens,
-            }
-
-        return GeneratedEmbeddings(embeddings, options=options, usage=usage_dict)
-
 
 class MistralEmbeddingClient(
     EmbeddingTelemetryLayer[str, list[float], MistralEmbeddingOptionsT],
@@ -340,7 +293,8 @@ class MistralEmbeddingClient(
         server_url: Optional server URL override. Defaults to ``MISTRAL_SERVER_URL``
             environment variable, or the Mistral default.
         http_client: Optional pre-configured ``httpx.AsyncClient``.
-        client: Deprecated. Accepts an ``httpx.AsyncClient`` or a ``mistralai.Mistral`` instance.
+        client: Optional pre-configured ``mistralai.client.Mistral``. Passing an HTTP client via
+            this parameter remains supported but is deprecated; use ``http_client`` instead.
         otel_provider_name: Optional telemetry provider name override.
         env_file_path: Path to ``.env`` file for settings.
         env_file_encoding: Encoding for ``.env`` file.
@@ -376,7 +330,7 @@ def __init__(
         api_key: str | SecretString | None = None,
         server_url: str | None = None,
         http_client: httpx.AsyncClient | None = None,
-        client: Any | None = None,
+        client: Mistral | None = None,
         otel_provider_name: str | None = None,
         additional_properties: dict[str, Any] | None = None,
         env_file_path: str | None = None,
diff --git a/python/packages/mistral/pyproject.toml b/python/packages/mistral/pyproject.toml
index 86ba721eea8..f200e69aadd 100644
--- a/python/packages/mistral/pyproject.toml
+++ b/python/packages/mistral/pyproject.toml
@@ -24,9 +24,7 @@ classifiers = [
 ]
 dependencies = [
     "agent-framework-core>=1.13.0,<2",
-    # Talks to the Mistral REST API directly; the mistralai SDK is not used because its
-    # pinned OpenTelemetry requirements conflict with the rest of the framework.
-    "httpx>=0.23.1,<1",
+    "mistralai>=2.9.2,<3",
 ]
 
 [tool.uv]
diff --git a/python/packages/mistral/tests/mistral/test_mistral_chat_client.py b/python/packages/mistral/tests/mistral/test_mistral_chat_client.py
index 2cda9397acd..67188a1d4ca 100644
--- a/python/packages/mistral/tests/mistral/test_mistral_chat_client.py
+++ b/python/packages/mistral/tests/mistral/test_mistral_chat_client.py
@@ -15,6 +15,8 @@
     ChatClientInvalidRequestException,
     ChatClientInvalidResponseException,
 )
+from mistralai.client import Mistral
+from mistralai.client.models import AssistantMessage, ThinkChunk
 from pydantic import BaseModel
 
 import agent_framework_mistral._chat_client as chat_client_module
@@ -105,7 +107,7 @@ def make_client(*responses: httpx.Response) -> tuple[MistralChatClient, MockMist
         base_url="https://api.mistral.ai",
         transport=httpx.MockTransport(server.handler),
     )
-    client = MistralChatClient(model="mistral-small-latest", client=http_client)
+    client = MistralChatClient(model="mistral-small-latest", http_client=http_client)
     return client, server
 
 
@@ -130,7 +132,7 @@ def test_mistral_chat_construction_env(monkeypatch: pytest.MonkeyPatch) -> None:
 def test_mistral_chat_construction_with_params() -> None:
     client = MistralChatClient(model="mistral-large-latest", api_key="test-key")
     assert client.model == "mistral-large-latest"
-    assert client.client.headers["Authorization"] == "Bearer test-key"
+    assert isinstance(client.client, Mistral)
 
 
 def test_mistral_chat_construction_with_server_url() -> None:
@@ -140,14 +142,26 @@ def test_mistral_chat_construction_with_server_url() -> None:
         server_url="https://custom.mistral.ai",
     )
     assert client.service_url() == "https://custom.mistral.ai"
-    assert str(client.client.base_url) == "https://custom.mistral.ai"
+    assert client.client.sdk_configuration.get_server_details()[0] == "https://custom.mistral.ai"
 
 
 def test_mistral_chat_construction_with_client_needs_no_api_key(monkeypatch: pytest.MonkeyPatch) -> None:
     monkeypatch.delenv("MISTRAL_API_KEY", raising=False)
     http_client = httpx.AsyncClient(base_url="https://api.mistral.ai")
-    client = MistralChatClient(model="mistral-large-latest", client=http_client)
-    assert client.client is http_client
+    client = MistralChatClient(model="mistral-large-latest", http_client=http_client)
+    assert client.client.sdk_configuration.async_client is http_client
+
+
+async def test_mistral_chat_deprecated_client_param_accepts_httpx() -> None:
+    http_client = httpx.AsyncClient(base_url="https://api.mistral.ai")
+    with pytest.deprecated_call(match="http_client"):
+        client = MistralChatClient(
+            model="mistral-large-latest",
+            client=http_client,  # type: ignore[arg-type]
+        )
+    assert client.client.sdk_configuration.async_client is http_client
+    await client.close()
+    await http_client.aclose()
 
 
 def test_mistral_chat_construction_missing_api_key_raises(monkeypatch: pytest.MonkeyPatch) -> None:
@@ -165,11 +179,13 @@ def test_mistral_chat_service_url_default() -> None:
 
 async def test_mistral_chat_close_only_closes_owned_client() -> None:
     owned = MistralChatClient(model="mistral-large-latest", api_key="test-key")
+    owned_http_client = owned.client.sdk_configuration.async_client
     await owned.close()
-    assert owned.client.is_closed
+    assert owned_http_client is not None
+    assert owned_http_client.is_closed
 
     http_client = httpx.AsyncClient(base_url="https://custom.mistral.ai")
-    injected = MistralChatClient(model="mistral-large-latest", client=http_client)
+    injected = MistralChatClient(model="mistral-large-latest", http_client=http_client)
     assert injected.service_url() == "https://custom.mistral.ai"
 
     await injected.close()
@@ -181,7 +197,7 @@ async def test_mistral_chat_close_only_closes_owned_client() -> None:
 async def test_mistral_chat_missing_model_raises_at_request(monkeypatch: pytest.MonkeyPatch) -> None:
     monkeypatch.delenv("MISTRAL_CHAT_MODEL", raising=False)
     http_client = httpx.AsyncClient(base_url="https://api.mistral.ai")
-    client = MistralChatClient(client=http_client, api_key="test-key")
+    client = MistralChatClient(http_client=http_client, api_key="test-key")
     with pytest.raises(ValueError, match="Mistral model is required"):
         await client.get_response([Message("user", ["hi"])])
 
@@ -241,6 +257,31 @@ async def test_get_response_includes_cached_input_tokens() -> None:
     assert response.usage_details["cache_read_input_token_count"] == 80
 
 
+async def test_get_response_includes_extended_usage() -> None:
+    client, _ = make_client(
+        json_response(
+            make_response_payload(
+                content="hello",
+                usage={
+                    "prompt_tokens": 100,
+                    "completion_tokens": 7,
+                    "total_tokens": 107,
+                    "prompt_audio_seconds": 4,
+                    "prompt_tokens_details": {"audio_tokens": 12},
+                    "service_tier": "priority",
+                },
+            )
+        )
+    )
+
+    response = await client.get_response([Message("user", ["hi"])])
+
+    assert response.usage_details is not None
+    assert response.usage_details["prompt_audio_seconds"] == 4
+    assert response.usage_details["prompt/audio_tokens"] == 12
+    assert response.additional_properties["service_tier"] == "priority"
+
+
 @pytest.mark.parametrize("cached_tokens", ["80", 80.5, True, False])
 async def test_get_response_ignores_invalid_cached_input_tokens(cached_tokens: Any) -> None:
     client, _ = make_client(
@@ -290,7 +331,7 @@ def raise_connect_error(request: httpx.Request) -> httpx.Response:
         base_url="https://api.mistral.ai",
         transport=httpx.MockTransport(raise_connect_error),
     )
-    client = MistralChatClient(model="mistral-small-latest", client=http_client)
+    client = MistralChatClient(model="mistral-small-latest", http_client=http_client)
 
     with pytest.raises(ChatClientException, match="Mistral chat request failed"):
         await client.get_response([Message("user", ["hi"])])
@@ -300,7 +341,7 @@ def raise_connect_error(request: httpx.Request) -> httpx.Response:
     ("response", "message"),
     [
         (httpx.Response(200, content=b"{"), "response was invalid"),
-        (json_response([]), "must be a JSON object"),
+        (json_response([]), "response was invalid"),
         (json_response({"choices": ["not-an-object"]}), "response was invalid"),
     ],
 )
@@ -321,9 +362,10 @@ async def test_get_response_option_mapping() -> None:
         "allow_multiple_tool_calls": False,
         "safe_prompt": True,
         "stop": ["END"],
-        "guardrails": [{"name": "test-guardrail"}],
+        "guardrails": [{"block_on_error": True}],
         "prompt_cache_key": "shared-prefix",
         "reasoning_effort": "high",
+        "service_tier": "standard_only",
     }
     await client.get_response([Message("user", ["hi"])], options=options)
 
@@ -335,9 +377,10 @@ async def test_get_response_option_mapping() -> None:
     assert request["safe_prompt"] is True
     assert request["stop"] == ["END"]
     assert "n" not in request
-    assert request["guardrails"] == [{"name": "test-guardrail"}]
+    assert request["guardrails"] == [{"block_on_error": True}]
     assert request["prompt_cache_key"] == "shared-prefix"
     assert request["reasoning_effort"] == "high"
+    assert request["service_tier"] == "standard_only"
     assert "seed" not in request
     assert "allow_multiple_tool_calls" not in request
 
@@ -382,7 +425,12 @@ async def test_message_conversion_roles() -> None:
     assert sent[2]["role"] == "assistant"
     assert sent[2]["content"] == "Let me check."
     assert sent[2]["tool_calls"] == [
-        {"id": "call123AB", "type": "function", "function": {"name": "lookup", "arguments": '{"q": "x"}'}}
+        {
+            "id": "call123AB",
+            "type": "function",
+            "function": {"name": "lookup", "arguments": '{"q": "x"}'},
+            "index": 0,
+        }
     ]
     assert sent[3]["role"] == "tool"
     assert sent[3]["tool_call_id"] == "call123AB"
@@ -665,23 +713,25 @@ async def test_parse_thinking_chunks() -> None:
 def test_response_content_edge_cases() -> None:
     client, _ = make_client()
     contents = client._parse_message_contents(  # pyright: ignore[reportPrivateUsage]
-        {
-            "content": [
-                {"type": "thinking", "thinking": "reasoning"},
-                {"type": "unsupported"},
-            ],
-            "tool_calls": [
-                tool_call_payload("mapping", {"value": 1}, call_id="abc123XYZ"),
-                tool_call_payload("missing", None, call_id="def456UVW"),
-            ],
-        }
+        AssistantMessage.model_validate(
+            {
+                "content": [
+                    {
+                        "type": "thinking",
+                        "thinking": [{"type": "text", "text": "reasoning"}],
+                    },
+                ],
+                "tool_calls": [
+                    tool_call_payload("mapping", {"value": 1}, call_id="abc123XYZ"),
+                ],
+            }
+        )
     )
     calls = [content for content in contents if content.type == "function_call"]
     assert contents[0].text == "reasoning"
     assert calls[0].arguments == {"value": 1}
-    assert calls[1].arguments == "None"
     assert client._format_created_at("invalid") is None  # pyright: ignore[reportPrivateUsage]
-    assert client._thinking_to_text({"thinking": object()}) == ""  # pyright: ignore[reportPrivateUsage]
+    assert client._thinking_to_text(ThinkChunk(thinking=[])) == ""  # pyright: ignore[reportPrivateUsage]
 
 
 async def test_parse_finish_reason_model_length() -> None:
@@ -701,12 +751,18 @@ async def test_parse_finish_reason_unmapped_is_preserved() -> None:
     assert response.finish_reason == "error"
 
 
-async def test_parse_finish_reason_absent() -> None:
-    """A response without a finish reason still reports no finish reason."""
+async def test_parse_finish_reason_null() -> None:
+    """A null finish reason still reports no finish reason."""
     client, _ = make_client(
         json_response(
             make_response_payload(
-                choices=[{"index": 0, "message": {"role": "assistant", "content": "x"}}],
+                choices=[
+                    {
+                        "index": 0,
+                        "finish_reason": None,
+                        "message": {"role": "assistant", "content": "x"},
+                    }
+                ],
             )
         )
     )
@@ -779,6 +835,23 @@ async def test_streaming_response() -> None:
     assert server.last_request["stream"] is True
 
 
+async def test_streaming_uses_first_choice_only() -> None:
+    chunk = make_chunk_payload(content="first", finish_reason="stop")
+    chunk["choices"].append(
+        {
+            "index": 1,
+            "finish_reason": "length",
+            "delta": {"role": "assistant", "content": "second"},
+        }
+    )
+    client, _ = make_client(stream_response(chunk))
+
+    response = await client.get_response([Message("user", ["hi"])], stream=True).get_final_response()
+
+    assert response.text == "first"
+    assert response.finish_reason == "stop"
+
+
 async def test_streaming_finish_reason_unmapped_is_preserved() -> None:
     """A streamed Mistral finish reason with no framework equivalent is passed through unchanged."""
     client, _ = make_client(stream_response(make_chunk_payload(content="partial", finish_reason="error")))
@@ -934,20 +1007,6 @@ async def test_streaming_http_error_wrapped() -> None:
             pass
 
 
-def test_parse_sse_line_variants() -> None:
-    payload = make_chunk_payload(content="hello")
-    assert MistralChatClient._parse_sse_line(f"data:{json.dumps(payload)}") == payload  # pyright: ignore[reportPrivateUsage]
-    assert MistralChatClient._parse_sse_line("") is None  # pyright: ignore[reportPrivateUsage]
-    assert MistralChatClient._parse_sse_line("event: message") is None  # pyright: ignore[reportPrivateUsage]
-    assert MistralChatClient._parse_sse_line("data:") is None  # pyright: ignore[reportPrivateUsage]
-    assert MistralChatClient._parse_sse_line("data: [DONE]") is None  # pyright: ignore[reportPrivateUsage]
-
-    with pytest.raises(ChatClientInvalidResponseException, match="malformed SSE"):
-        MistralChatClient._parse_sse_line("data: {")  # pyright: ignore[reportPrivateUsage]
-    with pytest.raises(ChatClientInvalidResponseException, match="must be a JSON object"):
-        MistralChatClient._parse_sse_line("data: []")  # pyright: ignore[reportPrivateUsage]
-
-
 async def test_streaming_tool_call_flushed_without_finish_chunk() -> None:
     """A stream that ends without a finish chunk still emits accumulated calls."""
     client, _ = make_client(
diff --git a/python/packages/mistral/tests/mistral/test_mistral_embedding_client.py b/python/packages/mistral/tests/mistral/test_mistral_embedding_client.py
index 1eee61cfe4d..be9166fa80e 100644
--- a/python/packages/mistral/tests/mistral/test_mistral_embedding_client.py
+++ b/python/packages/mistral/tests/mistral/test_mistral_embedding_client.py
@@ -15,6 +15,7 @@
     IntegrationInvalidRequestException,
     IntegrationInvalidResponseException,
 )
+from mistralai.client import Mistral
 
 from agent_framework_mistral import MistralEmbeddingClient, MistralEmbeddingOptions
 
@@ -27,6 +28,7 @@ def make_embeddings_payload(
     usage: dict[str, Any] | None = None,
 ) -> dict[str, Any]:
     return {
+        "id": "embed-id",
         "object": "list",
         "model": model,
         "data": [{"object": "embedding", "index": i, "embedding": list(vector)} for i, vector in enumerate(vectors)],
@@ -70,7 +72,7 @@ def test_mistral_embedding_construction_with_params() -> None:
     """Test construction with explicit parameters."""
     client = MistralEmbeddingClient(model="mistral-embed", api_key="test-key")
     assert client.model == "mistral-embed"
-    assert client.client.headers["Authorization"] == "Bearer test-key"
+    assert isinstance(client.client, Mistral)
 
 
 def test_mistral_embedding_construction_with_server_url() -> None:
@@ -82,21 +84,24 @@ def test_mistral_embedding_construction_with_server_url() -> None:
     )
     assert client.model == "mistral-embed"
     assert client.server_url == "https://custom.mistral.ai"
-    assert str(client.client.base_url) == "https://custom.mistral.ai"
+    assert client.client.sdk_configuration.get_server_details()[0] == "https://custom.mistral.ai"
 
 
 def test_mistral_embedding_construction_with_http_client() -> None:
     """Test construction with a pre-configured client."""
     http_client = httpx.AsyncClient(base_url="https://api.mistral.ai")
     client = MistralEmbeddingClient(model="mistral-embed", http_client=http_client)
-    assert client.client is http_client
+    assert client.client.sdk_configuration.async_client is http_client
 
 
-def test_mistral_embedding_deprecated_client_param_accepts_httpx() -> None:
+def test_mistral_embedding_client_param_accepts_httpx() -> None:
     http_client = httpx.AsyncClient(base_url="https://api.mistral.ai")
-    with pytest.deprecated_call():
-        client = MistralEmbeddingClient(model="mistral-embed", client=http_client)
-    assert client.client is http_client
+    with pytest.deprecated_call(match="http_client"):
+        client = MistralEmbeddingClient(
+            model="mistral-embed",
+            client=http_client,  # type: ignore[arg-type]
+        )
+    assert client.client.sdk_configuration.async_client is http_client
 
 
 class FakeMistralSDK:
@@ -116,31 +121,42 @@ async def _create_async(self, **kwargs: Any) -> Any:
         )
 
 
-async def test_mistral_embedding_deprecated_client_param_accepts_sdk_client() -> None:
-    """An injected mistralai.Mistral keeps working through the legacy SDK path."""
-    sdk = FakeMistralSDK()
-    with pytest.deprecated_call():
-        client = MistralEmbeddingClient(model="mistral-embed", client=sdk)
+async def test_mistral_embedding_client_param_accepts_sdk_client() -> None:
+    """An injected mistralai.Mistral is used directly."""
+    fake_embeddings = FakeMistralSDK()
+    sdk = Mistral(api_key="test-key")
+    sdk.embeddings = fake_embeddings.embeddings  # type: ignore[assignment]
+    client = MistralEmbeddingClient(model="mistral-embed", client=sdk)
 
     result = await client.get_embeddings(["hello"], options=MistralEmbeddingOptions(dimensions=2))
 
     assert [e.vector for e in result] == [[0.1, 0.2]]
     assert result.usage == {"input_token_count": 3, "total_token_count": 3}
-    assert sdk.requests == [{"model": "mistral-embed", "inputs": ["hello"], "output_dimension": 2}]
+    request = fake_embeddings.requests[0]
+    assert request["model"] == "mistral-embed"
+    assert request["inputs"] == ["hello"]
+    assert request["output_dimension"] == 2
+    assert "http_headers" in request
+    await sdk.__aexit__(None, None, None)
+    sdk.__exit__(None, None, None)
 
 
-def test_mistral_embedding_deprecated_client_param_rejects_unknown_client() -> None:
+def test_mistral_embedding_client_param_rejects_unknown_client() -> None:
     class NotAClient:
         pass
 
-    with pytest.deprecated_call(), pytest.raises(TypeError, match="httpx.AsyncClient"):
-        MistralEmbeddingClient(model="mistral-embed", client=NotAClient())
+    with pytest.raises(TypeError, match="mistralai.client.Mistral"):
+        MistralEmbeddingClient(model="mistral-embed", client=NotAClient())  # type: ignore[arg-type]
 
 
 def test_mistral_embedding_client_and_http_client_conflict() -> None:
     http_client = httpx.AsyncClient(base_url="https://api.mistral.ai")
-    with pytest.deprecated_call(), pytest.raises(ValueError, match="not both"):
-        MistralEmbeddingClient(model="mistral-embed", http_client=http_client, client=http_client)
+    with pytest.deprecated_call(match="http_client"), pytest.raises(ValueError, match="not both"):
+        MistralEmbeddingClient(
+            model="mistral-embed",
+            http_client=http_client,
+            client=http_client,  # type: ignore[arg-type]
+        )
 
 
 def test_mistral_embedding_construction_missing_model_raises(monkeypatch: pytest.MonkeyPatch) -> None:
@@ -181,8 +197,10 @@ def test_mistral_embedding_service_url_custom() -> None:
 
 async def test_mistral_embedding_close_only_closes_owned_client() -> None:
     owned = MistralEmbeddingClient(model="mistral-embed", api_key="test-key")
+    owned_http_client = owned.client.sdk_configuration.async_client
     await owned.close()
-    assert owned.client.is_closed
+    assert owned_http_client is not None
+    assert owned_http_client.is_closed
 
     http_client = httpx.AsyncClient(base_url="https://custom.mistral.ai")
     injected = MistralEmbeddingClient(model="mistral-embed", http_client=http_client)
@@ -323,7 +341,7 @@ def raise_connect_error(request: httpx.Request) -> httpx.Response:
     ("response", "message"),
     [
         (httpx.Response(200, content=b"{"), "response was invalid"),
-        (httpx.Response(200, json=[]), "must be a JSON object"),
+        (httpx.Response(200, json=[]), "response was invalid"),
         (httpx.Response(200, json={"data": ["not-an-object"]}), "response was invalid"),
     ],
 )
diff --git a/python/uv.lock b/python/uv.lock
index 5aa816032c6..52bb07990c0 100644
--- a/python/uv.lock
+++ b/python/uv.lock
@@ -911,13 +911,13 @@ version = "1.0.0b260813"
 source = { editable = "packages/mistral" }
 dependencies = [
     { name = "agent-framework-core" },
-    { name = "httpx" },
+    { name = "mistralai" },
 ]
 
 [package.metadata]
 requires-dist = [
     { name = "agent-framework-core", editable = "packages/core" },
-    { name = "httpx", specifier = ">=0.23.1,<1" },
+    { name = "mistralai", specifier = ">=2.9.2,<3" },
 ]
 
 [[package]]
@@ -2457,6 +2457,15 @@ wheels = [
     { url = "https://files.pythonhosted.org/packages/de/15/545e2b6cf2e3be84bc1ed85613edd75b8aea69807a71c26f4ca6a9258e82/email_validator-2.3.0-py3-none-any.whl", hash = "sha256:80f13f623413e6b197ae73bb10bf4eb0908faf509ad8362c5edeb0be7fd450b4", size = 35604, upload-time = "2025-08-26T13:09:05.858Z" },
 ]
 
+[[package]]
+name = "eval-type-backport"
+version = "0.4.0"
+source = { registry = "https://pypi.org/simple" }
+sdist = { url = "https://files.pythonhosted.org/packages/1c/15/273a4baf8248d6d76220723c3caf039d283774b31a7c46ba686120145b76/eval_type_backport-0.4.0.tar.gz", hash = "sha256:8397d25e6524c2e67b9576bb0636be27dea2192017711220c534ec2de921e9b0", size = 10260, upload-time = "2026-06-02T13:22:06.059Z" }
+wheels = [
+    { url = "https://files.pythonhosted.org/packages/50/a7/bb99bf5e6f78736ddb53480f2c3ff3702ffe2196a7c5e1661c03081d398e/eval_type_backport-0.4.0-py3-none-any.whl", hash = "sha256:ad5e2a8db71b6696a56eafb938b0f5a337d3217f256b8e158b469422b4772b20", size = 6432, upload-time = "2026-06-02T13:22:04.827Z" },
+]
+
 [[package]]
 name = "execnet"
 version = "2.1.2"
@@ -3575,6 +3584,15 @@ wheels = [
     { url = "https://files.pythonhosted.org/packages/03/99/33c7d78a3fb70d545fd5411ac67a651c81602cc09c9cf0df383733f068c5/jsonpath_ng-1.8.0-py3-none-any.whl", hash = "sha256:b8dde192f8af58d646fc031fac9c99fe4d00326afc4148f1f043c601a8cfe138", size = 67844, upload-time = "2026-02-28T00:53:19.637Z" },
 ]
 
+[[package]]
+name = "jsonpath-python"
+version = "1.1.6"
+source = { registry = "https://pypi.org/simple" }
+sdist = { url = "https://files.pythonhosted.org/packages/98/18/4ca8742534a5993ff383f7602e325ce2d5d7cc93d72ac5e1cdedbea8a458/jsonpath_python-1.1.6.tar.gz", hash = "sha256:dded9932b4ec41fb8726e09c83afa4e6be618f938c2db287cc2a81723c639671", size = 88178, upload-time = "2026-05-07T01:26:34.482Z" }
+wheels = [
+    { url = "https://files.pythonhosted.org/packages/55/8a/1270a6803bd821cbfcdda387eaa13cb41a7b1f7b9bd145979b3bfb9d6cb7/jsonpath_python-1.1.6-py3-none-any.whl", hash = "sha256:a1c50afd8d3fbbaf47a4873bc890dcb3c15da96f5c020327977d844d8731a2d4", size = 14453, upload-time = "2026-05-07T01:26:33.306Z" },
+]
+
 [[package]]
 name = "jsonschema"
 version = "4.26.0"
@@ -4228,6 +4246,25 @@ wheels = [
     { url = "https://files.pythonhosted.org/packages/f0/af/7ed341ee2040eed3c0fc738f84d4cbc21045f6d06d2461fa25cf5c866f93/microsoft_opentelemetry-1.3.7-py3-none-any.whl", hash = "sha256:babb480e2499016317b5898df490268ec41a669352f44b8ac34267e0ea7a2278", size = 210149, upload-time = "2026-08-05T15:38:57.018Z" },
 ]
 
+[[package]]
+name = "mistralai"
+version = "2.9.4"
+source = { registry = "https://pypi.org/simple" }
+dependencies = [
+    { name = "eval-type-backport" },
+    { name = "httpx" },
+    { name = "jsonpath-python" },
+    { name = "opentelemetry-api" },
+    { name = "opentelemetry-semantic-conventions" },
+    { name = "pydantic" },
+    { name = "python-dateutil" },
+    { name = "typing-inspection" },
+]
+sdist = { url = "https://files.pythonhosted.org/packages/34/52/51065b4b453b0cd410fbc9d659d9b43345d4cef74f610cf4cad3c2682e22/mistralai-2.9.4.tar.gz", hash = "sha256:e3607552d34cc38b6f81e80bf95201946eac119260259b0cc1094aac350dbb8e", size = 543545, upload-time = "2026-08-21T18:06:03.045Z" }
+wheels = [
+    { url = "https://files.pythonhosted.org/packages/ac/c6/ce6371bd7a56e8c01a8f59baf4bc16d134eed4f3e2c593b5733087012936/mistralai-2.9.4-py3-none-any.whl", hash = "sha256:184f8ba3cffee6ffc379481340a67c1d5f6db145cf690b2ffd06d6a4cec92d2d", size = 1298804, upload-time = "2026-08-21T18:06:00.898Z" },
+]
+
 [[package]]
 name = "ml-dtypes"
 version = "0.5.4"

From 34aabb1f912882e56445eff3de80e835b1e6021c Mon Sep 17 00:00:00 2001
From: eavanvalkenburg <github@vanvalkenburg.eu>
Date: Mon, 31 Aug 2026 11:42:27 +0200
Subject: [PATCH 2/4] Python: fix Mistral test typing

Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com>

Copilot-Session: 3cc4a2d5-640a-4c05-8fc9-abe08f8134ee
---
 .../tests/mistral/test_mistral_chat_client.py        | 11 ++++++-----
 .../tests/mistral/test_mistral_embedding_client.py   | 12 ++++++------
 2 files changed, 12 insertions(+), 11 deletions(-)

diff --git a/python/packages/mistral/tests/mistral/test_mistral_chat_client.py b/python/packages/mistral/tests/mistral/test_mistral_chat_client.py
index 67188a1d4ca..f49e659e63e 100644
--- a/python/packages/mistral/tests/mistral/test_mistral_chat_client.py
+++ b/python/packages/mistral/tests/mistral/test_mistral_chat_client.py
@@ -4,7 +4,7 @@
 import logging
 import os
 from collections.abc import AsyncIterator, Sequence
-from typing import Any
+from typing import Any, cast
 
 import httpx
 import pytest
@@ -157,7 +157,7 @@ async def test_mistral_chat_deprecated_client_param_accepts_httpx() -> None:
     with pytest.deprecated_call(match="http_client"):
         client = MistralChatClient(
             model="mistral-large-latest",
-            client=http_client,  # type: ignore[arg-type]
+            client=cast("Any", http_client),
         )
     assert client.client.sdk_configuration.async_client is http_client
     await client.close()
@@ -181,7 +181,7 @@ async def test_mistral_chat_close_only_closes_owned_client() -> None:
     owned = MistralChatClient(model="mistral-large-latest", api_key="test-key")
     owned_http_client = owned.client.sdk_configuration.async_client
     await owned.close()
-    assert owned_http_client is not None
+    assert isinstance(owned_http_client, httpx.AsyncClient)
     assert owned_http_client.is_closed
 
     http_client = httpx.AsyncClient(base_url="https://custom.mistral.ai")
@@ -277,8 +277,9 @@ async def test_get_response_includes_extended_usage() -> None:
     response = await client.get_response([Message("user", ["hi"])])
 
     assert response.usage_details is not None
-    assert response.usage_details["prompt_audio_seconds"] == 4
-    assert response.usage_details["prompt/audio_tokens"] == 12
+    usage_details = cast("dict[str, Any]", response.usage_details)
+    assert usage_details["prompt_audio_seconds"] == 4
+    assert usage_details["prompt/audio_tokens"] == 12
     assert response.additional_properties["service_tier"] == "priority"
 
 
diff --git a/python/packages/mistral/tests/mistral/test_mistral_embedding_client.py b/python/packages/mistral/tests/mistral/test_mistral_embedding_client.py
index be9166fa80e..be612f4fad8 100644
--- a/python/packages/mistral/tests/mistral/test_mistral_embedding_client.py
+++ b/python/packages/mistral/tests/mistral/test_mistral_embedding_client.py
@@ -4,7 +4,7 @@
 import os
 from collections.abc import Sequence
 from types import SimpleNamespace
-from typing import Any
+from typing import Any, cast
 
 import httpx
 import pytest
@@ -99,7 +99,7 @@ def test_mistral_embedding_client_param_accepts_httpx() -> None:
     with pytest.deprecated_call(match="http_client"):
         client = MistralEmbeddingClient(
             model="mistral-embed",
-            client=http_client,  # type: ignore[arg-type]
+            client=cast("Any", http_client),
         )
     assert client.client.sdk_configuration.async_client is http_client
 
@@ -125,7 +125,7 @@ async def test_mistral_embedding_client_param_accepts_sdk_client() -> None:
     """An injected mistralai.Mistral is used directly."""
     fake_embeddings = FakeMistralSDK()
     sdk = Mistral(api_key="test-key")
-    sdk.embeddings = fake_embeddings.embeddings  # type: ignore[assignment]
+    sdk.embeddings = cast("Any", fake_embeddings.embeddings)
     client = MistralEmbeddingClient(model="mistral-embed", client=sdk)
 
     result = await client.get_embeddings(["hello"], options=MistralEmbeddingOptions(dimensions=2))
@@ -146,7 +146,7 @@ class NotAClient:
         pass
 
     with pytest.raises(TypeError, match="mistralai.client.Mistral"):
-        MistralEmbeddingClient(model="mistral-embed", client=NotAClient())  # type: ignore[arg-type]
+        MistralEmbeddingClient(model="mistral-embed", client=cast("Any", NotAClient()))
 
 
 def test_mistral_embedding_client_and_http_client_conflict() -> None:
@@ -155,7 +155,7 @@ def test_mistral_embedding_client_and_http_client_conflict() -> None:
         MistralEmbeddingClient(
             model="mistral-embed",
             http_client=http_client,
-            client=http_client,  # type: ignore[arg-type]
+            client=cast("Any", http_client),
         )
 
 
@@ -199,7 +199,7 @@ async def test_mistral_embedding_close_only_closes_owned_client() -> None:
     owned = MistralEmbeddingClient(model="mistral-embed", api_key="test-key")
     owned_http_client = owned.client.sdk_configuration.async_client
     await owned.close()
-    assert owned_http_client is not None
+    assert isinstance(owned_http_client, httpx.AsyncClient)
     assert owned_http_client.is_closed
 
     http_client = httpx.AsyncClient(base_url="https://custom.mistral.ai")

From b4bb9411b3dae24444b3e87a1e7554fdad9d257d Mon Sep 17 00:00:00 2001
From: eavanvalkenburg <github@vanvalkenburg.eu>
Date: Mon, 31 Aug 2026 11:43:20 +0200
Subject: [PATCH 3/4] Python: namespace Mistral usage details

Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com>

Copilot-Session: 3cc4a2d5-640a-4c05-8fc9-abe08f8134ee
---
 .../packages/mistral/agent_framework_mistral/_chat_client.py  | 4 ++--
 .../mistral/tests/mistral/test_mistral_chat_client.py         | 4 ++--
 2 files changed, 4 insertions(+), 4 deletions(-)

diff --git a/python/packages/mistral/agent_framework_mistral/_chat_client.py b/python/packages/mistral/agent_framework_mistral/_chat_client.py
index 2e52b39cbf5..1739fa3c7e1 100644
--- a/python/packages/mistral/agent_framework_mistral/_chat_client.py
+++ b/python/packages/mistral/agent_framework_mistral/_chat_client.py
@@ -865,7 +865,7 @@ def _parse_usage(self, usage: UsageInfo | None) -> UsageDetails | None:
         if "total_tokens" in fields_set and usage.total_tokens is not None:
             details["total_token_count"] = usage.total_tokens
         if isinstance(usage.prompt_audio_seconds, int) and not isinstance(usage.prompt_audio_seconds, bool):
-            cast("dict[str, Any]", details)["prompt_audio_seconds"] = usage.prompt_audio_seconds
+            cast("dict[str, Any]", details)["mistral.prompt_audio_seconds"] = usage.prompt_audio_seconds
         prompt_tokens_details = getattr(usage, "prompt_tokens_details", None)
         if isinstance(prompt_tokens_details, Mapping):
             prompt_details = cast("Mapping[str, Any]", prompt_tokens_details)
@@ -875,7 +875,7 @@ def _parse_usage(self, usage: UsageInfo | None) -> UsageDetails | None:
                 details["cache_read_input_token_count"] = cached_tokens
             audio_tokens = prompt_details.get("audio_tokens")
             if isinstance(audio_tokens, int) and not isinstance(audio_tokens, bool):
-                cast("dict[str, Any]", details)["prompt/audio_tokens"] = audio_tokens
+                cast("dict[str, Any]", details)["mistral.prompt_audio_tokens"] = audio_tokens
         return details or None
 
     @staticmethod
diff --git a/python/packages/mistral/tests/mistral/test_mistral_chat_client.py b/python/packages/mistral/tests/mistral/test_mistral_chat_client.py
index f49e659e63e..e4d6e787f5f 100644
--- a/python/packages/mistral/tests/mistral/test_mistral_chat_client.py
+++ b/python/packages/mistral/tests/mistral/test_mistral_chat_client.py
@@ -278,8 +278,8 @@ async def test_get_response_includes_extended_usage() -> None:
 
     assert response.usage_details is not None
     usage_details = cast("dict[str, Any]", response.usage_details)
-    assert usage_details["prompt_audio_seconds"] == 4
-    assert usage_details["prompt/audio_tokens"] == 12
+    assert usage_details["mistral.prompt_audio_seconds"] == 4
+    assert usage_details["mistral.prompt_audio_tokens"] == 12
     assert response.additional_properties["service_tier"] == "priority"
 
 

From fd584cdc10b0a6635cdcdc713fbebf625da15bdb Mon Sep 17 00:00:00 2001
From: eavanvalkenburg <github@vanvalkenburg.eu>
Date: Mon, 31 Aug 2026 12:18:17 +0200
Subject: [PATCH 4/4] Python: address Mistral SDK review feedback

Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com>

Copilot-Session: 3cc4a2d5-640a-4c05-8fc9-abe08f8134ee
---
 .../agent_framework_mistral/_chat_client.py   | 26 +++--
 .../_embedding_client.py                      |  6 +-
 .../agent_framework_mistral/_http_client.py   | 23 +++++
 .../tests/mistral/test_mistral_chat_client.py | 99 ++++++++++++++++++-
 .../mistral/test_mistral_embedding_client.py  | 34 ++++++-
 5 files changed, 176 insertions(+), 12 deletions(-)
 create mode 100644 python/packages/mistral/agent_framework_mistral/_http_client.py

diff --git a/python/packages/mistral/agent_framework_mistral/_chat_client.py b/python/packages/mistral/agent_framework_mistral/_chat_client.py
index 1739fa3c7e1..597cbe60774 100644
--- a/python/packages/mistral/agent_framework_mistral/_chat_client.py
+++ b/python/packages/mistral/agent_framework_mistral/_chat_client.py
@@ -53,9 +53,10 @@
     ToolCall,
     UsageInfo,
 )
-from pydantic import BaseModel
+from pydantic import BaseModel, ValidationError
 
 from ._feature_usage import FeatureIndex
+from ._http_client import AsyncClientUsingConfiguredTimeout
 
 if sys.version_info >= (3, 13):
     from typing import TypeVar  # pragma: no cover
@@ -167,6 +168,8 @@ class MistralSettings(TypedDict, total=False):
 # endregion
 
 _MISTRAL_API_BASE_URL = "https://api.mistral.ai"
+_DEFAULT_TIMEOUT_MS = 60_000
+_SDK_TRANSPORT_ARGUMENTS = frozenset({"http_headers", "retries", "server_url", "timeout_ms"})
 
 # Keys mapping to a different Mistral chat-completion parameter name
 _OPTION_TRANSLATIONS: dict[str, str] = {
@@ -399,11 +402,11 @@ def __init__(
             if self.server_url is None:
                 self.server_url = client.sdk_configuration.get_server_details()[0]
         else:
-            client_kwargs: dict[str, Any] = {}
+            client_kwargs: dict[str, Any] = {"timeout_ms": _DEFAULT_TIMEOUT_MS}
             if resolved_api_key := mistral_settings.get("api_key"):
                 client_kwargs["api_key"] = resolved_api_key.get_secret_value()
             if http_client is not None:
-                client_kwargs["async_client"] = http_client
+                client_kwargs["async_client"] = AsyncClientUsingConfiguredTimeout(http_client)
                 if self.server_url is None:
                     client_base_url = str(http_client.base_url).rstrip("/")
                     self.server_url = client_base_url or None
@@ -442,13 +445,18 @@ async def _stream() -> AsyncIterable[ChatResponseUpdate]:
                 request.setdefault("http_headers", {"User-Agent": get_user_agent()})
                 tool_calls = _StreamedToolCalls()
                 try:
-                    response = await self.client.chat.stream_async(**request)
-                    async for event in response:
-                        yield self._parse_chunk(event.data, tool_calls)
+                    async with await self.client.chat.stream_async(**request) as response:
+                        async for event in response:
+                            yield self._parse_chunk(event.data, tool_calls)
                     if remaining := tool_calls.flush_all():
                         yield ChatResponseUpdate(contents=remaining, role="assistant")
                 except MistralError as ex:
                     self._raise_sdk_error(ex, streaming=True)
+                except ValidationError as ex:
+                    raise ChatClientInvalidResponseException(
+                        f"Mistral streaming chat response was invalid: {ex}",
+                        inner_exception=ex,
+                    ) from ex
                 except ChatClientException:
                     raise
                 except Exception as ex:
@@ -520,6 +528,12 @@ def _prepare_request(
         Raises:
             ValueError: If no model is set on the options or the client instance.
         """
+        if transport_arguments := _SDK_TRANSPORT_ARGUMENTS.intersection(kwargs):
+            arguments = ", ".join(sorted(transport_arguments))
+            raise ValueError(
+                f"Mistral transport arguments cannot be supplied per call: {arguments}. Configure the client instead."
+            )
+
         model = options.get("model") or self.model
         if not model:
             raise ValueError(
diff --git a/python/packages/mistral/agent_framework_mistral/_embedding_client.py b/python/packages/mistral/agent_framework_mistral/_embedding_client.py
index d335a6af132..29b0107c22e 100644
--- a/python/packages/mistral/agent_framework_mistral/_embedding_client.py
+++ b/python/packages/mistral/agent_framework_mistral/_embedding_client.py
@@ -30,6 +30,7 @@
 from mistralai.client.errors import MistralError
 
 from ._feature_usage import FeatureIndex
+from ._http_client import AsyncClientUsingConfiguredTimeout
 
 if sys.version_info >= (3, 13):
     from typing import TypeVar  # pragma: no cover
@@ -40,6 +41,7 @@
 logger = logging.getLogger("agent_framework.mistral")
 
 _MISTRAL_API_BASE_URL = "https://api.mistral.ai"
+_DEFAULT_TIMEOUT_MS = 60_000
 
 
 class MistralEmbeddingOptions(EmbeddingGenerationOptions, total=False):
@@ -158,11 +160,11 @@ def __init__(
             if self.server_url is None:
                 self.server_url = client.sdk_configuration.get_server_details()[0]
         else:
-            client_kwargs: dict[str, Any] = {}
+            client_kwargs: dict[str, Any] = {"timeout_ms": _DEFAULT_TIMEOUT_MS}
             if resolved_api_key := mistral_settings.get("api_key"):
                 client_kwargs["api_key"] = resolved_api_key.get_secret_value()
             if http_client is not None:
-                client_kwargs["async_client"] = http_client
+                client_kwargs["async_client"] = AsyncClientUsingConfiguredTimeout(http_client)
                 if self.server_url is None:
                     client_base_url = str(http_client.base_url).rstrip("/")
                     self.server_url = client_base_url or None
diff --git a/python/packages/mistral/agent_framework_mistral/_http_client.py b/python/packages/mistral/agent_framework_mistral/_http_client.py
new file mode 100644
index 00000000000..3b8924bab39
--- /dev/null
+++ b/python/packages/mistral/agent_framework_mistral/_http_client.py
@@ -0,0 +1,23 @@
+# Copyright (c) Microsoft. All rights reserved.
+
+from typing import Any
+
+import httpx
+
+
+class AsyncClientUsingConfiguredTimeout:
+    """Let an injected HTTPX client retain its configured per-phase timeouts."""
+
+    def __init__(self, client: httpx.AsyncClient) -> None:
+        self.client = client
+
+    def build_request(self, *args: Any, **kwargs: Any) -> httpx.Request:
+        kwargs["timeout"] = httpx.USE_CLIENT_DEFAULT
+        return self.client.build_request(*args, **kwargs)
+
+    async def send(self, request: httpx.Request, **kwargs: Any) -> httpx.Response:
+        return await self.client.send(request, **kwargs)
+
+    async def aclose(self) -> None:
+        # The caller owns the injected client.
+        return
diff --git a/python/packages/mistral/tests/mistral/test_mistral_chat_client.py b/python/packages/mistral/tests/mistral/test_mistral_chat_client.py
index e4d6e787f5f..8d2888bf041 100644
--- a/python/packages/mistral/tests/mistral/test_mistral_chat_client.py
+++ b/python/packages/mistral/tests/mistral/test_mistral_chat_client.py
@@ -22,6 +22,9 @@
 import agent_framework_mistral._chat_client as chat_client_module
 from agent_framework_mistral import MistralChatClient, MistralChatOptions
 from agent_framework_mistral._chat_client import _sanitize_tool_call_id  # pyright: ignore[reportPrivateUsage]
+from agent_framework_mistral._http_client import (  # pyright: ignore[reportPrivateUsage]
+    AsyncClientUsingConfiguredTimeout,
+)
 
 # region: Helpers
 
@@ -91,8 +94,10 @@ class MockMistral:
     def __init__(self, responses: Sequence[httpx.Response]) -> None:
         self._responses = list(responses)
         self.requests: list[dict[str, Any]] = []
+        self.http_requests: list[httpx.Request] = []
 
     def handler(self, request: httpx.Request) -> httpx.Response:
+        self.http_requests.append(request)
         self.requests.append(json.loads(request.content))
         return self._responses.pop(0)
 
@@ -133,6 +138,7 @@ def test_mistral_chat_construction_with_params() -> None:
     client = MistralChatClient(model="mistral-large-latest", api_key="test-key")
     assert client.model == "mistral-large-latest"
     assert isinstance(client.client, Mistral)
+    assert client.client.sdk_configuration.timeout_ms == 60_000
 
 
 def test_mistral_chat_construction_with_server_url() -> None:
@@ -149,7 +155,9 @@ def test_mistral_chat_construction_with_client_needs_no_api_key(monkeypatch: pyt
     monkeypatch.delenv("MISTRAL_API_KEY", raising=False)
     http_client = httpx.AsyncClient(base_url="https://api.mistral.ai")
     client = MistralChatClient(model="mistral-large-latest", http_client=http_client)
-    assert client.client.sdk_configuration.async_client is http_client
+    async_client = client.client.sdk_configuration.async_client
+    assert isinstance(async_client, AsyncClientUsingConfiguredTimeout)
+    assert async_client.client is http_client
 
 
 async def test_mistral_chat_deprecated_client_param_accepts_httpx() -> None:
@@ -159,7 +167,9 @@ async def test_mistral_chat_deprecated_client_param_accepts_httpx() -> None:
             model="mistral-large-latest",
             client=cast("Any", http_client),
         )
-    assert client.client.sdk_configuration.async_client is http_client
+    async_client = client.client.sdk_configuration.async_client
+    assert isinstance(async_client, AsyncClientUsingConfiguredTimeout)
+    assert async_client.client is http_client
     await client.close()
     await http_client.aclose()
 
@@ -236,6 +246,26 @@ async def test_get_response_basic() -> None:
     assert server.last_request["messages"] == [{"role": "user", "content": "hi"}]
 
 
+async def test_get_response_preserves_injected_http_client_timeout() -> None:
+    server = MockMistral([json_response(make_response_payload(content="hello"))])
+    timeout = httpx.Timeout(connect=1, read=2, write=3, pool=4)
+    http_client = httpx.AsyncClient(
+        base_url="https://api.mistral.ai",
+        transport=httpx.MockTransport(server.handler),
+        timeout=timeout,
+    )
+    client = MistralChatClient(model="mistral-small-latest", http_client=http_client)
+
+    await client.get_response([Message("user", ["hi"])])
+
+    assert server.http_requests[0].extensions["timeout"] == {
+        "connect": 1,
+        "read": 2,
+        "write": 3,
+        "pool": 4,
+    }
+
+
 async def test_get_response_includes_cached_input_tokens() -> None:
     client, _ = make_client(
         json_response(
@@ -395,6 +425,19 @@ async def test_get_response_instructions_prepended_as_system_message() -> None:
     assert "instructions" not in server.last_request
 
 
+@pytest.mark.parametrize("argument", ["http_headers", "retries", "server_url", "timeout_ms"])
+async def test_get_response_rejects_per_call_transport_arguments(argument: str) -> None:
+    client, server = make_client()
+
+    with pytest.raises(ValueError, match="cannot be supplied per call"):
+        await client.get_response(
+            [Message("user", ["hi"])],
+            client_kwargs={argument: "override"},
+        )
+
+    assert server.requests == []
+
+
 async def test_get_response_model_override() -> None:
     client, server = make_client(json_response(make_response_payload(content="ok")))
 
@@ -799,6 +842,43 @@ def get_weather(location: str) -> str:
     assert any(m["role"] == "tool" for m in server.requests[1]["messages"])
 
 
+async def test_streaming_function_invocation_loop() -> None:
+    client, server = make_client(
+        stream_response(
+            make_chunk_payload(tool_calls=[tool_call_payload("get_weather", '{"loc', call_id="abc123XYZ", index=0)]),
+            make_chunk_payload(
+                tool_calls=[tool_call_payload("", 'ation": "Paris"}', index=0)],
+                finish_reason="tool_calls",
+            ),
+        ),
+        stream_response(make_chunk_payload(content="It is sunny in Paris.", finish_reason="stop")),
+    )
+    executions = 0
+
+    @tool(approval_mode="never_require")
+    def get_weather(location: str) -> str:
+        """Get the weather."""
+        nonlocal executions
+        executions += 1
+        return f"sunny in {location}"
+
+    response = await client.get_response(
+        [Message("user", ["Weather in Paris?"])],
+        options={"tools": [get_weather]},
+        stream=True,
+    ).get_final_response()
+
+    assert response.text == "It is sunny in Paris."
+    assert executions == 1
+    assert len(server.requests) == 2
+    second_messages = server.requests[1]["messages"]
+    assistant_message = next(message for message in second_messages if message["role"] == "assistant")
+    tool_message = next(message for message in second_messages if message["role"] == "tool")
+    assert assistant_message["tool_calls"][0]["id"] == "abc123XYZ"
+    assert tool_message["tool_call_id"] == "abc123XYZ"
+    assert tool_message["content"] == "sunny in Paris"
+
+
 # region: Streaming
 
 
@@ -999,6 +1079,21 @@ async def __aiter__(self) -> AsyncIterator[bytes]:
             pass
 
 
+async def test_streaming_invalid_event_wrapped() -> None:
+    client, _ = make_client(
+        httpx.Response(
+            200,
+            content=b"data: {\n\n",
+            headers={"content-type": "text/event-stream"},
+        )
+    )
+
+    stream = client.get_response([Message("user", ["hi"])], stream=True)
+    with pytest.raises(ChatClientInvalidResponseException, match="response was invalid"):
+        async for _ in stream:
+            pass
+
+
 async def test_streaming_http_error_wrapped() -> None:
     client, _ = make_client(httpx.Response(429, json={"message": "rate limited"}))
 
diff --git a/python/packages/mistral/tests/mistral/test_mistral_embedding_client.py b/python/packages/mistral/tests/mistral/test_mistral_embedding_client.py
index be612f4fad8..c08a033b0c9 100644
--- a/python/packages/mistral/tests/mistral/test_mistral_embedding_client.py
+++ b/python/packages/mistral/tests/mistral/test_mistral_embedding_client.py
@@ -18,6 +18,9 @@
 from mistralai.client import Mistral
 
 from agent_framework_mistral import MistralEmbeddingClient, MistralEmbeddingOptions
+from agent_framework_mistral._http_client import (  # pyright: ignore[reportPrivateUsage]
+    AsyncClientUsingConfiguredTimeout,
+)
 
 # region: Unit Tests
 
@@ -40,8 +43,10 @@ class MockMistral:
     def __init__(self, responses: Sequence[httpx.Response]) -> None:
         self._responses = list(responses)
         self.requests: list[dict[str, Any]] = []
+        self.http_requests: list[httpx.Request] = []
 
     def handler(self, request: httpx.Request) -> httpx.Response:
+        self.http_requests.append(request)
         self.requests.append(json.loads(request.content))
         return self._responses.pop(0)
 
@@ -73,6 +78,7 @@ def test_mistral_embedding_construction_with_params() -> None:
     client = MistralEmbeddingClient(model="mistral-embed", api_key="test-key")
     assert client.model == "mistral-embed"
     assert isinstance(client.client, Mistral)
+    assert client.client.sdk_configuration.timeout_ms == 60_000
 
 
 def test_mistral_embedding_construction_with_server_url() -> None:
@@ -91,7 +97,9 @@ def test_mistral_embedding_construction_with_http_client() -> None:
     """Test construction with a pre-configured client."""
     http_client = httpx.AsyncClient(base_url="https://api.mistral.ai")
     client = MistralEmbeddingClient(model="mistral-embed", http_client=http_client)
-    assert client.client.sdk_configuration.async_client is http_client
+    async_client = client.client.sdk_configuration.async_client
+    assert isinstance(async_client, AsyncClientUsingConfiguredTimeout)
+    assert async_client.client is http_client
 
 
 def test_mistral_embedding_client_param_accepts_httpx() -> None:
@@ -101,7 +109,9 @@ def test_mistral_embedding_client_param_accepts_httpx() -> None:
             model="mistral-embed",
             client=cast("Any", http_client),
         )
-    assert client.client.sdk_configuration.async_client is http_client
+    async_client = client.client.sdk_configuration.async_client
+    assert isinstance(async_client, AsyncClientUsingConfiguredTimeout)
+    assert async_client.client is http_client
 
 
 class FakeMistralSDK:
@@ -242,6 +252,26 @@ async def test_mistral_embedding_get_embeddings() -> None:
     assert server.last_request == {"model": "mistral-embed", "input": ["hello", "world"]}
 
 
+async def test_mistral_embedding_preserves_injected_http_client_timeout() -> None:
+    server = MockMistral([httpx.Response(200, json=make_embeddings_payload([[0.1, 0.2]]))])
+    timeout = httpx.Timeout(connect=1, read=2, write=3, pool=4)
+    http_client = httpx.AsyncClient(
+        base_url="https://api.mistral.ai",
+        transport=httpx.MockTransport(server.handler),
+        timeout=timeout,
+    )
+    client = MistralEmbeddingClient(model="mistral-embed", http_client=http_client)
+
+    await client.get_embeddings(["hello"])
+
+    assert server.http_requests[0].extensions["timeout"] == {
+        "connect": 1,
+        "read": 2,
+        "write": 3,
+        "pool": 4,
+    }
+
+
 async def test_mistral_embedding_get_embeddings_empty_input() -> None:
     """Test generating embeddings with empty input."""
     client, server = make_client()