Python: migrate Mistral integration to official SDK

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
 

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,8 +26,11 @@
     IntegrationInvalidResponseException,
 )
 from agent_framework.observability import EmbeddingTelemetryLayer
+from mistralai.client import Mistral
+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
@@ -38,37 +41,7 @@
 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__}."
-    )
+_DEFAULT_TIMEOUT_MS = 60_000
 
 
 class MistralEmbeddingOptions(EmbeddingGenerationOptions, total=False):
@@ -105,7 +78,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 +89,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 +99,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 +115,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 +153,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] = {"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"] = 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
+            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 +232,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 +280,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 +295,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 +332,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,

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

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]