zeno-cli 0.3.4__py3-none-any.whl

zeno_cli/hud/zeno_attention.py

@@ -0,0 +1,288 @@
+#!/usr/bin/env python3
+"""zeno_attention - the single source of truth for the cognitive-attention model.
+
+This is the ONE attention formula. The HUD statusline (the writer) imports it to
+compute the score it persists into cognition_samples, and any other Python reader
+(exporter, dashboard SSR helpers) presents those persisted rows. Because the
+writer and the bar both derive their number from this module - and the bar reads
+the row this module produced - the in-terminal HUD and the tailnet dashboard can
+never diverge. See docs/COGNITION_UNIFIED_PLAN.md (Phase 2: shared signal).
+
+The attention signal is a novel cognitive read derived from the session
+transcript: how much effort/precision goes into prompts (effort), the
+deliberation rhythm between turns (deliberation), and whether engagement is
+rising or fading (trend). It nudges "keep going" when sharp and "rest" when
+fading.
+
+Design: stdlib-only, Python 3.9+ safe, never raises. The model is a documented,
+tunable v1 heuristic; env vars ZENO_HUD_* override the weights/window so the HUD
+and any consumer that imports this module read identical knobs.
+
+The TS mirror of this formula lives in the dashboard; keep the two in lockstep
+(weights, thresholds, label bands) when tuning - that is the whole contract.
+"""
+
+import json
+import math
+import os
+import re
+
+# ---------------------------------------------------------------------------
+# tunable model knobs (env-overridable; shared by every importer of this module)
+# ---------------------------------------------------------------------------
+# attention weights (effort / deliberation / trend), normalized internally
+W_EFFORT = float(os.environ.get("ZENO_HUD_W_EFFORT", "0.55"))
+W_DELIB = float(os.environ.get("ZENO_HUD_W_DELIB", "0.30"))
+W_TREND = float(os.environ.get("ZENO_HUD_W_TREND", "0.15"))
+WINDOW = int(os.environ.get("ZENO_HUD_WINDOW", "6"))  # recent prompts considered
+TAIL_BYTES = int(os.environ.get("ZENO_HUD_TAIL_BYTES", "1048576"))  # transcript tail to read
+
+# label bands (score 0..100 -> label/nudge): the single classifier the bar and
+# the dashboard both honor. Colors are ANSI codes for the HUD; the dashboard
+# maps the same bands to its theme.
+GREEN, YELLOW, RED, CYAN = "92", "93", "91", "96"
+
+_TAG_RE = re.compile(r"<(/?)(command-[\w-]+|local-command-[\w-]+|system-reminder)[^>]*>")
+
+
+# ---------------------------------------------------------------------------
+# prompt cleaning + per-prompt effort
+# ---------------------------------------------------------------------------
+def clean_prompt_text(content):
+    """Extract human-authored prompt text, dropping injected wrappers/tool-results."""
+    if isinstance(content, str):
+        text = content
+    elif isinstance(content, list):
+        parts = []
+        for b in content:
+            if not isinstance(b, dict):
+                continue
+            if b.get("type") == "tool_result":  # this is a tool return, not a prompt
+                return ""
+            if b.get("type") == "text" and isinstance(b.get("text"), str):
+                parts.append(b["text"])
+        text = "\n".join(parts)
+    else:
+        return ""
+    text = _TAG_RE.sub("", text)  # strip <command-*>/<system-reminder> tags
+    # drop lines that are pure injected blocks the user didn't type
+    keep = [
+        ln for ln in text.splitlines() if not ln.lstrip().startswith(("<", "Caveat:", "[SYSTEM"))
+    ]
+    return "\n".join(keep).strip()
+
+
+def prompt_effort(text):
+    """0..1 effort score for one prompt: longer + structured + code = more effort."""
+    n = len(text)
+    base = 1.0 - math.exp(-n / 180.0)  # ~0 tiny, ~0.8 @300 chars, ->1 long paragraph
+    bonus = 0.0
+    if "```" in text or text.count("\n") >= 2:
+        bonus += 0.15
+    if n <= 6 and text.lower() in ("y", "yes", "go", "ok", "continue", "do it", "ship it", "next"):
+        base *= 0.25  # one-word "let the AI lead" prompts
+    return max(0.0, min(1.0, base + bonus))
+
+
+def delib_score(gap_s):
+    """0..1 health of a between-turn gap: rapid-fire coasting low, thoughtful mid-high."""
+    if gap_s is None:
+        return 0.5
+    if gap_s < 4:
+        return 0.30  # rapid-fire / coasting
+    if gap_s <= 180:
+        return 1.0  # active thinking
+    if gap_s <= 900:
+        return 0.6  # slower, still engaged
+    return None  # >15min: away, exclude from mean
+
+
+def parse_ts(s):
+    if not isinstance(s, str):
+        return None
+    try:
+        s = s.replace("Z", "+0000").replace("-00:00", "+0000")
+        # tolerate fractional seconds + tz; parse epoch via time.strptime on the date part
+        m = re.match(r"(\d{4})-(\d{2})-(\d{2})T(\d{2}):(\d{2}):(\d{2})", s)
+        if not m:
+            return None
+        import calendar
+
+        return calendar.timegm(tuple(int(x) for x in m.groups()) + (0, 0, 0))
+    except Exception:
+        return None
+
+
+def parse_transcript(path):
+    """Read the transcript tail; return (prompts, asst_turns).
+    prompts: list of (ts_epoch, effort 0..1). asst_turns: list of (ts_epoch, out_tokens, tool_uses).
+    """
+    prompts, asst = [], []
+    if not path or not os.path.exists(path):
+        return prompts, asst
+    try:
+        size = os.path.getsize(path)
+        with open(path, "rb") as f:
+            if size > TAIL_BYTES:
+                f.seek(size - TAIL_BYTES)
+                f.readline()  # discard partial first line
+            data = f.read().decode("utf-8", "replace")
+    except Exception:
+        return prompts, asst
+    for line in data.splitlines():
+        line = line.strip()
+        if not line:
+            continue
+        try:
+            o = json.loads(line)
+        except Exception:
+            continue
+        t = o.get("type")
+        if t not in ("user", "assistant"):
+            continue
+        ts = parse_ts(o.get("timestamp"))
+        msg = o.get("message") or {}
+        if t == "user":
+            if o.get("isMeta") or o.get("isSidechain"):
+                continue
+            text = clean_prompt_text(msg.get("content"))
+            if not text:
+                continue
+            prompts.append((ts, prompt_effort(text)))
+        else:
+            usage = msg.get("usage") or {}
+            out = usage.get("output_tokens") or 0
+            content = msg.get("content")
+            tools = (
+                sum(1 for b in content if isinstance(b, dict) and b.get("type") == "tool_use")
+                if isinstance(content, list)
+                else 0
+            )
+            asst.append((ts, out, tools))
+    return prompts, asst
+
+
+# ---------------------------------------------------------------------------
+# the formula
+# ---------------------------------------------------------------------------
+def classify(score, trend, long_session):
+    """Map a score (+ trend, session length) to (label, nudge, glyph, color).
+
+    The single classifier the HUD bar and the dashboard both honor: the bands
+    here are the contract. Returns plain data so a reader presenting a persisted
+    row reproduces the exact same label/nudge/glyph the writer would have shown.
+    """
+    glyph = "▲" if trend > 0.08 else ("▼" if trend < -0.08 else "─")
+    if score >= 70:
+        label, nudge, col = "sharp", "keep going", GREEN
+    elif score >= 45:
+        label, nudge, col = "steady", "in rhythm", CYAN
+    elif score >= 25:
+        label, nudge, col = "easing", ("wrap soon" if long_session else "ride it out"), YELLOW
+    else:
+        label, nudge, col = "fading", "rest", RED
+    if long_session and trend < -0.2 and score < 70:
+        nudge = "rest soon"
+        col = RED if score < 45 else YELLOW
+    return label, nudge, glyph, col
+
+
+def attention(prompts, asst):
+    """Compute the attention signal from parsed transcript data.
+
+    Returns dict(ok, score 0..100, trend -1..1, effort, deliberation, label,
+    nudge, glyph, color) or {"ok": False} when there is too little data. This is
+    the canonical computation: the HUD writer persists score/effort/deliberation/
+    trend from here into cognition_samples, and present_sample() rebuilds the
+    label/nudge/glyph from the same classifier - so a row read back is identical
+    to a row freshly computed.
+    """
+    if len(prompts) < 2:
+        return {"ok": False}
+    recent = prompts[-WINDOW:]
+    efforts = [e for _, e in recent]
+    effort_now = sum(efforts) / len(efforts)
+
+    # deliberation: gap from each prompt back to the prior assistant turn (or prior prompt)
+    gaps = []
+    asst_ts = [a[0] for a in asst if a[0] is not None]
+    for ts, _ in recent:
+        if ts is None:
+            continue
+        prior = [x for x in asst_ts if x <= ts]
+        if prior:
+            gaps.append(ts - max(prior))
+    delib_scores = [s for s in (delib_score(g) for g in gaps) if s is not None]
+    delib = sum(delib_scores) / len(delib_scores) if delib_scores else 0.5
+
+    # trend: recent half vs earlier half of the window
+    if len(efforts) >= 4:
+        half = len(efforts) // 2
+        trend = max(
+            -1.0,
+            min(1.0, (sum(efforts[half:]) / (len(efforts) - half)) - (sum(efforts[:half]) / half)),
+        )
+    else:
+        trend = 0.0
+
+    wsum = W_EFFORT + W_DELIB + W_TREND
+    raw = (W_EFFORT * effort_now + W_DELIB * delib + W_TREND * (0.5 + 0.5 * trend)) / wsum
+    score = int(round(max(0.0, min(1.0, raw)) * 100))
+
+    # session length (minutes) for fatigue framing
+    ts_all = [t for t, _ in prompts if t is not None]
+    session_min = (max(ts_all) - min(ts_all)) / 60.0 if len(ts_all) >= 2 else 0.0
+    long_session = session_min >= 120
+
+    label, nudge, glyph, col = classify(score, trend, long_session)
+    return {
+        "ok": True,
+        "score": score,
+        "trend": trend,
+        "effort": effort_now,
+        "deliberation": delib,
+        "label": label,
+        "nudge": nudge,
+        "glyph": glyph,
+        "color": col,
+    }
+
+
+def present_sample(row):
+    """Rebuild the renderable attention dict from a persisted cognition_samples row.
+
+    `row` is a mapping with attention_score / attention_trend (+ optional
+    attention_effort, attention_deliberation, ts pairs for session length).
+    Re-derives label/nudge/glyph via the same `classify` the live path uses, so
+    the bar shown from a stored row equals the bar shown from a live compute.
+    Returns {"ok": False} when the row has no usable score.
+    """
+    if not row:
+        return {"ok": False}
+    score = row.get("attention_score")
+    if score is None:
+        return {"ok": False}
+    try:
+        score = int(round(float(score)))
+    except (TypeError, ValueError):
+        return {"ok": False}
+    trend = row.get("attention_trend")
+    trend = float(trend) if isinstance(trend, (int, float)) else 0.0
+    long_session = bool(row.get("long_session"))
+    label, nudge, glyph, col = classify(score, trend, long_session)
+    out = {
+        "ok": True,
+        "score": score,
+        "trend": trend,
+        "label": label,
+        "nudge": nudge,
+        "glyph": glyph,
+        "color": col,
+    }
+    eff = row.get("attention_effort")
+    if isinstance(eff, (int, float)):
+        out["effort"] = float(eff)
+    delib = row.get("attention_deliberation")
+    if isinstance(delib, (int, float)):
+        out["deliberation"] = float(delib)
+    return out