wehub-resource-sync
980 lines
29 KiB
Python
980 lines
29 KiB
Python
from __future__ import annotations
|
|
|
|
import asyncio
|
|
from uuid import uuid4
|
|
|
|
import pytest
|
|
from inline_snapshot import snapshot
|
|
from openai.types.responses.response_usage import InputTokensDetails
|
|
|
|
from agents import Agent, RunConfig, Runner, RunState, custom_span, function_tool, trace
|
|
from agents.sandbox.runtime import SandboxRuntime
|
|
from agents.usage import Usage
|
|
|
|
from .fake_model import FakeModel
|
|
from .test_responses import get_function_tool_call, get_text_message
|
|
from .testing_processor import (
|
|
assert_no_traces,
|
|
fetch_events,
|
|
fetch_normalized_spans,
|
|
fetch_ordered_spans,
|
|
fetch_traces,
|
|
)
|
|
|
|
|
|
def _make_approval_agent(model: FakeModel) -> Agent[None]:
|
|
@function_tool(name_override="approval_tool", needs_approval=True)
|
|
def approval_tool() -> str:
|
|
return "ok"
|
|
|
|
return Agent(name="test_agent", model=model, tools=[approval_tool])
|
|
|
|
|
|
def _usage_metadata(requests: int, input_tokens: int, output_tokens: int) -> dict[str, int]:
|
|
return {
|
|
"requests": requests,
|
|
"input_tokens": input_tokens,
|
|
"output_tokens": output_tokens,
|
|
"total_tokens": input_tokens + output_tokens,
|
|
}
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_single_run_is_single_trace():
|
|
agent = Agent(
|
|
name="test_agent",
|
|
model=FakeModel(
|
|
initial_output=[get_text_message("first_test")],
|
|
),
|
|
)
|
|
|
|
await Runner.run(agent, input="first_test")
|
|
|
|
assert fetch_normalized_spans() == snapshot(
|
|
[
|
|
{
|
|
"workflow_name": "Agent workflow",
|
|
"children": [
|
|
{
|
|
"type": "agent",
|
|
"data": {
|
|
"name": "test_agent",
|
|
"handoffs": [],
|
|
"tools": [],
|
|
"output_type": "str",
|
|
},
|
|
}
|
|
],
|
|
}
|
|
]
|
|
)
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_task_and_turn_spans_export_aggregate_usage():
|
|
@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")],
|
|
]
|
|
)
|
|
model.set_hardcoded_usage(
|
|
Usage(
|
|
requests=1,
|
|
input_tokens=10,
|
|
output_tokens=3,
|
|
total_tokens=13,
|
|
input_tokens_details=InputTokensDetails.model_validate(
|
|
{"cache_write_tokens": 3, "cached_tokens": 2}
|
|
),
|
|
)
|
|
)
|
|
agent = Agent(name="test_agent", model=model, tools=[foo_tool])
|
|
|
|
await Runner.run(agent, input="first_test")
|
|
|
|
spans = fetch_ordered_spans()
|
|
task_spans = [span.export() for span in spans if span.span_data.type == "task"]
|
|
turn_spans = [span.export() for span in spans if span.span_data.type == "turn"]
|
|
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 len(task_spans) == 1
|
|
assert task_spans[0]
|
|
assert task_spans[0]["span_data"] == {
|
|
"type": "custom",
|
|
"name": "task",
|
|
"data": {
|
|
"sdk_span_type": "task",
|
|
"name": "Agent workflow",
|
|
"usage": {
|
|
"requests": 2,
|
|
"input_tokens": 20,
|
|
"output_tokens": 6,
|
|
"total_tokens": 26,
|
|
"cached_input_tokens": 4,
|
|
"cache_write_input_tokens": 6,
|
|
},
|
|
},
|
|
}
|
|
assert "metadata" not in task_spans[0]
|
|
assert [span["span_data"]["data"]["usage"] for span in turn_spans if span] == [
|
|
{
|
|
"input_tokens": 10,
|
|
"output_tokens": 3,
|
|
"cached_input_tokens": 2,
|
|
"cache_write_input_tokens": 3,
|
|
},
|
|
{
|
|
"input_tokens": 10,
|
|
"output_tokens": 3,
|
|
"cached_input_tokens": 2,
|
|
"cache_write_input_tokens": 3,
|
|
},
|
|
]
|
|
assert [span["span_data"] for span in turn_spans if span] == [
|
|
{
|
|
"type": "custom",
|
|
"name": "turn",
|
|
"data": {
|
|
"sdk_span_type": "turn",
|
|
"turn": 1,
|
|
"agent_name": "test_agent",
|
|
"usage": {
|
|
"input_tokens": 10,
|
|
"output_tokens": 3,
|
|
"cached_input_tokens": 2,
|
|
"cache_write_input_tokens": 3,
|
|
},
|
|
},
|
|
},
|
|
{
|
|
"type": "custom",
|
|
"name": "turn",
|
|
"data": {
|
|
"sdk_span_type": "turn",
|
|
"turn": 2,
|
|
"agent_name": "test_agent",
|
|
"usage": {
|
|
"input_tokens": 10,
|
|
"output_tokens": 3,
|
|
"cached_input_tokens": 2,
|
|
"cache_write_input_tokens": 3,
|
|
},
|
|
},
|
|
},
|
|
]
|
|
assert task_spans[0]["span_data"]["data"]["usage"] == {
|
|
"requests": 2,
|
|
"input_tokens": 20,
|
|
"output_tokens": 6,
|
|
"total_tokens": 26,
|
|
"cached_input_tokens": 4,
|
|
"cache_write_input_tokens": 6,
|
|
}
|
|
|
|
assert len(agent_spans) == 1
|
|
assert len(generation_spans) == 2
|
|
assert task_spans[0]["parent_id"] is None
|
|
assert agent_spans[0].parent_id == task_spans[0]["id"]
|
|
assert turn_spans[0] and turn_spans[1]
|
|
assert [span["parent_id"] for span in turn_spans if span] == [
|
|
agent_spans[0].span_id,
|
|
agent_spans[0].span_id,
|
|
]
|
|
assert [span.parent_id for span in generation_spans] == [
|
|
turn_spans[0]["id"],
|
|
turn_spans[1]["id"],
|
|
]
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_task_span_resets_current_span_if_run_setup_fails(monkeypatch: pytest.MonkeyPatch):
|
|
agent = Agent(
|
|
name="test_agent",
|
|
model=FakeModel(
|
|
tracing_enabled=True,
|
|
initial_output=[get_text_message("first_test")],
|
|
),
|
|
)
|
|
|
|
def raise_setup_error(self: SandboxRuntime[None], agent: Agent[None]) -> None:
|
|
raise RuntimeError("setup failed")
|
|
|
|
monkeypatch.setattr(SandboxRuntime, "assert_agent_supported", raise_setup_error)
|
|
|
|
with trace(workflow_name="test_workflow"):
|
|
with pytest.raises(RuntimeError, match="setup failed"):
|
|
await Runner.run(agent, input="first_test")
|
|
|
|
with custom_span(name="after_setup_failure") as after_span:
|
|
pass
|
|
|
|
after_span_export = after_span.export()
|
|
assert after_span_export
|
|
assert after_span_export["parent_id"] is None
|
|
|
|
task_spans = [span.export() for span in fetch_ordered_spans() if span.span_data.type == "task"]
|
|
assert len(task_spans) == 1
|
|
assert task_spans[0]
|
|
assert task_spans[0]["parent_id"] is None
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_multiple_runs_are_multiple_traces():
|
|
model = FakeModel()
|
|
model.add_multiple_turn_outputs(
|
|
[
|
|
[get_text_message("first_test")],
|
|
[get_text_message("second_test")],
|
|
]
|
|
)
|
|
agent = Agent(
|
|
name="test_agent_1",
|
|
model=model,
|
|
)
|
|
|
|
await Runner.run(agent, input="first_test")
|
|
await Runner.run(agent, input="second_test")
|
|
|
|
assert fetch_normalized_spans() == snapshot(
|
|
[
|
|
{
|
|
"workflow_name": "Agent workflow",
|
|
"children": [
|
|
{
|
|
"type": "agent",
|
|
"data": {
|
|
"name": "test_agent_1",
|
|
"handoffs": [],
|
|
"tools": [],
|
|
"output_type": "str",
|
|
},
|
|
}
|
|
],
|
|
},
|
|
{
|
|
"workflow_name": "Agent workflow",
|
|
"children": [
|
|
{
|
|
"type": "agent",
|
|
"data": {
|
|
"name": "test_agent_1",
|
|
"handoffs": [],
|
|
"tools": [],
|
|
"output_type": "str",
|
|
},
|
|
}
|
|
],
|
|
},
|
|
]
|
|
)
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_resumed_run_reuses_original_trace_without_duplicate_trace_start():
|
|
model = FakeModel()
|
|
model.add_multiple_turn_outputs(
|
|
[
|
|
[get_function_tool_call("approval_tool", "{}", call_id="call-1")],
|
|
[get_text_message("done")],
|
|
]
|
|
)
|
|
agent = _make_approval_agent(model)
|
|
|
|
first = await Runner.run(agent, input="first_test")
|
|
assert first.interruptions
|
|
|
|
state = first.to_state()
|
|
state.approve(first.interruptions[0])
|
|
|
|
resumed = await Runner.run(agent, state)
|
|
|
|
assert resumed.final_output == "done"
|
|
traces = fetch_traces()
|
|
assert len(traces) == 1
|
|
assert fetch_events().count("trace_start") == 1
|
|
assert fetch_events().count("trace_end") == 1
|
|
assert all(span.trace_id == traces[0].trace_id for span in fetch_ordered_spans())
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_resumed_run_task_span_usage_is_run_local_delta():
|
|
model = FakeModel()
|
|
model.add_multiple_turn_outputs(
|
|
[
|
|
[get_function_tool_call("approval_tool", "{}", call_id="call-1")],
|
|
[get_text_message("done")],
|
|
]
|
|
)
|
|
model.set_hardcoded_usage(Usage(requests=1, input_tokens=10, output_tokens=3, total_tokens=13))
|
|
agent = _make_approval_agent(model)
|
|
|
|
first = await Runner.run(agent, input="first_test")
|
|
assert first.interruptions
|
|
|
|
state = first.to_state()
|
|
state.approve(first.interruptions[0])
|
|
|
|
resumed = await Runner.run(agent, state)
|
|
|
|
assert resumed.final_output == "done"
|
|
task_spans = [span.export() for span in fetch_ordered_spans() if span.span_data.type == "task"]
|
|
assert [span["span_data"]["data"]["usage"] for span in task_spans if span] == [
|
|
{
|
|
**_usage_metadata(requests=1, input_tokens=10, output_tokens=3),
|
|
"cached_input_tokens": 0,
|
|
"cache_write_input_tokens": 0,
|
|
},
|
|
{
|
|
**_usage_metadata(requests=1, input_tokens=10, output_tokens=3),
|
|
"cached_input_tokens": 0,
|
|
"cache_write_input_tokens": 0,
|
|
},
|
|
]
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_resumed_run_from_serialized_state_reuses_original_trace():
|
|
model = FakeModel()
|
|
model.add_multiple_turn_outputs(
|
|
[
|
|
[get_function_tool_call("approval_tool", "{}", call_id="call-1")],
|
|
[get_text_message("done")],
|
|
]
|
|
)
|
|
agent = _make_approval_agent(model)
|
|
|
|
first = await Runner.run(agent, input="first_test")
|
|
assert first.interruptions
|
|
|
|
restored_state = await RunState.from_string(agent, first.to_state().to_string())
|
|
restored_interruptions = restored_state.get_interruptions()
|
|
assert len(restored_interruptions) == 1
|
|
restored_state.approve(restored_interruptions[0])
|
|
|
|
resumed = await Runner.run(agent, restored_state)
|
|
|
|
assert resumed.final_output == "done"
|
|
traces = fetch_traces()
|
|
assert len(traces) == 1
|
|
assert fetch_events().count("trace_start") == 1
|
|
assert fetch_events().count("trace_end") == 1
|
|
assert all(span.trace_id == traces[0].trace_id for span in fetch_ordered_spans())
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_resumed_run_from_serialized_state_preserves_explicit_trace_key():
|
|
model = FakeModel()
|
|
model.add_multiple_turn_outputs(
|
|
[
|
|
[get_function_tool_call("approval_tool", "{}", call_id="call-1")],
|
|
[get_text_message("done")],
|
|
]
|
|
)
|
|
agent = _make_approval_agent(model)
|
|
|
|
first = await Runner.run(
|
|
agent,
|
|
input="first_test",
|
|
run_config=RunConfig(tracing={"api_key": "trace-key"}),
|
|
)
|
|
assert first.interruptions
|
|
|
|
restored_state = await RunState.from_string(agent, first.to_state().to_string())
|
|
restored_interruptions = restored_state.get_interruptions()
|
|
assert len(restored_interruptions) == 1
|
|
restored_state.approve(restored_interruptions[0])
|
|
|
|
resumed = await Runner.run(
|
|
agent,
|
|
restored_state,
|
|
run_config=RunConfig(tracing={"api_key": "trace-key"}),
|
|
)
|
|
|
|
assert resumed.final_output == "done"
|
|
traces = fetch_traces()
|
|
assert len(traces) == 1
|
|
assert traces[0].tracing_api_key == "trace-key"
|
|
assert fetch_events().count("trace_start") == 1
|
|
assert fetch_events().count("trace_end") == 1
|
|
assert all(span.trace_id == traces[0].trace_id for span in fetch_ordered_spans())
|
|
assert all(span.tracing_api_key == "trace-key" for span in fetch_ordered_spans())
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_resumed_run_with_workflow_override_starts_new_trace() -> None:
|
|
trace_id = f"trace_{uuid4().hex}"
|
|
model = FakeModel()
|
|
model.add_multiple_turn_outputs(
|
|
[
|
|
[get_function_tool_call("approval_tool", "{}", call_id="call-1")],
|
|
[get_text_message("done")],
|
|
]
|
|
)
|
|
agent = _make_approval_agent(model)
|
|
|
|
first = await Runner.run(
|
|
agent,
|
|
input="first_test",
|
|
run_config=RunConfig(
|
|
workflow_name="original_workflow",
|
|
trace_id=trace_id,
|
|
group_id="group-1",
|
|
),
|
|
)
|
|
assert first.interruptions
|
|
|
|
state = first.to_state()
|
|
state.approve(first.interruptions[0])
|
|
|
|
resumed = await Runner.run(
|
|
agent,
|
|
state,
|
|
run_config=RunConfig(workflow_name="override_workflow"),
|
|
)
|
|
|
|
assert resumed.final_output == "done"
|
|
traces = fetch_traces()
|
|
assert len(traces) == 2
|
|
assert fetch_events().count("trace_start") == 2
|
|
assert fetch_events().count("trace_end") == 2
|
|
assert [trace.trace_id for trace in traces] == [trace_id, trace_id]
|
|
assert [trace.name for trace in traces] == ["original_workflow", "override_workflow"]
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_wrapped_trace_is_single_trace():
|
|
model = FakeModel()
|
|
model.add_multiple_turn_outputs(
|
|
[
|
|
[get_text_message("first_test")],
|
|
[get_text_message("second_test")],
|
|
[get_text_message("third_test")],
|
|
]
|
|
)
|
|
with trace(workflow_name="test_workflow"):
|
|
agent = Agent(
|
|
name="test_agent_1",
|
|
model=model,
|
|
)
|
|
|
|
await Runner.run(agent, input="first_test")
|
|
await Runner.run(agent, input="second_test")
|
|
await Runner.run(agent, input="third_test")
|
|
|
|
assert fetch_normalized_spans() == snapshot(
|
|
[
|
|
{
|
|
"workflow_name": "test_workflow",
|
|
"children": [
|
|
{
|
|
"type": "agent",
|
|
"data": {
|
|
"name": "test_agent_1",
|
|
"handoffs": [],
|
|
"tools": [],
|
|
"output_type": "str",
|
|
},
|
|
},
|
|
{
|
|
"type": "agent",
|
|
"data": {
|
|
"name": "test_agent_1",
|
|
"handoffs": [],
|
|
"tools": [],
|
|
"output_type": "str",
|
|
},
|
|
},
|
|
{
|
|
"type": "agent",
|
|
"data": {
|
|
"name": "test_agent_1",
|
|
"handoffs": [],
|
|
"tools": [],
|
|
"output_type": "str",
|
|
},
|
|
},
|
|
],
|
|
}
|
|
]
|
|
)
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_parent_disabled_trace_disabled_agent_trace():
|
|
with trace(workflow_name="test_workflow", disabled=True):
|
|
agent = Agent(
|
|
name="test_agent",
|
|
model=FakeModel(
|
|
initial_output=[get_text_message("first_test")],
|
|
),
|
|
)
|
|
|
|
await Runner.run(agent, input="first_test")
|
|
|
|
assert_no_traces()
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_manual_disabling_works():
|
|
agent = Agent(
|
|
name="test_agent",
|
|
model=FakeModel(
|
|
initial_output=[get_text_message("first_test")],
|
|
),
|
|
)
|
|
|
|
await Runner.run(agent, input="first_test", run_config=RunConfig(tracing_disabled=True))
|
|
|
|
assert_no_traces()
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_trace_config_works():
|
|
agent = Agent(
|
|
name="test_agent",
|
|
model=FakeModel(
|
|
initial_output=[get_text_message("first_test")],
|
|
),
|
|
)
|
|
|
|
await Runner.run(
|
|
agent,
|
|
input="first_test",
|
|
run_config=RunConfig(workflow_name="Foo bar", group_id="123", trace_id="trace_456"),
|
|
)
|
|
|
|
assert fetch_normalized_spans(keep_trace_id=True) == snapshot(
|
|
[
|
|
{
|
|
"id": "trace_456",
|
|
"workflow_name": "Foo bar",
|
|
"group_id": "123",
|
|
"children": [
|
|
{
|
|
"type": "agent",
|
|
"data": {
|
|
"name": "test_agent",
|
|
"handoffs": [],
|
|
"tools": [],
|
|
"output_type": "str",
|
|
},
|
|
}
|
|
],
|
|
}
|
|
]
|
|
)
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_not_starting_streaming_creates_trace():
|
|
agent = Agent(
|
|
name="test_agent",
|
|
model=FakeModel(
|
|
initial_output=[get_text_message("first_test")],
|
|
),
|
|
)
|
|
|
|
result = Runner.run_streamed(agent, input="first_test")
|
|
|
|
# Purposely don't await the stream
|
|
while True:
|
|
if result.is_complete:
|
|
break
|
|
await asyncio.sleep(0.1)
|
|
|
|
assert fetch_normalized_spans() == snapshot(
|
|
[
|
|
{
|
|
"workflow_name": "Agent workflow",
|
|
"children": [
|
|
{
|
|
"type": "agent",
|
|
"data": {
|
|
"name": "test_agent",
|
|
"handoffs": [],
|
|
"tools": [],
|
|
"output_type": "str",
|
|
},
|
|
}
|
|
],
|
|
}
|
|
]
|
|
)
|
|
|
|
# Await the stream to avoid warnings about it not being awaited
|
|
async for _ in result.stream_events():
|
|
pass
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_streaming_single_run_is_single_trace():
|
|
agent = Agent(
|
|
name="test_agent",
|
|
model=FakeModel(
|
|
initial_output=[get_text_message("first_test")],
|
|
),
|
|
)
|
|
|
|
x = Runner.run_streamed(agent, input="first_test")
|
|
async for _ in x.stream_events():
|
|
pass
|
|
|
|
assert fetch_normalized_spans() == snapshot(
|
|
[
|
|
{
|
|
"workflow_name": "Agent workflow",
|
|
"children": [
|
|
{
|
|
"type": "agent",
|
|
"data": {
|
|
"name": "test_agent",
|
|
"handoffs": [],
|
|
"tools": [],
|
|
"output_type": "str",
|
|
},
|
|
}
|
|
],
|
|
}
|
|
]
|
|
)
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_multiple_streamed_runs_are_multiple_traces():
|
|
model = FakeModel()
|
|
model.add_multiple_turn_outputs(
|
|
[
|
|
[get_text_message("first_test")],
|
|
[get_text_message("second_test")],
|
|
]
|
|
)
|
|
agent = Agent(
|
|
name="test_agent_1",
|
|
model=model,
|
|
)
|
|
|
|
x = Runner.run_streamed(agent, input="first_test")
|
|
async for _ in x.stream_events():
|
|
pass
|
|
|
|
x = Runner.run_streamed(agent, input="second_test")
|
|
async for _ in x.stream_events():
|
|
pass
|
|
|
|
assert fetch_normalized_spans() == snapshot(
|
|
[
|
|
{
|
|
"workflow_name": "Agent workflow",
|
|
"children": [
|
|
{
|
|
"type": "agent",
|
|
"data": {
|
|
"name": "test_agent_1",
|
|
"handoffs": [],
|
|
"tools": [],
|
|
"output_type": "str",
|
|
},
|
|
}
|
|
],
|
|
},
|
|
{
|
|
"workflow_name": "Agent workflow",
|
|
"children": [
|
|
{
|
|
"type": "agent",
|
|
"data": {
|
|
"name": "test_agent_1",
|
|
"handoffs": [],
|
|
"tools": [],
|
|
"output_type": "str",
|
|
},
|
|
}
|
|
],
|
|
},
|
|
]
|
|
)
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_resumed_streaming_run_reuses_original_trace_without_duplicate_trace_start():
|
|
model = FakeModel()
|
|
model.add_multiple_turn_outputs(
|
|
[
|
|
[get_function_tool_call("approval_tool", "{}", call_id="call-1")],
|
|
[get_text_message("done")],
|
|
]
|
|
)
|
|
agent = _make_approval_agent(model)
|
|
|
|
first = Runner.run_streamed(agent, input="first_test")
|
|
async for _ in first.stream_events():
|
|
pass
|
|
assert first.interruptions
|
|
|
|
state = first.to_state()
|
|
state.approve(first.interruptions[0])
|
|
|
|
resumed = Runner.run_streamed(agent, state)
|
|
async for _ in resumed.stream_events():
|
|
pass
|
|
|
|
assert resumed.final_output == "done"
|
|
traces = fetch_traces()
|
|
assert len(traces) == 1
|
|
assert fetch_events().count("trace_start") == 1
|
|
assert fetch_events().count("trace_end") == 1
|
|
assert all(span.trace_id == traces[0].trace_id for span in fetch_ordered_spans())
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_resumed_streaming_run_task_span_usage_is_run_local_delta():
|
|
model = FakeModel()
|
|
model.add_multiple_turn_outputs(
|
|
[
|
|
[get_function_tool_call("approval_tool", "{}", call_id="call-1")],
|
|
[get_text_message("done")],
|
|
]
|
|
)
|
|
model.set_hardcoded_usage(Usage(requests=1, input_tokens=11, output_tokens=4, total_tokens=15))
|
|
agent = _make_approval_agent(model)
|
|
|
|
first = Runner.run_streamed(agent, input="first_test")
|
|
async for _ in first.stream_events():
|
|
pass
|
|
assert first.interruptions
|
|
|
|
state = first.to_state()
|
|
state.approve(first.interruptions[0])
|
|
|
|
resumed = Runner.run_streamed(agent, state)
|
|
async for _ in resumed.stream_events():
|
|
pass
|
|
|
|
assert resumed.final_output == "done"
|
|
task_spans = [span.export() for span in fetch_ordered_spans() if span.span_data.type == "task"]
|
|
assert [span["span_data"]["data"]["usage"] for span in task_spans if span] == [
|
|
{
|
|
**_usage_metadata(requests=1, input_tokens=11, output_tokens=4),
|
|
"cached_input_tokens": 0,
|
|
"cache_write_input_tokens": 0,
|
|
},
|
|
{
|
|
**_usage_metadata(requests=1, input_tokens=11, output_tokens=4),
|
|
"cached_input_tokens": 0,
|
|
"cache_write_input_tokens": 0,
|
|
},
|
|
]
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_wrapped_streaming_trace_is_single_trace():
|
|
model = FakeModel()
|
|
model.add_multiple_turn_outputs(
|
|
[
|
|
[get_text_message("first_test")],
|
|
[get_text_message("second_test")],
|
|
[get_text_message("third_test")],
|
|
]
|
|
)
|
|
with trace(workflow_name="test_workflow"):
|
|
agent = Agent(
|
|
name="test_agent_1",
|
|
model=model,
|
|
)
|
|
|
|
x = Runner.run_streamed(agent, input="first_test")
|
|
async for _ in x.stream_events():
|
|
pass
|
|
|
|
x = Runner.run_streamed(agent, input="second_test")
|
|
async for _ in x.stream_events():
|
|
pass
|
|
|
|
x = Runner.run_streamed(agent, input="third_test")
|
|
async for _ in x.stream_events():
|
|
pass
|
|
|
|
assert fetch_normalized_spans() == snapshot(
|
|
[
|
|
{
|
|
"workflow_name": "test_workflow",
|
|
"children": [
|
|
{
|
|
"type": "agent",
|
|
"data": {
|
|
"name": "test_agent_1",
|
|
"handoffs": [],
|
|
"tools": [],
|
|
"output_type": "str",
|
|
},
|
|
},
|
|
{
|
|
"type": "agent",
|
|
"data": {
|
|
"name": "test_agent_1",
|
|
"handoffs": [],
|
|
"tools": [],
|
|
"output_type": "str",
|
|
},
|
|
},
|
|
{
|
|
"type": "agent",
|
|
"data": {
|
|
"name": "test_agent_1",
|
|
"handoffs": [],
|
|
"tools": [],
|
|
"output_type": "str",
|
|
},
|
|
},
|
|
],
|
|
}
|
|
]
|
|
)
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_wrapped_streaming_run_creates_root_task_span():
|
|
agent = Agent(
|
|
name="test_agent",
|
|
model=FakeModel(
|
|
tracing_enabled=True,
|
|
initial_output=[get_text_message("first_test")],
|
|
),
|
|
)
|
|
|
|
with trace(workflow_name="test_workflow"):
|
|
result = Runner.run_streamed(agent, input="first_test")
|
|
async for _ in result.stream_events():
|
|
pass
|
|
|
|
spans = fetch_ordered_spans()
|
|
task_spans = [span.export() for span in spans if span.span_data.type == "task"]
|
|
agent_spans = [span for span in spans if span.span_data.type == "agent"]
|
|
turn_spans = [span.export() for span in spans if span.span_data.type == "turn"]
|
|
generation_spans = [span for span in spans if span.span_data.type == "generation"]
|
|
|
|
assert len(task_spans) == 1
|
|
assert task_spans[0]
|
|
assert task_spans[0]["parent_id"] is None
|
|
assert len(agent_spans) == 1
|
|
assert agent_spans[0].parent_id == task_spans[0]["id"]
|
|
assert len(turn_spans) == 1
|
|
assert turn_spans[0]
|
|
assert turn_spans[0]["parent_id"] == agent_spans[0].span_id
|
|
assert len(generation_spans) == 1
|
|
assert generation_spans[0].parent_id == turn_spans[0]["id"]
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_wrapped_mixed_trace_is_single_trace():
|
|
model = FakeModel()
|
|
model.add_multiple_turn_outputs(
|
|
[
|
|
[get_text_message("first_test")],
|
|
[get_text_message("second_test")],
|
|
[get_text_message("third_test")],
|
|
]
|
|
)
|
|
with trace(workflow_name="test_workflow"):
|
|
agent = Agent(
|
|
name="test_agent_1",
|
|
model=model,
|
|
)
|
|
|
|
x = Runner.run_streamed(agent, input="first_test")
|
|
async for _ in x.stream_events():
|
|
pass
|
|
|
|
await Runner.run(agent, input="second_test")
|
|
|
|
x = Runner.run_streamed(agent, input="third_test")
|
|
async for _ in x.stream_events():
|
|
pass
|
|
|
|
assert fetch_normalized_spans() == snapshot(
|
|
[
|
|
{
|
|
"workflow_name": "test_workflow",
|
|
"children": [
|
|
{
|
|
"type": "agent",
|
|
"data": {
|
|
"name": "test_agent_1",
|
|
"handoffs": [],
|
|
"tools": [],
|
|
"output_type": "str",
|
|
},
|
|
},
|
|
{
|
|
"type": "agent",
|
|
"data": {
|
|
"name": "test_agent_1",
|
|
"handoffs": [],
|
|
"tools": [],
|
|
"output_type": "str",
|
|
},
|
|
},
|
|
{
|
|
"type": "agent",
|
|
"data": {
|
|
"name": "test_agent_1",
|
|
"handoffs": [],
|
|
"tools": [],
|
|
"output_type": "str",
|
|
},
|
|
},
|
|
],
|
|
}
|
|
]
|
|
)
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_parent_disabled_trace_disables_streaming_agent_trace():
|
|
model = FakeModel()
|
|
model.add_multiple_turn_outputs(
|
|
[
|
|
[get_text_message("first_test")],
|
|
[get_text_message("second_test")],
|
|
]
|
|
)
|
|
with trace(workflow_name="test_workflow", disabled=True):
|
|
agent = Agent(
|
|
name="test_agent",
|
|
model=model,
|
|
)
|
|
|
|
x = Runner.run_streamed(agent, input="first_test")
|
|
async for _ in x.stream_events():
|
|
pass
|
|
|
|
assert_no_traces()
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_manual_streaming_disabling_works():
|
|
model = FakeModel()
|
|
model.add_multiple_turn_outputs(
|
|
[
|
|
[get_text_message("first_test")],
|
|
[get_text_message("second_test")],
|
|
]
|
|
)
|
|
agent = Agent(
|
|
name="test_agent",
|
|
model=model,
|
|
)
|
|
|
|
x = Runner.run_streamed(agent, input="first_test", run_config=RunConfig(tracing_disabled=True))
|
|
async for _ in x.stream_events():
|
|
pass
|
|
|
|
assert_no_traces()
|