Skip to content

Commit 944254d

Browse files
rohitprasad15rohitprasad15
authored
Merge pull request #294 from drp8226/drp/viewer-and-cli
Add agent runs viewer and aisuite-code CLI
2 parents 4034290 + eb160aa commit 944254d

35 files changed

Lines changed: 10269 additions & 85 deletions

‎aisuite/tracing/__init__.py‎

Lines changed: 6 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1,23 +1,28 @@
11
from .sinks import (
22
InMemoryTraceSink,
3+
HttpTraceSink,
34
LocalTraceSink,
45
TRACE_SCHEMA_VERSION,
56
TraceEvent,
67
TraceSink,
8+
TraceStoreSink,
79
configure,
810
get_configured_sinks,
911
)
10-
from .store import JsonlTraceStore, TraceStore
12+
from .store import InMemoryTraceStore, JsonlTraceStore, TraceStore
1113
from .viewer import ViewerServer, read_trace_file, start_viewer
1214

1315
__all__ = [
1416
"InMemoryTraceSink",
17+
"HttpTraceSink",
18+
"InMemoryTraceStore",
1519
"JsonlTraceStore",
1620
"LocalTraceSink",
1721
"TRACE_SCHEMA_VERSION",
1822
"TraceEvent",
1923
"TraceSink",
2024
"TraceStore",
25+
"TraceStoreSink",
2126
"ViewerServer",
2227
"configure",
2328
"get_configured_sinks",

‎aisuite/tracing/sinks.py‎

Lines changed: 54 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -1,13 +1,16 @@
11
from __future__ import annotations
22

33
import copy
4+
import json
45
import uuid
56
from dataclasses import dataclass, field
67
from datetime import datetime, timezone
78
from pathlib import Path
89
from typing import Any, Literal, Optional, Protocol
10+
from urllib.error import URLError
11+
from urllib.request import Request, urlopen
912

10-
from .store import JsonlTraceStore
13+
from .store import JsonlTraceStore, TraceStore
1114

1215

1316
TRACE_SCHEMA_VERSION = "2026-05-15"
@@ -17,11 +20,14 @@
1720
"run.started",
1821
"run.completed",
1922
"run.failed",
20-
"model.started",
21-
"model.completed",
23+
"model.send",
24+
"model.response",
25+
"model.error",
2226
"tool.allowed",
2327
"tool.denied",
28+
"tool.started",
2429
"tool.completed",
30+
"tool.failed",
2531
]
2632

2733

@@ -83,6 +89,51 @@ def emit(self, event: TraceEvent) -> None:
8389
self.store.write_event(event)
8490

8591

92+
class TraceStoreSink:
93+
def __init__(self, store: TraceStore):
94+
self.store = store
95+
96+
def emit(self, event: TraceEvent) -> None:
97+
self.store.append_event(event)
98+
99+
100+
class HttpTraceSink:
101+
def __init__(
102+
self,
103+
endpoint: str,
104+
*,
105+
timeout: float = 2.0,
106+
headers: Optional[dict[str, str]] = None,
107+
fail_silently: bool = True,
108+
):
109+
self.endpoint = endpoint
110+
self.timeout = timeout
111+
self.headers = headers or {}
112+
self.fail_silently = fail_silently
113+
114+
def emit(self, event: TraceEvent) -> None:
115+
body = json.dumps(event.to_dict()).encode("utf-8")
116+
headers = {
117+
"Content-Type": "application/json",
118+
**self.headers,
119+
}
120+
request = Request(
121+
self.endpoint,
122+
data=body,
123+
headers=headers,
124+
method="POST",
125+
)
126+
try:
127+
with urlopen(request, timeout=self.timeout) as response:
128+
if response.status >= 400:
129+
raise RuntimeError(
130+
f"Trace HTTP sink failed with status {response.status}"
131+
)
132+
except (OSError, RuntimeError, URLError):
133+
if not self.fail_silently:
134+
raise
135+
136+
86137
class InMemoryTraceSink:
87138
def __init__(self):
88139
self.events: list[TraceEvent] = []

‎aisuite/tracing/static/viewer/assets/index-C29aAR1d.js‎

Lines changed: 10 additions & 0 deletions
Some generated files are not rendered by default. Learn more about customizing how changed files appear on GitHub.

‎aisuite/tracing/static/viewer/assets/index-C2bMCWqj.css‎

Lines changed: 1 addition & 0 deletions
Some generated files are not rendered by default. Learn more about customizing how changed files appear on GitHub.

‎aisuite/tracing/static/viewer/index.html‎

Lines changed: 13 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,13 @@
1+
<!doctype html>
2+
<html lang="en">
3+
<head>
4+
<meta charset="UTF-8" />
5+
<meta name="viewport" content="width=device-width, initial-scale=1.0" />
6+
<title>aisuite runs</title>
7+
<script type="module" crossorigin src="/assets/index-C29aAR1d.js"></script>
8+
<link rel="stylesheet" crossorigin href="/assets/index-C2bMCWqj.css">
9+
</head>
10+
<body>
11+
<div id="root"></div>
12+
</body>
13+
</html>

