ai-workflows/openinference-reasoning.patch at main · packit/ai-workflows · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
diff --git a/python/instrumentation/openinference-instrumentation-beeai/src/openinference/instrumentation/beeai/processors/chat.py b/python/instrumentation/openinference-instrumentation-beeai/src/openinference/instrumentation/beeai/processors/chat.py
index a135b4b7..a8a7a0f4 100644
--- a/python/instrumentation/openinference-instrumentation-beeai/src/openinference/instrumentation/beeai/processors/chat.py
+++ b/python/instrumentation/openinference-instrumentation-beeai/src/openinference/instrumentation/beeai/processors/chat.py
@@ -5,6 +5,7 @@ from beeai_framework.backend import (
     AnyMessage,
     ChatModel,
     MessageImageContent,
+    MessageReasoningContent,
     MessageTextContent,
     MessageToolCallContent,
     MessageToolResultContent,
@@ -218,6 +219,11 @@ def _process_messages(
                                 ),
                             }
                             if isinstance(content, MessageToolResultContent)
+                            else {
+                                MessageContentAttributes.MESSAGE_CONTENT_TYPE: "reasoning",
+                                MessageContentAttributes.MESSAGE_CONTENT_TEXT: content.text,
+                            }
+                            if isinstance(content, MessageReasoningContent)
                             else None
                         )
                         for content in msg.content
@@ -244,7 +250,7 @@ def _process_messages(


 def _aggregate_msg_content(message: "AnyMessage") -> None:
-    from beeai_framework.backend import MessageTextContent, MessageToolCallContent
+    from beeai_framework.backend import MessageReasoningContent, MessageTextContent, MessageToolCallContent

     contents = message.content.copy()
     aggregated_content: list[Any] = []
@@ -257,6 +263,10 @@ def _aggregate_msg_content(message: "AnyMessage") -> None:
             content, MessageToolCallContent
         ):
             last_content.args += content.args
+        elif isinstance(last_content, MessageReasoningContent) and isinstance(
+            content, MessageReasoningContent
+        ):
+            last_content.text += content.text
         else:
             aggregated_content.append(content)