ss-tools/scripts/pretty_cot.py

#!/usr/bin/env python3
# #region Scripts.PrettyCoT [C:2] [TYPE Script] [SEMANTICS logging,pretty,agent,cli]
# @BRIEF Agent-first pretty printer and trace visualizer for Molecular CoT JSON logs.
# @RELATION IMPLEMENTS -> [molecular-cot-logging:CLI Reader]
# @USAGE
#   python scripts/pretty_cot.py backend/logs/app.log --last 80
#   tail -f logs/app.log | python scripts/pretty_cot.py
#   python scripts/pretty_cot.py backend/logs/app.log --follow
# @INVARIANT Output is always a readable narrative grouped by trace_id.
"""
pretty_cot.py — Agent-centric pretty printer for Molecular CoT logs.

Usage:
    python scripts/pretty_cot.py backend/logs/app.log --last 100
    python scripts/pretty_cot.py --trace <trace_id> < log.txt
    python scripts/pretty_cot.py backend/logs/app.log --follow
    tail -f backend/logs/app.log | python scripts/pretty_cot.py

Makes raw JSON CoT lines human- and agent-readable with icons, grouping,
and optional filtering. Designed so an agent can quickly understand a trace.

This is the reference tool for "agent view" of logs (see molecular-cot-logging skill).
"""
# #endregion Scripts.PrettyCoT

import argparse
import json
import os
import sys
import time
from collections import defaultdict
from datetime import datetime
from pathlib import Path
from typing import Any, Iterable, Optional

# ── Central suppression (mirrors shared/ss_tools/shared/cot_logger.py) ──
# This is the canonical list. If you update this, update shared too.
_ROUTINE_PHRASES = (
    "Reusing cached Superset auth tokens",
    "Resolve authenticated user principal",
    "User principal resolved",
    "Resolving current user preference",
    "Loading current user's dashboard preference",
    "Validated ENCRYPTION_KEY",
    "Superset client ready",
    "Superset client initialized",
    "SupersetClientRegistry.get_client",
    "Initialized ResourceService",
    "ResourceService initialized",
    "Created shared HTTP client",
    "Ensured directory",
)


def _is_routine(intent: str) -> bool:
    """Check if intent matches known infrastructure noise (defense-in-depth)."""
    return any(p in intent for p in _ROUTINE_PHRASES)


ICONS = {
    "REASON": "→",
    "REFLECT": "✓",
    "EXPLORE": "⚠",
}

LEVEL_COLOR = {
    "INFO": "",
    "WARNING": "⚡ ",
    "ERROR": "🔥 ",
    "DEBUG": "… ",
}


def parse_line(line: str) -> Optional[dict[str, Any]]:
    line = line.strip()
    if not line:
        return None
    try:
        # Handle lines that may be prefixed by docker timestamps etc.
        if line.startswith("{"):
            return json.loads(line)
        # Try to find embedded JSON
        start = line.find("{")
        if start != -1:
            candidate = line[start:]
            end = candidate.rfind("}") + 1
            if end > 1:
                return json.loads(candidate[:end])
    except Exception:
        pass
    return None


def format_record(rec: dict[str, Any], compact: bool = False) -> str:
    ts = rec.get("ts", "")
    if isinstance(ts, str) and "T" in ts:
        try:
            dt = datetime.fromisoformat(ts.replace("Z", "+00:00").split(".")[0])
            ts = dt.strftime("%H:%M:%S")
        except Exception:
            ts = ts[-12:]

    level = rec.get("level", "INFO")
    marker = rec.get("marker", "REASON")
    src = rec.get("src", "?")
    intent = rec.get("intent", "")
    payload = rec.get("payload")
    error = rec.get("error")
    elapsed = rec.get("elapsed_ms")

    # Suppress routine infra noise (defense-in-depth; call-site cleanup is primary)
    if _is_routine(intent):
        return ""

    icon = ICONS.get(marker, "·")
    lvl = LEVEL_COLOR.get(level, "") + level[:4]

    trace = rec.get("trace_id", "")
    trace_short = trace[:8] if trace and trace != "no-trace" else ""

    base = f"{ts} {icon} {lvl:5} {src:30.30}  {intent}"

    if elapsed is not None:
        base += f"  ⏱{elapsed}ms"

    if error:
        base += f"  | error={error}"

    if payload and not compact:
        try:
            p = json.dumps(payload, ensure_ascii=False, default=str)[:180]
            base += f"  | {p}"
        except Exception:
            pass

    if trace_short and not compact:
        base = f"[{trace_short}] {base}"

    return base


