tokenjam 0.2.0__py3-none-any.whl

tokenjam/cli/cmd_cost.py

@@ -0,0 +1,90 @@
+import click
+import json
+from tokenjam.core.models import CostFilters
+from tokenjam.utils.formatting import console, make_table, format_cost, format_tokens
+from tokenjam.utils.time_parse import parse_since
+
+
+@click.command("cost")
+@click.option("--agent", default=None, help="Filter to specific agent_id")
+@click.option("--since", default="7d", help="Time window (e.g. 1h, 7d, 2026-03-01)")
+@click.option("--group-by", "group_by",
+              type=click.Choice(["agent", "model", "day", "tool"]),
+              default="day")
+@click.option("--json", "output_json", is_flag=True)
+@click.pass_context
+def cmd_cost(ctx: click.Context, agent: str | None, since: str,
+             group_by: str, output_json: bool) -> None:
+    """Show cost breakdown by agent, model, day, or tool."""
+    db = ctx.obj["db"]
+    try:
+        since_dt = parse_since(since)
+    except ValueError as exc:
+        raise click.BadParameter(str(exc), param_hint="'--since'") from exc
+    filters = CostFilters(
+        agent_id=agent,
+        since=since_dt,
+        group_by=group_by,
+    )
+    rows = db.get_cost_summary(filters)
+    total = sum(r.cost_usd for r in rows)
+
+    if output_json:
+        click.echo(json.dumps({
+            "rows": [vars(r) for r in rows],
+            "total_cost_usd": total,
+        }, default=str))
+        return
+
+    if not rows:
+        console.print("[dim]No cost data found for the given filters.[/dim]")
+        return
+
+    if group_by == "day":
+        table = make_table("DATE", "AGENT", "MODEL", "TOKENS IN", "TOKENS OUT", "COST")
+        for r in rows:
+            table.add_row(
+                r.group,
+                r.agent_id or "-",
+                r.model or "-",
+                format_tokens(r.input_tokens),
+                format_tokens(r.output_tokens),
+                format_cost(r.cost_usd),
+            )
+    elif group_by == "agent":
+        table = make_table("AGENT", "MODEL", "TOKENS IN", "TOKENS OUT", "COST")
+        for r in rows:
+            table.add_row(
+                r.group,
+                r.model or "-",
+                format_tokens(r.input_tokens),
+                format_tokens(r.output_tokens),
+                format_cost(r.cost_usd),
+            )
+    elif group_by == "model":
+        table = make_table("MODEL", "TOKENS IN", "TOKENS OUT", "COST")
+        for r in rows:
+            table.add_row(
+                r.group,
+                format_tokens(r.input_tokens),
+                format_tokens(r.output_tokens),
+                format_cost(r.cost_usd),
+            )
+    elif group_by == "tool":
+        table = make_table("TOOL", "COST")
+        for r in rows:
+            table.add_row(
+                r.group,
+                format_cost(r.cost_usd),
+            )
+
+    if group_by == "day":
+        table.add_row("", "", "", "", "[bold]TOTAL[/bold]", f"[bold]{format_cost(total)}[/bold]")
+    elif group_by == "agent":
+        table.add_row("", "", "", "[bold]TOTAL[/bold]", f"[bold]{format_cost(total)}[/bold]")
+    elif group_by == "model":
+        table.add_row("", "", "[bold]TOTAL[/bold]", f"[bold]{format_cost(total)}[/bold]")
+    elif group_by == "tool":
+        table.add_row("[bold]TOTAL[/bold]", f"[bold]{format_cost(total)}[/bold]")
+
+    console.print(table)

tokenjam/cli/cmd_demo.py

