Skip to content

Commit 2d24550

Browse files
nikhilNavanikhilc-microsoftCopilot
authored
OpenAI auto instrumentation fix attributes (#122)
* fix input and output messages for invoke span and fix tool call messages * add test for invoke agent auto instrumented * Add bounded size limit to _pending_tool_calls dictionary (#125) * Initial plan * Implement bounded size for _pending_tool_calls to prevent memory growth Co-authored-by: nikhilNava <211831449+nikhilNava@users.noreply.github.com> * Fix trailing empty line in test file Co-authored-by: nikhilNava <211831449+nikhilNava@users.noreply.github.com> --------- Co-authored-by: copilot-swe-agent[bot] <198982749+Copilot@users.noreply.github.com> Co-authored-by: nikhilNava <211831449+nikhilNava@users.noreply.github.com> * Fix code formatting with ruff (#126) * Initial plan * Fix formatting with ruff format Co-authored-by: nikhilNava <211831449+nikhilNava@users.noreply.github.com> --------- Co-authored-by: copilot-swe-agent[bot] <198982749+Copilot@users.noreply.github.com> Co-authored-by: nikhilNava <211831449+nikhilNava@users.noreply.github.com> --------- Co-authored-by: Nikhil Chitlur Navakiran (from Dev Box) <nikhilc@microsoft.com> Co-authored-by: Copilot <198982749+Copilot@users.noreply.github.com> Co-authored-by: nikhilNava <211831449+nikhilNava@users.noreply.github.com>
1 parent 67b8726 commit 2d24550

File tree

5 files changed

+417
-5
lines changed

5 files changed

+417
-5
lines changed

‎libraries/microsoft-agents-a365-observability-core/microsoft_agents_a365/observability/core/trace_processor/util.py‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -23,6 +23,7 @@
2323
consts.SESSION_ID_KEY,
2424
consts.SESSION_DESCRIPTION_KEY,
2525
consts.HIRING_MANAGER_ID_KEY,
26+
consts.GEN_AI_CALLER_CLIENT_IP_KEY, # gen_ai.caller.client.ip
2627
# Execution context
2728
consts.GEN_AI_EXECUTION_SOURCE_NAME_KEY, # gen_ai.channel.name
2829
consts.GEN_AI_EXECUTION_SOURCE_DESCRIPTION_KEY, # gen_ai.channel.link

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

Lines changed: 56 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -21,13 +21,17 @@
2121
from microsoft_agents_a365.observability.core.constants import (
2222
CUSTOM_PARENT_SPAN_ID_KEY,
2323
EXECUTE_TOOL_OPERATION_NAME,
24+
GEN_AI_EXECUTION_TYPE_KEY,
2425
GEN_AI_INPUT_MESSAGES_KEY,
2526
GEN_AI_OPERATION_NAME_KEY,
2627
GEN_AI_OUTPUT_MESSAGES_KEY,
2728
GEN_AI_REQUEST_MODEL_KEY,
2829
GEN_AI_SYSTEM_KEY,
30+
GEN_AI_TOOL_CALL_ID_KEY,
31+
GEN_AI_TOOL_TYPE_KEY,
2932
INVOKE_AGENT_OPERATION_NAME,
3033
)
34+
from microsoft_agents_a365.observability.core.execution_type import ExecutionType
3135
from microsoft_agents_a365.observability.core.utils import as_utc_nano, safe_json_dumps
3236
from opentelemetry import trace as ot_trace
3337
from opentelemetry.context import attach, detach
@@ -44,10 +48,13 @@
4448
)
4549

4650
from .constants import (
47-
GEN_AI_GRAPH_NODE_ID,
4851
GEN_AI_GRAPH_NODE_PARENT_ID,
4952
)
5053
from .utils import (
54+
capture_input_message,
55+
capture_output_message,
56+
capture_tool_call_ids,
57+
find_ancestor_agent_span_id,
5158
get_attributes_from_function_span_data,
5259
get_attributes_from_generation_span_data,
5360
get_attributes_from_input,
@@ -56,6 +63,7 @@
5663
get_span_kind,
5764
get_span_name,
5865
get_span_status,
66+
get_tool_call_id,
5967
)
6068

6169
logger = logging.getLogger(__name__)
@@ -68,6 +76,7 @@
6876

6977
class OpenAIAgentsTraceProcessor(TracingProcessor):
7078
_MAX_HANDOFFS_IN_FLIGHT = 1000
79+
_MAX_PENDING_TOOL_CALLS = 1000
7180

7281
def __init__(self, tracer: Tracer) -> None:
7382
self._tracer = tracer
@@ -79,6 +88,17 @@ def __init__(self, tracer: Tracer) -> None:
7988
# Use an OrderedDict and _MAX_HANDOFFS_IN_FLIGHT to cap the size of the dict
8089
# in case there are large numbers of orphaned handoffs
8190
self._reverse_handoffs_dict: OrderedDict[str, str] = OrderedDict()
91+
# Track input/output messages for agent spans (keyed by agent span_id)
92+
self._agent_inputs: dict[str, str] = {}
93+
self._agent_outputs: dict[str, str] = {}
94+
# Track agent span IDs to find nearest ancestor
95+
self._agent_span_ids: set[str] = set()
96+
# Track parent-child relationships: child_span_id -> parent_span_id
97+
self._span_parents: dict[str, str] = {}
98+
# Track tool_call_ids from GenerationSpan: (function_name, trace_id) -> call_id
99+
# Use an OrderedDict and _MAX_PENDING_TOOL_CALLS to cap the size of the dict
100+
# in case tool calls are captured but never consumed
101+
self._pending_tool_calls: OrderedDict[str, str] = OrderedDict()
82102

83103
# helper
84104
def _stamp_custom_parent(self, otel_span: OtelSpan, trace_id: str) -> None:
@@ -133,6 +153,12 @@ def on_span_start(self, span: Span[Any]) -> None:
133153
)
134154
self._otel_spans[span.span_id] = otel_span
135155
self._tokens[span.span_id] = attach(set_span_in_context(otel_span))
156+
# Track parent-child relationship for ancestor lookup
157+
if span.parent_id:
158+
self._span_parents[span.span_id] = span.parent_id
159+
# Track AgentSpan IDs
160+
if isinstance(span.span_data, AgentSpanData):
161+
self._agent_span_ids.add(span.span_id)
136162

