chore: import upstream snapshot with attribution
This commit is contained in:
wehub-resource-sync
@@ -0,0 +1,530 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import asyncio
|
||||
from typing import Any
|
||||
|
||||
import pytest
|
||||
from inline_snapshot import snapshot
|
||||
|
||||
from agents.tracing import (
|
||||
Span,
|
||||
Trace,
|
||||
TracingProcessor,
|
||||
agent_span,
|
||||
custom_span,
|
||||
function_span,
|
||||
generation_span,
|
||||
handoff_span,
|
||||
set_trace_processors,
|
||||
trace,
|
||||
)
|
||||
from agents.tracing.spans import SpanError
|
||||
|
||||
from .testing_processor import (
|
||||
SPAN_PROCESSOR_TESTING,
|
||||
assert_no_traces,
|
||||
fetch_events,
|
||||
fetch_normalized_spans,
|
||||
)
|
||||
|
||||
### HELPERS
|
||||
|
||||
|
||||
def standard_span_checks(
|
||||
span: Span[Any], trace_id: str, parent_id: str | None, span_type: str
|
||||
) -> None:
|
||||
assert span.span_id is not None
|
||||
assert span.trace_id == trace_id
|
||||
assert span.parent_id == parent_id
|
||||
assert span.started_at is not None
|
||||
assert span.ended_at is not None
|
||||
assert span.span_data.type == span_type
|
||||
|
||||
|
||||
def standard_trace_checks(trace: Trace, name_check: str | None = None) -> None:
|
||||
assert trace.trace_id is not None
|
||||
|
||||
if name_check:
|
||||
assert trace.name == name_check
|
||||
|
||||
|
||||
### TESTS
|
||||
|
||||
|
||||
def simple_tracing():
|
||||
x = trace("test")
|
||||
x.start()
|
||||
|
||||
span_1 = agent_span(name="agent_1", span_id="span_1", parent=x)
|
||||
span_1.start()
|
||||
span_1.finish()
|
||||
|
||||
span_2 = custom_span(name="custom_1", span_id="span_2", parent=x)
|
||||
span_2.start()
|
||||
|
||||
span_3 = custom_span(name="custom_2", span_id="span_3", parent=span_2)
|
||||
span_3.start()
|
||||
span_3.finish()
|
||||
|
||||
span_2.finish()
|
||||
|
||||
x.finish()
|
||||
|
||||
|
||||
def test_simple_tracing() -> None:
|
||||
simple_tracing()
|
||||
|
||||
assert fetch_normalized_spans(keep_span_id=True) == snapshot(
|
||||
[
|
||||
{
|
||||
"workflow_name": "test",
|
||||
"children": [
|
||||
{
|
||||
"type": "agent",
|
||||
"id": "span_1",
|
||||
"data": {"name": "agent_1"},
|
||||
},
|
||||
{
|
||||
"type": "custom",
|
||||
"id": "span_2",
|
||||
"data": {"name": "custom_1", "data": {}},
|
||||
"children": [
|
||||
{
|
||||
"type": "custom",
|
||||
"id": "span_3",
|
||||
"data": {"name": "custom_2", "data": {}},
|
||||
}
|
||||
],
|
||||
},
|
||||
],
|
||||
}
|
||||
]
|
||||
)
|
||||
|
||||
|
||||
def ctxmanager_spans():
|
||||
with trace(workflow_name="test", trace_id="trace_123", group_id="456"):
|
||||
with custom_span(name="custom_1", span_id="span_1"):
|
||||
with custom_span(name="custom_2", span_id="span_1_inner"):
|
||||
pass
|
||||
|
||||
with custom_span(name="custom_2", span_id="span_2"):
|
||||
pass
|
||||
|
||||
|
||||
def test_ctxmanager_spans() -> None:
|
||||
ctxmanager_spans()
|
||||
|
||||
assert fetch_normalized_spans(keep_span_id=True) == snapshot(
|
||||
[
|
||||
{
|
||||
"workflow_name": "test",
|
||||
"group_id": "456",
|
||||
"children": [
|
||||
{
|
||||
"type": "custom",
|
||||
"id": "span_1",
|
||||
"data": {"name": "custom_1", "data": {}},
|
||||
"children": [
|
||||
{
|
||||
"type": "custom",
|
||||
"id": "span_1_inner",
|
||||
"data": {"name": "custom_2", "data": {}},
|
||||
}
|
||||
],
|
||||
},
|
||||
{"type": "custom", "id": "span_2", "data": {"name": "custom_2", "data": {}}},
|
||||
],
|
||||
}
|
||||
]
|
||||
)
|
||||
|
||||
|
||||
async def run_subtask(span_id: str | None = None) -> None:
|
||||
with generation_span(span_id=span_id):
|
||||
await asyncio.sleep(0.0001)
|
||||
|
||||
|
||||
async def simple_async_tracing():
|
||||
with trace(workflow_name="test", trace_id="trace_123", group_id="group_456"):
|
||||
await run_subtask(span_id="span_1")
|
||||
await run_subtask(span_id="span_2")
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_async_tracing() -> None:
|
||||
await simple_async_tracing()
|
||||
|
||||
assert fetch_normalized_spans(keep_span_id=True) == snapshot(
|
||||
[
|
||||
{
|
||||
"workflow_name": "test",
|
||||
"group_id": "group_456",
|
||||
"children": [
|
||||
{"type": "generation", "id": "span_1"},
|
||||
{"type": "generation", "id": "span_2"},
|
||||
],
|
||||
}
|
||||
]
|
||||
)
|
||||
|
||||
|
||||
async def run_tasks_parallel(span_ids: list[str]) -> None:
|
||||
await asyncio.gather(
|
||||
*[run_subtask(span_id=span_id) for span_id in span_ids],
|
||||
)
|
||||
|
||||
|
||||
async def run_tasks_as_children(first_span_id: str, second_span_id: str) -> None:
|
||||
with generation_span(span_id=first_span_id):
|
||||
await run_subtask(span_id=second_span_id)
|
||||
|
||||
|
||||
async def complex_async_tracing():
|
||||
with trace(workflow_name="test", trace_id="trace_123", group_id="456"):
|
||||
await asyncio.gather(
|
||||
run_tasks_parallel(["span_1", "span_2"]),
|
||||
run_tasks_parallel(["span_3", "span_4"]),
|
||||
)
|
||||
await asyncio.gather(
|
||||
run_tasks_as_children("span_5", "span_6"),
|
||||
run_tasks_as_children("span_7", "span_8"),
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_complex_async_tracing() -> None:
|
||||
for _ in range(300):
|
||||
SPAN_PROCESSOR_TESTING.clear()
|
||||
await complex_async_tracing()
|
||||
|
||||
assert fetch_normalized_spans(keep_span_id=True) == (
|
||||
[
|
||||
{
|
||||
"workflow_name": "test",
|
||||
"group_id": "456",
|
||||
"children": [
|
||||
{"type": "generation", "id": "span_1"},
|
||||
{"type": "generation", "id": "span_2"},
|
||||
{"type": "generation", "id": "span_3"},
|
||||
{"type": "generation", "id": "span_4"},
|
||||
{
|
||||
"type": "generation",
|
||||
"id": "span_5",
|
||||
"children": [{"type": "generation", "id": "span_6"}],
|
||||
},
|
||||
{
|
||||
"type": "generation",
|
||||
"id": "span_7",
|
||||
"children": [{"type": "generation", "id": "span_8"}],
|
||||
},
|
||||
],
|
||||
}
|
||||
]
|
||||
)
|
||||
|
||||
|
||||
def spans_with_setters():
|
||||
with trace(workflow_name="test", trace_id="trace_123", group_id="456"):
|
||||
with agent_span(name="agent_1") as span_a:
|
||||
span_a.span_data.name = "agent_2"
|
||||
|
||||
with function_span(name="function_1") as span_b:
|
||||
span_b.span_data.input = "i"
|
||||
span_b.span_data.output = "o"
|
||||
|
||||
with generation_span() as span_c:
|
||||
span_c.span_data.input = [{"foo": "bar"}]
|
||||
|
||||
with handoff_span(from_agent="agent_1", to_agent="agent_2"):
|
||||
pass
|
||||
|
||||
|
||||
def test_spans_with_setters() -> None:
|
||||
spans_with_setters()
|
||||
|
||||
assert fetch_normalized_spans() == snapshot(
|
||||
[
|
||||
{
|
||||
"workflow_name": "test",
|
||||
"group_id": "456",
|
||||
"children": [
|
||||
{
|
||||
"type": "agent",
|
||||
"data": {"name": "agent_2"},
|
||||
"children": [
|
||||
{
|
||||
"type": "function",
|
||||
"data": {"name": "function_1", "input": "i", "output": "o"},
|
||||
},
|
||||
{
|
||||
"type": "generation",
|
||||
"data": {"input": [{"foo": "bar"}]},
|
||||
},
|
||||
{
|
||||
"type": "handoff",
|
||||
"data": {"from_agent": "agent_1", "to_agent": "agent_2"},
|
||||
},
|
||||
],
|
||||
}
|
||||
],
|
||||
}
|
||||
]
|
||||
)
|
||||
|
||||
|
||||
def disabled_tracing():
|
||||
with trace(workflow_name="test", trace_id="123", group_id="456", disabled=True):
|
||||
with agent_span(name="agent_1"):
|
||||
with function_span(name="function_1"):
|
||||
pass
|
||||
|
||||
|
||||
def test_disabled_tracing():
|
||||
disabled_tracing()
|
||||
assert_no_traces()
|
||||
|
||||
|
||||
def enabled_trace_disabled_span():
|
||||
with trace(workflow_name="test", trace_id="trace_123"):
|
||||
with agent_span(name="agent_1"):
|
||||
with function_span(name="function_1", disabled=True):
|
||||
with generation_span():
|
||||
pass
|
||||
|
||||
|
||||
def test_enabled_trace_disabled_span():
|
||||
enabled_trace_disabled_span()
|
||||
|
||||
assert fetch_normalized_spans() == snapshot(
|
||||
[
|
||||
{
|
||||
"workflow_name": "test",
|
||||
"children": [
|
||||
{
|
||||
"type": "agent",
|
||||
"data": {"name": "agent_1"},
|
||||
}
|
||||
],
|
||||
}
|
||||
]
|
||||
)
|
||||
|
||||
|
||||
def test_start_and_end_called_manual():
|
||||
simple_tracing()
|
||||
|
||||
events = fetch_events()
|
||||
|
||||
assert events == [
|
||||
"trace_start",
|
||||
"span_start", # span_1
|
||||
"span_end", # span_1
|
||||
"span_start", # span_2
|
||||
"span_start", # span_3
|
||||
"span_end", # span_3
|
||||
"span_end", # span_2
|
||||
"trace_end",
|
||||
]
|
||||
|
||||
|
||||
def test_start_and_end_called_ctxmanager():
|
||||
with trace(workflow_name="test", trace_id="123", group_id="456"):
|
||||
with custom_span(name="custom_1", span_id="span_1"):
|
||||
with custom_span(name="custom_2", span_id="span_1_inner"):
|
||||
pass
|
||||
|
||||
with custom_span(name="custom_2", span_id="span_2"):
|
||||
pass
|
||||
|
||||
events = fetch_events()
|
||||
|
||||
assert events == [
|
||||
"trace_start",
|
||||
"span_start", # span_1
|
||||
"span_start", # span_1_inner
|
||||
"span_end", # span_1_inner
|
||||
"span_end", # span_1
|
||||
"span_start", # span_2
|
||||
"span_end", # span_2
|
||||
"trace_end",
|
||||
]
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_start_and_end_called_async_ctxmanager():
|
||||
await simple_async_tracing()
|
||||
|
||||
events = fetch_events()
|
||||
|
||||
assert events == [
|
||||
"trace_start",
|
||||
"span_start", # span_1
|
||||
"span_end", # span_1
|
||||
"span_start", # span_2
|
||||
"span_end", # span_2
|
||||
"trace_end",
|
||||
]
|
||||
|
||||
|
||||
async def test_noop_span_doesnt_record():
|
||||
with trace(workflow_name="test", disabled=True) as t:
|
||||
with custom_span(name="span_1") as span:
|
||||
span.set_error(SpanError(message="test", data={}))
|
||||
|
||||
assert_no_traces()
|
||||
|
||||
assert t.export() is None
|
||||
assert span.export() is None
|
||||
assert span.started_at is None
|
||||
assert span.ended_at is None
|
||||
assert span.error is None
|
||||
|
||||
|
||||
async def test_multiple_span_start_finish_doesnt_crash():
|
||||
with trace(workflow_name="test", trace_id="123", group_id="456"):
|
||||
with custom_span(name="span_1") as span:
|
||||
span.start()
|
||||
|
||||
span.finish()
|
||||
|
||||
|
||||
async def test_noop_parent_is_noop_child():
|
||||
tr = trace(workflow_name="test", disabled=True)
|
||||
|
||||
span = custom_span(name="span_1", parent=tr)
|
||||
span.start()
|
||||
span.finish()
|
||||
|
||||
assert span.export() is None
|
||||
|
||||
span_2 = custom_span(name="span_2", parent=span)
|
||||
span_2.start()
|
||||
span_2.finish()
|
||||
|
||||
assert span_2.export() is None
|
||||
|
||||
|
||||
def test_trace_and_spans_use_tracing_config_key():
|
||||
with trace(workflow_name="test", tracing={"api_key": "tracing-key"}) as tr:
|
||||
assert tr.tracing_api_key == "tracing-key"
|
||||
with custom_span(name="span_with_key") as span:
|
||||
assert span.tracing_api_key == "tracing-key"
|
||||
|
||||
|
||||
def test_trace_metadata_propagates_to_spans():
|
||||
metadata = {"source": "run"}
|
||||
with trace(workflow_name="test", metadata=metadata) as current_trace:
|
||||
with custom_span(name="direct_child", parent=current_trace) as direct_child:
|
||||
assert direct_child.trace_metadata == metadata
|
||||
direct_child_export = direct_child.export()
|
||||
assert direct_child_export is not None
|
||||
assert "metadata" not in direct_child_export
|
||||
with custom_span(name="parent") as parent:
|
||||
assert parent.trace_metadata == metadata
|
||||
parent_export = parent.export()
|
||||
assert parent_export is not None
|
||||
assert "metadata" not in parent_export
|
||||
with custom_span(name="child", parent=parent) as child:
|
||||
assert child.trace_metadata == metadata
|
||||
child_export = child.export()
|
||||
assert child_export is not None
|
||||
assert "metadata" not in child_export
|
||||
|
||||
|
||||
def test_agent_span_metadata_exports_with_routing_metadata():
|
||||
routing_metadata = {
|
||||
"agent_harness_id": "harness_123",
|
||||
}
|
||||
with trace(
|
||||
workflow_name="test",
|
||||
metadata={
|
||||
**routing_metadata,
|
||||
"agent_id": "agent_123",
|
||||
"agent_task_id": "task_123",
|
||||
"tenant_id": "tenant_123",
|
||||
"user_id": "user_123",
|
||||
},
|
||||
):
|
||||
with agent_span(name="agent") as span:
|
||||
span.span_data.metadata = {
|
||||
"usage": {
|
||||
"requests": 1,
|
||||
"input_tokens": 10,
|
||||
"output_tokens": 4,
|
||||
"total_tokens": 14,
|
||||
"cached_input_tokens": 3,
|
||||
}
|
||||
}
|
||||
|
||||
span_export = span.export()
|
||||
|
||||
assert span_export is not None
|
||||
assert span_export["metadata"] == {
|
||||
**routing_metadata,
|
||||
"usage": {
|
||||
"requests": 1,
|
||||
"input_tokens": 10,
|
||||
"output_tokens": 4,
|
||||
"total_tokens": 14,
|
||||
"cached_input_tokens": 3,
|
||||
},
|
||||
}
|
||||
|
||||
|
||||
def test_processor_can_lookup_trace_metadata_by_span_trace_id():
|
||||
class MetadataPropagatingProcessor(TracingProcessor):
|
||||
def __init__(self) -> None:
|
||||
self.trace_metadata_by_id: dict[str, dict[str, Any]] = {}
|
||||
self.looked_up_metadata: dict[str, Any] | None = None
|
||||
self.span_trace_metadata: dict[str, Any] | None = None
|
||||
|
||||
def on_trace_start(self, trace: Trace) -> None:
|
||||
trace_metadata = getattr(trace, "metadata", None)
|
||||
if trace_metadata:
|
||||
self.trace_metadata_by_id[trace.trace_id] = dict(trace_metadata)
|
||||
|
||||
def on_trace_end(self, trace: Trace) -> None:
|
||||
return None
|
||||
|
||||
def on_span_start(self, span: Span[Any]) -> None:
|
||||
return None
|
||||
|
||||
def on_span_end(self, span: Span[Any]) -> None:
|
||||
if span.span_data.type != "agent":
|
||||
return
|
||||
self.looked_up_metadata = self.trace_metadata_by_id.get(span.trace_id)
|
||||
self.span_trace_metadata = span.trace_metadata
|
||||
|
||||
def shutdown(self) -> None:
|
||||
return None
|
||||
|
||||
def force_flush(self) -> None:
|
||||
return None
|
||||
|
||||
metadata = {
|
||||
"user_id": "u_123",
|
||||
"chat_type": "support",
|
||||
}
|
||||
processor = MetadataPropagatingProcessor()
|
||||
set_trace_processors([processor])
|
||||
try:
|
||||
with trace(workflow_name="workflow", metadata=metadata):
|
||||
with agent_span(name="agent"):
|
||||
pass
|
||||
finally:
|
||||
set_trace_processors([SPAN_PROCESSOR_TESTING])
|
||||
|
||||
assert processor.looked_up_metadata == metadata
|
||||
assert processor.span_trace_metadata == metadata
|
||||
|
||||
|
||||
def test_trace_to_json_only_includes_tracing_api_key_when_requested():
|
||||
with trace(workflow_name="test", tracing={"api_key": "secret-key"}) as tr:
|
||||
default_json = tr.to_json()
|
||||
assert default_json is not None
|
||||
assert "tracing_api_key" not in default_json
|
||||
|
||||
with_key = tr.to_json(include_tracing_api_key=True)
|
||||
assert with_key is not None
|
||||
assert with_key["tracing_api_key"] == "secret-key"
|
||||
Reference in New Issue
Block a user