diff --git a/sentry_sdk/_types.py b/sentry_sdk/_types.py index fa6cf656df..2eb00fcbd9 100644 --- a/sentry_sdk/_types.py +++ b/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] diff --git a/sentry_sdk/integrations/mistral.py b/sentry_sdk/integrations/mistral.py index 042982fafc..bc3fe9ede0 100644 --- a/sentry_sdk/integrations/mistral.py +++ b/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) diff --git a/tests/integrations/mistral/test_mistral.py b/tests/integrations/mistral/test_mistral.py index 74fa37f251..f679825e2b 100644 --- a/tests/integrations/mistral/test_mistral.py +++ b/tests/integrations/mistral/test_mistral.py @@ -268,7 +268,7 @@ async def test_nonstreaming_chat_async( @pytest.mark.parametrize( - "messages,expected_system_instructions", + "messages,expected_system_instructions,expected_input_messages", ( ( [ @@ -276,6 +276,9 @@ async def test_nonstreaming_chat_async( content="You are a helpful math tutor. You will be provided with a math problem, and your goal will be to output a step by step solution, along with a final answer. For each step, just provide the output as an equation use the explanation field to detail the reasoning." ), UserMessage(content="How can I solve 8x + 7 = -23"), + AssistantMessage( + content="Subtract 7 from both sides to isolate the term with x ..." + ), ], [ { @@ -283,6 +286,26 @@ async def test_nonstreaming_chat_async( "content": "You are a helpful math tutor. You will be provided with a math problem, and your goal will be to output a step by step solution, along with a final answer. For each step, just provide the output as an equation use the explanation field to detail the reasoning.", } ], + [ + { + "role": "user", + "parts": [ + { + "type": "text", + "content": "How can I solve 8x + 7 = -23", + } + ], + }, + { + "role": "assistant", + "parts": [ + { + "type": "text", + "content": "Subtract 7 from both sides to isolate the term with x ...", + } + ], + }, + ], ), ( [ @@ -298,11 +321,45 @@ async def test_nonstreaming_chat_async( TextChunk(text="give the best 50"), ] ), + AssistantMessage( + content=[ + TextChunk(text="Camembert de Normandie"), + TextChunk(text="Brie de Meaux"), + ] + ), ], [ {"type": "text", "content": "You are a helpful assistant."}, {"type": "text", "content": "Be concise and clear."}, ], + [ + { + "role": "user", + "parts": [ + { + "type": "text", + "content": "What is the best French cheese?", + }, + { + "type": "text", + "content": "give the best 50", + }, + ], + }, + { + "role": "assistant", + "parts": [ + { + "type": "text", + "content": "Camembert de Normandie", + }, + { + "type": "text", + "content": "Brie de Meaux", + }, + ], + }, + ], ), ( [ @@ -311,6 +368,10 @@ async def test_nonstreaming_chat_async( "content": "You are a helpful math tutor. You will be provided with a math problem, and your goal will be to output a step by step solution, along with a final answer. For each step, just provide the output as an equation use the explanation field to detail the reasoning.", }, {"role": "user", "content": "How can I solve 8x + 7 = -23"}, + { + "role": "assistant", + "content": "Subtract 7 from both sides to isolate the term with x ...", + }, ], [ { @@ -318,6 +379,26 @@ async def test_nonstreaming_chat_async( "content": "You are a helpful math tutor. You will be provided with a math problem, and your goal will be to output a step by step solution, along with a final answer. For each step, just provide the output as an equation use the explanation field to detail the reasoning.", } ], + [ + { + "role": "user", + "parts": [ + { + "type": "text", + "content": "How can I solve 8x + 7 = -23", + } + ], + }, + { + "role": "assistant", + "parts": [ + { + "type": "text", + "content": "Subtract 7 from both sides to isolate the term with x ...", + } + ], + }, + ], ), ( [ @@ -335,11 +416,46 @@ async def test_nonstreaming_chat_async( {"type": "text", "text": "give the best 50"}, ], }, + { + "role": "assistant", + "content": [ + {"type": "text", "text": "Camembert de Normandie"}, + {"type": "text", "text": "Brie de Meaux"}, + ], + }, ], [ {"type": "text", "content": "You are a helpful assistant."}, {"type": "text", "content": "Be concise and clear."}, ], + [ + { + "role": "user", + "parts": [ + { + "type": "text", + "content": "What is the best French cheese?", + }, + { + "type": "text", + "content": "give the best 50", + }, + ], + }, + { + "role": "assistant", + "parts": [ + { + "type": "text", + "content": "Camembert de Normandie", + }, + { + "type": "text", + "content": "Brie de Meaux", + }, + ], + }, + ], ), ), ) @@ -353,6 +469,7 @@ def test_input_attributes_nonstreaming_chat( mistral_response, messages, expected_system_instructions, + expected_input_messages, data_collection, stream_gen_ai_spans, span_streaming, @@ -407,6 +524,10 @@ def test_input_attributes_nonstreaming_chat( json.loads(span["attributes"][SPANDATA.GEN_AI_SYSTEM_INSTRUCTIONS]) == expected_system_instructions ) + assert ( + json.loads(span["attributes"][SPANDATA.GEN_AI_INPUT_MESSAGES]) + == expected_input_messages + ) else: items = capture_items("transaction") @@ -428,10 +549,14 @@ def test_input_attributes_nonstreaming_chat( json.loads(span["data"][SPANDATA.GEN_AI_SYSTEM_INSTRUCTIONS]) == expected_system_instructions ) + assert ( + json.loads(span["data"][SPANDATA.GEN_AI_INPUT_MESSAGES]) + == expected_input_messages + ) @pytest.mark.parametrize( - "messages,expected_system_instructions", + "messages,expected_system_instructions,expected_input_messages", ( ( [ @@ -439,6 +564,9 @@ def test_input_attributes_nonstreaming_chat( content="You are a helpful math tutor. You will be provided with a math problem, and your goal will be to output a step by step solution, along with a final answer. For each step, just provide the output as an equation use the explanation field to detail the reasoning." ), UserMessage(content="How can I solve 8x + 7 = -23"), + AssistantMessage( + content="Subtract 7 from both sides to isolate the term with x ..." + ), ], [ { @@ -446,6 +574,26 @@ def test_input_attributes_nonstreaming_chat( "content": "You are a helpful math tutor. You will be provided with a math problem, and your goal will be to output a step by step solution, along with a final answer. For each step, just provide the output as an equation use the explanation field to detail the reasoning.", } ], + [ + { + "role": "user", + "parts": [ + { + "type": "text", + "content": "How can I solve 8x + 7 = -23", + } + ], + }, + { + "role": "assistant", + "parts": [ + { + "type": "text", + "content": "Subtract 7 from both sides to isolate the term with x ...", + } + ], + }, + ], ), ( [ @@ -461,11 +609,45 @@ def test_input_attributes_nonstreaming_chat( TextChunk(text="give the best 50"), ] ), + AssistantMessage( + content=[ + TextChunk(text="Camembert de Normandie"), + TextChunk(text="Brie de Meaux"), + ] + ), ], [ {"type": "text", "content": "You are a helpful assistant."}, {"type": "text", "content": "Be concise and clear."}, ], + [ + { + "role": "user", + "parts": [ + { + "type": "text", + "content": "What is the best French cheese?", + }, + { + "type": "text", + "content": "give the best 50", + }, + ], + }, + { + "role": "assistant", + "parts": [ + { + "type": "text", + "content": "Camembert de Normandie", + }, + { + "type": "text", + "content": "Brie de Meaux", + }, + ], + }, + ], ), ( [ @@ -474,6 +656,10 @@ def test_input_attributes_nonstreaming_chat( "content": "You are a helpful math tutor. You will be provided with a math problem, and your goal will be to output a step by step solution, along with a final answer. For each step, just provide the output as an equation use the explanation field to detail the reasoning.", }, {"role": "user", "content": "How can I solve 8x + 7 = -23"}, + { + "role": "assistant", + "content": "Subtract 7 from both sides to isolate the term with x ...", + }, ], [ { @@ -481,6 +667,26 @@ def test_input_attributes_nonstreaming_chat( "content": "You are a helpful math tutor. You will be provided with a math problem, and your goal will be to output a step by step solution, along with a final answer. For each step, just provide the output as an equation use the explanation field to detail the reasoning.", } ], + [ + { + "role": "user", + "parts": [ + { + "type": "text", + "content": "How can I solve 8x + 7 = -23", + } + ], + }, + { + "role": "assistant", + "parts": [ + { + "type": "text", + "content": "Subtract 7 from both sides to isolate the term with x ...", + } + ], + }, + ], ), ( [ @@ -498,11 +704,46 @@ def test_input_attributes_nonstreaming_chat( {"type": "text", "text": "give the best 50"}, ], }, + { + "role": "assistant", + "content": [ + {"type": "text", "text": "Camembert de Normandie"}, + {"type": "text", "text": "Brie de Meaux"}, + ], + }, ], [ {"type": "text", "content": "You are a helpful assistant."}, {"type": "text", "content": "Be concise and clear."}, ], + [ + { + "role": "user", + "parts": [ + { + "type": "text", + "content": "What is the best French cheese?", + }, + { + "type": "text", + "content": "give the best 50", + }, + ], + }, + { + "role": "assistant", + "parts": [ + { + "type": "text", + "content": "Camembert de Normandie", + }, + { + "type": "text", + "content": "Brie de Meaux", + }, + ], + }, + ], ), ), ) @@ -517,6 +758,7 @@ async def test_input_attributes_nonstreaming_chat_async( mistral_response, messages, expected_system_instructions, + expected_input_messages, data_collection, stream_gen_ai_spans, span_streaming, @@ -570,6 +812,10 @@ async def test_input_attributes_nonstreaming_chat_async( json.loads(span["attributes"][SPANDATA.GEN_AI_SYSTEM_INSTRUCTIONS]) == expected_system_instructions ) + assert ( + json.loads(span["attributes"][SPANDATA.GEN_AI_INPUT_MESSAGES]) + == expected_input_messages + ) else: items = capture_items("transaction") @@ -591,3 +837,7 @@ async def test_input_attributes_nonstreaming_chat_async( json.loads(span["data"][SPANDATA.GEN_AI_SYSTEM_INSTRUCTIONS]) == expected_system_instructions ) + assert ( + json.loads(span["data"][SPANDATA.GEN_AI_INPUT_MESSAGES]) + == expected_input_messages + )