feat(mistral): Add integration with Chat.complete and Chat.complete_async patches

.github/workflows/test-integrations-ai.yml

@@ -17,7 +17,7 @@ jobs:
     strategy:
       fail-fast: false
       matrix:
-        python-version: ["3.8","3.9","3.10","3.11","3.12","3.13","3.14","3.14t"]
+        python-version: ["3.8","3.9","3.10","3.11","3.12","3.13","3.14","3.14t","3.15","3.15t"]
         os: [ubuntu-22.04]
     env:
       # 3.6/3.7 run in the python:X.Y container; tell uv to use that system Python.
@@ -57,6 +57,10 @@ jobs:
         run: |
           set -x # print commands that are executed
           ./scripts/runtox.sh "py${{ matrix.python-version }}-litellm"
+      - name: Test mistral
+        run: |
+          set -x # print commands that are executed
+          ./scripts/runtox.sh "py${{ matrix.python-version }}-mistral"
       - name: Test openai-base
         run: |
           set -x # print commands that are executed

pyproject.toml

@@ -80,6 +80,7 @@ typing = [
     "tiktoken>=0.14.0",
     "aiohttp>=3.14.3",
     "openai-agents>=0.22.0",
+    "mistralai>=2.10.1",
 ]
 test = [
     "dataclasses ; python_full_version < '3.7'",

scripts/populate_tox/config.py

@@ -312,6 +312,12 @@
             "*": ["pytest-asyncio", "httpx"],
         },
     },