def pretty_print(
    records: Iterable[dict[str, Any]],
    group_by_trace: bool = True,
    compact: bool = False,
    max_per_trace: int = 60,
) -> None:
    if not group_by_trace:
        for r in records:
            line = format_record(r, compact=compact)
            if line:
                print(line)
        return

    # Group and limit per trace to keep signal high for agent
    by_trace: dict[str, list[dict]] = defaultdict(list)
    for r in records:
        tid = r.get("trace_id") or "no-trace"
        by_trace[tid].append(r)

    for tid, items in by_trace.items():
        filtered = [r for r in items if format_record(r, compact=compact)]
        if not filtered:
            continue
        if tid != "no-trace":
            print(f"\n=== TRACE {tid} ({len(filtered)} events) ===")
        shown = 0
        for r in filtered:
            if shown >= max_per_trace:
                print(f"   ... ({len(filtered) - shown} more events truncated for agent focus)")
                break
            line = format_record(r, compact=compact)
            if line:
                print("  " + line)
            shown += 1


def follow_file(file_path: str) -> None:
    """Tail -f equivalent for a single log file, processing new lines in real time."""
    with open(file_path, encoding="utf-8", errors="replace") as f:
        # Seek to end
        f.seek(0, os.SEEK_END)
        print(f"Following {file_path}... (Ctrl-C to stop)", file=sys.stderr)
        buf = ""
        try:
            while True:
                chunk = f.read()
                if chunk:
                    buf += chunk
                    lines = buf.split("\n")
                    buf = lines[-1]  # keep incomplete last line
                    for line in lines[:-1]:
                        rec = parse_line(line)
                        if rec:
                            line_out = format_record(rec)
                            if line_out:
                                print(line_out)
                else:
                    time.sleep(0.25)
        except KeyboardInterrupt:
            pass


def main() -> None:
    parser = argparse.ArgumentParser(description="Pretty-print Molecular CoT logs for agents")
    parser.add_argument("files", nargs="*", help="Log files (JSON lines). Use - for stdin.")
    parser.add_argument("--last", type=int, default=200, help="Only process last N lines")
    parser.add_argument("--trace", help="Filter to a specific trace_id")
    parser.add_argument("--compact", action="store_true", help="Less verbose output")
    parser.add_argument("--no-group", action="store_true", help="Do not group by trace_id")
    parser.add_argument("--follow", action="store_true", help="Follow (like tail -f)")

    args = parser.parse_args()

    # ── Follow mode ────────────────────────────────────────────────────────
    if args.follow and args.files:
        for f in args.files:
            follow_file(f)
        return

    # ── Batch mode ─────────────────────────────────────────────────────────
    sources: list[Any] = []
    if not args.files:
        sources.append(sys.stdin)
    else:
        for f in args.files:
            if f == "-":
                sources.append(sys.stdin)
            else:
                sources.append(Path(f).open(encoding="utf-8", errors="replace"))

    all_records: list[dict] = []
    for src in sources:
        try:
            lines = src if hasattr(src, "readline") else src.read_text(encoding="utf-8", errors="replace").splitlines()
            if args.last and hasattr(lines, "__iter__") and not args.follow:
                if isinstance(lines, list):
                    lines = lines[-args.last:]
            for line in lines:
                rec = parse_line(line if isinstance(line, str) else line)
                if rec:
                    if args.trace and args.trace not in (rec.get("trace_id") or ""):
                        continue
                    all_records.append(rec)
        except Exception as e:
            print(f"Error reading: {e}", file=sys.stderr)

    pretty_print(all_records, group_by_trace=not args.no_group, compact=args.compact)

    # Close files
    for s in sources:
        if hasattr(s, "close") and s is not sys.stdin:
            try:
                s.close()
            except Exception:
                pass


if __name__ == "__main__":
    main()