Skip to content

Commit f3d1d5f

Browse files
nikhilNavanikhilc-microsoft
nikhilNava
and
nikhilc-microsoft
authored
update mapper for open ai (#236)
Co-authored-by: Nikhil Chitlur Navakiran (from Dev Box) <nikhilc@microsoft.com>
1 parent cd747b5 commit f3d1d5f

5 files changed

Lines changed: 952 additions & 1 deletion

File tree

‎libraries/microsoft-agents-a365-observability-extensions-openai/microsoft_agents_a365/observability/extensions/openai/message_mapper.py‎

Lines changed: 381 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,381 @@
1+
# Copyright (c) Microsoft Corporation.
2+
# Licensed under the MIT License.
3+
4+
"""Maps OpenAI span tag messages to A365 versioned message format.
5+
6+
Handles three input shapes produced by the OpenAI trace processor:
7+
8+
1. **Chat-completions format** (from ``GenerationSpanData``):
9+
``[{"role":"system","content":"..."}, ...]``
10+
2. **Response API format** (from ``ResponseSpanData``):
11+
- Input: ``[{"type":"message","role":"user","content":"..."}, ...]``
12+
- Output: ``{"id":"...","model":"...","output":[...], ...}`` (full Response JSON)
13+
3. **Plain string** (from ``AgentSpanData``):
14+
A bare user/assistant message captured from child generation spans.
15+
"""
16+
17+
from __future__ import annotations
18+
19+
import json
20+
import logging
21+
from collections.abc import Mapping
22+
from typing import Any
23+
24+
from microsoft_agents_a365.observability.core.message_utils import serialize_messages
25+
from microsoft_agents_a365.observability.core.models.messages import (
26+
ChatMessage,
27+
InputMessages,
28+
MessagePart,
29+
MessageRole,
30+
OutputMessage,
31+
OutputMessages,
32+
TextPart,
33+
ToolCallRequestPart,
34+
ToolCallResponsePart,
35+
)
36+
37+
logger = logging.getLogger(__name__)
38+
39+
_ROLE_MAP: dict[str, MessageRole] = {
40+
"system": MessageRole.SYSTEM,
41+
"user": MessageRole.USER,
42+
"assistant": MessageRole.ASSISTANT,
43+
"tool": MessageRole.TOOL,
44+
}
45+
46+
47+
# ---------------------------------------------------------------------------
48+
# Public API
49+
# ---------------------------------------------------------------------------
50+
51+
52+
def map_input_messages(messages_json: str) -> str | None:
53+
"""Map a ``gen_ai.input.messages`` tag value to a serialized A365 JSON string.
54+
55+
Args:
56+
messages_json: The raw JSON string from the span attribute.
57+
58+
Returns:
59+
Serialized :class:`InputMessages` JSON string, or ``None`` if the
60+
input is empty or cannot be parsed.
61+
"""
62+
if not messages_json:
63+
return None
64+
65+
# Plain string (AgentSpanData captures bare user text)
66+
try:
67+
raw = json.loads(messages_json)
68+
except (json.JSONDecodeError, TypeError):
69+
return _wrap_plain_input(messages_json)
70+
71+
if isinstance(raw, list):
72+
return _map_input_list(raw)
73+
74+
# Unexpected shape
75+
return _wrap_plain_input(messages_json)
76+
77+
78+
def map_output_messages(messages_json: str) -> str | None:
79+
"""Map a ``gen_ai.output.messages`` tag value to a serialized A365 JSON string.
80+
81+
Args:
82+
messages_json: The raw JSON string from the span attribute.
83+
84+
Returns:
85+
Serialized :class:`OutputMessages` JSON string, or ``None`` if the
86+
input is empty or cannot be parsed.
87+
"""
88+
if not messages_json:
89+
return None
90+
91+
try:
92+
raw = json.loads(messages_json)
93+
except (json.JSONDecodeError, TypeError):
94+
return _wrap_plain_output(messages_json)
95+
96+
if isinstance(raw, list):
97+
return _map_output_list(raw)
98+
99+
if isinstance(raw, dict):
100+
# Full Response JSON from ResponseSpanData (model_dump_json)
101+
return _map_response_output(raw)
102+
103+
return _wrap_plain_output(messages_json)
104+
105+
106+
# ---------------------------------------------------------------------------
107+
# Input mapping
108+
# ---------------------------------------------------------------------------
109+
110+
111+
def _map_input_list(items: list[Any]) -> str | None:
112+
"""Map a list of input items (chat completions or ResponseInputItemParam)."""
113+
chat_messages: list[ChatMessage] = []
114+
115+
for item in items:
116+
if not isinstance(item, dict):
117+
continue
118+
119+
item_type = item.get("type")
120+
121+
if item_type == "function_call":
122+
# ResponseInputItemParam: function_call → assistant tool call request
123+
name = item.get("name", "")
124+
if name:
125+
parts: list[MessagePart] = [
126+
ToolCallRequestPart(
127+
name=name,
128+
id=item.get("call_id"),
129+
arguments=item.get("arguments"),
130+
)
131+
]
132+
chat_messages.append(ChatMessage(role=MessageRole.ASSISTANT, parts=parts))
133+
134+
elif item_type == "function_call_output":
135+
# ResponseInputItemParam: function_call_output → tool response
136+
parts = [
137+
ToolCallResponsePart(
138+
id=item.get("call_id"),
139+
response=item.get("output"),
140+
)
141+
]
142+
chat_messages.append(ChatMessage(role=MessageRole.TOOL, parts=parts))
143+
144+
elif item_type == "custom_tool_call":
145+
name = item.get("name", "")
146+
if name:
147+
input_data = item.get("input")
148+
args = json.dumps({"input": input_data}) if input_data is not None else None
149+
parts = [ToolCallRequestPart(name=name, id=item.get("call_id"), arguments=args)]
150+
chat_messages.append(ChatMessage(role=MessageRole.ASSISTANT, parts=parts))
151+
152+
elif item_type == "custom_tool_call_output":
153+
parts = [
154+
ToolCallResponsePart(
155+
id=item.get("call_id"),
156+
response=item.get("output"),
157+
)
158+
]
159+
chat_messages.append(ChatMessage(role=MessageRole.TOOL, parts=parts))
160+
161+
elif item_type == "message" or "role" in item:
162+
# Standard message (ResponseInputItemParam or chat completions)
163+
mapped = _map_chat_completions_message(item)
164+
if mapped is not None:
165+
chat_messages.append(mapped)
166+
167+
else:
168+
# Unknown type, try as generic message
169+
mapped = _map_chat_completions_message(item)
170+
if mapped is not None:
171+
chat_messages.append(mapped)
172+
173+
if not chat_messages:
174+
return None
175+
return serialize_messages(InputMessages(messages=chat_messages))
176+
177+
178+
def _map_chat_completions_message(msg: dict[str, Any]) -> ChatMessage | None:
179+
"""Map a single chat-completions-style message dict."""
180+
role_str = msg.get("role", "")
181+
role = _ROLE_MAP.get(str(role_str).lower(), MessageRole.USER)
182+
parts: list[MessagePart] = []
183+
184+
# Tool response message
185+
if role == MessageRole.TOOL:
186+
content = msg.get("content", "")
187+
tool_call_id = msg.get("tool_call_id")
188+
response = str(content) if content else ""
189+
if response or tool_call_id:
190+
parts.append(ToolCallResponsePart(id=tool_call_id, response=response))
191+
return ChatMessage(role=role, parts=parts) if parts else None
192+
193+
# Text content (string or list)
194+
content = msg.get("content")
195+
if isinstance(content, str) and content.strip():
196+
parts.append(TextPart(content=content))
197+
elif isinstance(content, list):
198+
for item in content:
199+
if isinstance(item, dict):
200+
if item.get("type") in ("input_text", "text"):
201+
text = item.get("text", "")
202+
if text:
203+
parts.append(TextPart(content=text))
204+
elif item.get("type") == "output_text":
205+
text = item.get("text", "")
206+
if text:
207+
parts.append(TextPart(content=text))
208+
209+
# Tool calls
210+
tool_calls = msg.get("tool_calls")
211+
if isinstance(tool_calls, list):
212+
for tc in tool_calls:
213+
if not isinstance(tc, dict):
214+
continue
215+
func = tc.get("function", {})
216+
if isinstance(func, dict):
217+
name = func.get("name")
218+
if name:
219+
parts.append(
220+
ToolCallRequestPart(
221+
name=name,
222+
id=tc.get("id"),
223+
arguments=func.get("arguments"),
224+
)
225+
)
226+
227+
if not parts:
228+
return None
229+
return ChatMessage(role=role, parts=parts, name=msg.get("name"))
230+
231+
232+
# ---------------------------------------------------------------------------
233+
# Output mapping
234+
# ---------------------------------------------------------------------------
235+
236+
237+
def _map_output_list(items: list[Any]) -> str | None:
238+
"""Map a list of chat-completions-style output messages."""
239+
output_messages: list[OutputMessage] = []
240+
241+
for item in items:
242+
if not isinstance(item, dict):
243+
continue
244+
role_str = item.get("role", "assistant")
245+
role = _ROLE_MAP.get(str(role_str).lower(), MessageRole.ASSISTANT)
246+
parts: list[MessagePart] = []
247+
248+
# Tool response
249+
if role == MessageRole.TOOL:
250+
content = item.get("content", "")
251+
tool_call_id = item.get("tool_call_id")
252+
response = str(content) if content else ""
253+
if response or tool_call_id:
254+
parts.append(ToolCallResponsePart(id=tool_call_id, response=response))
255+
else:
256+
# Text content
257+
content = item.get("content")
258+
if isinstance(content, str) and content.strip():
259+
parts.append(TextPart(content=content))
260+
elif isinstance(content, list):
261+
for c in content:
262+
if isinstance(c, dict):
263+
text = c.get("text", "")
264+
if text:
265+
parts.append(TextPart(content=text))
266+
267+
# Tool calls
268+
tool_calls = item.get("tool_calls")
269+
if isinstance(tool_calls, list):
270+
for tc in tool_calls:
271+
if not isinstance(tc, dict):
272+
continue
273+
func = tc.get("function", {})
274+
if isinstance(func, dict):
275+
name = func.get("name")
276+
if name:
277+
parts.append(
278+
ToolCallRequestPart(
279+
name=name,
280+
id=tc.get("id"),
281+
arguments=func.get("arguments"),
282+
)
283+
)
284+
285+
finish_reason = item.get("finish_reason")
286+
if parts:
287+
output_messages.append(
288+
OutputMessage(role=role, parts=parts, finish_reason=finish_reason)
289+
)
290+
291+
if not output_messages:
292+
return None
293+
return serialize_messages(OutputMessages(messages=output_messages))
294+
295+
296+
def _map_response_output(response: dict[str, Any]) -> str | None:
297+
"""Map a full OpenAI Response JSON to A365 OutputMessages.
298+
299+
The Response object has ``output: [...]`` containing items with
300+
``type`` of ``message`` or ``function_call``.
301+
"""
302+
output_items = response.get("output")
303+
if not isinstance(output_items, list):
304+
return None
305+
306+
output_messages: list[OutputMessage] = []
307+
308+
for item in output_items:
309+
if not isinstance(item, Mapping):
310+
continue
311+
item_type = item.get("type")
312+
313+
if item_type == "message":
314+
parts: list[MessagePart] = []
315+
role_str = item.get("role", "assistant")
316+
role = _ROLE_MAP.get(str(role_str).lower(), MessageRole.ASSISTANT)
317+
318+
for content_item in item.get("content", []):
319+
if isinstance(content_item, Mapping):
320+
content_type = content_item.get("type")
321+
if content_type == "output_text":
322+
text = content_item.get("text", "")
323+
if text:
324+
parts.append(TextPart(content=text))
325+
elif content_type == "refusal":
326+
text = content_item.get("refusal", "")
327+
if text:
328+
parts.append(TextPart(content=text))
329+
330+
if parts:
331+
finish_reason = item.get("status")
332+
output_messages.append(
333+
OutputMessage(role=role, parts=parts, finish_reason=finish_reason)
334+
)
335+
336+
elif item_type == "function_call":
337+
name = item.get("name", "")
338+
if name:
339+
parts = [
340+
ToolCallRequestPart(
341+
name=name,
342+
id=item.get("call_id"),
343+
arguments=item.get("arguments"),
344+
)
345+
]
346+
output_messages.append(
347+
OutputMessage(
348+
role=MessageRole.ASSISTANT,
349+
parts=parts,
350+
finish_reason="tool_call",
351+
)
352+
)
353+
354+
if not output_messages:
355+
return None
356+
return serialize_messages(OutputMessages(messages=output_messages))
357+
358+
359+
# ---------------------------------------------------------------------------
360+
# Plain-string wrappers
361+
# ---------------------------------------------------------------------------
362+
363+
364+
def _wrap_plain_input(text: str) -> str | None:
365+
"""Wrap a plain text string as a versioned InputMessages."""
366+
if not text or not text.strip():
367+
return None
368+
return serialize_messages(
369+
InputMessages(messages=[ChatMessage(role=MessageRole.USER, parts=[TextPart(content=text)])])
370+
)
371+
372+
373+
def _wrap_plain_output(text: str) -> str | None:
374+
"""Wrap a plain text string as a versioned OutputMessages."""
375+
if not text or not text.strip():
376+
return None
377+
return serialize_messages(
378+
OutputMessages(
379+
messages=[OutputMessage(role=MessageRole.ASSISTANT, parts=[TextPart(content=text)])]
380+
)
381+
)

‎libraries/microsoft-agents-a365-observability-extensions-openai/microsoft_agents_a365/observability/extensions/openai/trace_instrumentor.py‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -66,4 +66,4 @@ def _instrument(self, **kwargs: Any) -> None:
6666
set_trace_processors([OpenAIAgentsTraceProcessor(agent365_tracer)])
6767

6868
def _uninstrument(self, **kwargs: Any) -> None:
69-
pass
69+
set_trace_processors([])

0 commit comments

Comments
 (0)