137163
def on_span_end(self, span: Span[Any]) -> None:
138164
"""Called when a span is finished. Should not block or raise exceptions.
@@ -142,6 +168,8 @@ def on_span_end(self, span: Span[Any]) -> None:
142168
"""
143169
if token := self._tokens.pop(span.span_id, None):
144170
detach(token) # type: ignore[arg-type]
171+
# Clean up parent tracking
172+
self._span_parents.pop(span.span_id, None)
145173
if not (otel_span := self._otel_spans.pop(span.span_id, None)):
146174
return
147175
otel_span.update_name(get_span_name(span))
@@ -167,14 +195,32 @@ def on_span_end(self, span: Span[Any]) -> None:
167195
for k, v in get_attributes_from_generation_span_data(data):
168196
otel_span.set_attribute(k, v)
169197
self._stamp_custom_parent(otel_span, span.trace_id)
198+
# Capture input/output messages for nearest ancestor agent span
199+
if agent_span_id := find_ancestor_agent_span_id(
200+
span.parent_id, self._agent_span_ids, self._span_parents
201+
):
202+
if data.input:
203+
capture_input_message(agent_span_id, data.input, self._agent_inputs)
204+
if data.output:
205+
capture_output_message(agent_span_id, data.output, self._agent_outputs)
206+
# Capture tool_call_ids for later use by FunctionSpan
207+
if data.output:
208+
capture_tool_call_ids(
209+
data.output, self._pending_tool_calls, self._MAX_PENDING_TOOL_CALLS
210+
)
170211
otel_span.update_name(
171212
f"{otel_span.attributes[GEN_AI_OPERATION_NAME_KEY]} {otel_span.attributes[GEN_AI_REQUEST_MODEL_KEY]}"
172213
)
173214
elif isinstance(data, FunctionSpanData):
174215
for k, v in get_attributes_from_function_span_data(data):
175216
otel_span.set_attribute(k, v)
176217
self._stamp_custom_parent(otel_span, span.trace_id)
177-
otel_span.update_name(f"{EXECUTE_TOOL_OPERATION_NAME} {data.function_name}")
218+
otel_span.set_attribute(GEN_AI_TOOL_TYPE_KEY, data.type)
219+
# Set tool_call_id if available from preceding GenerationSpan
220+
func_args = data.input if data.input else ""
221+
if tool_call_id := get_tool_call_id(data.name, func_args, self._pending_tool_calls):
222+
otel_span.set_attribute(GEN_AI_TOOL_CALL_ID_KEY, tool_call_id)
223+
otel_span.update_name(f"{EXECUTE_TOOL_OPERATION_NAME} {data.name}")
178224
elif isinstance(data, MCPListToolsSpanData):
179225
for k, v in get_attributes_from_mcp_list_tool_span_data(data):
180226
otel_span.set_attribute(k, v)
@@ -187,12 +233,19 @@ def on_span_end(self, span: Span[Any]) -> None:
187233
while len(self._reverse_handoffs_dict) > self._MAX_HANDOFFS_IN_FLIGHT:
188234
self._reverse_handoffs_dict.popitem(last=False)
189235
elif isinstance(data, AgentSpanData):
190-
otel_span.set_attribute(GEN_AI_GRAPH_NODE_ID, data.name)
236+
otel_span.set_attribute(GEN_AI_EXECUTION_TYPE_KEY, ExecutionType.HUMAN_TO_AGENT.value)
191237
# Lookup the parent node if exists
192238
key = f"{data.name}:{span.trace_id}"
193239
if parent_node := self._reverse_handoffs_dict.pop(key, None):
194240
otel_span.set_attribute(GEN_AI_GRAPH_NODE_PARENT_ID, parent_node)
241+
# Apply captured input/output messages from child spans
242+
if input_msg := self._agent_inputs.pop(span.span_id, None):
243+
otel_span.set_attribute(GEN_AI_INPUT_MESSAGES_KEY, input_msg)
244+
if output_msg := self._agent_outputs.pop(span.span_id, None):
245+
otel_span.set_attribute(GEN_AI_OUTPUT_MESSAGES_KEY, output_msg)
195246
otel_span.update_name(f"{INVOKE_AGENT_OPERATION_NAME} {get_span_name(span)}")
247+
# Clean up tracking
248+
self._agent_span_ids.discard(span.span_id)
196249