‎aisuite/tracing/store.py‎

Lines changed: 108 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -7,7 +7,29 @@
77

88

99
class TraceStore(Protocol):
10-
def write_event(self, event: Any) -> None:
10+
"""Storage contract used by the local viewer and trace sinks.
11+
12+
Implementations may store raw trace_event records, legacy run snapshots,
13+
or both. Read methods return reconstructed run dictionaries and raw event
14+
records so the viewer API does not depend on the physical storage backend.
15+
"""
16+
17+
def append_event(self, event: Any) -> None:
18+
...
19+
20+
def append_events(self, events: list[Any]) -> None:
21+
...
22+
23+
def append_record(self, record: dict[str, Any]) -> None:
24+
...
25+
26+
def append_records(self, records: list[dict[str, Any]]) -> None:
27+
...
28+
29+
def import_jsonl(self, content: str) -> int:
30+
...
31+
32+
def list_records(self) -> list[dict[str, Any]]:
1133
...
1234

1335
def list_runs(self) -> list[dict[str, Any]]:
@@ -21,13 +43,38 @@ def list_events(self, trace_id: str) -> list[dict[str, Any]]:
2143

2244

2345
class JsonlTraceStore:
46+
"""TraceStore implementation backed by a local JSONL file."""
47+
2448
def __init__(self, path: str | Path = ".aisuite/events.jsonl"):
2549
self.path = Path(path)
2650

27-
def write_event(self, event: Any) -> None:
51+
def append_event(self, event: Any) -> None:
52+
self.append_record(event.to_dict())
53+
54+
def append_events(self, events: list[Any]) -> None:
55+
self.append_records([event.to_dict() for event in events])
56+
57+
def append_record(self, record: dict[str, Any]) -> None:
58+
self.append_records([record])
59+
60+
def append_records(self, records: list[dict[str, Any]]) -> None:
61+
if not records:
62+
return
2863
self.path.parent.mkdir(parents=True, exist_ok=True)
2964
with self.path.open("a", encoding="utf-8") as handle:
30-
handle.write(json.dumps(event.to_dict()) + "\n")
65+
for record in records:
66+
handle.write(json.dumps(record) + "\n")
67+
68+
def write_event(self, event: Any) -> None:
69+
self.append_event(event)
70+
71+
def write_record(self, record: dict[str, Any]) -> None:
72+
self.append_record(record)
73+
74+
def import_jsonl(self, content: str) -> int:
75+
records = parse_jsonl_records(content)
76+
self.append_records(records)
77+
return len(records)
3178

3279
def list_records(self) -> list[dict[str, Any]]:
3380
if not self.path.exists():
@@ -63,6 +110,64 @@ def list_events(self, trace_id: str) -> list[dict[str, Any]]:
63110
]
64111

65112

113+
class InMemoryTraceStore:
114+
"""TraceStore implementation backed by an in-memory record list."""
115+
116+
def __init__(self, records: Optional[list[dict[str, Any]]] = None):
117+
self.records = list(records or [])
118+
119+
def append_event(self, event: Any) -> None:
120+
self.append_record(event.to_dict())
121+
122+
def append_events(self, events: list[Any]) -> None:
123+
self.append_records([event.to_dict() for event in events])
124+
125+
def append_record(self, record: dict[str, Any]) -> None:
126+
self.append_records([record])
127+
128+
def append_records(self, records: list[dict[str, Any]]) -> None:
129+
self.records.extend(copy.deepcopy(records))
130+
131+
def import_jsonl(self, content: str) -> int:
132+
records = parse_jsonl_records(content)
133+
self.append_records(records)
134+
return len(records)
135+
136+
def list_records(self) -> list[dict[str, Any]]:
137+
return copy.deepcopy(self.records)
138+
139+
def list_runs(self) -> list[dict[str, Any]]:
140+
return reconstruct_runs(self.list_records())
141+
142+
def get_run(self, trace_id: str) -> Optional[dict[str, Any]]:
143+
for run in self.list_runs():
144+
if run.get("trace_id") == trace_id:
145+
return run
146+
return None
147+
148+
def list_events(self, trace_id: str) -> list[dict[str, Any]]:
149+
return [
150+
record
151+
for record in self.list_records()
152+
if record.get("record_type") == "trace_event"
153+
and record.get("trace_id") == trace_id
154+
]
155+
156+
157+
def parse_jsonl_records(content: str) -> list[dict[str, Any]]:
158+
records = []
159+
for line in content.splitlines():
160+
line = line.strip()
161+
if not line:
162+
continue
163+
try:
164+
record = json.loads(line)
165+
except json.JSONDecodeError:
166+
continue
167+
records.append(record)
168+
return records
169+
170+
66171
def reconstruct_runs(records: list[dict[str, Any]]) -> list[dict[str, Any]]:
67172
snapshots = []
68173
events_by_trace_id: dict[str, list[dict[str, Any]]] = {}

0 commit comments

Comments
 (0)
⚡