diff --git a/src/agents/run.py b/src/agents/run.py index f28462d461..2ac1189663 100644 --- a/src/agents/run.py +++ b/src/agents/run.py @@ -121,6 +121,7 @@ from .tool import dispose_resolved_computers from .tool_guardrails import ToolInputGuardrailResult, ToolOutputGuardrailResult from .tracing import Span, SpanError, agent_span, get_current_trace, task_span, turn_span +from .tracing.config import include_task_and_turn_spans from .tracing.context import TraceCtxManager, create_trace_for_run from .tracing.span_data import AgentSpanData, TaskSpanData from .util import _error_tracing @@ -643,8 +644,12 @@ async def run( run_state._reasoning_item_id_policy = resolved_reasoning_item_id_policy run_state.set_trace(get_current_trace()) - current_task_span: Span[TaskSpanData] = task_span(name=trace_workflow_name) - current_task_span.start(mark_as_current=True) + use_task_and_turn_spans = include_task_and_turn_spans(run_config.tracing) + current_task_span: Span[TaskSpanData] | None = ( + task_span(name=trace_workflow_name) if use_task_and_turn_spans else None + ) + if current_task_span: + current_task_span.start(mark_as_current=True) task_usage_start = snapshot_usage(context_wrapper.usage) try: @@ -758,11 +763,12 @@ def _finalize_result(result: RunResult) -> RunResult: ) session_input_items_for_persistence = [] except BaseException: - attach_usage_to_span( - current_task_span, - usage_delta(task_usage_start, context_wrapper.usage), - ) - current_task_span.finish(reset_current=True) + if current_task_span: + attach_usage_to_span( + current_task_span, + usage_delta(task_usage_start, context_wrapper.usage), + ) + current_task_span.finish(reset_current=True) raise try: @@ -1165,11 +1171,16 @@ def _finalize_result(result: RunResult) -> RunResult: ) turn_usage_start = snapshot_usage(context_wrapper.usage) - current_turn_span = turn_span( - turn=current_turn, - agent_name=current_agent.name, + current_turn_span = ( + turn_span( + turn=current_turn, + agent_name=current_agent.name, + ) + if use_task_and_turn_spans + else None ) - current_turn_span.start(mark_as_current=True) + if current_turn_span: + current_turn_span.start(mark_as_current=True) try: if current_turn <= 1: try: @@ -1273,11 +1284,12 @@ def _finalize_result(result: RunResult) -> RunResult: error_handlers=error_handlers, ) finally: - attach_usage_to_span( - current_turn_span, - usage_delta(turn_usage_start, context_wrapper.usage), - ) - current_turn_span.finish(reset_current=True) + if current_turn_span: + attach_usage_to_span( + current_turn_span, + usage_delta(turn_usage_start, context_wrapper.usage), + ) + current_turn_span.finish(reset_current=True) # Start hooks should only run on the first turn unless reset by a handoff. last_saved_input_snapshot_for_rewind = None diff --git a/src/agents/run_internal/run_loop.py b/src/agents/run_internal/run_loop.py index b85347f6bf..6476519b2e 100644 --- a/src/agents/run_internal/run_loop.py +++ b/src/agents/run_internal/run_loop.py @@ -83,6 +83,7 @@ get_function_tool_origin, ) from ..tracing import Span, SpanError, agent_span, get_current_trace, task_span, turn_span +from ..tracing.config import include_task_and_turn_spans from ..tracing.model_tracing import get_model_tracing_impl from ..tracing.span_data import AgentSpanData, TaskSpanData from ..usage import Usage, _response_usage_to_usage @@ -486,8 +487,9 @@ async def start_streaming( ) current_trace = streamed_result.trace or get_current_trace() + use_task_and_turn_spans = include_task_and_turn_spans(run_config.tracing) current_task_span: Span[TaskSpanData] | None = ( - task_span(name=current_trace.name) if current_trace else None + task_span(name=current_trace.name) if current_trace and use_task_and_turn_spans else None ) if current_task_span: current_task_span.start(mark_as_current=True) @@ -1012,11 +1014,16 @@ async def _save_stream_items_without_count( current_agent.name, ) turn_usage_start = snapshot_usage(context_wrapper.usage) - current_turn_span = turn_span( - turn=current_turn, - agent_name=current_agent.name, + current_turn_span = ( + turn_span( + turn=current_turn, + agent_name=current_agent.name, + ) + if use_task_and_turn_spans + else None ) - current_turn_span.start(mark_as_current=True) + if current_turn_span: + current_turn_span.start(mark_as_current=True) try: if ( session is not None @@ -1050,11 +1057,12 @@ async def _save_stream_items_without_count( error_handlers=error_handlers, ) finally: - attach_usage_to_span( - current_turn_span, - usage_delta(turn_usage_start, context_wrapper.usage), - ) - current_turn_span.finish(reset_current=True) + if current_turn_span: + attach_usage_to_span( + current_turn_span, + usage_delta(turn_usage_start, context_wrapper.usage), + ) + current_turn_span.finish(reset_current=True) logger.debug( "Turn %s complete, next_step type=%s", current_turn, diff --git a/src/agents/tracing/config.py b/src/agents/tracing/config.py index 24aaa9a5fd..fb6ee013d1 100644 --- a/src/agents/tracing/config.py +++ b/src/agents/tracing/config.py @@ -4,6 +4,15 @@ class TracingConfig(TypedDict, total=False): - """Configuration for tracing export.""" + """Configuration for tracing behavior and export.""" api_key: str + """Optional API key used to export traces.""" + + include_task_and_turn_spans: bool + """Whether the runner creates task and turn spans. Defaults to True when omitted.""" + + +def include_task_and_turn_spans(config: TracingConfig | None) -> bool: + """Resolve whether automatic runner task and turn spans are enabled.""" + return config is None or config.get("include_task_and_turn_spans", True) diff --git a/tests/test_agent_tracing.py b/tests/test_agent_tracing.py index b6990dcadc..2b63a0ef41 100644 --- a/tests/test_agent_tracing.py +++ b/tests/test_agent_tracing.py @@ -193,6 +193,66 @@ def foo_tool() -> str: ] +@pytest.mark.asyncio +async def test_task_and_turn_spans_can_be_disabled(): + @function_tool + def foo_tool() -> str: + return "foo result" + + model = FakeModel(tracing_enabled=True) + model.add_multiple_turn_outputs( + [ + [get_function_tool_call("foo_tool", "{}", call_id="call-1")], + [get_text_message("done")], + ] + ) + agent = Agent(name="test_agent", model=model, tools=[foo_tool]) + + result = await Runner.run( + agent, + input="first_test", + run_config=RunConfig(tracing={"include_task_and_turn_spans": False}), + ) + + spans = fetch_ordered_spans() + agent_spans = [span for span in spans if span.span_data.type == "agent"] + generation_spans = [span for span in spans if span.span_data.type == "generation"] + function_spans = [span for span in spans if span.span_data.type == "function"] + + assert result.final_output == "done" + assert not [span for span in spans if span.span_data.type in {"task", "turn"}] + assert len(agent_spans) == 1 + assert agent_spans[0].parent_id is None + assert len(generation_spans) == 2 + assert len(function_spans) == 1 + assert [span.parent_id for span in generation_spans + function_spans] == [ + agent_spans[0].span_id, + agent_spans[0].span_id, + agent_spans[0].span_id, + ] + + +@pytest.mark.asyncio +async def test_task_and_turn_spans_can_be_explicitly_enabled(): + agent = Agent( + name="test_agent", + model=FakeModel( + tracing_enabled=True, + initial_output=[get_text_message("done")], + ), + ) + + await Runner.run( + agent, + input="first_test", + run_config=RunConfig(tracing={"include_task_and_turn_spans": True}), + ) + + span_types = [span.span_data.type for span in fetch_ordered_spans()] + assert span_types.count("task") == 1 + assert span_types.count("turn") == 1 + + @pytest.mark.asyncio async def test_task_span_resets_current_span_if_run_setup_fails(monkeypatch: pytest.MonkeyPatch): agent = Agent( @@ -872,6 +932,37 @@ async def test_wrapped_streaming_run_creates_root_task_span(): assert generation_spans[0].parent_id == turn_spans[0]["id"] +@pytest.mark.asyncio +async def test_wrapped_streaming_run_can_disable_task_and_turn_spans(): + agent = Agent( + name="test_agent", + model=FakeModel( + tracing_enabled=True, + initial_output=[get_text_message("done")], + ), + ) + + with trace(workflow_name="test_workflow"): + result = Runner.run_streamed( + agent, + input="first_test", + run_config=RunConfig(tracing={"include_task_and_turn_spans": False}), + ) + async for _ in result.stream_events(): + pass + + spans = fetch_ordered_spans() + agent_spans = [span for span in spans if span.span_data.type == "agent"] + generation_spans = [span for span in spans if span.span_data.type == "generation"] + + assert result.final_output == "done" + assert not [span for span in spans if span.span_data.type in {"task", "turn"}] + assert len(agent_spans) == 1 + assert agent_spans[0].parent_id is None + assert len(generation_spans) == 1 + assert generation_spans[0].parent_id == agent_spans[0].span_id + + @pytest.mark.asyncio async def test_wrapped_mixed_trace_is_single_trace(): model = FakeModel()