197250
end_time: int | None = None
198251
if span.ended_at:

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

Lines changed: 100 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -22,6 +22,7 @@
2222
)
2323
from microsoft_agents_a365.observability.core.constants import (
2424
GEN_AI_CHOICE,
25+
GEN_AI_EVENT_CONTENT,
2526
GEN_AI_EXECUTION_PAYLOAD_KEY,
2627
GEN_AI_INPUT_MESSAGES_KEY,
2728
GEN_AI_OUTPUT_MESSAGES_KEY,
@@ -364,9 +365,9 @@ def get_attributes_from_function_span_data(
364365
) -> Iterator[tuple[str, AttributeValue]]:
365366
yield GEN_AI_TOOL_NAME_KEY, obj.name
366367
if obj.input:
367-
yield GEN_AI_INPUT_MESSAGES_KEY, obj.input
368+
yield GEN_AI_TOOL_ARGS_KEY, obj.input
368369
if obj.output is not None:
369-
yield GEN_AI_OUTPUT_MESSAGES_KEY, _convert_to_primitive(obj.output)
370+
yield GEN_AI_EVENT_CONTENT, _convert_to_primitive(obj.output)
370371

371372

372373
def get_attributes_from_message_content_list(
@@ -534,3 +535,100 @@ def get_span_status(obj: Span[Any]) -> Status:
534535
)
535536
else:
536537
return Status(StatusCode.OK)
538+
539+
540+
def capture_tool_call_ids(
541+
output_list: Any, pending_tool_calls: dict[str, str], max_size: int = 1000
542+
) -> None:
543+
"""Extract and store tool_call_ids from generation output for later use by FunctionSpan.
544+
545+
Args:
546+
output_list: The generation output containing tool calls
547+
pending_tool_calls: OrderedDict to store pending tool calls
548+
max_size: Maximum number of pending tool calls to keep in memory
549+
"""
550+
if not output_list:
551+
return
552+
try:
553+
for msg in output_list:
554+
if isinstance(msg, dict) and msg.get("role") == "assistant":
555+
tool_calls = msg.get("tool_calls")
556+
if tool_calls:
557+
for tc in tool_calls:
558+
if isinstance(tc, dict):
559+
call_id = tc.get("id")
560+
func = tc.get("function", {})
561+
func_name = func.get("name") if isinstance(func, dict) else None
562+
func_args = func.get("arguments", "") if isinstance(func, dict) else ""
563+
if call_id and func_name:
564+
# Key by (function_name, arguments) to uniquely identify each call
565+
key = f"{func_name}:{func_args}"
566+
pending_tool_calls[key] = call_id
567+
# Cap the size of the dict to prevent unbounded growth
568+
while len(pending_tool_calls) > max_size:
569+
pending_tool_calls.popitem(last=False)
570+
except Exception:
571+
pass
572+
573+
574+
def get_tool_call_id(
575+
function_name: str, function_args: str, pending_tool_calls: dict[str, str]
576+
) -> str | None:
577+
"""Get and remove the tool_call_id for a function with specific arguments."""
578+
key = f"{function_name}:{function_args}"
579+
return pending_tool_calls.pop(key, None)
580+
581+
582+
def capture_input_message(
583+
parent_span_id: str, input_list: Any, agent_inputs: dict[str, str]
584+
) -> None:
585+
"""Extract and store the first user message from input list for parent agent span."""
586+
if parent_span_id in agent_inputs:
587+
return # Already captured
588+
if not input_list:
589+
return
590+
try:
591+
for msg in input_list:
592+
if isinstance(msg, dict) and msg.get("role") == "user":
593+
content = msg.get("content", "")
594+
if content:
595+
agent_inputs[parent_span_id] = str(content)
596+
return
597+
except Exception:
598+
pass
599+
600+
601+
def capture_output_message(
602+
parent_span_id: str, output_list: Any, agent_outputs: dict[str, str]
603+
) -> None:
604+
"""Extract and store the last assistant message with actual content (no tool calls) for parent agent span."""
605+
if not output_list:
606+
return
607+
try:
608+
# Iterate in reverse to get the last assistant message with content (not a tool call)
609+
output_items = list(output_list) if not isinstance(output_list, list) else output_list
610+
for msg in reversed(output_items):
611+
if isinstance(msg, dict) and msg.get("role") == "assistant":
612+
content = msg.get("content")
613+
tool_calls = msg.get("tool_calls")
614+
# Only capture if there's actual content and no tool_calls
615+
# (tool_calls means this is an intermediate step, not the final response)
616+
if content and not tool_calls:
617+
agent_outputs[parent_span_id] = str(content)
618+
return
619+
except Exception:
620+
pass
621+
622+
623+
def find_ancestor_agent_span_id(
624+
span_id: str | None, agent_span_ids: set[str], span_parents: dict[str, str]
625+
) -> str | None:
626+
"""Walk up the parent chain to find the nearest ancestor AgentSpan."""
627+
current = span_id
628+
visited: set[str] = set() # Prevent infinite loops
629+
while current and current not in visited:
630+
if current in agent_span_ids:
631+
return current
632+
visited.add(current)
633+
current = span_parents.get(current)
634+
return None

0 commit comments

Comments
 (0)