thinkless 0.2.0__py3-none-any.whl

thinkless/tracing/summary.py

@@ -0,0 +1,143 @@
+"""Roll-ups computed from the spans of one trace."""
+
+from __future__ import annotations
+
+import time
+from collections import Counter, defaultdict
+from collections.abc import Iterable, Mapping
+from typing import Any
+
+from pydantic import BaseModel, Field
+
+__all__ = ["TraceSummary", "summarize"]
+
+PLANES = ("rule", "model", "llm", "tool", "code")
+
+
+class TraceSummary(BaseModel):
+    """What a run cost and where its time went.
+
+    Attributes:
+        llm_calls: Every call to a generative model, whether it generated text
+            or answered decision questions.
+        generation_calls: LLM calls made through ``Engine.generate``.
+        decisions: Questions resolved in the run, plus deterministic checks
+            recorded with ``Engine.rule``.
+        decisions_by_plane: How many decisions each plane resolved.
+            ``unresolved`` counts abstentions.
+        uncertain: Decisions that ended below threshold on every provider.
+        escalations: Decisions where a provider answered below threshold (or
+            failed) and the next provider was asked.
+        llm_input_tokens, llm_output_tokens: Tokens spent on generative models.
+        model_input_tokens: Tokens processed by small decision models.
+        cost_usd: Estimated spend, from the configured price table.
+        time_by_plane_ms: Time inside leaf work (provider attempts, generations,
+            tools and rules) grouped by plane.
+    """
+
+    trace_id: str
+    name: str
+    status: str
+    started_at: str
+    duration_ms: float
+    attributes: dict[str, Any] = Field(default_factory=dict)
+    llm_calls: int = 0
+    generation_calls: int = 0
+    decisions: int = 0
+    decisions_by_plane: dict[str, int] = Field(default_factory=dict)
+    uncertain: int = 0
+    escalations: int = 0
+    tool_calls: int = 0
+    errors: int = 0
+    llm_input_tokens: int = 0
+    llm_output_tokens: int = 0
+    model_input_tokens: int = 0
+    cost_usd: float = 0.0
+    time_by_plane_ms: dict[str, float] = Field(default_factory=dict)
+
+    @property
+    def llm_tokens(self) -> int:
+        return self.llm_input_tokens + self.llm_output_tokens
+
+
+def _usage(attrs: Mapping[str, Any]) -> tuple[int, int]:
+    usage = attrs.get("usage") or {}
+    return int(usage.get("input_tokens") or 0), int(usage.get("output_tokens") or 0)
+
+
+def summarize(spans: Iterable[Mapping[str, Any]]) -> TraceSummary:
+    """Summarize one trace.
+
+    Args:
+        spans: Span dictionaries of a single trace, in any order.
+
+    Raises:
+        ValueError: If there are no spans.
+    """
+    spans = list(spans)
+    if not spans:
+        raise ValueError("cannot summarize an empty trace")
+    root = next((s for s in spans if s.get("parent_id") is None), None)
+    if root is None:
+        root = min(spans, key=lambda s: s.get("start_ms") or 0.0)
+
+    by_plane: Counter[str] = Counter()
+    time_by_plane: defaultdict[str, float] = defaultdict(float)
+    summary = TraceSummary(
+        trace_id=root["trace_id"],
+        name=root["name"],
+        status=root.get("status", "ok"),
+        started_at=time.strftime(
+            "%Y-%m-%dT%H:%M:%SZ", time.gmtime((root.get("start_ms") or 0.0) / 1000.0)
+        ),
+        duration_ms=float(root.get("duration_ms") or 0.0),
+        attributes={k: v for k, v in (root.get("attributes") or {}).items() if k != "input"},
+    )
+
+    for span in spans:
+        kind = span.get("kind")
+        plane = span.get("plane")
+        attrs = span.get("attributes") or {}
+        duration = float(span.get("duration_ms") or 0.0)
+        if span.get("status") == "error":
+            summary.errors += 1
+        summary.cost_usd += float(attrs.get("cost_usd") or 0.0)
+
+        if kind in ("attempt", "llm", "tool", "rule") and plane:
+            time_by_plane[plane] += duration
+
+        if kind == "llm":
+            summary.llm_calls += 1
+            summary.generation_calls += 1
+            inp, out = _usage(attrs)
+            summary.llm_input_tokens += inp
+            summary.llm_output_tokens += out
+        elif kind == "attempt":
+            inp, out = _usage(attrs)
+            if plane == "llm":
+                summary.llm_calls += 1
+                summary.llm_input_tokens += inp
+                summary.llm_output_tokens += out
+            elif plane == "model":
+                summary.model_input_tokens += inp
+        elif kind == "tool":
+            summary.tool_calls += 1
+        elif kind == "rule":
+            summary.decisions += 1
+            by_plane["rule"] += 1
+        elif kind == "decide":
+            for decision in attrs.get("decisions") or []:
+                summary.decisions += 1
+                by_plane[decision.get("plane") or "unresolved"] += 1
+                if decision.get("status") == "uncertain":
+                    summary.uncertain += 1
+                escalated = decision.get("escalated")
+                if escalated is None:
+                    escalated = len(decision.get("attempts") or []) > 1
+                if escalated:
+                    summary.escalations += 1
+
+    summary.cost_usd = round(summary.cost_usd, 8)
+    summary.decisions_by_plane = dict(by_plane)
+    summary.time_by_plane_ms = {k: round(v, 3) for k, v in time_by_plane.items()}
+    return summary

