tanglebrain 0.16.0__py3-none-any.whl

tanglebrain/measurement.py

@@ -0,0 +1,548 @@
+"""Measurement / "spend avoided" rollup.
+
+Every routed task is logged as one JSON line in an append-only usage log, and ``tanglebrain
+--stats`` rolls those records up into a "spend avoided" figure: what the routed work *would* have
+cost on a paid frontier API, had it not gone to the free local tier or a subscription CLI. This
+makes the cloud-equivalent cost avoided by routing visible.
+
+Design notes:
+
+- **Tokens are estimated, not measured.** Authenticated CLIs expose no usable token counts, and a
+  local reasoning model's real ``usage`` is inflated by dropped reasoning tokens. So a
+  single ``chars/4`` heuristic over the visible prompt + response is applied *uniformly* across all
+  tiers — one consistent, if approximate, methodology (see :func:`estimate_tokens`).
+- **Pricing is config-driven** (``config/pricing.yaml``): a reference frontier price the operator
+  tunes (the knob GUI edits it).
+- **All I/O is fault-tolerant.** A logging failure must never break the user's actual answer, and a
+  corrupt log line must never break the rollup. Reads return sensible defaults; the writer swallows
+  every exception. This mirrors the router's state-file idiom (:mod:`tanglebrain.router`).
+"""
+from __future__ import annotations
+
+import json
+import os
+import shutil
+import threading
+from dataclasses import dataclass
+from datetime import datetime, timezone
+from pathlib import Path
+
+import yaml
+
+from tanglebrain.router import DEFAULT_STATE_SUBDIR, STATE_DIR_ENV
+
+LOG_FILENAME = "usage.jsonl"
+
+#: Env var carrying the top-level task id from the orchestrator down to a delegated sub-call. The
+#: CLI mints a task id per routed task and the orchestrator-CLI adapter injects it into the
+#: orchestrator subprocess env (only when the delegate tool is injected); the orchestrator forwards
+#: its env to the MCP delegate child it spawns, where :func:`tanglebrain.delegate.run_delegate` reads
+#: it back and stamps each delegate record's ``parent_task_id``. This is what links a delegated
+#: sub-call to the specific top-level task that spawned it, across the process boundary.
+PARENT_TASK_ID_ENV = "TANGLEBRAIN_TASK_ID"
+
+# Serializes appends to the usage log so concurrent writers in one process (delegate_many fans
+# sub-tasks out across threads) can't interleave bytes mid-line. Per-process only; cross-process
+# appends rely on the OS's O_APPEND atomicity for short lines, as before.
+_LOG_LOCK = threading.Lock()
+
+# Chars per token for the uniform estimation heuristic. ~4 chars/token is the standard rough
+# approximation for English-ish text across modern BPE tokenizers; good enough for an *estimate*.
+_CHARS_PER_TOKEN = 4
+
+
+@dataclass(frozen=True)
+class Pricing:
+    """Cloud-equivalent reference pricing for the rollup (loaded from ``config/pricing.yaml``).
+
+    Attributes:
+        reference_model: Human-readable label for the frontier model these rates represent.
+        input_per_mtok: US dollars per 1,000,000 input (prompt) tokens.
+        output_per_mtok: US dollars per 1,000,000 output (completion) tokens.
+        is_placeholder: ``True`` while the rates are rough/illustrative; the rollup renders a
+            PLACEHOLDER caveat so no figure is mistaken for a precise cost.
+    """
+
+    reference_model: str
+    input_per_mtok: float
+    output_per_mtok: float
+    is_placeholder: bool
+
+
+# Fallback used when ``config/pricing.yaml`` is missing or unreadable — always flagged placeholder.
+PLACEHOLDER_PRICING = Pricing(
+    reference_model="unconfigured (PLACEHOLDER — pricing.yaml unreadable)",
+    input_per_mtok=3.00,
+    output_per_mtok=15.00,
+    is_placeholder=True,
+)
+
+
+def default_log_path() -> Path:
+    """Return the usage-log file path.
+
+    Honors ``TANGLEBRAIN_STATE_DIR`` (``~`` expanded); otherwise ``~/.cache/tanglebrain/``. The log
+    lives alongside the router's state file (same dir, same env override).
+
+    Returns:
+        The absolute path to the append-only usage JSONL file.
+    """
+    base = os.environ.get(STATE_DIR_ENV)
+    root = Path(base).expanduser() if base else Path.home() / DEFAULT_STATE_SUBDIR
+    return root / LOG_FILENAME
+
+
+def default_pricing_path() -> Path:
+    """Return the path to the pricing YAML shipped with the package.
+
+    Returns:
+        The absolute path to ``tanglebrain/config/pricing.yaml``.
+    """
+    return Path(__file__).resolve().parent / "config" / "pricing.yaml"
+
+
+def load_pricing(path: str | os.PathLike[str] | None = None) -> Pricing:
+    """Load cloud-equivalent reference pricing, tolerating a missing/corrupt file.
+
+    Args:
+        path: Path to a pricing YAML. Defaults to the packaged ``config/pricing.yaml``.
+
+    Returns:
+        The parsed :class:`Pricing`, or :data:`PLACEHOLDER_PRICING` if the file is absent,
+        unreadable, or malformed — bad config must never crash the rollup.
+    """
+    pricing_path = Path(path) if path is not None else default_pricing_path()
+    try:
+        raw = yaml.safe_load(pricing_path.read_text())
+        return Pricing(
+            reference_model=str(raw.get("reference_model", "unknown")),
+            input_per_mtok=float(raw["input_per_mtok"]),
+            output_per_mtok=float(raw["output_per_mtok"]),
+            is_placeholder=bool(raw.get("placeholder", False)),
+        )
+    except (OSError, yaml.YAMLError, ValueError, TypeError, KeyError, AttributeError):
+        return PLACEHOLDER_PRICING
+
+
+# Fallback header, used ONLY when the target file is absent (e.g. a fresh write to a new path).
+# A normal save preserves the existing file's own leading comment block verbatim (see
+# :func:`_leading_comment_block`), so the curated methodology note is never replaced or drifted.
+PRICING_HEADER = """\
+# Cloud-equivalent reference pricing for the "spend avoided" rollup.
+#
+# Methodology: for each routed task, estimate what it WOULD have cost on a paid frontier API, valued
+# at a reference model's per-million-token price. The rollup multiplies estimated tokens (a chars/4
+# heuristic over the visible prompt + response) by these rates. Values are US dollars per 1,000,000
+# tokens.
+#
+# `placeholder: true` makes `tanglebrain --stats` flag every figure as PLACEHOLDER (use it if you
+# fork these reference rates before re-checking them). Edited via `tanglebrain-gui` or by hand.
+"""
+
+
+def _leading_comment_block(text: str) -> str:
+    """Return the file's leading run of comment/blank lines (its header), or ``""`` if none.
+
+    Used to preserve a pricing file's curated header verbatim across a save, so no documentation
+    is lost or replaced. Stops at the first non-comment, non-blank line (the first YAML key).
+    """
+    out: list[str] = []
+    for line in text.splitlines():
+        if line.startswith("#") or not line.strip():
+            out.append(line)
+        else:
+            break
+    while out and not out[-1].strip():  # drop trailing blank lines before the keys
+        out.pop()
+    return "\n".join(out) + "\n" if out else ""
+
+
+def validate_pricing(data: dict) -> Pricing:
+    """Strictly validate raw pricing fields and build a :class:`Pricing`.
+
+    Unlike :func:`load_pricing` (lenient — bad reads fall back to a placeholder), this rejects
+    invalid input so the panel never persists garbage.
+
+    Args:
+        data: ``{reference_model, input_per_mtok, output_per_mtok, placeholder}``.
+
+    Returns:
+        A validated :class:`Pricing`.
+
+    Raises:
+        ValueError: If a field is missing, the wrong type, a non-finite/negative rate, or an
+            empty ``reference_model``.
+    """
+    if not isinstance(data, dict):
+        raise ValueError("pricing must be an object")
+
+    model = data.get("reference_model")
+    if not isinstance(model, str) or not model.strip():
+        raise ValueError("reference_model must be a non-empty string")
+
+    rates = {}
+    for key in ("input_per_mtok", "output_per_mtok"):
+        value = data.get(key)
+        if isinstance(value, bool) or not isinstance(value, (int, float)):
+            raise ValueError(f"{key} must be a number")
+        value = float(value)
+        if value != value or value in (float("inf"), float("-inf")):  # NaN / inf guard
+            raise ValueError(f"{key} must be finite")
+        if value < 0:
+            raise ValueError(f"{key} must be >= 0")
+        rates[key] = value
+
+    placeholder = data.get("placeholder", False)
+    if not isinstance(placeholder, bool):
+        raise ValueError("placeholder must be a boolean")
+
+    return Pricing(
+        reference_model=model.strip(),
+        input_per_mtok=rates["input_per_mtok"],
+        output_per_mtok=rates["output_per_mtok"],
+        is_placeholder=placeholder,
+    )
+
+
+def _backup_dir() -> Path:
+    """Return the directory for config backups (under the state dir, never the repo config dir)."""
+    base = os.environ.get(STATE_DIR_ENV)
+    root = Path(base).expanduser() if base else Path.home() / DEFAULT_STATE_SUBDIR
+    return root / "backups"
+
+
+def _atomic_write(path: Path, text: str) -> None:
+    """Write ``text`` to ``path`` atomically (temp file in the same dir, then ``os.replace``)."""
+    path.parent.mkdir(parents=True, exist_ok=True)
+    tmp = path.with_name(path.name + ".tmp")
+    tmp.write_text(text, encoding="utf-8")
+    os.replace(tmp, path)
+
+
+def _render_pricing(pricing: Pricing, header: str) -> str:
+    """Render a :class:`Pricing` to YAML text beneath ``header``.
+
+    ``reference_model`` is emitted via ``json.dumps`` — a JSON string is valid YAML and safely
+    quotes/escapes any colons, quotes, unicode, or backslashes. Float rates use Python's ``repr``
+    (valid YAML), which round-trips exactly through :func:`load_pricing`.
+    """
+    return (
+        header
+        + f"placeholder: {str(pricing.is_placeholder).lower()}\n"
+        + f"reference_model: {json.dumps(pricing.reference_model)}\n"
+        + f"input_per_mtok: {pricing.input_per_mtok}\n"
+        + f"output_per_mtok: {pricing.output_per_mtok}\n"
+    )
+
+
+def save_pricing(pricing: Pricing, path: str | os.PathLike[str] | None = None) -> None:
+    """Persist pricing to the config YAML — header-preserving, with a backup, written atomically.
+
+    Preserves the target's existing leading comment block verbatim (falling back to
+    :data:`PRICING_HEADER` only when the file is absent), backs up any existing file to
+    ``<state_dir>/backups/pricing-<ts>.yaml``, then atomically replaces the target.
+
+    Args:
+        pricing: The validated pricing to write (see :func:`validate_pricing`).
+        path: Target YAML path. Defaults to the packaged ``config/pricing.yaml``.
+    """
+    target = Path(path) if path is not None else default_pricing_path()
+    header = PRICING_HEADER
+    if target.exists():
+        existing = target.read_text(encoding="utf-8")
+        block = _leading_comment_block(existing)
+        if block.strip():
+            header = block  # keep the curated header verbatim — no drift, no doc loss
+        backup_dir = _backup_dir()
+        backup_dir.mkdir(parents=True, exist_ok=True)
+        stamp = datetime.now(timezone.utc).strftime("%Y%m%dT%H%M%S_%fZ")  # sub-second: no same-second collision
+        shutil.copy2(target, backup_dir / f"pricing-{stamp}.yaml")
+    _atomic_write(target, _render_pricing(pricing, header))
+
+
+def estimate_tokens(text: str) -> int:
+    """Estimate the token count of ``text`` via the uniform ``chars/4`` heuristic.
+
+    This is an approximation, applied identically to every tier (CLI subs expose no real counts).
+    Empty/falsy text counts as 0; any non-empty text is at least 1 token.
+
+    Args:
+        text: The prompt or response text.
+
+    Returns:
+        The estimated token count (``>= 0``).
+    """
+    if not text:
+        return 0
+    return max(1, len(text) // _CHARS_PER_TOKEN)
+
+
+def cloud_equiv_usd(in_tokens: int, out_tokens: int, pricing: Pricing) -> float:
+    """Compute the cloud-equivalent cost of a task at the reference frontier price.
+
+    Args:
+        in_tokens: Estimated input (prompt) tokens.
+        out_tokens: Estimated output (completion) tokens.
+        pricing: The reference pricing to apply.
+
+    Returns:
+        The estimated US-dollar cost on the reference frontier API.
+    """
+    return (
+        in_tokens / 1_000_000 * pricing.input_per_mtok
+        + out_tokens / 1_000_000 * pricing.output_per_mtok
+    )
+
+
+def record_task(
+    *,
+    path: str,
+    entry: object,
+    prompt: str,
+    response: str,
+    kind: str = "task",
+    task_id: str | None = None,
+    parent_task_id: str | None = None,
+    log_path: str | os.PathLike[str] | None = None,
+    pricing: Pricing | None = None,
+) -> None:
+    """Append one usage record for a routed task or a delegated sub-call. Never raises.
+
+    A logging failure is dropped — measurement is a side-effect that must never affect the returned
+    answer.
+
+    Args:
+        path: Which execution path served the work — ``router`` | ``local`` | ``model`` |
+            ``delegate``.
+        entry: The served :class:`~tanglebrain.roster.RosterEntry` (read for ``tier``/``id``); may
+            be ``None`` (e.g. the router didn't surface one), in which case both are ``"unknown"``.
+        prompt: The prompt (for input-token estimation).
+        response: The returned response text (for output-token estimation).
+        kind: ``"task"`` for a top-level routed task (the default; what the spend-avoided headline
+            counts) or ``"delegate"`` for a delegated sub-call. Delegate records are rolled up
+            **separately** so a sub-call's saving is never double-counted against its parent task.
+        task_id: For a top-level task, the id minted for this routed task (so its delegated sub-calls
+            can be linked back to it). Omitted from the record when ``None``.
+        parent_task_id: For a delegated sub-call, the id of the top-level task that spawned it (read
+            from :data:`PARENT_TASK_ID_ENV`). Omitted from the record when ``None`` — e.g. a delegate
+            invoked outside a propagated task, which rolls up as ``unlinked``.
+        log_path: Override the usage-log path (tests inject a temp path). Defaults to
+            :func:`default_log_path`.
+        pricing: Override the pricing. Defaults to :func:`load_pricing`.
+    """
+    try:
+        if pricing is None:
+            pricing = load_pricing()
+        tier = getattr(entry, "tier", None) or "unknown"
+        model = getattr(entry, "id", None) or "unknown"
+        in_tok = estimate_tokens(prompt)
+        out_tok = estimate_tokens(response)
+        equiv = cloud_equiv_usd(in_tok, out_tok, pricing)
+        # A paid `api` task incurs real spend, so it avoids nothing (avoided = 0); every other tier
+        # routes work off a paid frontier API, so it avoids the full cloud-equivalent.
+        avoided = 0.0 if tier == "api" else equiv
+        record = {
+            "ts": datetime.now(timezone.utc).isoformat(timespec="seconds"),
+            "kind": str(kind),
+            "path": str(path),
+            "tier": str(tier),
+            "model": str(model),
+            "in_tokens_est": in_tok,
+            "out_tokens_est": out_tok,
+            "cloud_equiv_usd": round(equiv, 6),
+            "spend_avoided_usd": round(avoided, 6),
+            "pricing_ref": pricing.reference_model,
+        }
+        # Optional linkage fields — only written when present, so existing records/readers that
+        # never set them are unaffected (a missing field reads as "no linkage").
+        if task_id is not None:
+            record["task_id"] = str(task_id)
+        if parent_task_id is not None:
+            record["parent_task_id"] = str(parent_task_id)
+        target = Path(log_path) if log_path is not None else default_log_path()
+        target.parent.mkdir(parents=True, exist_ok=True)
+        with _LOG_LOCK:
+            with target.open("a", encoding="utf-8") as fh:
+                fh.write(json.dumps(record) + "\n")
+    except Exception:
+        # Measurement is a side-effect: a failure here must never affect the returned answer.
+        return
+
+
+def _as_int(value: object) -> int:
+    """Coerce a stored numeric field to int, defaulting to 0 on any bad value."""
+    try:
+        return int(value)  # type: ignore[arg-type]
+    except (ValueError, TypeError):
+        return 0
+
+
+def _as_float(value: object) -> float:
+    """Coerce a stored numeric field to float, defaulting to 0.0 on any bad value."""
+    try:
+        return float(value)  # type: ignore[arg-type]
+    except (ValueError, TypeError):
+        return 0.0
+
+
+def read_records(log_path: str | os.PathLike[str] | None = None) -> list[dict]:
+    """Read all usage records from the log, skipping malformed lines.
+
+    Args:
+        log_path: Override the usage-log path. Defaults to :func:`default_log_path`.
+
+    Returns:
+        The parsed records in file (chronological) order. An absent log yields ``[]``.
+    """
+    target = Path(log_path) if log_path is not None else default_log_path()
+    records: list[dict] = []
+    try:
+        text = target.read_text(encoding="utf-8")
+    except OSError:
+        return records
+    for line in text.splitlines():
+        line = line.strip()
+        if not line:
+            continue
+        try:
+            obj = json.loads(line)
+        except json.JSONDecodeError:
+            continue
+        if isinstance(obj, dict):
+            records.append(obj)
+    return records
+
+
+def rollup(records: list[dict]) -> dict:
+    """Aggregate usage records into a summary.
+
+    Args:
+        records: The records from :func:`read_records`.
+
+    Returns:
+        A dict with: ``tasks`` (int), ``by_tier`` (tier → count), ``in_tokens_est`` /
+        ``out_tokens_est`` (summed estimates), and ``cloud_equiv_usd`` / ``spend_avoided_usd``
+        (summed dollars) — all over **top-level tasks only** — plus ``delegates``, a separate
+        sub-rollup of delegated sub-calls ``{count, by_backend: {model: {count, in_tokens_est,
+        out_tokens_est}}, by_parent: {parent_task_id: {count, by_backend: {model: count}}},
+        in_tokens_est, out_tokens_est, cloud_equiv_usd}``. ``by_parent`` groups each delegate under
+        the top-level task that spawned it (via ``parent_task_id``); delegates with no
+        ``parent_task_id`` are grouped under the sentinel ``"unlinked"``. Delegate records are kept
+        out of the headline so a sub-call's saving is never double-counted against its parent task;
+        their cloud-equiv is informational. A record without a ``kind`` field counts as a task.
+    """
+    summary: dict = {
+        "tasks": 0,
+        "by_tier": {},
+        "in_tokens_est": 0,
+        "out_tokens_est": 0,
+        "cloud_equiv_usd": 0.0,
+        "spend_avoided_usd": 0.0,
+    }
+    delegates: dict = {
+        "count": 0,
+        "by_backend": {},
+        "by_parent": {},
+        "in_tokens_est": 0,
+        "out_tokens_est": 0,
+        "cloud_equiv_usd": 0.0,
+    }
+    for r in records:
+        in_tok = _as_int(r.get("in_tokens_est"))
+        out_tok = _as_int(r.get("out_tokens_est"))
+        if str(r.get("kind", "task")) == "delegate":
+            delegates["count"] += 1
+            model = str(r.get("model", "unknown"))
+            backend = delegates["by_backend"].setdefault(
+                model, {"count": 0, "in_tokens_est": 0, "out_tokens_est": 0}
+            )
+            backend["count"] += 1
+            backend["in_tokens_est"] += in_tok
+            backend["out_tokens_est"] += out_tok
+            # Per-parent tree: link this sub-call to the top-level task that spawned it. A delegate
+            # with no parent_task_id (run outside a propagated task) groups under "unlinked".
+            parent_id = r.get("parent_task_id")
+            parent_key = str(parent_id) if parent_id not in (None, "") else "unlinked"
+            parent = delegates["by_parent"].setdefault(parent_key, {"count": 0, "by_backend": {}})
+            parent["count"] += 1
+            parent["by_backend"][model] = parent["by_backend"].get(model, 0) + 1
+            delegates["in_tokens_est"] += in_tok
+            delegates["out_tokens_est"] += out_tok
+            delegates["cloud_equiv_usd"] += _as_float(r.get("cloud_equiv_usd"))
+            continue
+        summary["tasks"] += 1
+        tier = str(r.get("tier", "unknown"))
+        summary["by_tier"][tier] = summary["by_tier"].get(tier, 0) + 1
+        summary["in_tokens_est"] += in_tok
+        summary["out_tokens_est"] += out_tok
+        summary["cloud_equiv_usd"] += _as_float(r.get("cloud_equiv_usd"))
+        summary["spend_avoided_usd"] += _as_float(r.get("spend_avoided_usd"))
+    summary["cloud_equiv_usd"] = round(summary["cloud_equiv_usd"], 4)
+    summary["spend_avoided_usd"] = round(summary["spend_avoided_usd"], 4)
+    delegates["cloud_equiv_usd"] = round(delegates["cloud_equiv_usd"], 4)
+    summary["delegates"] = delegates
+    return summary
+
+
+def format_rollup(summary: dict, pricing: Pricing) -> str:
+    """Render a rollup summary as a human-readable block for the CLI.
+
+    Args:
+        summary: The aggregate from :func:`rollup`.
+        pricing: The currently-configured pricing (for the reference-model label + placeholder
+            caveat). Per-record costs were computed when each task ran; this only labels the figure.
+
+    Returns:
+        A multi-line string suitable for printing.
+    """
+    lines = [
+        "TangleBrain — spend avoided (cloud-equivalent)",
+        f"  Tasks routed:   {summary.get('tasks', 0)}",
+    ]
+    by_tier = summary.get("by_tier") or {}
+    if by_tier:
+        tiers = ", ".join(f"{k} {v}" for k, v in sorted(by_tier.items()))
+        lines.append(f"  By tier:        {tiers}")
+    lines.append(
+        f"  Est. tokens:    in {summary.get('in_tokens_est', 0):,} / "
+        f"out {summary.get('out_tokens_est', 0):,}"
+    )
+    lines.append(f"  Spend avoided:  ${summary.get('spend_avoided_usd', 0.0):,.2f}")
+    lines.append(f"  Pricing ref:    {pricing.reference_model}")
+    if pricing.is_placeholder:
+        lines.append(
+            "  ⚠ pricing: PLACEHOLDER — figures are illustrative; set real rates in "
+            "config/pricing.yaml and flip placeholder to false."
+        )
+
+    delegates = summary.get("delegates") or {}
+    if delegates.get("count"):
+        lines.append("")
+        lines.append("  Delegated sub-tasks (offloaded by orchestrators)")
+        lines.append(f"    Count:        {delegates.get('count', 0)}")
+        by_backend = delegates.get("by_backend") or {}
+        if by_backend:
+            backends = ", ".join(
+                f"{model} {info.get('count', 0)}" for model, info in sorted(by_backend.items())
+            )
+            lines.append(f"    By backend:   {backends}")
+        by_parent = delegates.get("by_parent") or {}
+        if by_parent:
+            linked = [k for k in by_parent if k != "unlinked"]
+            unlinked = (by_parent.get("unlinked") or {}).get("count", 0)
+            if linked:
+                tree = f"{len(linked)} parent task(s)"
+                if unlinked:
+                    tree += f", {unlinked} unlinked"
+            else:
+                tree = f"{unlinked} unlinked"  # all sub-calls ran outside a propagated task
+            lines.append(f"    Linked to:    {tree}")
+        lines.append(
+            f"    Est. tokens:  in {delegates.get('in_tokens_est', 0):,} / "
+            f"out {delegates.get('out_tokens_est', 0):,}"
+        )
+        lines.append(
+            f"    Cloud-equiv:  ${delegates.get('cloud_equiv_usd', 0.0):,.2f} "
+            "(informational — already credited within parent tasks)"
+        )
+    return "\n".join(lines)