diff --git a/sentry_sdk/_types.py b/sentry_sdk/_types.py index 2eb00fcbd9..39e0ca43d7 100644 --- a/sentry_sdk/_types.py +++ b/sentry_sdk/_types.py @@ -485,6 +485,10 @@ class InputMessage(TypedDict): role: Literal["user", "assistant", "tool"] parts: list[Union[TextPart, ReasoningPart, ToolCallPart]] + class OutputMessage(TypedDict): + role: Literal["assistant"] + parts: list[Union[TextPart, ReasoningPart]] + class ToolDefinition(TypedDict): type: str name: NotRequired[str] diff --git a/sentry_sdk/integrations/mistral.py b/sentry_sdk/integrations/mistral.py index bc3fe9ede0..1ef4d42259 100644 --- a/sentry_sdk/integrations/mistral.py +++ b/sentry_sdk/integrations/mistral.py @@ -19,13 +19,14 @@ from typing import Any, Callable, Iterable, Optional, TypeGuard, Union from mistralai.client.models import ( + AssistantMessageTypedDict, ChatCompletionRequestMessage, ChatCompletionRequestMessageTypedDict, SystemMessageTypedDict, TextChunkTypedDict, ) - from sentry_sdk._types import InputMessage, TextPart + from sentry_sdk._types import InputMessage, OutputMessage, TextPart try: from mistralai.client.chat import Chat @@ -148,6 +149,46 @@ def _transform_input_messages( return input_messages +def _transform_output_message( + message: "Union[AssistantMessage, AssistantMessageTypedDict]", +) -> "Optional[OutputMessage]": + if isinstance(message, AssistantMessage): + if message.content is None: + return None + + if isinstance(message.content, str): + return { + "role": "assistant", + "parts": [{"type": "text", "content": message.content}], + } + + parts = [part for part in message.content if isinstance(part, TextChunk)] + return { + "role": "assistant", + "parts": [{"type": "text", "content": part.text} for part in parts], + } + + content = message.get("content") + if content is None: + return None + + if isinstance(content, str): + return {"role": "assistant", "parts": [{"type": "text", "content": content}]} + + text_parts = [ + part + for part in content + if isinstance(part, dict) and part.get("type") == "text" + ] + return { + "role": "assistant", + "parts": [ + {"type": "text", "content": cast("TextChunkTypedDict", part)["text"]} + for part in text_parts + ], + } + + def _transform_system_instructions( messages: "list[Union[SystemMessage, SystemMessageTypedDict]]", ) -> "list[TextPart]": @@ -296,6 +337,29 @@ def wrap_complete(self: "Chat", *args: "Any", **kwargs: "Any") -> "Any": SPANDATA.GEN_AI_USAGE_TOTAL_TOKENS, response.usage.total_tokens ) + if ( + has_data_collection_enabled(client.options) + and client.options["data_collection"]["gen_ai"]["outputs"] + ) or ( + not has_data_collection_enabled(client.options) + and should_send_default_pii() + ): + output_messages: "list[OutputMessage]" = [] + for choice in response.choices: + if choice.message is None: + continue + + transformed_message = _transform_output_message(choice.message) + if transformed_message is None: + continue + + output_messages.append(transformed_message) + + set_on_span( + SPANDATA.GEN_AI_OUTPUT_MESSAGES, + json.dumps(output_messages), + ) + return response return wrap_complete @@ -414,6 +478,29 @@ async def wrap_complete_async(self: "Chat", *args: "Any", **kwargs: "Any") -> "A SPANDATA.GEN_AI_USAGE_TOTAL_TOKENS, response.usage.total_tokens ) + if ( + has_data_collection_enabled(client.options) + and client.options["data_collection"]["gen_ai"]["outputs"] + ) or ( + not has_data_collection_enabled(client.options) + and should_send_default_pii() + ): + output_messages: "list[OutputMessage]" = [] + for choice in response.choices: + if choice.message is None: + continue + + transformed_message = _transform_output_message(choice.message) + if transformed_message is None: + continue + + output_messages.append(transformed_message) + + set_on_span( + SPANDATA.GEN_AI_OUTPUT_MESSAGES, + json.dumps(output_messages), + ) + return response return wrap_complete_async diff --git a/tests/integrations/mistral/test_mistral.py b/tests/integrations/mistral/test_mistral.py index f679825e2b..0854683f4d 100644 --- a/tests/integrations/mistral/test_mistral.py +++ b/tests/integrations/mistral/test_mistral.py @@ -36,7 +36,17 @@ def mistral_response(): index=0, finish_reason="stop", message=AssistantMessage(content="Hello, how can I help you?"), - ) + ), + ChatCompletionChoice( + index=1, + finish_reason="stop", + message=AssistantMessage( + content=[ + TextChunk(text="Response 1"), + TextChunk(text="Response 2"), + ] + ), + ), ], ) @@ -841,3 +851,222 @@ async def test_input_attributes_nonstreaming_chat_async( json.loads(span["data"][SPANDATA.GEN_AI_INPUT_MESSAGES]) == expected_input_messages ) + + +@pytest.mark.parametrize("data_collection", [True, False]) +@pytest.mark.parametrize("span_streaming", [True, False]) +@pytest.mark.parametrize("stream_gen_ai_spans", [True, False]) +def test_output_attributes_nonstreaming_chat( + sentry_init, + capture_items, + get_model_response, + mistral_response, + data_collection, + stream_gen_ai_spans, + span_streaming, +): + if data_collection: + 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", + _experiments={"data_collection": {}}, + ) + else: + 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", + send_default_pii=True, + ) + + 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 json.loads(span["attributes"][SPANDATA.GEN_AI_OUTPUT_MESSAGES]) == [ + { + "role": "assistant", + "parts": [ + {"type": "text", "content": "Hello, how can I help you?"}, + ], + }, + { + "role": "assistant", + "parts": [ + {"type": "text", "content": "Response 1"}, + {"type": "text", "content": "Response 2"}, + ], + }, + ] + 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 json.loads(span["data"][SPANDATA.GEN_AI_OUTPUT_MESSAGES]) == [ + { + "role": "assistant", + "parts": [ + {"type": "text", "content": "Hello, how can I help you?"}, + ], + }, + { + "role": "assistant", + "parts": [ + {"type": "text", "content": "Response 1"}, + {"type": "text", "content": "Response 2"}, + ], + }, + ] + + +@pytest.mark.asyncio +@pytest.mark.parametrize("data_collection", [True, False]) +@pytest.mark.parametrize("span_streaming", [True, False]) +@pytest.mark.parametrize("stream_gen_ai_spans", [True, False]) +async def test_output_attributes_nonstreaming_chat_async( + sentry_init, + capture_items, + get_model_response, + mistral_response, + data_collection, + stream_gen_ai_spans, + span_streaming, +): + if data_collection: + 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", + _experiments={"data_collection": {}}, + ) + else: + 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", + send_default_pii=True, + ) + + 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 json.loads(span["attributes"][SPANDATA.GEN_AI_OUTPUT_MESSAGES]) == [ + { + "role": "assistant", + "parts": [ + {"type": "text", "content": "Hello, how can I help you?"}, + ], + }, + { + "role": "assistant", + "parts": [ + {"type": "text", "content": "Response 1"}, + {"type": "text", "content": "Response 2"}, + ], + }, + ] + 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 json.loads(span["data"][SPANDATA.GEN_AI_OUTPUT_MESSAGES]) == [ + { + "role": "assistant", + "parts": [ + {"type": "text", "content": "Hello, how can I help you?"}, + ], + }, + { + "role": "assistant", + "parts": [ + {"type": "text", "content": "Response 1"}, + {"type": "text", "content": "Response 2"}, + ], + }, + ]