andrewyng--aisuite
Devika
eb160aae74
Two developer surfaces on top of the agent runtime: a local web viewer for inspecting agent runs and traces, and a coding-agent CLI. Viewer: - aisuite/tracing: a local ViewerServer that serves a runs/trace UI, plus HttpTraceSink, TraceStoreSink, and InMemoryTraceStore - viewer-ui: React/Vite single-page app for browsing runs, with the prebuilt bundle shipped under aisuite/tracing/static/viewer CLI: - cli/py/aisuite-code-cli: the aisuite-code coding-agent CLI (app, agent loop, approval flow, config, rendering) as a standalone package - scripts/aisuite-code launcher and examples/cli usage - tests for the viewer server, trace sinks, and the CLI Co-authored-by: Rohit <rohit.prasad15@gmail.com>
251 行
7.5 KiB
Python
251 行
7.5 KiB
Python
import json
|
|
from types import SimpleNamespace
|
|
from urllib.error import URLError
|
|
from urllib.request import Request
|
|
from unittest.mock import Mock
|
|
|
|
import aisuite as ai
|
|
from aisuite.tracing.normalize import normalize_usage
|
|
from tests.agents.helpers import chat_response
|
|
|
|
|
|
def test_runner_emits_trace_events_to_memory_sink():
|
|
client = ai.Client()
|
|
client.chat.completions.create = Mock(return_value=chat_response("ok"))
|
|
sink = ai.tracing.InMemoryTraceSink()
|
|
agent = ai.Agent(name="assistant", model="openai:gpt-4o")
|
|
|
|
result = ai.Runner.run_sync(
|
|
agent,
|
|
"Hello",
|
|
client=client,
|
|
run_name="run",
|
|
group_id="group",
|
|
trace_sinks=[sink],
|
|
)
|
|
|
|
event_types = [event.event_type for event in sink.events]
|
|
assert event_types == [
|
|
"run.started",
|
|
"model.send",
|
|
"model.response",
|
|
"run.completed",
|
|
]
|
|
assert sink.events[0].trace_id == result.trace_id
|
|
assert sink.events[0].group_id == "group"
|
|
assert sink.events[1].data["input"]["items"][-1]["text_preview"] == "Hello"
|
|
assert sink.events[2].data["response"]["text_preview"] == "ok"
|
|
assert sink.events[-1].data["run"]["final_output"] == "ok"
|
|
|
|
|
|
def test_trace_store_sink_writes_events_to_store():
|
|
store = ai.tracing.InMemoryTraceStore()
|
|
sink = ai.tracing.TraceStoreSink(store)
|
|
event = ai.tracing.TraceEvent(
|
|
event_type="run.started",
|
|
trace_id="trace_1",
|
|
agent_name="assistant",
|
|
run_name="run",
|
|
)
|
|
|
|
sink.emit(event)
|
|
|
|
records = store.list_records()
|
|
assert records[0]["event_type"] == "run.started"
|
|
assert store.list_events("trace_1")[0]["event_id"] == event.event_id
|
|
|
|
|
|
def test_runner_can_emit_to_trace_store_sink_and_viewer_reads_same_store():
|
|
client = ai.Client()
|
|
client.chat.completions.create = Mock(return_value=chat_response("ok"))
|
|
store = ai.tracing.InMemoryTraceStore()
|
|
sink = ai.tracing.TraceStoreSink(store)
|
|
|
|
result = ai.Runner.run_sync(
|
|
ai.Agent(name="assistant", model="openai:gpt-4o"),
|
|
"Hello",
|
|
client=client,
|
|
run_name="store_run",
|
|
trace_sinks=[sink],
|
|
)
|
|
|
|
run = store.get_run(result.trace_id)
|
|
assert run["run_name"] == "store_run"
|
|
assert run["final_output"] == "ok"
|
|
assert [event["event_type"] for event in store.list_events(result.trace_id)] == [
|
|
"run.started",
|
|
"model.send",
|
|
"model.response",
|
|
"run.completed",
|
|
]
|
|
|
|
|
|
def test_normalize_usage_preserves_aisuite_and_provider_fields():
|
|
usage = normalize_usage(
|
|
SimpleNamespace(
|
|
prompt_tokens=1200,
|
|
completion_tokens=34,
|
|
total_tokens=None,
|
|
prompt_tokens_details={"cached_tokens": 100},
|
|
)
|
|
)
|
|
|
|
assert usage["input_tokens"] == 1200
|
|
assert usage["output_tokens"] == 34
|
|
assert usage["total_tokens"] == 1234
|
|
assert usage["prompt_tokens"] == 1200
|
|
assert usage["completion_tokens"] == 34
|
|
assert usage["provider_raw"]["prompt_tokens_details"] == {"cached_tokens": 100}
|
|
|
|
|
|
def test_local_trace_sink_writes_event_jsonl(tmp_path):
|
|
sink = ai.tracing.LocalTraceSink(tmp_path / "events.jsonl")
|
|
event = ai.tracing.TraceEvent(
|
|
event_type="run.started",
|
|
trace_id="trace_1",
|
|
agent_name="assistant",
|
|
group_id="group",
|
|
)
|
|
|
|
sink.emit(event)
|
|
|
|
line = (tmp_path / "events.jsonl").read_text(encoding="utf-8").splitlines()[0]
|
|
payload = json.loads(line)
|
|
assert payload["record_type"] == "trace_event"
|
|
assert payload["schema_version"] == ai.tracing.TRACE_SCHEMA_VERSION
|
|
assert payload["event_type"] == "run.started"
|
|
assert payload["trace_id"] == "trace_1"
|
|
|
|
|
|
def test_http_trace_sink_posts_event(monkeypatch):
|
|
calls = []
|
|
|
|
class Response:
|
|
status = 200
|
|
|
|
def __enter__(self):
|
|
return self
|
|
|
|
def __exit__(self, *_args):
|
|
return False
|
|
|
|
def fake_urlopen(request: Request, timeout):
|
|
calls.append((request, timeout))
|
|
return Response()
|
|
|
|
monkeypatch.setattr("aisuite.tracing.sinks.urlopen", fake_urlopen)
|
|
sink = ai.tracing.HttpTraceSink("http://127.0.0.1:8780/api/events", timeout=3)
|
|
event = ai.tracing.TraceEvent(
|
|
event_type="run.started",
|
|
trace_id="trace_1",
|
|
agent_name="assistant",
|
|
)
|
|
|
|
sink.emit(event)
|
|
|
|
request, timeout = calls[0]
|
|
assert request.full_url == "http://127.0.0.1:8780/api/events"
|
|
assert timeout == 3
|
|
payload = json.loads(request.data.decode("utf-8"))
|
|
assert payload["event_type"] == "run.started"
|
|
assert payload["trace_id"] == "trace_1"
|
|
|
|
|
|
def test_http_trace_sink_can_fail_silently(monkeypatch):
|
|
def fake_urlopen(_request, timeout=None):
|
|
raise URLError("offline")
|
|
|
|
monkeypatch.setattr("aisuite.tracing.sinks.urlopen", fake_urlopen)
|
|
sink = ai.tracing.HttpTraceSink("http://127.0.0.1:8780/api/events")
|
|
|
|
sink.emit(
|
|
ai.tracing.TraceEvent(
|
|
event_type="run.started",
|
|
trace_id="trace_1",
|
|
agent_name="assistant",
|
|
)
|
|
)
|
|
|
|
|
|
def test_jsonl_trace_store_lists_runs_and_events(tmp_path):
|
|
store = ai.tracing.JsonlTraceStore(tmp_path / "events.jsonl")
|
|
started = ai.tracing.TraceEvent(
|
|
event_type="run.started",
|
|
trace_id="trace_1",
|
|
agent_name="assistant",
|
|
group_id="group",
|
|
run_name="run",
|
|
)
|
|
completed = ai.tracing.TraceEvent(
|
|
event_type="run.completed",
|
|
trace_id="trace_1",
|
|
agent_name="assistant",
|
|
group_id="group",
|
|
run_name="run",
|
|
data={
|
|
"run": {
|
|
"trace_id": "trace_1",
|
|
"agent_name": "assistant",
|
|
"run_name": "run",
|
|
"group_id": "group",
|
|
"status": "completed",
|
|
"messages": [],
|
|
"steps": [],
|
|
"tags": [],
|
|
"metadata": {},
|
|
"final_output": "ok",
|
|
}
|
|
},
|
|
)
|
|
|
|
store.append_event(started)
|
|
store.append_events([completed])
|
|
|
|
records = store.list_records()
|
|
assert [record["event_type"] for record in records] == [
|
|
"run.started",
|
|
"run.completed",
|
|
]
|
|
|
|
run = store.get_run("trace_1")
|
|
assert run["status"] == "completed"
|
|
assert run["final_output"] == "ok"
|
|
assert [event["event_type"] for event in store.list_events("trace_1")] == [
|
|
"run.started",
|
|
"run.completed",
|
|
]
|
|
|
|
|
|
def test_trace_stores_import_jsonl_and_ignore_invalid_lines(tmp_path):
|
|
valid_event = ai.tracing.TraceEvent(
|
|
event_type="run.started",
|
|
trace_id="trace_1",
|
|
agent_name="assistant",
|
|
).to_dict()
|
|
content = "\n".join([json.dumps(valid_event), "not json", ""])
|
|
|
|
jsonl_store = ai.tracing.JsonlTraceStore(tmp_path / "events.jsonl")
|
|
memory_store = ai.tracing.InMemoryTraceStore()
|
|
|
|
assert jsonl_store.import_jsonl(content) == 1
|
|
assert memory_store.import_jsonl(content) == 1
|
|
assert jsonl_store.list_events("trace_1")[0]["event_type"] == "run.started"
|
|
assert memory_store.list_events("trace_1")[0]["event_type"] == "run.started"
|
|
|
|
|
|
def test_global_trace_configuration_is_used_by_runner():
|
|
client = ai.Client()
|
|
client.chat.completions.create = Mock(return_value=chat_response("ok"))
|
|
sink = ai.tracing.InMemoryTraceSink()
|
|
ai.tracing.configure(sink)
|
|
try:
|
|
ai.Runner.run_sync(
|
|
ai.Agent(name="assistant", model="openai:gpt-4o"),
|
|
"Hello",
|
|
client=client,
|
|
)
|
|
finally:
|
|
ai.tracing.configure()
|
|
|
|
assert [event.event_type for event in sink.events][-1] == "run.completed"
|