Skip to content

Commit e063b12

Browse files
danielmillerpdanielmillerp
committed
Add better tracing for sync_provider
1 parent 7046381 commit e063b12

File tree

1 file changed

+254
-19
lines changed

1 file changed

+254
-19
lines changed

‎src/agentex/lib/adk/providers/_modules/sync_provider.py‎

Lines changed: 254 additions & 19 deletions
Original file line numberDiff line numberDiff line change
@@ -109,10 +109,94 @@ async def get_response(
109109

110110
response = await self.original_model.get_response(**kwargs)
111111

112-
# Set span output
113-
if span:
112+
# Set span output with structured data
113+
if span and response:
114+
new_items = []
115+
final_output = None
116+
117+
# Extract final output text from response
118+
response_final_output = getattr(response, 'final_output', None)
119+
if response_final_output:
120+
final_output = response_final_output
121+
122+
# Extract items from the response output
123+
response_output = getattr(response, 'output', None)
124+
if response_output:
125+
output_items = response_output if isinstance(response_output, list) else [response_output]
126+
127+
for item in output_items:
128+
item_type = getattr(item, 'type', None)
129+
130+
# Handle reasoning items
131+
if item_type == 'reasoning':
132+
reasoning_summary = []
133+
summary = getattr(item, 'summary', None)
134+
if summary:
135+
for summary_part in summary:
136+
text = getattr(summary_part, 'text', None)
137+
if text:
138+
reasoning_summary.append({
139+
"text": text,
140+
"type": "summary_text"
141+
})
142+
143+
new_items.append({
144+
"id": getattr(item, 'id', None),
145+
"type": "reasoning",
146+
"status": getattr(item, 'status', None),
147+
"content": None,
148+
"summary": reasoning_summary if reasoning_summary else None,
149+
})
150+
151+
# Handle tool call items
152+
elif item_type == 'function_call':
153+
new_items.append({
154+
"id": getattr(item, 'id', None),
155+
"name": getattr(item, 'name', None),
156+
"type": "function_call",
157+
"status": getattr(item, 'status', 'completed'),
158+
"call_id": getattr(item, 'call_id', None),
159+
"arguments": getattr(item, 'arguments', None),
160+
})
161+
162+
# Handle tool output items
163+
elif item_type == 'function_call_output':
164+
new_items.append({
165+
"type": "function_call_output",
166+
"output": getattr(item, 'output', None),
167+
"call_id": getattr(item, 'call_id', None),
168+
})
169+
170+
# Handle message items
171+
elif item_type == 'message':
172+
content = []
173+
message_text = ""
174+
item_content = getattr(item, 'content', None)
175+
if item_content:
176+
for content_part in item_content:
177+
text = getattr(content_part, 'text', None)
178+
if text:
179+
content.append({
180+
"text": text,
181+
"type": "output_text",
182+
})
183+
message_text = text
184+
185+
new_items.append({
186+
"id": getattr(item, 'id', None),
187+
"role": getattr(item, 'role', 'assistant'),
188+
"type": "message",
189+
"status": getattr(item, 'status', 'completed'),
190+
"content": content,
191+
})
192+
193+
# Use message text as final output if we have it
194+
if message_text and not final_output:
195+
final_output = message_text
196+
114197
span.output = {
115-
"response": str(response) if response else None,
198+
"new_items": new_items,
199+
"final_output": final_output,
116200
}
117201

118202
return response
@@ -160,7 +244,9 @@ async def stream_response(
160244
# Wrap the streaming in a tracing span if tracer is available
161245
if self.tracer and self.trace_id:
162246
trace = self.tracer.trace(self.trace_id)
163-
async with trace.span(
247+
248+
# Manually start the span instead of using context manager
249+
span = await trace.start_span(
164250
parent_id=self.parent_span_id,
165251
name="run_agent_streamed",
166252
input={
@@ -172,7 +258,9 @@ async def stream_response(
172258
"handoffs": [str(h) for h in handoffs] if handoffs else [],
173259
"previous_response_id": previous_response_id,
174260
},
175-
) as span:
261+
)
262+
263+
try:
176264
# Get the stream from the original model
177265
stream_kwargs = {
178266
"system_instructions": system_instructions,
@@ -193,23 +281,170 @@ async def stream_response(
193281
# Get the stream response from the original model and yield each event
194282
stream_response = self.original_model.stream_response(**stream_kwargs)
195283

196-
# Pass through each event from the original stream
197-
event_count = 0
198-
final_output = None
284+
# Pass through each event from the original stream and track items
285+
new_items = []
286+
final_response_text = ""
287+
current_text_item = None
288+
tool_call_map = {} # Map call_id to tool name
289+
199290
async for event in stream_response:
200-
event_count += 1
201-
# Track the final output if available
202-
if hasattr(event, 'type') and event.type == 'raw_response_event':
203-
if hasattr(event.data, 'output'):
204-
final_output = event.data.output
291+
event_type = getattr(event, 'type', 'no-type')
292+
293+
# Handle response.output_item.done events which contain completed items
294+
if event_type == 'response.output_item.done':
295+
item = getattr(event, 'item', None)
296+
if item is not None:
297+
item_type = getattr(item, 'type', None)
298+
299+
# Handle function call (tool request)
300+
if item_type == 'function_call':
301+
call_id = getattr(item, 'call_id', None)
302+
name = getattr(item, 'name', None)
303+
new_items.append({
304+
"id": getattr(item, 'id', None),
305+
"name": name,
306+
"type": "function_call",
307+
"status": getattr(item, 'status', 'completed'),
308+
"call_id": call_id,
309+
"arguments": getattr(item, 'arguments', None),
310+
})
311+
if call_id and name:
312+
tool_call_map[call_id] = name
313+
314+
# Handle completed message items (final text output)
315+
elif item_type == 'message':
316+
content = []
317+
message_text = ""
318+
item_content = getattr(item, 'content', None)
319+
if item_content:
320+
for content_part in item_content:
321+
text = getattr(content_part, 'text', None)
322+
if text:
323+
content.append({
324+
"text": text,
325+
"type": "output_text",
326+
})
327+
# Use the complete text from the message as final_output
328+
message_text = text
329+
330+
new_items.append({
331+
"id": getattr(item, 'id', None),
332+
"role": getattr(item, 'role', 'assistant'),
333+
"type": "message",
334+
"status": getattr(item, 'status', 'completed'),
335+
"content": content,
336+
})
337+
338+
# Update final_response_text with the complete message text
339+
if message_text:
340+
final_response_text = message_text
341+
342+
# Track reasoning, tool calls, and responses from run_item_stream_event (kept for compatibility)
343+
if event_type == 'run_item_stream_event':
344+
item = getattr(event, 'item', None)
345+
if item is not None:
346+
item_type = getattr(item, 'type', None)
347+
348+
# Handle reasoning items
349+
if item_type == 'reasoning_item':
350+
reasoning_summary = []
351+
raw_item = getattr(item, 'raw_item', None)
352+
if raw_item is not None:
353+
summary = getattr(raw_item, 'summary', None)
354+
if summary:
355+
for summary_part in summary:
356+
text = getattr(summary_part, 'text', None)
357+
if text:
358+
reasoning_summary.append({
359+
"text": text,
360+
"type": "summary_text"
361+
})
362+
363+
new_items.append({
364+
"id": getattr(raw_item, 'id', None),
365+
"type": "reasoning",
366+
"status": getattr(raw_item, 'status', None),
367+
"content": None,
368+
"summary": reasoning_summary if reasoning_summary else None,
369+
})
370+
371+
# Handle tool call items
372+
elif item_type == 'tool_call_item':
373+
raw_item = getattr(item, 'raw_item', None)
374+
if raw_item is not None:
375+
call_id, tool_name, tool_arguments = _extract_tool_call_info(raw_item)
376+
tool_call_map[call_id] = tool_name
377+
378+
new_items.append({
379+
"id": getattr(raw_item, 'id', None),
380+
"name": tool_name,
381+
"type": "function_call",
382+
"status": getattr(raw_item, 'status', 'completed'),
383+
"call_id": call_id,
384+
"arguments": str(tool_arguments) if isinstance(tool_arguments, dict) else tool_arguments,
385+
})
386+
387+
# Handle tool output items
388+
elif item_type == 'tool_call_output_item':
389+
raw_item = getattr(item, 'raw_item', None)
390+
if raw_item is not None:
391+
call_id, tool_name, content = _extract_tool_response_info(tool_call_map, raw_item)
392+
393+
new_items.append({
394+
"type": "function_call_output",
395+
"output": content,
396+
"call_id": call_id,
397+
})
398+
399+
# Accumulate text deltas to build final response
400+
# Note: OpenAI Agents SDK can emit events in different formats
401+
if hasattr(event, 'type') and event.type == 'response.output_text.delta':
402+
# Direct event type from OpenAI Agents SDK (observed in practice)
403+
if hasattr(event, 'delta'):
404+
final_response_text += event.delta
405+
406+
# Handle raw_response_event wrapper (alternative event format, kept for compatibility)
407+
elif hasattr(event, 'type') and event.type == 'raw_response_event':
408+
if hasattr(event, 'data'):
409+
raw_event = event.data
410+
411+
# Track when output items are added
412+
if isinstance(raw_event, ResponseOutputItemAddedEvent):
413+
if hasattr(raw_event, 'item') and raw_event.item.type == 'message':
414+
current_text_item = {
415+
"id": getattr(raw_event.item, 'id', None),
416+
"role": getattr(raw_event.item, 'role', 'assistant'),
417+
"type": "message",
418+
"status": "in_progress",
419+
"content": []
420+
}
421+
422+
# Check if this is a text delta event
423+
elif isinstance(raw_event, ResponseTextDeltaEvent):
424+
if hasattr(raw_event, 'delta') and raw_event.delta:
425+
final_response_text += raw_event.delta
426+
427+
# Track when output items are done
428+
elif isinstance(raw_event, ResponseOutputItemDoneEvent):
429+
if current_text_item and final_response_text:
430+
current_text_item["status"] = "completed"
431+
current_text_item["content"] = [{
432+
"text": final_response_text,
433+
"type": "output_text",
434+
}]
435+
new_items.append(current_text_item)
436+
current_text_item = None
437+
205438
yield event
206439

207-
# Set span output
208-
if span:
209-
span.output = {
210-
"event_count": event_count,
211-
"final_output": str(final_output) if final_output else None,
212-
}
440+
# Set span output with structured data including tool calls and final response
441+
span.output = {
442+
"new_items": new_items,
443+
"final_output": final_response_text if final_response_text else None,
444+
}
445+
finally:
446+
# End the span after all events have been yielded
447+
await trace.end_span(span)
213448
else:
214449
# No tracing, just stream normally
215450
# Get the stream from the original model

0 commit comments

Comments
 (0)