|
85 | 85 | ToolUnionParam, |
86 | 86 | ) |
87 | 87 |
|
88 | | - from sentry_sdk._types import TextPart |
| 88 | + from sentry_sdk._types import ReasoningPart, TextPart |
89 | 89 |
|
90 | 90 | class _PatchedRawMessageStream(Stream[Any]): |
91 | 91 | _span: Union[Span, StreamedSpan] |
@@ -731,30 +731,30 @@ def _set_output_data( |
731 | 731 | record_outputs = True |
732 | 732 |
|
733 | 733 | if record_outputs: |
734 | | - output_messages: "dict[str, list[Any]]" = { |
735 | | - "response": [], |
736 | | - "tool": [], |
737 | | - } |
| 734 | + parts: "list[Union[TextPart, ReasoningPart]]" = [] |
| 735 | + tool_calls = [] |
738 | 736 |
|
739 | 737 | for output in content_blocks: |
740 | 738 | if output["type"] == "text": |
741 | | - output_messages["response"].append(output["text"]) |
| 739 | + parts.append({"type": "text", "content": output["text"]}) |
| 740 | + elif output["type"] == "thinking": |
| 741 | + parts.append({"type": "reasoning", "content": output.get("thinking")}) |
742 | 742 | elif output["type"] == "tool_use": |
743 | | - output_messages["tool"].append(output) |
| 743 | + tool_calls.append(output) |
744 | 744 |
|
745 | | - if len(output_messages["tool"]) > 0: |
| 745 | + if len(tool_calls) > 0: |
746 | 746 | set_data_normalized( |
747 | 747 | span, |
748 | 748 | SPANDATA.GEN_AI_RESPONSE_TOOL_CALLS, |
749 | | - output_messages["tool"], |
| 749 | + tool_calls, |
750 | 750 | unpack=False, |
751 | 751 | ) |
752 | 752 |
|
753 | | - if len(output_messages["response"]) > 0: |
| 753 | + if len(parts) > 0: |
754 | 754 | set_data_normalized( |
755 | 755 | span, |
756 | | - SPANDATA.GEN_AI_RESPONSE_TEXT, |
757 | | - output_messages["response"], |
| 756 | + SPANDATA.GEN_AI_OUTPUT_MESSAGES, |
| 757 | + [{"role": "assistant", "parts": parts}], |
758 | 758 | unpack=False, |
759 | 759 | ) |
760 | 760 |
|
|
0 commit comments