thinkless/tracing/tracer.py

@@ -0,0 +1,159 @@
+"""The tracer: opens spans, links them and hands them to sinks."""
+
+from __future__ import annotations
+
+import functools
+import inspect
+import logging
+from collections.abc import Callable, Iterator, Sequence
+from contextlib import contextmanager
+from typing import Any, TypeVar, overload
+
+from .._json import jsonable
+from .span import Span, _current_span, _current_tracer, current_tracer, new_span_id, new_trace_id
+
+__all__ = ["REDACTED", "Tracer", "jsonable", "tool"]
+
+logger = logging.getLogger("thinkless.tracing")
+
+REDACTED = "[redacted]"
+
+F = TypeVar("F", bound=Callable[..., Any])
+
+
+class Tracer:
+    """Creates spans and forwards them to sinks.
+
+    Args:
+        sinks: Where finished spans go. See :mod:`thinkless.tracing.sinks`.
+        capture_content: Record inputs, outputs, prompts and tool arguments.
+            Turn it off where traces must not hold user content; structure,
+            timings, confidences and token counts are still recorded.
+
+    Sink failures are logged and never propagate into application code.
+    """
+
+    def __init__(self, sinks: Sequence[Any] | None = None, *, capture_content: bool = True) -> None:
+        self.sinks: list[Any] = list(sinks or [])
+        self.capture_content = capture_content
+
+    def add_sink(self, sink: Any) -> None:
+        self.sinks.append(sink)
+
+    def content(self, value: Any) -> Any:
+        """Return ``value`` as JSON-safe data, or a marker when content capture is off."""
+        return jsonable(value) if self.capture_content else REDACTED
+
+    @contextmanager
+    def span(
+        self, kind: str, name: str, *, plane: str | None = None, **attributes: Any
+    ) -> Iterator[Span]:
+        parent = _current_span.get()
+        span = Span(
+            trace_id=parent.trace_id if parent else new_trace_id(),
+            span_id=new_span_id(),
+            parent_id=parent.span_id if parent else None,
+            kind=kind,
+            name=name,
+            plane=plane,
+            attributes={k: jsonable(v) for k, v in attributes.items()},
+        )
+        span_token = _current_span.set(span)
+        tracer_token = _current_tracer.set(self)
+        self._emit("on_start", span)
+        try:
+            yield span
+        except BaseException as exc:
+            span.fail(exc)
+            raise
+        finally:
+            span.finish()
+            _current_span.reset(span_token)
+            _current_tracer.reset(tracer_token)
+            self._emit("on_end", span)
+
+    def flush(self) -> None:
+        for sink in self.sinks:
+            flush = getattr(sink, "flush", None)
+            if flush is not None:
+                self._call(flush)
+
+    def close(self) -> None:
+        for sink in self.sinks:
+            close = getattr(sink, "close", None)
+            if close is not None:
+                self._call(close)
+
+    def _emit(self, hook: str, span: Span) -> None:
+        for sink in tuple(self.sinks):
+            method = getattr(sink, hook, None)
+            if method is not None:
+                self._call(method, span)
+
+    @staticmethod
+    def _call(fn: Callable[..., Any], *args: Any) -> None:
+        try:
+            fn(*args)
+        except Exception:
+            logger.exception("trace sink %r failed", fn)
+
+
+@overload
+def tool(fn: F, /) -> F: ...
+@overload
+def tool(*, name: str | None = None) -> Callable[[F], F]: ...
+
+
+def tool(fn: F | None = None, /, *, name: str | None = None) -> F | Callable[[F], F]:
+    """Record calls to a function as ``tool`` spans.
+
+    Outside a traced run the function runs untouched, so decorated tools stay
+    usable in tests and scripts.
+
+    Example:
+        >>> @tool
+        ... def refund_payment(payment_id: str, amount: float) -> dict: ...
+    """
+
+    def decorate(func: F) -> F:
+        span_name = name or func.__name__
+        signature = inspect.signature(func)
+
+        def arguments(args: tuple[Any, ...], kwargs: dict[str, Any]) -> dict[str, Any]:
+            try:
+                bound = signature.bind_partial(*args, **kwargs)
+            except TypeError:
+                return {"args": args, "kwargs": kwargs}
+            return {k: v for k, v in bound.arguments.items() if k not in ("self", "cls")}
+
+        if inspect.iscoroutinefunction(func):
+
+            @functools.wraps(func)
+            async def async_wrapper(*args: Any, **kwargs: Any) -> Any:
+                tracer = current_tracer()
+                if tracer is None:
+                    return await func(*args, **kwargs)
+                with tracer.span("tool", span_name, plane="tool") as span:
+                    span.set(arguments=tracer.content(arguments(args, kwargs)))
+                    result = await func(*args, **kwargs)
+                    span.set(result=tracer.content(result))
+                    return result
+
+            return async_wrapper  # type: ignore[return-value]
+
+        @functools.wraps(func)
+        def wrapper(*args: Any, **kwargs: Any) -> Any:
+            tracer = current_tracer()
+            if tracer is None:
+                return func(*args, **kwargs)
+            with tracer.span("tool", span_name, plane="tool") as span:
+                span.set(arguments=tracer.content(arguments(args, kwargs)))
+                result = func(*args, **kwargs)
+                span.set(result=tracer.content(result))
+                return result
+
+        return wrapper  # type: ignore[return-value]
+
+    if fn is not None:
+        return decorate(fn)
+    return decorate

