404 lines
14 KiB
Python
404 lines
14 KiB
Python
"""Structures for data captured from LLM calls via telemetry.
|
|
|
|
All capture is best-effort: `dispatch_captured_call` swallows exporter exceptions
|
|
so telemetry can never break the LLM call path.
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
import base64
|
|
import contextlib
|
|
import hashlib
|
|
import json
|
|
import logging
|
|
from dataclasses import dataclass, field
|
|
from typing import Any, Protocol, cast, runtime_checkable
|
|
|
|
from src.config import settings
|
|
|
|
from .backend import CompletionResult as BackendCompletionResult
|
|
from .backend import ToolCallResult
|
|
from .types import LLMTelemetryContext
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
# Sentinel roles for non-message content stored in the shared content store so
|
|
# the same hash+dedup machinery covers them. They never collide with real
|
|
# conversation roles ("user"/"assistant"/"system"/"tool").
|
|
ROLE_OUTPUT = "assistant"
|
|
ROLE_TOOL_SCHEMA = "__tool_schema__"
|
|
ROLE_THINKING = "__thinking__"
|
|
|
|
|
|
def canonical_json(obj: Any) -> str:
|
|
"""Deterministic JSON encoding used for every content hash."""
|
|
return json.dumps(
|
|
obj, sort_keys=True, separators=(",", ":"), ensure_ascii=False, default=str
|
|
)
|
|
|
|
|
|
def compute_content_hash(
|
|
role: str,
|
|
content: Any,
|
|
tool_call_id: str | None,
|
|
tool_calls: list[dict[str, Any]] | None = None,
|
|
) -> str:
|
|
"""Content hash covering the FULL message identity, not just the text.
|
|
|
|
Includes `tool_calls` so two assistant turns with identical (often empty)
|
|
content but different tool calls don't collide in the dedup store.
|
|
"""
|
|
digest = hashlib.sha256(
|
|
canonical_json(
|
|
{
|
|
"role": role,
|
|
"content": content,
|
|
"tool_call_id": tool_call_id,
|
|
"tool_calls": tool_calls or [],
|
|
}
|
|
).encode("utf-8")
|
|
).hexdigest()
|
|
return f"sha256:{digest}"
|
|
|
|
|
|
def clip_for_trace(content: Any) -> tuple[Any, bool]:
|
|
"""Clip a content value to `TELEMETRY.TRACE_MAX_BYTES`, returning (content, truncated).
|
|
|
|
Only oversized string content is clipped (with a marker); non-string
|
|
structured content is left intact. Returns the input unchanged when it
|
|
fits or when the cap is non-positive.
|
|
"""
|
|
max_bytes = settings.TELEMETRY.TRACE_MAX_BYTES
|
|
if max_bytes <= 0 or not isinstance(content, str):
|
|
return content, False
|
|
encoded = content.encode("utf-8")
|
|
if len(encoded) <= max_bytes:
|
|
return content, False
|
|
marker = "…[truncated]"
|
|
keep = max(0, max_bytes - len(marker.encode("utf-8")))
|
|
clipped = encoded[:keep].decode("utf-8", errors="ignore") + marker
|
|
return clipped, True
|
|
|
|
|
|
@dataclass(slots=True)
|
|
class CapturedMessage:
|
|
"""One input message, normalized to a provider-agnostic shape.
|
|
|
|
`content` is the message text; `tool_calls` holds any tool calls in a
|
|
unified `{id, name, input}` shape regardless of provider. `content_hash`
|
|
covers all identity fields so the ref and the shipped `trace.content` agree.
|
|
"""
|
|
|
|
role: str
|
|
content: Any
|
|
tool_call_id: str | None
|
|
content_hash: str
|
|
truncated: bool = False
|
|
tool_calls: list[dict[str, Any]] = field(default_factory=list)
|
|
|
|
|
|
@dataclass(slots=True)
|
|
class CapturedLLMCall:
|
|
"""Everything one LLM call needs to be reconstructed, captured once."""
|
|
|
|
# Correlation (span tree)
|
|
trace_id: str | None
|
|
span_id: str | None
|
|
parent_span_id: str | None
|
|
iteration: int | None
|
|
step_seq: int
|
|
attempt: int
|
|
was_fallback: bool
|
|
run_id: str | None
|
|
# Path identity
|
|
workspace_name: str | None
|
|
call_purpose: str | None
|
|
parent_category: str | None
|
|
agent_type: str | None
|
|
# unique session ID for grouping traces
|
|
session_id: str | None
|
|
observer: str | None
|
|
observed: str | None
|
|
peer_name: str | None
|
|
track_name: str | None
|
|
transport: str
|
|
provider_label: str | None
|
|
model: str
|
|
# Context window
|
|
input_messages: list[CapturedMessage]
|
|
tool_schemas: list[dict[str, Any]]
|
|
tool_choice: Any
|
|
# Output (replay-grade)
|
|
output_content: Any
|
|
output_tool_calls: list[dict[str, Any]]
|
|
thinking_content: str | None
|
|
thinking_blocks: list[dict[str, Any]]
|
|
reasoning_details: list[dict[str, Any]]
|
|
finish_reason: str | None
|
|
# Accounting copy (so the trace stream stands alone)
|
|
input_tokens: int
|
|
output_tokens: int
|
|
cache_read_tokens: int
|
|
cache_creation_tokens: int
|
|
was_stream: bool
|
|
# True when any input message was clipped to TRACE_MAX_BYTES.
|
|
input_truncated: bool = False
|
|
|
|
|
|
def _normalize_message(
|
|
message: dict[str, Any], transport: str | None
|
|
) -> tuple[Any, str | None, list[dict[str, Any]]]:
|
|
"""Normalize a provider-native message to (content, tool_call_id, tool_calls).
|
|
|
|
Providers stash tool calls and results outside `content` (openai's
|
|
`tool_calls`, gemini's `parts`), so a naive `content` read loses them. This
|
|
lifts them into a unified shape: `content` becomes text, `tool_calls` is a
|
|
list of `{id, name, input}`, and tool results surface as `content` keyed by
|
|
`tool_call_id`.
|
|
"""
|
|
content: Any = message.get("content")
|
|
tool_call_id: str | None = message.get("tool_call_id")
|
|
tool_calls: list[dict[str, Any]] = []
|
|
|
|
if transport == "openai":
|
|
for tc in cast("list[dict[str, Any]]", message.get("tool_calls") or []):
|
|
fn = cast("dict[str, Any]", tc.get("function") or {})
|
|
args = fn.get("arguments")
|
|
if isinstance(args, str):
|
|
with contextlib.suppress(json.JSONDecodeError):
|
|
args = json.loads(args)
|
|
tool_calls.append(
|
|
{"id": tc.get("id"), "name": fn.get("name"), "input": args}
|
|
)
|
|
|
|
elif transport == "gemini":
|
|
parts = message.get("parts")
|
|
if isinstance(parts, list):
|
|
texts: list[str] = []
|
|
results: list[Any] = []
|
|
for raw_part in cast("list[Any]", parts):
|
|
if not isinstance(raw_part, dict):
|
|
continue
|
|
part = cast("dict[str, Any]", raw_part)
|
|
text = part.get("text")
|
|
if isinstance(text, str):
|
|
texts.append(text)
|
|
elif "function_call" in part:
|
|
fc = cast("dict[str, Any]", part["function_call"] or {})
|
|
tool_calls.append(
|
|
{"id": None, "name": fc.get("name"), "input": fc.get("args")}
|
|
)
|
|
elif "function_response" in part:
|
|
fr = cast("dict[str, Any]", part["function_response"] or {})
|
|
resp = fr.get("response")
|
|
if isinstance(resp, dict):
|
|
results.append(cast("dict[str, Any]", resp).get("result"))
|
|
else:
|
|
results.append(resp)
|
|
if tool_call_id is None:
|
|
tool_call_id = fr.get("name")
|
|
content = "\n".join(texts) if texts else (results[0] if results else None)
|
|
|
|
elif transport == "anthropic" and isinstance(content, list):
|
|
texts = []
|
|
for raw_block in cast("list[Any]", content):
|
|
if not isinstance(raw_block, dict):
|
|
continue
|
|
block = cast("dict[str, Any]", raw_block)
|
|
btype = block.get("type")
|
|
text = block.get("text")
|
|
if btype == "text" and isinstance(text, str):
|
|
texts.append(text)
|
|
elif btype == "tool_use":
|
|
tool_calls.append(
|
|
{
|
|
"id": block.get("id"),
|
|
"name": block.get("name"),
|
|
"input": block.get("input"),
|
|
}
|
|
)
|
|
elif btype == "tool_result":
|
|
if tool_call_id is None:
|
|
tool_call_id = block.get("tool_use_id")
|
|
inner = block.get("content")
|
|
texts.append(inner if isinstance(inner, str) else canonical_json(inner))
|
|
content = "\n".join(texts) if texts else None
|
|
|
|
return content, tool_call_id, tool_calls
|
|
|
|
|
|
def build_captured_messages(
|
|
messages: list[dict[str, Any]],
|
|
memo: dict[int, CapturedMessage] | None,
|
|
transport: str | None = None,
|
|
) -> tuple[list[CapturedMessage], bool]:
|
|
"""Create a list of CapturedMessage from LLM response messages.
|
|
|
|
Conversation is append-only. Uses hashed message content to deduplicate
|
|
across turns. Messages are normalized per provider, then content is
|
|
truncated and hashed.
|
|
"""
|
|
captured: list[CapturedMessage] = []
|
|
any_truncated = False
|
|
for message in messages:
|
|
key = id(message)
|
|
cached = memo.get(key) if memo is not None else None
|
|
if cached is not None:
|
|
captured.append(cached)
|
|
any_truncated = any_truncated or cached.truncated
|
|
continue
|
|
role = str(message.get("role", ""))
|
|
raw_content, tool_call_id, tool_calls = _normalize_message(message, transport)
|
|
content, truncated = clip_for_trace(raw_content)
|
|
any_truncated = any_truncated or truncated
|
|
captured_message = CapturedMessage(
|
|
role=role,
|
|
content=content,
|
|
tool_call_id=tool_call_id,
|
|
content_hash=compute_content_hash(role, content, tool_call_id, tool_calls),
|
|
truncated=truncated,
|
|
tool_calls=tool_calls,
|
|
)
|
|
if memo is not None:
|
|
memo[key] = captured_message
|
|
captured.append(captured_message)
|
|
return captured, any_truncated
|
|
|
|
|
|
def build_captured_call(
|
|
*,
|
|
telemetry: LLMTelemetryContext | None,
|
|
transport: str,
|
|
provider_label: str | None,
|
|
model: str,
|
|
messages: list[dict[str, Any]],
|
|
tools: list[dict[str, Any]] | None,
|
|
tool_choice: Any,
|
|
result: BackendCompletionResult | None,
|
|
attempt: int,
|
|
was_fallback: bool,
|
|
was_stream: bool,
|
|
finish_reason: str | None,
|
|
) -> CapturedLLMCall:
|
|
"""Assemble a `CapturedLLMCall` from telemetry + the provider result."""
|
|
memo = telemetry.hash_memo if telemetry is not None else None
|
|
captured_messages, input_truncated = build_captured_messages(
|
|
messages, memo, transport
|
|
)
|
|
|
|
output_tool_calls = [
|
|
_tool_call_to_dict(tc) for tc in (result.tool_calls if result else [])
|
|
]
|
|
|
|
return CapturedLLMCall(
|
|
trace_id=telemetry.trace_id if telemetry else None,
|
|
span_id=telemetry.span_id if telemetry else None,
|
|
parent_span_id=telemetry.exported_parent_span_id() if telemetry else None,
|
|
iteration=telemetry.iteration if telemetry else None,
|
|
step_seq=telemetry.step_seq if telemetry else 0,
|
|
attempt=attempt,
|
|
was_fallback=was_fallback,
|
|
run_id=telemetry.run_id if telemetry else None,
|
|
workspace_name=telemetry.workspace_name if telemetry else None,
|
|
call_purpose=telemetry.call_purpose if telemetry else None,
|
|
parent_category=telemetry.parent_category if telemetry else None,
|
|
agent_type=telemetry.agent_type if telemetry else None,
|
|
session_id=telemetry.session_id if telemetry else None,
|
|
observer=telemetry.observer if telemetry else None,
|
|
observed=telemetry.observed if telemetry else None,
|
|
peer_name=telemetry.peer_name if telemetry else None,
|
|
track_name=telemetry.track_name if telemetry else None,
|
|
transport=transport,
|
|
provider_label=provider_label,
|
|
model=model,
|
|
input_messages=captured_messages,
|
|
tool_schemas=list(tools) if tools else [],
|
|
tool_choice=tool_choice,
|
|
output_content=result.content if result else None,
|
|
output_tool_calls=output_tool_calls,
|
|
thinking_content=result.thinking_content if result else None,
|
|
thinking_blocks=result.thinking_blocks if result else [],
|
|
reasoning_details=result.reasoning_details if result else [],
|
|
finish_reason=finish_reason,
|
|
input_tokens=result.input_tokens if result else 0,
|
|
output_tokens=result.output_tokens if result else 0,
|
|
cache_read_tokens=result.cache_read_input_tokens if result else 0,
|
|
cache_creation_tokens=result.cache_creation_input_tokens if result else 0,
|
|
was_stream=was_stream,
|
|
input_truncated=input_truncated,
|
|
)
|
|
|
|
|
|
def _tool_call_to_dict(tool_call: ToolCallResult) -> dict[str, Any]:
|
|
"""Normalize a ToolCallResult to a JSON-safe dict for the trace stream.
|
|
|
|
`thought_signature` arrives as raw bytes from Gemini; base64-encode it so
|
|
CloudEvents JSON serialization can't choke on non-UTF8 bytes (which would
|
|
silently drop the whole event via the best-effort emit path).
|
|
"""
|
|
out: dict[str, Any] = {
|
|
"id": tool_call.id,
|
|
"name": tool_call.name,
|
|
"input": tool_call.input,
|
|
}
|
|
sig = tool_call.thought_signature
|
|
if sig is not None:
|
|
out["thought_signature"] = (
|
|
base64.b64encode(sig).decode("ascii") if isinstance(sig, bytes) else sig
|
|
)
|
|
return out
|
|
|
|
|
|
@runtime_checkable
|
|
class LLMCallExporter(Protocol):
|
|
"""A sink that consumes a `CapturedLLMCall`"""
|
|
|
|
def export(self, call: CapturedLLMCall) -> None: ...
|
|
|
|
|
|
_EXPORTERS: list[LLMCallExporter] = []
|
|
|
|
|
|
def register_exporter(exporter: LLMCallExporter) -> None:
|
|
"""Register an exporter (idempotent on identity). Called at startup."""
|
|
if exporter not in _EXPORTERS:
|
|
_EXPORTERS.append(exporter)
|
|
|
|
|
|
def clear_exporters() -> None:
|
|
"""Drop all exporters — used on shutdown and in tests."""
|
|
_EXPORTERS.clear()
|
|
|
|
|
|
def has_exporters() -> bool:
|
|
"""True when at least one exporter is registered."""
|
|
return bool(_EXPORTERS)
|
|
|
|
|
|
def dispatch_captured_call(call: CapturedLLMCall) -> None:
|
|
"""Fan a captured call out to every exporter."""
|
|
for exporter in _EXPORTERS:
|
|
try:
|
|
exporter.export(call)
|
|
except Exception: # pragma: no cover - best-effort telemetry
|
|
logger.debug("LLM call exporter failed", exc_info=True)
|
|
|
|
|
|
__all__ = [
|
|
"ROLE_OUTPUT",
|
|
"ROLE_THINKING",
|
|
"ROLE_TOOL_SCHEMA",
|
|
"CapturedLLMCall",
|
|
"CapturedMessage",
|
|
"LLMCallExporter",
|
|
"build_captured_call",
|
|
"build_captured_messages",
|
|
"canonical_json",
|
|
"clear_exporters",
|
|
"clip_for_trace",
|
|
"compute_content_hash",
|
|
"dispatch_captured_call",
|
|
"has_exporters",
|
|
"register_exporter",
|
|
]
|