66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257 | class GraphObserver:
"""Assemble a :class:`Run` from graph events and forward it to an exporter."""
def __init__(
self,
name: str,
*,
exporter: TraceExporter | None = None,
topology: GraphTopology | None = None,
owner: str | None = None,
checkpoint_id: str | None = None,
redact: bool = True,
redact_fn: Callable[[Any], Any] = _default_redact,
):
self.name = name
self.exporter = exporter
self.topology = topology or GraphTopology()
self.owner = owner
self.checkpoint_id = checkpoint_id
self._redact_fn = redact_fn if redact else (lambda value: value)
self.tracer = RunTracer()
self._start = time.monotonic()
self._wall_start = time.time()
self._spans: dict[str, NodeSpan] = {}
self._active: list[NodeSpan] = []
self._tool_seen: dict[str, dict[str, ToolCall]] = {}
self._status = "ok"
self._error: str | None = None
self._total_ms = 0.0
self._wire_tracer()
async def on_llm_payload(
self,
provider: str,
model: str,
messages: list[dict[str, Any]],
response: str,
usage: dict[str, Any],
latency_ms: float,
cached: bool,
) -> None:
"""Sink for the run's ``on_llm_payload`` channel."""
redact = self._redact_fn
call = LLMCall(
node_id=self._active[-1].node_id if self._active else None,
provider=provider,
model=model,
messages=redact(messages),
response=str(redact(response) or ""),
prompt_tokens=int(usage.get("prompt", 0) or 0),
completion_tokens=int(usage.get("completion", 0) or 0),
latency_ms=latency_ms,
cached=cached,
)
if self._active:
span = self._active[-1]
# Tool calls are discovered in the *next* call's request payload
# (they ran after the previous reply), so capture them first to
# keep the event list in real chronological order.
self._capture_tool_calls(span, messages)
span.llm_calls.append(call)
span.events.append(SpanEvent(kind="llm", index=len(span.llm_calls) - 1))
def _capture_tool_calls(
self, span: NodeSpan, messages: list[dict[str, Any]]
) -> None:
"""Extract tool calls from an LLM payload into *span*.
An assistant ``tool_calls`` block and its matching ``role: tool``
result can arrive in *different* payloads (the result is appended
before the next model call), so already-seen calls are backfilled
with their result instead of duplicated.
"""
results: dict[str, str] = {}
for msg in messages:
if msg.get("role") == "tool":
results[str(msg.get("tool_call_id") or "")] = str(
msg.get("content") or ""
)
seen = self._tool_seen.setdefault(span.node_id, {})
for msg in messages:
if msg.get("role") != "assistant":
continue
for tc in msg.get("tool_calls") or []:
name, raw_args, call_id = _tool_call_parts(tc)
result = results.get(call_id)
existing = seen.get(call_id)
if existing is not None:
if result is not None and not existing.result:
existing.result = self._redact_fn(result)
existing.ok = not _is_tool_error(result)
continue
call = ToolCall(
name=name,
args=self._redact_fn(raw_args),
result=self._redact_fn(result) if result else "",
ok=not (result and _is_tool_error(result)),
)
seen[call_id] = call
span.tool_calls.append(call)
span.events.append(
SpanEvent(kind="tool", index=len(span.tool_calls) - 1)
)
def _start_node(self, node_id: str, node_type: str) -> None:
span = self._spans.get(node_id)
if span is None:
span = NodeSpan(
node_id=node_id,
node_type=node_type,
start_ms=(time.monotonic() - self._start) * 1000.0,
)
self._spans[node_id] = span
# A node can be visited many times in one run (react loops, retries);
# reuse the span so its LLM calls, tool calls and events accumulate
# in chronological order instead of keeping only the last visit.
self._active.append(span)
def _end_node(
self, node_id: str, status: str = "ok", error: str | None = None
) -> None:
span = self._spans.get(node_id)
if span is None:
return
span.end_ms = (time.monotonic() - self._start) * 1000.0
span.status = status
span.error = error
if self._active and self._active[-1] is span:
self._active.pop()
def _wire_tracer(self) -> None:
tracer = self.tracer
node_start = tracer.node_start
node_end = tracer.node_end
node_error = tracer.node_error
run_end = tracer.run_end
def _node_start(node_id, node_type):
node_start(node_id, node_type)
self._start_node(node_id, node_type)
def _node_end(node_id, node_type, duration_ms):
node_end(node_id, node_type, duration_ms)
self._end_node(node_id)
def _node_error(node_id, node_type, duration_ms, error):
node_error(node_id, node_type, duration_ms, error)
self._end_node(node_id, status="error", error=str(error))
def _run_end(status, total_ms, error=None):
run_end(status, total_ms, error)
self._status = status
self._total_ms = total_ms
if error is not None:
self._error = str(error)
for span in list(self._active):
self._end_node(span.node_id)
tracer.node_start = _node_start # type: ignore[method-assign]
tracer.node_end = _node_end # type: ignore[method-assign]
tracer.node_error = _node_error # type: ignore[method-assign]
tracer.run_end = _run_end # type: ignore[method-assign]
def build(self) -> Run:
return Run(
name=self.name,
status=self._status,
total_ms=self._total_ms,
owner=self.owner,
checkpoint_id=self.checkpoint_id,
created_at=self._wall_start,
topology=self.topology,
nodes=list(self._spans.values()),
)
def export(self) -> str | None:
"""Persist the collected :class:`Run` and return its backend run id.
Returns ``None`` when no exporter is attached (tracing disabled) or
the backend does not expose ids.
"""
if self.exporter is None:
return None
return self.exporter.export(self.build())
def close(self) -> None:
if self.exporter is not None:
self.exporter.close()
|