Skip to content

feat: Add output function tracing #2191

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
wants to merge 8 commits into
base: main
Choose a base branch
from
+1,494 −18
Open

feat: Add output function tracing #2191

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
8 commits
Select commit Hold shift + click to select a range
0ae3c18
Add span generation to ObjectOutputProcessor
bitnahian Jul 13, 2025
539cfd3
Add working version with non-ToolCallPart call
bitnahian Jul 13, 2025
0bb1ba3
Add working version
bitnahian Jul 13, 2025
16b9f4d
Working tests
bitnahian Jul 13, 2025
491c856
Fix tool_response serialisation in tracing
bitnahian Jul 13, 2025
e2b96a2
Add tests for TraceContext
bitnahian Jul 13, 2025
2d736b1
Add tracing for TextOutputSchema
bitnahian Jul 13, 2025
61b17de
Raise errors in unreachable code
bitnahian Jul 13, 2025
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
8 changes: 6 additions & 2 deletions pydantic_ai_slim/pydantic_ai/_agent_graph.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -341,6 +341,7 @@ async def stream(
ctx.deps.output_schema,
ctx.deps.output_validators,
build_run_context(ctx),
_output.build_trace_context(ctx),
ctx.deps.usage_limits,
)
yield agent_stream
Expand Down Expand Up @@ -529,7 +530,9 @@ async def _handle_tool_calls(
if isinstance(output_schema, _output.ToolOutputSchema):
for call, output_tool in output_schema.find_tool(tool_calls):
try:
result_data = await output_tool.process(call, run_context)
trace_context = _output.build_trace_context(ctx)
with trace_context.with_call(call):
result_data = await output_tool.process(call, run_context, trace_context)
result_data = await _validate_output(result_data, ctx, call)
except _output.ToolRetryError as e:
# TODO: Should only increment retry stuff once per node execution, not for each tool call
Expand Down Expand Up @@ -586,7 +589,8 @@ async def _handle_text_response(
try:
if isinstance(output_schema, _output.TextOutputSchema):
run_context = build_run_context(ctx)
result_data = await output_schema.process(text, run_context)
trace_context = _output.build_trace_context(ctx)
result_data = await output_schema.process(text, run_context, trace_context)
else:
m = _messages.RetryPromptPart(
content='Plain text responses are not permitted, please include your response in a tool call',
Expand Down
Loading