feat(mistral): Record gen_ai.input.messages

sentry_sdk/_types.py

@@ -481,6 +481,10 @@ class ToolCallPart(TypedDict):
         name: NotRequired[str]
         arguments: NotRequired[Any]
 
+    class InputMessage(TypedDict):
+        role: Literal["user", "assistant", "tool"]
+        parts: list[Union[TextPart, ReasoningPart, ToolCallPart]]
+
     class ToolDefinition(TypedDict):
         type: str
         name: NotRequired[str]

sentry_sdk/integrations/mistral.py

@@ -25,14 +25,16 @@
         TextChunkTypedDict,
     )
 
-    from sentry_sdk._types import TextPart
+    from sentry_sdk._types import InputMessage, TextPart
 
 try:
     from mistralai.client.chat import Chat
     from mistralai.client.models import (
+        AssistantMessage,
         ChatCompletionResponse,
         SystemMessage,
         TextChunk,
+        UserMessage,
     )
 except ImportError:
     raise DidNotEnable("mistralai not installed")
@@ -61,6 +63,91 @@ def _is_system_instruction(
     return False
 
 
+def _transform_input_messages(
+    messages: "Sequence[Union[ChatCompletionRequestMessage, ChatCompletionRequestMessageTypedDict]]",
+) -> "list[InputMessage]":
+    input_messages: "list[InputMessage]" = []
+
+    for message in messages:
+        if isinstance(message, UserMessage) and isinstance(message.content, str):
+            input_messages.append(
+                {
+                    "role": "user",
+                    "parts": [{"type": "text", "content": message.content}],
+                }
+            )
+        elif isinstance(message, UserMessage) and isinstance(message.content, list):
+            text_parts = [
+                part for part in message.content if isinstance(part, TextChunk)
+            ]
+            input_messages.append(
+                {
+                    "role": "user",
+                    "parts": [
+                        {"type": "text", "content": part.text} for part in text_parts
+                    ],
+                }
+            )
+
+        if isinstance(message, AssistantMessage) and isinstance(message.content, str):
+            input_messages.append(
+                {
+                    "role": "assistant",
+                    "parts": [{"type": "text", "content": message.content}],
+                }
+            )
+        elif isinstance(message, AssistantMessage) and isinstance(
+            message.content, list
+        ):
+            text_parts = [
+                part for part in message.content if isinstance(part, TextChunk)
+            ]
+            input_messages.append(
+                {
+                    "role": "assistant",
+                    "parts": [
+                        {"type": "text", "content": part.text} for part in text_parts
+                    ],
+                }
+            )
+
+        if not isinstance(message, dict):
+            continue
+
+        role = message.get("role")
+        if role != "user" and role != "assistant":
+            continue
+
+        content = message.get("content")
+        if isinstance(content, str):
+            input_messages.append(
+                {"role": role, "parts": [{"type": "text", "content": content}]}
+            )
+
+        if not isinstance(content, list):
+            continue
+
+        text_parts = [
+            part
+            for part in content
+            if isinstance(part, dict) and part.get("type") == "text" and "text" in part
+        ]
+        input_messages.append(
+            {
+                "role": role,
+                "parts": [
+                    {
+                        "type": "text",
+                        "content": cast("TextChunkTypedDict", part)["text"],
+                    }
+                    for part in text_parts
+                ],
+            }
+        )
+
+    return input_messages
+
+
 def _transform_system_instructions(
     messages: "list[Union[SystemMessage, SystemMessageTypedDict]]",
 ) -> "list[TextPart]":
@@ -157,6 +244,11 @@ def wrap_complete(self: "Chat", *args: "Any", **kwargs: "Any") -> "Any":
                         json.dumps(_transform_system_instructions(system_instructions)),
                     )
 
+                set_on_span(
+                    SPANDATA.GEN_AI_INPUT_MESSAGES,
+                    json.dumps(_transform_input_messages(messages)),
+                )
+
             max_tokens = kwargs.get("max_tokens")
             if max_tokens is not None:
                 set_on_span(SPANDATA.GEN_AI_REQUEST_MAX_TOKENS, max_tokens)
@@ -270,6 +362,11 @@ async def wrap_complete_async(self: "Chat", *args: "Any", **kwargs: "Any") -> "A
                         json.dumps(_transform_system_instructions(system_instructions)),
                     )
 
+                set_on_span(
+                    SPANDATA.GEN_AI_INPUT_MESSAGES,
+                    json.dumps(_transform_input_messages(messages)),
+                )
+
             max_tokens = kwargs.get("max_tokens")
             if max_tokens is not None:
                 set_on_span(SPANDATA.GEN_AI_REQUEST_MAX_TOKENS, max_tokens)