feat: Add output function tracing #2191
Conversation
|
@bitnahian I didn't understand many of the design decisions behind |
| @@ -529,7 +530,8 @@ 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) | |||
| result_data = await output_tool.process(call, run_context, trace_context.with_call(call)) | |||
There was a problem hiding this comment.
looks like output_tool.process can do the .with_call(call)
| @dataclass(frozen=True) | ||
| class TraceContext: | ||
| """A context for tracing output processing.""" | ||
|
|
||
| tracer: Tracer | ||
| include_content: bool | ||
| call: _messages.ToolCallPart | None = None | ||
|
|
There was a problem hiding this comment.
I've noticed a preference towards dataclasses in the code base. Just curious as to why this is the preferred choice? Is it primarily because of serialisation semantics?
There was a problem hiding this comment.
a dataclass works nicely here and makes more sense than a 'plain' class.
a pydantic BaseModel would be better suited for data intended to be (de)serialized.
| if trace_context.call: | ||
| span_manager = trace_context.span(trace_context.call, include_tool_call_id=True) | ||
| else: | ||
| function_name = getattr(self._function_schema.function, '__name__', 'output_function') |
There was a problem hiding this comment.
Is there a better way to extract the function name from the schema?
|
@alexmojaki thanks for the review. I've made your suggested changes. |
| @@ -529,7 +530,8 @@ 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) | |||
| result_data = await output_tool.process(call, run_context, trace_context.with_call(call)) | |||
| @dataclass(frozen=True) | ||
| class TraceContext: | ||
| """A context for tracing output processing.""" | ||
|
|
||
| tracer: Tracer | ||
| include_content: bool | ||
| call: _messages.ToolCallPart | None = None | ||
|
|
There was a problem hiding this comment.
a dataclass works nicely here and makes more sense than a 'plain' class.
a pydantic BaseModel would be better suited for data intended to be (de)serialized.
Co-authored-by: Alex Hall <[email protected]>
|
Have addressed comments in places I missed. Lmk if any more changes are required. |
| try: | ||
| output = await self._function_schema.call(output, run_context) | ||
| output = await trace_context.execute_function_with_span( |
There was a problem hiding this comment.
Possible followup: dedupe this whole try/except
|
Thanks! |
Fixes #2108