@@ -0,0 +1,82 @@
+"""tj demo — Agent Incident Library CLI command."""
+from __future__ import annotations
+
+import importlib.util
+from pathlib import Path
+from types import ModuleType
+
+import click
+
+from tokenjam.utils.formatting import console
+
+# incidents/ lives two levels above this file (tj/cli/ -> tj/ -> repo/site-packages root)
+_INCIDENTS_DIR = Path(__file__).parent.parent.parent / "incidents"
+
+
+def _discover_scenarios() -> dict[str, ModuleType]:
+    """
+    Scan incidents/*/scenario.py for modules exposing a `run` callable.
+    Returns a dict mapping scenario slug to loaded module.
+    """
+    scenarios: dict[str, ModuleType] = {}
+    if not _INCIDENTS_DIR.exists():
+        return scenarios
+    for scenario_file in sorted(_INCIDENTS_DIR.glob("*/scenario.py")):
+        slug = scenario_file.parent.name
+        spec = importlib.util.spec_from_file_location(
+            f"incidents.{slug}.scenario", scenario_file
+        )
+        if spec is None or spec.loader is None:
+            continue
+        mod = importlib.util.module_from_spec(spec)
+        spec.loader.exec_module(mod)  # type: ignore[union-attr]
+        if callable(getattr(mod, "run", None)):
+            scenarios[slug] = mod
+    return scenarios
+
+
+@click.command("demo")
+@click.argument("scenario", required=False, default=None)
+@click.option("--json", "output_json", is_flag=True, help="Output JSON instead of Rich panels")
+@click.pass_context
+def cmd_demo(ctx: click.Context, scenario: str | None, output_json: bool) -> None:
+    """Run a reproducible AI agent incident scenario.
+
+    \b
+    tj demo                     List available scenarios
+    tj demo retry-loop          Run a specific scenario
+    tj demo retry-loop --json   Machine-readable output
+    """
+    scenarios = _discover_scenarios()
+
+    if scenario is None:
+        _list_scenarios(scenarios)
+        return
+
+    if scenario not in scenarios:
+        click.echo(
+            f"Unknown scenario '{scenario}'. Run `tj demo` to see available scenarios.",
+            err=True,
+        )
+        raise SystemExit(1)
+
+    scenarios[scenario].run()
+
+
+def _list_scenarios(scenarios: dict[str, ModuleType]) -> None:
+    from rich import box
+    from rich.table import Table
+
+    console.print()
+    console.print(
+        "[bold]OCW Agent Incident Library[/bold]\n"
+        "Reproducible AI agent failures — no API keys, no config needed.\n"
+    )
+    table = Table(box=box.SIMPLE, show_header=True, header_style="bold")
+    table.add_column("Scenario", style="cyan", no_wrap=True)
+    table.add_column("Description")
+    for slug, mod in scenarios.items():
+        table.add_row(slug, getattr(mod, "DESCRIPTION", ""))
+    console.print(table)
+    console.print("[dim]Usage:[/dim] tj demo <scenario>  [dim]|[/dim]  tj demo <scenario> --json")
+    console.print()

tokenjam/cli/cmd_doctor.py

