primarypath-server 1.0.0__py3-none-any.whl

app/traces.py

@@ -0,0 +1,309 @@
+"""Run traces + cost, folded from the audit spine. No new store.
+
+The drangue bridge/gateway ingests a run's whole log as `run_event`s (each
+wrapping a drangue Event: model decisions with reasoning and usage, tool
+calls with args/result, timing). A run's trace is a fold of those plus the
+run's `approval_*` events: an ordered timeline, the approval pause/decision,
+per-step token usage, and a cost estimate. This is a read over what the
+audit already needs — not a second observability backend.
+
+Tokens are exact (summed from the ingested usage). Cost is an estimate from
+an optional per-model price map (PRIMARYPATH_MODEL_PRICES); with no prices
+configured, cost is null and only tokens are reported.
+"""
+from __future__ import annotations
+
+from collections import defaultdict
+
+from .ledger.store import DEFAULT_TENANT, Ledger
+
+
+def parse_prices(spec: str) -> dict[str, tuple[float, float]]:
+    """"model:in/out,..." USD per 1M tokens -> {model: (in, out)}."""
+    out: dict[str, tuple[float, float]] = {}
+    for part in spec.split(","):
+        part = part.strip()
+        if not part or ":" not in part:
+            continue
+        model, _, rates = part.partition(":")
+        try:
+            inp, _, outp = rates.partition("/")
+            out[model.strip()] = (float(inp), float(outp))
+        except ValueError:
+            continue
+    return out
+
+
+def _cost(prices: dict, model: str, usage: dict) -> float | None:
+    rate = prices.get(model)
+    if rate is None or not usage:
+        return None
+    return (usage.get("input_tokens", 0) * rate[0]
+            + usage.get("output_tokens", 0) * rate[1]) / 1_000_000
+
+
+class TraceStore:
+    def __init__(self, pool, ledger: Ledger, settings):
+        self.pool = pool
+        self.ledger = ledger
+        self.settings = settings
+        self.prices = parse_prices(settings.model_prices)
+
+    # -- one run's trace -----------------------------------------------------
+
+    async def run_trace(self, run_id: str,
+                        tenant: str = DEFAULT_TENANT,
+                        visible=None) -> dict | None:
+        """`visible(agent_id, unit)` narrows the fold to the caller's scope
+        event by event: a run_id written by two agents must never fold one
+        agent's steps into a trace the other can read."""
+        events = await self.ledger.events_for_run(run_id, tenant)
+        if visible is not None:
+            events = [e for e in events
+                      if visible(e.get("agent_id", "default"), e.get("unit"))]
+        if not events:
+            return None
+        return self._fold_run(run_id, events)
+
+    def _fold_run(self, run_id: str, events: list[dict]) -> dict:
+        import json
+
+        steps: list[dict] = []
+        tools: list[str] = []
+        approval = None
+        status = "running"
+        agent = "default"
+        unit = None
+        started_ts = None
+        tokens = {"input_tokens": 0, "output_tokens": 0}
+        by_model: dict[str, dict] = defaultdict(
+            lambda: {"input_tokens": 0, "output_tokens": 0, "decisions": 0})
+        # a run can hold several approvals; it is paused while ANY of them
+        # is undecided, not merely until the first decision lands
+        pending_ids: set[str] = set()
+
+        for e in events:
+            payload = json.loads(e["payload"])
+            ts = e["ts_ms"]
+            if started_ts is None:
+                started_ts = ts
+            if e.get("agent_id"):
+                agent = e["agent_id"]
+            if e.get("unit"):
+                unit = e["unit"]
+            t = e["type"]
+
+            if t == "run_event":
+                inner = payload.get("event") or {}
+                kind = inner.get("type", "event")
+                ip = inner.get("payload") or {}
+                step = {"seq": e["seq"], "ts_ms": ts, "kind": kind,
+                        "duration_ms": inner.get("duration_ms")}
+                if kind == "model_decision":
+                    step["reasoning"] = ip.get("reasoning") or ip.get("text")
+                    step["tool_calls"] = ip.get("tool_calls") or []
+                    usage = ip.get("usage") or {}
+                    model = ip.get("model") or "unknown"
+                    step["model"] = model
+                    step["usage"] = usage
+                    tokens["input_tokens"] += usage.get("input_tokens", 0)
+                    tokens["output_tokens"] += usage.get("output_tokens", 0)
+                    m = by_model[model]
+                    m["input_tokens"] += usage.get("input_tokens", 0)
+                    m["output_tokens"] += usage.get("output_tokens", 0)
+                    m["decisions"] += 1
+                elif kind == "tool_result":
+                    step["tool"] = ip.get("name")
+                    step["result"] = ip.get("content")
+                    step["is_error"] = ip.get("is_error", False)
+                if kind == "run_finished":
+                    status = "completed"
+                steps.append(step)
+
+            elif t == "approval_requested":
+                tools.append(payload.get("tool"))
+                pending_ids.add(payload.get("approval_id"))
+                steps.append({"seq": e["seq"], "ts_ms": ts,
+                              "kind": "approval_requested",
+                              "tool": payload.get("tool"),
+                              "arguments": payload.get("arguments"),
+                              "reasoning": payload.get("reasoning"),
+                              "approval_id": payload.get("approval_id")})
+            elif t == "approval_endorsed":
+                # dual control: first approval; the run stays paused
+                steps.append({"seq": e["seq"], "ts_ms": ts,
+                              "kind": "approval_endorsed",
+                              "reviewer": payload.get("reviewer"),
+                              "reason": payload.get("reason")})
+            elif t in ("approval_granted", "approval_denied"):
+                pending_ids.discard(payload.get("approval_id"))
+                approval = {"status": ("granted" if t == "approval_granted"
+                                       else "denied"),
+                            "reviewer": payload.get("reviewer"),
+                            "reason": payload.get("reason"),
+                            "reviewer_role": payload.get("reviewer_role"),
+                            "ts_ms": ts}
+                steps.append({"seq": e["seq"], "ts_ms": ts, "kind": t,
+                              "reviewer": payload.get("reviewer"),
+                              "reason": payload.get("reason")})
+            elif t == "approval_annotated":
+                steps.append({"seq": e["seq"], "ts_ms": ts,
+                              "kind": "approval_annotated",
+                              "reviewer": payload.get("reviewer"),
+                              "note": payload.get("note")})
+
+        if pending_ids:
+            status = "paused"
+
+        cost = 0.0
+        cost_seen = False
+        cost_by_model = {}
+        for model, u in by_model.items():
+            c = _cost(self.prices, model, u)
+            cost_by_model[model] = {**u, "cost": c}
+            if c is not None:
+                cost += c
+                cost_seen = True
+
+        return {
+            "run_id": run_id, "status": status,
+            "agent": agent, "unit": unit,
+            "tools": [t for t in dict.fromkeys(tools) if t],
+            "step_count": sum(1 for s in steps if s["kind"] == "model_decision"),
+            "tokens": tokens,
+            "cost": round(cost, 6) if cost_seen else None,
+            "cost_by_model": cost_by_model,
+            "started_ts_ms": started_ts,
+            "last_ts_ms": events[-1]["ts_ms"],
+            "approval": approval,
+            "steps": steps,
+        }
+
+    # -- run list (summaries) ------------------------------------------------
+
+    async def run_ids(self, tenant: str = DEFAULT_TENANT,
+                      limit: int | None = 50, offset: int = 0) -> list[str]:
+        # order by rowid (global ingest order), not seq: F24 restarts seq
+        # per (tenant, chain), so seqs on different agents' chains are not
+        # comparable and a new agent's fresh run would sort last
+        q = ("SELECT run_id, max(rowid) AS s FROM events "
+             "WHERE tenant_id = %s AND run_id IS NOT NULL "
+             "GROUP BY run_id ORDER BY s DESC")
+        params: list = [tenant]
+        if limit is not None:
+            q += " LIMIT %s OFFSET %s"
+            params += [limit, offset]
+        async with self.pool.connection() as conn:
+            cur = await conn.execute(q, params)
+            return [r[0] for r in await cur.fetchall()]
+
+    async def run_list(self, tenant: str = DEFAULT_TENANT,
+                       status: str | None = None, tool: str | None = None,
+                       limit: int = 50, offset: int = 0,
+                       units: set[str] | None = None,
+                       agent: str | None = None) -> list[dict]:
+        """Paginate AFTER filtering: fold + filter the full ordered run set,
+        then slice. Slicing run_ids before the fold would drop matching runs
+        that fell outside the id page. With no filter the DB LIMIT/OFFSET is
+        the fast path; the fold-everything path only runs when a filter is
+        active (bounded to the console's own use).
+
+        `units`/`agent` are the F24 need-to-know gate (None = unrestricted):
+        a run whose agent sits in another unit is filtered out server-side,
+        never sent for the client to hide."""
+        filtered = (status is not None or tool is not None
+                    or units is not None or agent is not None)
+        ids = await self.run_ids(tenant, limit=None if filtered else limit,
+                                 offset=0 if filtered else offset)
+        # the same per-event gate as run_detail: a summary (tools, steps,
+        # tokens) folded over another agent's events would leak them
+        visible = None
+        if units is not None or agent is not None:
+            def visible(a, u):
+                if agent is not None:
+                    return a == agent
+                return u is None or units is None or u in units
+        summaries = []
+        for run_id in ids:
+            trace = await self.run_trace(run_id, tenant, visible=visible)
+            if trace is None:
+                continue
+            if units is not None and trace["unit"] is not None \
+                    and trace["unit"] not in units:
+                continue
+            if agent is not None and trace["agent"] != agent:
+                continue
+            if status and trace["status"] != status:
+                continue
+            if tool and tool not in trace["tools"]:
+                continue
+            summaries.append({k: trace[k] for k in (
+                "run_id", "status", "agent", "unit", "tools", "step_count",
+                "tokens", "cost", "started_ts_ms", "last_ts_ms", "approval")})
+        return summaries[offset:offset + limit] if filtered else summaries
+
+    # -- cost aggregation over a period --------------------------------------
+
+    async def cost_summary(self, tenant: str = DEFAULT_TENANT,
+                           since_ms: int | None = None,
+                           until_ms: int | None = None,
+                           units: set[str] | None = None) -> dict:
+        import json
+
+        q = ("SELECT payload FROM events WHERE tenant_id = %s "
+             "AND type = 'run_event'")
+        params: list = [tenant]
+        if units is not None:
+            # spend is volume, volume is information: the unit gate applies
+            # to aggregates too (F24)
+            q += " AND (unit IS NULL OR unit = ANY(%s))"
+            params.append(sorted(units))
+        if since_ms is not None:
+            q += " AND ts_ms >= %s"
+            params.append(since_ms)
+        if until_ms is not None:
+            q += " AND ts_ms <= %s"
+            params.append(until_ms)
+        async with self.pool.connection() as conn:
+            cur = await conn.execute(q, params)
+            rows = await cur.fetchall()
+
+        by_tool: dict[str, dict] = defaultdict(
+            lambda: {"input_tokens": 0, "output_tokens": 0, "cost": 0.0,
+                     "cost_seen": False})
+        by_model: dict[str, dict] = defaultdict(
+            lambda: {"input_tokens": 0, "output_tokens": 0, "cost": 0.0,
+                     "cost_seen": False})
+        totals = {"input_tokens": 0, "output_tokens": 0, "cost": 0.0,
+                  "cost_seen": False}
+
+        for (raw,) in rows:
+            payload = json.loads(raw)
+            inner = payload.get("event") or {}
+            if inner.get("type") != "model_decision":
+                continue
+            ip = inner.get("payload") or {}
+            usage = ip.get("usage") or {}
+            model = ip.get("model") or "unknown"
+            calls = ip.get("tool_calls") or []
+            tool = (calls[0].get("name") if calls else None) or "(no tool)"
+            c = _cost(self.prices, model, usage)
+            for bucket, key in ((by_tool, tool), (by_model, model), (None, None)):
+                b = totals if bucket is None else bucket[str(key)]
+                b["input_tokens"] += usage.get("input_tokens", 0)
+                b["output_tokens"] += usage.get("output_tokens", 0)
+                if c is not None:
+                    b["cost"] += c
+                    b["cost_seen"] = True
+
+        def _clean(bucket: dict) -> dict:
+            return {"input_tokens": bucket["input_tokens"],
+                    "output_tokens": bucket["output_tokens"],
+                    "cost": round(bucket["cost"], 6) if bucket["cost_seen"]
+                    else None}
+
+        return {
+            "by_tool": {k: _clean(v) for k, v in by_tool.items()},
+            "by_model": {k: _clean(v) for k, v in by_model.items()},
+            "totals": _clean(totals),
+        }