thinkless/tracing/viewer.py

@@ -0,0 +1,87 @@
+"""A self-contained HTML viewer for traces.
+
+The output is one HTML file with the trace data embedded and no external
+requests, so it opens offline and can be attached to an issue or a pull
+request as-is.
+"""
+
+from __future__ import annotations
+
+import json
+import time
+from collections.abc import Iterable, Sequence
+from importlib import resources
+from pathlib import Path
+from typing import Any
+
+from .sinks import iter_trace_files, read_trace
+from .summary import summarize
+
+__all__ = ["build_viewer", "collect_traces", "write_viewer"]
+
+PLACEHOLDER = "__THINKLESS_DATA__"
+
+
+def collect_traces(
+    paths: Iterable[str | Path], *, limit: int | None = None
+) -> list[dict[str, Any]]:
+    """Load trace files, expanding directories, newest first."""
+    files: list[Path] = []
+    for raw in paths:
+        path = Path(raw)
+        files.extend(iter_trace_files(path) if path.is_dir() else [path])
+    traces = []
+    for file in files[:limit] if limit else files:
+        spans = read_trace(file)
+        if not spans:
+            continue
+        traces.append({"file": file.name, "spans": spans})
+    return traces
+
+
+def build_viewer(
+    traces: Sequence[dict[str, Any]] | Sequence[Sequence[dict[str, Any]]], *, title: str = "Traces"
+) -> str:
+    """Render traces into the viewer HTML.
+
+    Args:
+        traces: Either ``{"file", "spans"}`` records from :func:`collect_traces`
+            or plain lists of span dictionaries.
+        title: Shown in the header.
+    """
+    records = []
+    for item in traces:
+        if isinstance(item, dict):
+            file, spans = item.get("file"), list(item["spans"])
+        else:
+            file, spans = None, list(item)
+        if not spans:
+            continue
+        records.append(
+            {"file": file, "spans": spans, "summary": summarize(spans).model_dump(mode="json")}
+        )
+    payload = {
+        "title": title,
+        "generated_at": time.strftime("%Y-%m-%d %H:%M UTC", time.gmtime()),
+        "traces": records,
+    }
+    # Keep the JSON inert inside the script tag.
+    embedded = json.dumps(payload, ensure_ascii=False, default=str).replace("</", "<\\/")
+    template = resources.files("thinkless").joinpath("data/viewer.html").read_text(encoding="utf-8")
+    return template.replace(PLACEHOLDER, embedded)
+
+
+def write_viewer(
+    paths: Iterable[str | Path],
+    output: str | Path,
+    *,
+    title: str = "Traces",
+    limit: int | None = None,
+) -> Path:
+    """Collect traces from ``paths`` and write the viewer to ``output``."""
+    target = Path(output)
+    target.parent.mkdir(parents=True, exist_ok=True)
+    target.write_text(
+        build_viewer(collect_traces(paths, limit=limit), title=title), encoding="utf-8"
+    )
+    return target