+    "mistral": {
+        "package": "mistralai",
+        "deps": {
+            "*": ["pytest-asyncio", "httpx"],
+        },
+    },
     "fastmcp": {
         "package": "fastmcp",
         "deps": {

scripts/split_tox_gh_actions/split_tox_gh_actions.py

@@ -92,6 +92,7 @@
         "google_genai",
         "huggingface_hub",
         "litellm",
+        "mistral",
         "openai-base",
         "openai-notiktoken",
     ],

sentry_sdk/integrations/__init__.py

@@ -149,6 +149,7 @@ def iter_default_integrations(
     "litellm": (1, 77, 5),
     "loguru": (0, 7, 0),
     "mcp": (1, 15, 0),
+    "mistral": (2, 0, 5),
     "openai": (1, 0, 0),
     "openai_agents": (0, 0, 19),
     "openfeature": (0, 7, 1),

sentry_sdk/integrations/mistral.py

@@ -0,0 +1,117 @@
+from functools import wraps
+from typing import TYPE_CHECKING
+
+import sentry_sdk
+from sentry_sdk.ai.utils import (
+    get_start_span_function,
+)
+from sentry_sdk.consts import OP, SPANDATA
+from sentry_sdk.integrations import DidNotEnable, Integration
+from sentry_sdk.tracing_utils import (
+    has_span_streaming_enabled,
+)
+
+if TYPE_CHECKING:
+    from typing import Any, Callable
+
+try:
+    from mistralai.client.chat import Chat
+except ImportError:
+    raise DidNotEnable("mistralai not installed")
+
+
+class MistralIntegration(Integration):
+    identifier = "mistral"
+    origin = f"auto.ai.{identifier}"
+
+    @staticmethod
+    def setup_once() -> None:
+        Chat.complete = _wrap_complete(Chat.complete)  # type: ignore[method-assign]
+
+        Chat.complete_async = _wrap_complete_async(Chat.complete_async)  # type: ignore[method-assign]
+
+
+def _wrap_complete(f: "Callable[..., Any]") -> "Callable[..., Any]":
+    @wraps(f)
+    def wrap_complete(self: "Chat", *args: "Any", **kwargs: "Any") -> "Any":
+        client = sentry_sdk.get_client()
+        integration = client.get_integration(MistralIntegration)
+        if integration is None or kwargs.get("stream"):
+            return f(self, *args, **kwargs)
+
+        model = kwargs.get("model")
+
+        if has_span_streaming_enabled(client.options):
+            span = sentry_sdk.traces.start_span(
+                name=f"chat {model}" if model is not None else "chat",
+                attributes={
+                    "sentry.op": OP.GEN_AI_CHAT,
+                    "sentry.origin": MistralIntegration.origin,
+                    SPANDATA.GEN_AI_PROVIDER_NAME: "mistral",
+                    SPANDATA.GEN_AI_OPERATION_NAME: "chat",
+                },
+            )
+
+            set_on_span = span.set_attribute
+        else:
+            span = get_start_span_function()(
+                op=OP.GEN_AI_CHAT,
+                name=f"chat {model}" if model is not None else "chat",
+                origin=MistralIntegration.origin,
+            )
+            span.set_data(SPANDATA.GEN_AI_PROVIDER_NAME, "mistral")
+            span.set_data(SPANDATA.GEN_AI_OPERATION_NAME, "chat")
+
+            set_on_span = span.set_data
+
+        with span:
+            if model is not None:
+                set_on_span(SPANDATA.GEN_AI_REQUEST_MODEL, model)
+
+            set_on_span(SPANDATA.GEN_AI_RESPONSE_STREAMING, False)
+            return f(self, *args, **kwargs)
+
+    return wrap_complete
+
+
+def _wrap_complete_async(f: "Callable[..., Any]") -> "Callable[..., Any]":
+    @wraps(f)
+    async def wrap_complete_async(self: "Chat", *args: "Any", **kwargs: "Any") -> "Any":
+        client = sentry_sdk.get_client()
+        integration = client.get_integration(MistralIntegration)
+        if integration is None or kwargs.get("stream"):
+            return await f(self, *args, **kwargs)
+
+        model = kwargs.get("model")
+
+        if has_span_streaming_enabled(client.options):
+            span = sentry_sdk.traces.start_span(
+                name=f"chat {model}" if model is not None else "chat",
+                attributes={
+                    "sentry.op": OP.GEN_AI_CHAT,
+                    "sentry.origin": MistralIntegration.origin,
+                    SPANDATA.GEN_AI_PROVIDER_NAME: "mistral",
+                    SPANDATA.GEN_AI_OPERATION_NAME: "chat",
+                },
+            )
+
+            set_on_span = span.set_attribute
+        else:
+            span = get_start_span_function()(
+                op=OP.GEN_AI_CHAT,
+                name=f"chat {model}" if model is not None else "chat",
+                origin=MistralIntegration.origin,
+            )
+            span.set_data(SPANDATA.GEN_AI_PROVIDER_NAME, "mistral")
+            span.set_data(SPANDATA.GEN_AI_OPERATION_NAME, "chat")
+
+            set_on_span = span.set_data
+
+        with span:
+            if model is not None:
+                set_on_span(SPANDATA.GEN_AI_REQUEST_MODEL, model)
+
+            set_on_span(SPANDATA.GEN_AI_RESPONSE_STREAMING, False)
+            return await f(self, *args, **kwargs)
+
+    return wrap_complete_async

tests/conftest.py

@@ -1281,6 +1281,7 @@ def inner(response_content, serialize_pydantic=False, request_headers=None):
             200,
             request=model_request,
             content=response_content,
+            headers={"Content-Type": "application/json"},
         )
 
         return response

tests/integrations/mistral/__init__.py

@@ -0,0 +1,3 @@
+import pytest
+
+pytest.importorskip("mistralai")

tests/integrations/mistral/test_mistral.py

@@ -0,0 +1,196 @@
+from unittest import mock
+
+import pytest
+from mistralai.client import Mistral
+from mistralai.client.chat import Chat
+from mistralai.client.models import (
+    AssistantMessage,
+    ChatCompletionChoice,
+    ChatCompletionResponse,
+    UsageInfo,
+)
+
+import sentry_sdk
+from sentry_sdk.consts import OP, SPANDATA
+from sentry_sdk.integrations.mistral import MistralIntegration
+
+
+@pytest.fixture
+def mistral_response():
+    return ChatCompletionResponse(
+        id="chat-id",
+        object="chat.completion",
+        model="mistral-medium-3-5",
+        created=10000000,
+        usage=UsageInfo(
+            prompt_tokens=10,
+            completion_tokens=20,
+            total_tokens=30,
+        ),
+        choices=[
+            ChatCompletionChoice(
+                index=0,
+                finish_reason="stop",
+                message=AssistantMessage(content="Hello, how can I help you?"),
+            )
+        ],
+    )
+
+
+@pytest.mark.parametrize("span_streaming", [True, False])
+@pytest.mark.parametrize("stream_gen_ai_spans", [True, False])
+def test_nonstreaming_chat(
+    sentry_init,
+    capture_items,
+    get_model_response,
+    mistral_response,
+    stream_gen_ai_spans,
+    span_streaming,
+):
+    sentry_init(
+        integrations=[MistralIntegration()],
+        traces_sample_rate=1.0,
+        stream_gen_ai_spans=stream_gen_ai_spans,
+        trace_lifecycle="stream" if span_streaming else "static",
+    )
+
+    client = Mistral(api_key="z")
+
+    model_response = get_model_response(
+        mistral_response,
+        serialize_pydantic=True,
+    )
+
+    if span_streaming or stream_gen_ai_spans:
+        items = capture_items("span")
+
+        with mock.patch.object(
+            Chat,
+            "do_request",
+            return_value=model_response,
+        ), sentry_sdk.start_transaction(name="mistral"):
+            client.chat.complete(
+                model="mistral-medium-latest",
+                messages=[
+                    {"role": "user", "content": "What is the best French cheese?"}
+                ],
+            )
+
+        sentry_sdk.flush()
+        spans = [item.payload for item in items]
+        (span,) = (
+            span
+            for span in spans
+            if span["attributes"].get("sentry.op") == OP.GEN_AI_CHAT
+        )
+
+        assert span["name"] == "chat mistral-medium-latest"
+        assert span["attributes"][SPANDATA.GEN_AI_PROVIDER_NAME] == "mistral"
+        assert span["attributes"][SPANDATA.GEN_AI_OPERATION_NAME] == "chat"
+
+        assert (
+            span["attributes"][SPANDATA.GEN_AI_REQUEST_MODEL] == "mistral-medium-latest"
+        )
+        assert span["attributes"][SPANDATA.GEN_AI_RESPONSE_STREAMING] is False
+    else:
+        items = capture_items("transaction")
+
+        with mock.patch.object(
+            Chat,
+            "do_request",
+            return_value=model_response,
+        ), sentry_sdk.start_transaction(name="mistral"):
+            client.chat.complete(
+                model="open-mistral",
+                messages=[{"role": "user", "content": "Hello, Mistral"}],
+            )
+
+        (transaction,) = [item.payload for item in items]
+        (span,) = transaction["spans"]
+
+        assert span["description"] == "chat open-mistral"
+        assert span["data"][SPANDATA.GEN_AI_PROVIDER_NAME] == "mistral"
+        assert span["data"][SPANDATA.GEN_AI_OPERATION_NAME] == "chat"
+
+        assert span["data"][SPANDATA.GEN_AI_REQUEST_MODEL] == "open-mistral"
+        assert span["data"][SPANDATA.GEN_AI_RESPONSE_STREAMING] is False
+
+
+@pytest.mark.asyncio
+@pytest.mark.parametrize("span_streaming", [True, False])
+@pytest.mark.parametrize("stream_gen_ai_spans", [True, False])
+async def test_nonstreaming_chat_async(
+    sentry_init,
+    capture_items,
+    get_model_response,
+    mistral_response,
+    stream_gen_ai_spans,
+    span_streaming,
+):
+    sentry_init(
+        integrations=[MistralIntegration()],
+        traces_sample_rate=1.0,
+        stream_gen_ai_spans=stream_gen_ai_spans,
+        trace_lifecycle="stream" if span_streaming else "static",
+    )
+
+    client = Mistral(api_key="z")
+
+    model_response = get_model_response(
+        mistral_response,
+        serialize_pydantic=True,
+    )
+
+    if span_streaming or stream_gen_ai_spans:
+        items = capture_items("span")
+
+        with mock.patch.object(
+            Chat,
+            "do_request_async",
+            return_value=model_response,
+        ), sentry_sdk.start_transaction(name="mistral"):
+            await client.chat.complete_async(
+                model="mistral-medium-latest",
+                messages=[
+                    {"role": "user", "content": "What is the best French cheese?"}
+                ],
+            )
+
+        sentry_sdk.flush()
+        spans = [item.payload for item in items]
+        (span,) = (
+            span
+            for span in spans
+            if span["attributes"].get("sentry.op") == OP.GEN_AI_CHAT
+        )
+
+        assert span["name"] == "chat mistral-medium-latest"
+        assert span["attributes"][SPANDATA.GEN_AI_PROVIDER_NAME] == "mistral"
+        assert span["attributes"][SPANDATA.GEN_AI_OPERATION_NAME] == "chat"
+
+        assert (
+            span["attributes"][SPANDATA.GEN_AI_REQUEST_MODEL] == "mistral-medium-latest"
+        )
+        assert span["attributes"][SPANDATA.GEN_AI_RESPONSE_STREAMING] is False
+    else:
+        items = capture_items("transaction")
+
+        with mock.patch.object(
+            Chat,
+            "do_request_async",
+            return_value=model_response,
+        ), sentry_sdk.start_transaction(name="mistral"):
+            await client.chat.complete_async(
+                model="mistral-medium-latest",
+                messages=[{"role": "user", "content": "Hello, Mistral"}],
+            )
+
+        (transaction,) = [item.payload for item in items]
+        (span,) = transaction["spans"]
+
+        assert span["description"] == "chat mistral-medium-latest"
+        assert span["data"][SPANDATA.GEN_AI_PROVIDER_NAME] == "mistral"
+        assert span["data"][SPANDATA.GEN_AI_OPERATION_NAME] == "chat"
+
+        assert span["data"][SPANDATA.GEN_AI_REQUEST_MODEL] == "mistral-medium-latest"
+        assert span["data"][SPANDATA.GEN_AI_RESPONSE_STREAMING] is False