@@ -0,0 +1,173 @@
+from __future__ import annotations
+
+import json
+
+import click
+import duckdb
+
+from tokenjam.core.config import find_config_file, load_config
+from tokenjam.utils.formatting import console
+
+
+@click.command("doctor")
+@click.option("--json", "output_json", is_flag=True)
+@click.pass_context
+def cmd_doctor(ctx: click.Context, output_json: bool) -> None:
+    """Run health checks on tj configuration and environment."""
+    config = ctx.obj["config"]
+    checks: list[dict] = []
+
+    # 1. Config file found and valid
+    checks.append(_check_config())
+
+    # 2. DuckDB file writable
+    checks.append(_check_db(config))
+
+    # 3. Ingest secret set
+    checks.append(_check_ingest_secret(config))
+
+    # 4. Prometheus configured
+    checks.append(_check_prometheus(config))
+
+    # 5. Schema validation vs capture
+    checks.append(_check_schema_vs_capture(config))
+
+    # 6. Drift configured but inactive
+    checks.append(_check_drift_inactive(config, ctx.obj["db"]))
+
+    # 7. Webhook URL security
+    checks.extend(_check_webhook_security(config))
+
+    # 8. Webhook domain allowlist
+    checks.extend(_check_webhook_allowlist(config))
+
+    if output_json:
+        click.echo(json.dumps(checks, default=str))
+    else:
+        for c in checks:
+            _print_check(c)
+
+    has_errors = any(c["level"] == "error" for c in checks)
+    has_warnings = any(c["level"] == "warning" for c in checks)
+    if has_errors:
+        ctx.exit(2)
+    elif has_warnings:
+        ctx.exit(1)
+    else:
+        ctx.exit(0)
+
+
+def _check_config() -> dict:
+    try:
+        path = find_config_file()
+        if path is None:
+            return {"name": "Config file", "level": "error",
+                    "message": "No config file found. Run `tj onboard` to create one."}
+        load_config(str(path))
+        return {"name": "Config file", "level": "ok",
+                "message": f"Found and valid: {path}"}
+    except Exception as e:
+        return {"name": "Config file", "level": "error",
+                "message": f"Config parse error: {e}"}
+
+
+def _check_db(config: object) -> dict:
+    try:
+        from pathlib import Path
+        db_path = Path(config.storage.path).expanduser()
+        conn = duckdb.connect(str(db_path))
+        conn.close()
+        return {"name": "DuckDB writable", "level": "ok",
+                "message": f"Database accessible: {db_path}"}
+    except Exception as e:
+        return {"name": "DuckDB writable", "level": "error",
+                "message": f"Cannot open database: {e}"}
+
+
+def _check_ingest_secret(config: object) -> dict:
+    if config.security.ingest_secret:
+        return {"name": "Ingest secret", "level": "ok",
+                "message": "Ingest secret is configured."}
+    return {"name": "Ingest secret", "level": "warning",
+            "message": "No ingest secret set. API ingest endpoint is unprotected."}
+
+
+def _check_prometheus(config: object) -> dict:
+    if config.export.prometheus.enabled:
+        return {"name": "Prometheus", "level": "ok",
+                "message": f"Enabled on port {config.export.prometheus.port}"}
+    return {"name": "Prometheus", "level": "info",
+            "message": "Prometheus export disabled."}
+
+
+def _check_schema_vs_capture(config: object) -> dict:
+    has_schema = any(
+        ac.output_schema for ac in config.agents.values()
+    )
+    if has_schema and not config.capture.tool_outputs:
+        return {"name": "Schema vs capture", "level": "warning",
+                "message": "Agent has output_schema but capture.tool_outputs is false. "
+                           "Schema validation will have no data to validate."}
+    return {"name": "Schema vs capture", "level": "ok",
+            "message": "Schema and capture settings are consistent."}
+
+
+def _check_drift_inactive(config: object, db: object) -> dict:
+    for agent_id, ac in config.agents.items():
+        if ac.drift.enabled:
+            count = db.get_completed_session_count(agent_id)
+            if count < ac.drift.baseline_sessions:
+                return {"name": "Drift detection", "level": "warning",
+                        "message": f"Agent '{agent_id}' has drift enabled but only "
+                                   f"{count}/{ac.drift.baseline_sessions} baseline sessions."}
+    return {"name": "Drift detection", "level": "ok",
+            "message": "Drift detection status is consistent."}
+
+
+def _check_webhook_security(config: object) -> list[dict]:
+    results = []
+    for ch in config.alerts.channels:
+        url = ch.url or ch.webhook_url
+        if url and not url.startswith("https://") and not _is_local_url(url):
+            results.append({
+                "name": "Webhook security",
+                "level": "warning",
+                "message": f"Non-HTTPS, non-local webhook URL: {url}",
+            })
+    if not results:
+        results.append({"name": "Webhook security", "level": "ok",
+                        "message": "All webhook URLs are secure or local."})
+    return results
+
+
+def _check_webhook_allowlist(config: object) -> list[dict]:
+    allowed = config.security.webhook_allowed_domains
+    if not allowed:
+        return []
+    results = []
+    for ch in config.alerts.channels:
+        url = ch.url or ch.webhook_url
+        if url:
+            from urllib.parse import urlparse
+            domain = urlparse(url).hostname
+            if domain and domain not in allowed:
+                results.append({
+                    "name": "Webhook allowlist",
+                    "level": "error",
+                    "message": f"Webhook domain '{domain}' not in allowed list.",
+                })
+    return results
+
+
+def _is_local_url(url: str) -> bool:
+    from urllib.parse import urlparse
+    hostname = urlparse(url).hostname
+    return hostname in ("localhost", "127.0.0.1", "::1", "0.0.0.0") if hostname else False
+
+
+def _print_check(check: dict) -> None:
+    level = check["level"]
+    icons = {"ok": "[green]\u2713[/green]", "warning": "[yellow]\u26a0[/yellow]",
+             "error": "[red]\u2717[/red]", "info": "[blue]i[/blue]"}
+    icon = icons.get(level, "?")
+    console.print(f"  {icon}  {check['name']}: {check['message']}")

tokenjam/cli/cmd_drift.py

@@ -0,0 +1,238 @@
+"""tj drift — show behavioral drift baselines and Z-scores."""
+from __future__ import annotations
+
+import json as json_mod
+
+import click
+from rich.table import Table
+
+from tokenjam.core.drift import evaluate_drift
+from tokenjam.utils.formatting import console
+
+
+@click.command("drift")
+@click.option("--agent", default=None, help="Filter to specific agent_id")
+@click.option("--json", "output_json", is_flag=True, help="JSON output")
+@click.pass_context
+def cmd_drift(ctx: click.Context, agent: str | None, output_json: bool) -> None:
+    """Show drift baselines and Z-scores for recent sessions."""
+    db = ctx.obj["db"]
+    config = ctx.obj["config"]
+    agent_filter = agent or ctx.obj.get("agent")
+
+    # Discover agents with baselines
+    if agent_filter:
+        agent_ids = [agent_filter]
+    elif hasattr(db, "conn"):
+        rows = db.conn.execute(
+            "SELECT DISTINCT agent_id FROM drift_baselines ORDER BY agent_id"
+        ).fetchall()
+        agent_ids = [r[0] for r in rows]
+    else:
+        agent_ids = []
+
+    if not agent_ids:
+        if output_json:
+            click.echo(json_mod.dumps({"agents": [], "drifted": False}))
+        else:
+            console.print(
+                "[dim]No drift baselines found. "
+                "Need at least 10 completed sessions to build a baseline.[/dim]"
+            )
+        ctx.exit(0)
+        return
+
+    all_results = []
+    any_drifted = False
+
+    for aid in agent_ids:
+        baseline = db.get_baseline(aid)
+        if baseline is None:
+            continue
+
+        sessions = db.get_completed_sessions(aid, limit=1)
+        if not sessions:
+            continue
+        latest = sessions[0]
+
+        agent_cfg = config.agents.get(aid)
+        threshold = agent_cfg.drift.token_threshold if agent_cfg else 2.0
+        seq_threshold = agent_cfg.drift.tool_sequence_diff if agent_cfg else 0.4
+
+        result = evaluate_drift(
+            session=latest,
+            baseline=baseline,
+            config_threshold=threshold,
+            sequence_diff_threshold=seq_threshold,
+            db=db,
+        )
+
+        if result.drifted:
+            any_drifted = True
+
+        agent_data = {
+            "agent_id": aid,
+            "baseline_sessions": baseline.sessions_sampled,
+            "drifted": result.drifted,
+            "violations": [
+                {
+                    "dimension": v.dimension,
+                    "z_score": v.z_score,
+                    "expected": v.expected,
+                    "observed": v.observed,
+                    "detail": v.detail,
+                }
+                for v in result.violations
+            ],
+            "metrics": _build_metrics(baseline, latest, result, threshold),
+        }
+        all_results.append(agent_data)
+
+        if not output_json:
+            _print_drift_table(aid, baseline, latest, result, threshold, seq_threshold)
+
+    if output_json:
+        click.echo(json_mod.dumps(
+            {"agents": all_results, "drifted": any_drifted},
+            default=str,
+        ))
+
+    ctx.exit(1 if any_drifted else 0)
+
+
+def _build_metrics(baseline, session, result, threshold: float) -> list[dict]:
+    """Return per-dimension metric dicts for JSON output."""
+    from tokenjam.core.drift import z_score
+
+    violated_dims = {v.dimension for v in result.violations}
+    metrics = []
+
+    def _add(dimension: str, mean, stddev, current) -> None:
+        if mean is None or stddev is None:
+            return
+        z = z_score(float(current), float(mean), float(stddev))
+        metrics.append({
+            "dimension": dimension,
+            "baseline_mean": mean,
+            "baseline_stddev": stddev,
+            "current_value": current,
+            "z_score": z,
+            "status": "DRIFT" if dimension in violated_dims else "ok",
+        })
+
+    _add("input_tokens", baseline.avg_input_tokens, baseline.stddev_input_tokens,
+         session.input_tokens)
+    _add("output_tokens", baseline.avg_output_tokens, baseline.stddev_output_tokens,
+         session.output_tokens)
+    if session.duration_seconds is not None:
+        _add("session_duration", baseline.avg_session_duration_s,
+             baseline.stddev_session_duration, session.duration_seconds)
+    _add("tool_call_count", baseline.avg_tool_call_count, baseline.stddev_tool_call_count,
+         session.tool_call_count)
+
+    # tool_sequence is special (Jaccard, no z-score)
+    if "tool_sequence" in violated_dims:
+        seq_viol = next((v for v in result.violations if v.dimension == "tool_sequence"), None)
+        if seq_viol:
+            metrics.append({
+                "dimension": "tool_sequence",
+                "baseline_mean": None,
+                "baseline_stddev": None,
+                "current_value": seq_viol.observed,
+                "z_score": None,
+                "status": "DRIFT",
+            })
+
+    return metrics
+
+
+def _print_drift_table(aid, baseline, session, result, threshold: float, seq_threshold: float = 0.4) -> None:
+    """Render a Rich table for a single agent's drift state."""
+    from tokenjam.core.drift import z_score
+
+    violated_dims = {v.dimension for v in result.violations}
+    status_label = "[bold red]DRIFTED[/bold red]" if result.drifted else "[green]ok[/green]"
+
+    console.print()
+    console.print(
+        f"[bold]Agent:[/bold] {aid}  |  "
+        f"[bold]Baseline:[/bold] {baseline.sessions_sampled} sessions  |  "
+        f"[bold]Status:[/bold] {status_label}"
+    )
+    console.print()
+
+    table = Table(show_header=True, header_style="bold")
+    table.add_column("Dimension", style="dim")
+    table.add_column("Baseline")
+    table.add_column("Current")
+    table.add_column("Z-Score", justify="right")
+    table.add_column("Status")
+
+    def _z_color(z: float | None) -> str:
+        if z is None:
+            return "--"
+        az = abs(z)
+        if az < 1.0:
+            return f"[green]{z:.2f}[/green]"
+        if az <= threshold:
+            return f"[yellow]{z:.2f}[/yellow]"
+        return f"[red]{z:.2f}[/red]"
+
+    def _status_cell(dimension: str) -> str:
+        if dimension in violated_dims:
+            return "[bold red]DRIFT[/bold red]"
+        return "[green]ok[/green]"
+
+    def _add_row(dimension: str, mean, stddev, current, fmt_baseline: str, fmt_current: str) -> None:
+        z = z_score(float(current), float(mean), float(stddev)) if (mean is not None and stddev is not None) else None
+        table.add_row(dimension, fmt_baseline, fmt_current, _z_color(z), _status_cell(dimension))
+
+    if baseline.avg_input_tokens is not None and baseline.stddev_input_tokens is not None:
+        _add_row(
+            "input_tokens",
+            baseline.avg_input_tokens, baseline.stddev_input_tokens, session.input_tokens,
+            f"{baseline.avg_input_tokens:,.0f} +/- {baseline.stddev_input_tokens:,.0f}",
+            f"{session.input_tokens:,}",
+        )
+    if baseline.avg_output_tokens is not None and baseline.stddev_output_tokens is not None:
+        _add_row(
+            "output_tokens",
+            baseline.avg_output_tokens, baseline.stddev_output_tokens, session.output_tokens,
+            f"{baseline.avg_output_tokens:,.0f} +/- {baseline.stddev_output_tokens:,.0f}",
+            f"{session.output_tokens:,}",
+        )
+    if (
+        session.duration_seconds is not None
+        and baseline.avg_session_duration_s is not None
+        and baseline.stddev_session_duration is not None
+    ):
+        _add_row(
+            "session_duration",
+            baseline.avg_session_duration_s, baseline.stddev_session_duration,
+            session.duration_seconds,
+            f"{baseline.avg_session_duration_s:.1f}s +/- {baseline.stddev_session_duration:.1f}s",
+            f"{session.duration_seconds:.1f}s",
+        )
+    if baseline.avg_tool_call_count is not None and baseline.stddev_tool_call_count is not None:
+        _add_row(
+            "tool_call_count",
+            baseline.avg_tool_call_count, baseline.stddev_tool_call_count, session.tool_call_count,
+            f"{baseline.avg_tool_call_count:.0f} +/- {baseline.stddev_tool_call_count:.0f}",
+            str(session.tool_call_count),
+        )
+
+    # Tool sequence row (Jaccard, no z-score)
+    seq_viol = next((v for v in result.violations if v.dimension == "tool_sequence"), None)
+    if seq_viol:
+        table.add_row(
+            "tool_sequence",
+            seq_viol.expected or "",
+            seq_viol.observed or "",
+            "--",
+            "[bold red]DRIFT[/bold red]",
+        )
+    elif baseline.common_tool_sequences:
+        min_sim = 1.0 - seq_threshold
+        table.add_row("tool_sequence", f"similarity >= {min_sim:.2f}", "--", "--", "[green]ok[/green]")
+
+    console.print(table)