diff --git a/pyproject.toml b/pyproject.toml index 0d5bb39..c488e76 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -43,16 +43,13 @@ thirdeye-claude-stop-failure = "thirdeye.platforms.claude.hooks:stop_failure" thirdeye-claude-notification = "thirdeye.platforms.claude.hooks:notification" thirdeye-claude-permission-request = "thirdeye.platforms.claude.hooks:permission_request" thirdeye-claude-session-end = "thirdeye.platforms.claude.hooks:session_end" +thirdeye-claude-post-tool-use-failure = "thirdeye.platforms.claude.hooks:post_tool_use_failure" +thirdeye-claude-subagent-start = "thirdeye.platforms.claude.hooks:subagent_start" +thirdeye-claude-user-prompt-expansion = "thirdeye.platforms.claude.hooks:user_prompt_expansion" +thirdeye-claude-pre-compact = "thirdeye.platforms.claude.hooks:pre_compact" +thirdeye-claude-post-compact = "thirdeye.platforms.claude.hooks:post_compact" +thirdeye-claude-permission-denied = "thirdeye.platforms.claude.hooks:permission_denied" thirdeye-codex-notify = "thirdeye.platforms.codex.hooks:notify" -thirdeye-gemini-session-start = "thirdeye.platforms.gemini.hooks:session_start" -thirdeye-gemini-session-end = "thirdeye.platforms.gemini.hooks:session_end" -thirdeye-gemini-before-agent = "thirdeye.platforms.gemini.hooks:before_agent" -thirdeye-gemini-after-agent = "thirdeye.platforms.gemini.hooks:after_agent" -thirdeye-gemini-before-model = "thirdeye.platforms.gemini.hooks:before_model" -thirdeye-gemini-after-model = "thirdeye.platforms.gemini.hooks:after_model" -thirdeye-gemini-before-tool = "thirdeye.platforms.gemini.hooks:before_tool" -thirdeye-gemini-after-tool = "thirdeye.platforms.gemini.hooks:after_tool" -thirdeye-cursor-hook = "thirdeye.platforms.cursor.hook:main" [build-system] requires = ["setuptools>=64", "setuptools-scm>=8", "wheel"] diff --git a/src/thirdeye/commands/add.py b/src/thirdeye/commands/add.py index e02f54c..fc8badb 100644 --- a/src/thirdeye/commands/add.py +++ b/src/thirdeye/commands/add.py @@ -1,38 +1,100 @@ from __future__ import annotations +import json +from collections.abc import Iterable +from pathlib import Path + import click from thirdeye.platforms.base import Platform from thirdeye.platforms.claude.install import ClaudePlatform from thirdeye.platforms.codex.install import CodexPlatform -from thirdeye.platforms.cursor.install import CursorPlatform -from thirdeye.platforms.gemini.install import GeminiPlatform PLATFORMS: dict[str, type[Platform]] = { "claude": ClaudePlatform, - "gemini": GeminiPlatform, "codex": CodexPlatform, - "cursor": CursorPlatform, } +# Config files that may still reference console scripts for platforms this +# version no longer installs (Gemini and Cursor). Deleting those platforms +# leaves their hook entries orphaned in other tools' config, firing a missing +# binary on every event, so `thirdeye add --list` warns about them. +ORPHAN_CONFIG_PATHS: tuple[Path, ...] = ( + Path.home() / ".gemini" / "settings.json", + Path.home() / ".cursor" / "hooks.json", +) + def _platform_options(fn): - fn = click.option("--cursor", "platform_flag", flag_value="cursor", help="Cursor.")(fn) fn = click.option("--codex", "platform_flag", flag_value="codex", help="Codex CLI.")(fn) - fn = click.option("--gemini", "platform_flag", flag_value="gemini", help="Gemini CLI.")(fn) fn = click.option("--claude", "platform_flag", flag_value="claude", help="Claude Code.")(fn) return fn def _resolve_platform(platform_flag: str | None) -> Platform: if not platform_flag: - raise click.UsageError("Pick a platform: --claude, --gemini, --codex, --cursor") + raise click.UsageError("Pick a platform: --claude, --codex") return PLATFORMS[platform_flag]() +def _is_stale_command(command: str) -> bool: + """True if a hook command targets a removed-platform console script.""" + name = Path(command).name + return name.startswith("thirdeye-gemini-") or name == "thirdeye-cursor-hook" + + +def find_orphaned_hooks( + config_paths: Iterable[Path] = ORPHAN_CONFIG_PATHS, +) -> list[tuple[Path, str]]: + """Return (config_file, stale_command) for each removed-platform hook found. + + A command is stale when its basename starts with "thirdeye-gemini-" or + equals "thirdeye-cursor-hook". Missing or malformed files yield nothing. + """ + found: list[tuple[Path, str]] = [] + for path in config_paths: + try: + data = json.loads(path.read_text()) + except (OSError, json.JSONDecodeError): + continue + for command in _iter_commands(data): + if _is_stale_command(command): + found.append((path, command)) + return found + + +def _iter_commands(node: object) -> Iterable[str]: + """Yield every string under a "command" key anywhere in a nested structure.""" + if isinstance(node, dict): + for key, value in node.items(): + if key == "command" and isinstance(value, str): + yield value + else: + yield from _iter_commands(value) + elif isinstance(node, list): + for item in node: + yield from _iter_commands(item) + + @click.command(help="Install tracing hooks for an agentic platform.") +@click.option( + "--list", + "list_platforms", + is_flag=True, + help="List supported platforms and warn about orphaned hooks from removed platforms.", +) @_platform_options -def add(platform_flag: str | None) -> None: +def add(platform_flag: str | None, list_platforms: bool) -> None: + if list_platforms: + click.echo("Supported platforms:") + for name in PLATFORMS: + click.echo(f" {name}") + for path, command in find_orphaned_hooks(ORPHAN_CONFIG_PATHS): + click.echo( + f"Warning: {path} still references removed hook {command!r}. " + "Remove it from that tool's config to stop the missing-binary error.", + ) + return platform = _resolve_platform(platform_flag) platform.install() click.echo(f"Installed tracing for {platform.display_name}") diff --git a/src/thirdeye/commands/eval.py b/src/thirdeye/commands/eval.py index 8ba13dc..b1f27bd 100644 --- a/src/thirdeye/commands/eval.py +++ b/src/thirdeye/commands/eval.py @@ -49,8 +49,7 @@ def run_cmd(session_prefix, using, agent, background, as_json, save): if agent not in list_agent_names(config.root): raise click.ClickException( - f"unknown agent {agent!r} — choose one of " - f"{', '.join(list_agent_names(config.root))}" + f"unknown agent {agent!r} — choose one of {', '.join(list_agent_names(config.root))}" ) if background: @@ -267,7 +266,7 @@ def status_cmd(session_prefix, as_json): click.echo("No background eval jobs.") return click.echo( - f"{'JOB':<28} {'SESSION':<14} {'USING':<18} {'AGENT':<8} " f"{'STATUS':<10} {'STARTED':<26}" + f"{'JOB':<28} {'SESSION':<14} {'USING':<18} {'AGENT':<8} {'STATUS':<10} {'STARTED':<26}" ) for j in jobs: click.echo( diff --git a/src/thirdeye/commands/ingest.py b/src/thirdeye/commands/ingest.py index 56a6d65..dc16d8b 100644 --- a/src/thirdeye/commands/ingest.py +++ b/src/thirdeye/commands/ingest.py @@ -11,7 +11,7 @@ @click.command(help="Read newline-delimited JSON events from stdin and append them to a session.") -@click.option("--platform", required=True, help="Platform name (e.g. claude, cursor).") +@click.option("--platform", required=True, help="Platform name (e.g. claude, codex).") @click.option("--session-id", default=None, help="Session ID. Generated if omitted.") @click.option("--cwd", default=None, help="Working directory for the session.") def ingest(platform: str, session_id: str | None, cwd: str | None) -> None: diff --git a/src/thirdeye/commands/ui.py b/src/thirdeye/commands/ui.py index d5222e0..cb71033 100644 --- a/src/thirdeye/commands/ui.py +++ b/src/thirdeye/commands/ui.py @@ -30,7 +30,7 @@ def ui(host: str, port: int, no_browser: bool) -> None: msg = str(e) if "starlette" in msg or "uvicorn" in msg or "jinja2" in msg: raise click.ClickException( - "The UI requires the 'ui' extra. Install with:\n" " pip install 'thrdi[ui]'" + "The UI requires the 'ui' extra. Install with:\n pip install 'thrdi[ui]'" ) from e raise diff --git a/src/thirdeye/commands/usage.py b/src/thirdeye/commands/usage.py index 5ebf15b..269b55b 100644 --- a/src/thirdeye/commands/usage.py +++ b/src/thirdeye/commands/usage.py @@ -1,15 +1,27 @@ from __future__ import annotations import json +import sqlite3 import time -from datetime import datetime +from datetime import UTC, datetime +from pathlib import Path import click +from thirdeye.commands.add import find_orphaned_hooks from thirdeye.config import Config -from thirdeye.paths import session_dir, usage_log_path +from thirdeye.paths import ( + session_dir, + sessions_root, + usage_db_path, + usage_jsonl_path, + usage_log_path, + usage_state_path, +) from thirdeye.timeparse import parse_when from thirdeye.usage.index import UsageIndex +from thirdeye.usage.read import iter_calls +from thirdeye.usage.types import UsageRow def _parse_window(value: str | None, flag: str) -> datetime | None: @@ -30,13 +42,71 @@ def _resolve_session(config: Config, prefix: str) -> tuple[str, str]: raise click.ClickException(str(e)) from e +def _row_ts(row: UsageRow) -> datetime | None: + s = row.ts + if s.endswith("Z"): + s = s[:-1] + "+00:00" + try: + dt = datetime.fromisoformat(s) + except (TypeError, ValueError): + return None + # Normalize to UTC-aware so comparisons against the aware --since/--until + # bounds never raise on a naive timestamp lacking an offset. + if dt.tzinfo is None: + dt = dt.replace(tzinfo=UTC) + return dt + + +def _keep_row( + row: UsageRow, + *, + platform_filter: str | None, + model_filter: str | None, + since_dt: datetime | None, + until_dt: datetime | None, +) -> bool: + if platform_filter and row.platform != platform_filter: + return False + if model_filter and model_filter not in row.response_model: + return False + if since_dt or until_dt: + dt = _row_ts(row) + if dt is None: + return False + if since_dt and dt < since_dt: + return False + if until_dt and dt > until_dt: + return False + return True + + +def _iter_session_dirs(root: Path): + """Yield (platform, session_id, session_dir) for every captured session.""" + troot = sessions_root(root) + if not troot.exists(): + return + for platform_dir_ in sorted(troot.iterdir()): + if not platform_dir_.is_dir(): + continue + for sd in sorted(platform_dir_.iterdir()): + if not sd.is_dir(): + continue + yield platform_dir_.name, sd.name, sd + + +def _fmt_cache(value: int | None) -> str: + """Render an absent cache attribute as '-', distinct from a reported 0.""" + return "-" if value is None else f"{value:,}" + + class _UsageGroup(click.Group): """Route non-subcommand args to the default `show` subcommand. Lets `thirdeye usage [SESSION_PREFIX] [OPTIONS]` work alongside - `thirdeye usage reindex` and `thirdeye usage errors` without the - positional-vs-subcommand parsing collision that arises from putting - a positional argument directly on an `invoke_without_command` group. + `thirdeye usage reindex`, `thirdeye usage reset`, and `thirdeye usage + errors` without the positional-vs-subcommand parsing collision that arises + from putting a positional argument directly on an `invoke_without_command` + group. """ def parse_args(self, ctx, args): @@ -50,7 +120,7 @@ def parse_args(self, ctx, args): @click.group( cls=_UsageGroup, name="usage", - help="Per-event model and token usage.", + help="Per-call model and token usage.", ) def usage(): pass @@ -70,7 +140,7 @@ def usage(): "--model", "model_filter", default=None, - help="Filter rows where `model` contains this substring.", + help="Filter rows where the response model contains this substring.", ) @click.option("--since", default=None, help="Time window lower bound.") @click.option("--until", default=None, help="Time window upper bound.") @@ -78,7 +148,7 @@ def usage(): "--top", type=int, default=None, - help="Rollup mode: keep top N sessions by total_tokens.", + help="Rollup mode: keep top N sessions by total tokens.", ) @click.option( "--sort", @@ -136,6 +206,18 @@ def reindex_cmd(session_prefix): click.echo(f"Indexed {n} rows from {sessions} sessions in {elapsed_ms} ms") +@usage.command(name="reset") +@click.option("--yes", is_flag=True, help="Required to actually delete anything.") +def reset_cmd(yes): + """Delete all captured usage data (sidecars + usage.db). + + Leaves events.alog, events.idx, tags.jsonl, meta.yaml, and every upstream + transcript untouched. + """ + config = Config.load() + _run_reset(config.root, yes=yes) + + @usage.command(name="errors") @click.option("-n", "n", type=int, default=20, help="Last N entries.") @click.option("--json", "as_json", is_flag=True) @@ -191,9 +273,9 @@ def errors_cmd(n, as_json, platform_filter, phase, since, until): return for e in entries: click.echo( - f"{e.get('ts','')} {e.get('level','?'):<5} " - f"{e.get('platform','?'):<7} {e.get('phase','?'):<20} " - f"{e.get('session_id','')[:12]:<12} {e.get('message','')}" + f"{e.get('ts', '')} {e.get('level', '?'):<5} " + f"{e.get('platform', '?'):<7} {e.get('phase', '?'):<20} " + f"{e.get('session_id', '')[:12]:<12} {e.get('message', '')}" ) @@ -209,17 +291,13 @@ def _run_show( sort, ): config = Config.load() - idx = UsageIndex(config.root) - conn = idx.connect() - idx.refresh(conn) - since_dt = _parse_window(since, "--since") until_dt = _parse_window(until, "--until") if session_prefix: platform, sid = _resolve_session(config, session_prefix) _render_session( - conn, + session_dir(config.root, platform, sid), sid, platform_filter, model_filter, @@ -230,7 +308,7 @@ def _run_show( ) else: _render_rollup( - conn, + config.root, platform_filter, model_filter, since_dt, @@ -242,7 +320,7 @@ def _run_show( def _render_session( - conn, + session_dir_, sid, platform_filter, model_filter, @@ -251,66 +329,54 @@ def _render_session( sort, as_json, ): - sql = [ - "SELECT seq, ts, platform, model, input_tokens, output_tokens, " - "total_tokens FROM usage WHERE session_id = ?" + rows = [ + r + for r in iter_calls(session_dir_) + if _keep_row( + r, + platform_filter=platform_filter, + model_filter=model_filter, + since_dt=since_dt, + until_dt=until_dt, + ) ] - params: list = [sid] - if platform_filter: - sql.append("AND platform = ?") - params.append(platform_filter) - if model_filter: - sql.append("AND model LIKE ?") - params.append(f"%{model_filter}%") - if since_dt: - sql.append("AND ts >= ?") - params.append(since_dt.isoformat()) - if until_dt: - sql.append("AND ts <= ?") - params.append(until_dt.isoformat()) - sort_col = { - "total": "total_tokens DESC", - "input": "input_tokens DESC", - "output": "output_tokens DESC", - "ts": "ts ASC", + sort_key = { + "total": lambda r: -r.total_tokens, + "input": lambda r: -r.input_tokens, + "output": lambda r: -r.output_tokens, + "ts": lambda r: r.ts, }[sort] - sql.append(f"ORDER BY {sort_col}") - rows = conn.execute(" ".join(sql), params).fetchall() + rows.sort(key=sort_key) if as_json: for r in rows: - click.echo( - json.dumps( - { - "session_id": sid, - "seq": r[0], - "ts": r[1], - "platform": r[2], - "model": r[3], - "input_tokens": r[4], - "output_tokens": r[5], - "total_tokens": r[6], - }, - separators=(",", ":"), - ) - ) + click.echo(json.dumps(r.to_dict(), separators=(",", ":"))) return if not rows: click.echo(f"No usage data for session {sid}.") return - click.echo(f"{'SEQ':<5} {'TS':<26} {'MODEL':<25} {'INPUT':>10} {'OUTPUT':>8} {'TOTAL':>10}") + click.echo( + f"{'SEQ':<5} {'TS':<26} {'MODEL':<25} {'INPUT':>10} {'OUTPUT':>8} " + f"{'CACHE_R':>12} {'CACHE_C':>12} {'TOTAL':>10}" + ) tot_in = tot_out = tot = 0 for r in rows: - click.echo(f"{r[0]:<5} {r[1]:<26} {r[3][:25]:<25} " f"{r[4]:>10,} {r[5]:>8,} {r[6]:>10,}") - tot_in += r[4] - tot_out += r[5] - tot += r[6] - click.echo(f"\n{len(rows)} turns · {tot_in:,} input · {tot_out:,} output · {tot:,} total") + click.echo( + f"{r.seq:<5} {r.ts:<26} {r.response_model[:25]:<25} " + f"{r.input_tokens:>10,} {r.output_tokens:>8,} " + f"{_fmt_cache(r.cache_read_input_tokens):>12} " + f"{_fmt_cache(r.cache_creation_input_tokens):>12} " + f"{r.total_tokens:>10,}" + ) + tot_in += r.input_tokens + tot_out += r.output_tokens + tot += r.total_tokens + click.echo(f"\n{len(rows)} calls · {tot_in:,} input · {tot_out:,} output · {tot:,} total") def _render_rollup( - conn, + root, platform_filter, model_filter, since_dt, @@ -319,51 +385,47 @@ def _render_rollup( sort, as_json, ): - sql = [ - "SELECT session_id, platform, " - "COUNT(*) AS turns, " - "SUM(input_tokens) AS in_tok, " - "SUM(output_tokens) AS out_tok, " - "SUM(total_tokens) AS total_tok " - "FROM usage WHERE 1=1" + # (session_id, platform) -> [turns, input_tokens, output_tokens] + agg: dict[tuple[str, str], list[int]] = {} + for platform, sid, sd in _iter_session_dirs(root): + for r in iter_calls(sd): + if not _keep_row( + r, + platform_filter=platform_filter, + model_filter=model_filter, + since_dt=since_dt, + until_dt=until_dt, + ): + continue + bucket = agg.setdefault((sid, platform), [0, 0, 0]) + bucket[0] += 1 + bucket[1] += r.input_tokens + bucket[2] += r.output_tokens + + rows = [ + (sid, platform, turns, in_tok, out_tok, in_tok + out_tok) + for (sid, platform), (turns, in_tok, out_tok) in agg.items() ] - params: list = [] - if platform_filter: - sql.append("AND platform = ?") - params.append(platform_filter) - if model_filter: - sql.append("AND model LIKE ?") - params.append(f"%{model_filter}%") - if since_dt: - sql.append("AND ts >= ?") - params.append(since_dt.isoformat()) - if until_dt: - sql.append("AND ts <= ?") - params.append(until_dt.isoformat()) - sql.append("GROUP BY session_id, platform") - sort_col = { - "total": "total_tok DESC", - "input": "in_tok DESC", - "output": "out_tok DESC", - "ts": "session_id ASC", + sort_key = { + "total": lambda r: -r[5], + "input": lambda r: -r[3], + "output": lambda r: -r[4], + "ts": lambda r: r[0], }[sort] - sql.append(f"ORDER BY {sort_col}") + rows.sort(key=sort_key) if top is not None: - sql.append("LIMIT ?") - params.append(top) - rows = conn.execute(" ".join(sql), params).fetchall() + rows = rows[:top] if as_json: - for r in rows: + for sid, platform, turns, in_tok, out_tok, _total in rows: click.echo( json.dumps( { - "session_id": r[0], - "platform": r[1], - "turns": r[2], - "input_tokens": r[3], - "output_tokens": r[4], - "total_tokens": r[5], + "session_id": sid, + "platform": platform, + "calls": turns, + "gen_ai.usage.input_tokens": in_tok, + "gen_ai.usage.output_tokens": out_tok, }, separators=(",", ":"), ) @@ -374,15 +436,81 @@ def _render_rollup( click.echo("No usage data.") return click.echo( - f"{'SESSION':<14} {'PLATFORM':<9} {'TURNS':>5} {'INPUT':>12} {'OUTPUT':>10} {'TOTAL':>12}" + f"{'SESSION':<14} {'PLATFORM':<9} {'CALLS':>5} {'INPUT':>12} {'OUTPUT':>10} {'TOTAL':>12}" ) tot_in = tot_out = tot = 0 - for r in rows: - click.echo(f"{r[0][:14]:<14} {r[1]:<9} {r[2]:>5} " f"{r[3]:>12,} {r[4]:>10,} {r[5]:>12,}") - tot_in += r[3] - tot_out += r[4] - tot += r[5] + for sid, platform, turns, in_tok, out_tok, total in rows: + click.echo( + f"{sid[:14]:<14} {platform:<9} {turns:>5} {in_tok:>12,} {out_tok:>10,} {total:>12,}" + ) + tot_in += in_tok + tot_out += out_tok + tot += total click.echo(f"\n{len(rows)} sessions · {tot_in:,} input · {tot_out:,} output · {tot:,} total") +def _run_reset(root: Path, *, yes: bool) -> None: + """Destroy every usage sidecar and usage.db under `root`. + + Counts are computed and reported before anything is removed. Without + `--yes` the command refuses and exits non-zero, deleting nothing. + """ + sidecar_files: list[Path] = [] + sessions_affected: set[Path] = set() + for _platform, _sid, sd in _iter_session_dirs(root): + jsonl = usage_jsonl_path(sd) + state = usage_state_path(sd) + if jsonl.exists(): + sidecar_files.append(jsonl) + sessions_affected.add(sd) + if state.exists(): + sidecar_files.append(state) + sessions_affected.add(sd) + + db = usage_db_path(root) + db_rows = 0 + if db.exists(): + conn = None + try: + conn = sqlite3.connect(db) + db_rows = conn.execute("SELECT COUNT(*) FROM usage").fetchone()[0] + except sqlite3.Error: + db_rows = 0 + finally: + if conn is not None: + conn.close() + + click.echo("usage reset will destroy:") + click.echo(f" sidecar files: {len(sidecar_files)}") + click.echo(f" sessions affected: {len(sessions_affected)}") + click.echo(f" usage.db rows: {db_rows}") + + if not yes: + raise click.ClickException("refusing to delete without --yes; nothing was removed.") + + for f in sidecar_files: + try: + f.unlink() + except FileNotFoundError: + pass + # Remove usage.db and its WAL/SHM companions if present. + for p in (db, db.with_name(db.name + "-wal"), db.with_name(db.name + "-shm")): + try: + p.unlink() + except FileNotFoundError: + pass + + click.echo( + f"Deleted {len(sidecar_files)} sidecar file(s) across " + f"{len(sessions_affected)} session(s) and {db_rows} usage.db row(s)." + ) + + orphans = find_orphaned_hooks() + for path, command in orphans: + click.echo( + f"Warning: {path} still references removed hook {command!r}. " + "Remove it from that tool's config (detection only — not edited here)." + ) + + __all__ = ["usage"] diff --git a/src/thirdeye/eval/definition.py b/src/thirdeye/eval/definition.py index e614814..fb9b510 100644 --- a/src/thirdeye/eval/definition.py +++ b/src/thirdeye/eval/definition.py @@ -58,7 +58,7 @@ def load_definition(thirdeye_home: Path, name: str) -> EvalDefinition: shipped = _shipped_path(name) if shipped is None: raise FileNotFoundError( - f"no eval definition named '{name}' " f"(checked {user_path} and shipped defaults)" + f"no eval definition named '{name}' (checked {user_path} and shipped defaults)" ) user_path.parent.mkdir(parents=True, exist_ok=True) user_path.write_text(shipped.read_text(encoding="utf-8"), encoding="utf-8") diff --git a/src/thirdeye/eval/prompt.py b/src/thirdeye/eval/prompt.py index 5cfc8ee..51f2551 100644 --- a/src/thirdeye/eval/prompt.py +++ b/src/thirdeye/eval/prompt.py @@ -74,7 +74,7 @@ def build_prompt( blocks.append("=== Tool inventory ===") blocks.append("You have read-only access to:") blocks.append( - "- `thirdeye` CLI: list, events, show, tail, event, search, " "tag, tags, stats, usage" + "- `thirdeye` CLI: list, events, show, tail, event, search, tag, tags, stats, usage" ) blocks.append(f"- `sqlite3 {db}`") blocks.append("- `jq`, `Read`") diff --git a/src/thirdeye/platforms/claude/constants.py b/src/thirdeye/platforms/claude/constants.py index 299c9fe..c63b665 100644 --- a/src/thirdeye/platforms/claude/constants.py +++ b/src/thirdeye/platforms/claude/constants.py @@ -17,4 +17,10 @@ "Notification": "thirdeye-claude-notification", "PermissionRequest": "thirdeye-claude-permission-request", "SessionEnd": "thirdeye-claude-session-end", + "PostToolUseFailure": "thirdeye-claude-post-tool-use-failure", + "SubagentStart": "thirdeye-claude-subagent-start", + "UserPromptExpansion": "thirdeye-claude-user-prompt-expansion", + "PreCompact": "thirdeye-claude-pre-compact", + "PostCompact": "thirdeye-claude-post-compact", + "PermissionDenied": "thirdeye-claude-permission-denied", } diff --git a/src/thirdeye/platforms/claude/hooks.py b/src/thirdeye/platforms/claude/hooks.py index 9682be4..04e80bc 100644 --- a/src/thirdeye/platforms/claude/hooks.py +++ b/src/thirdeye/platforms/claude/hooks.py @@ -136,6 +136,10 @@ def stop() -> None: def subagent_stop() -> None: + # SubagentStop keeps its historical "subagent_message" type on purpose: + # renaming it would orphan the 966 sessions already recorded under it. + # SubagentStart (below) uses a distinct "subagent_start" type; the + # start/stop asymmetry is intentional, not an oversight. _emit("subagent_message", _read_stdin()) @@ -155,3 +159,30 @@ def session_end() -> None: payload = _read_stdin() if _emit("session_end", payload) is not None: Store(Config.load()).close_session(payload["session_id"], platform=_PLATFORM) + + +def post_tool_use_failure() -> None: + # Emits "tool_result", not "error", on purpose: web/routes/sessions.py pairs + # each "tool_call" with a "tool_result", so a distinct type would leave failed + # tool calls rendering as dangling. The failure is evident from the payload. + _emit("tool_result", _read_stdin()) + + +def subagent_start() -> None: + _emit("subagent_start", _read_stdin()) + + +def user_prompt_expansion() -> None: + _emit("user_prompt_expansion", _read_stdin()) + + +def pre_compact() -> None: + _emit("compact_start", _read_stdin()) + + +def post_compact() -> None: + _emit("compact_end", _read_stdin()) + + +def permission_denied() -> None: + _emit("permission_denied", _read_stdin()) diff --git a/src/thirdeye/platforms/claude/usage.py b/src/thirdeye/platforms/claude/usage.py index 6abfa94..e64c3b7 100644 --- a/src/thirdeye/platforms/claude/usage.py +++ b/src/thirdeye/platforms/claude/usage.py @@ -22,6 +22,10 @@ def capture_usage_claude( Returns the number of rows appended. Wrapped in @safe_capture so any error is logged to usage-errors.jsonl and the function returns None instead of raising. + + One row is appended per assistant frame - Claude writes one frame per content + block, all carrying the identical ``message.usage``. Collapsing those + duplicates is ``usage/read.py``'s job, never this writer's. """ if not transcript_path: return 0 @@ -40,20 +44,6 @@ def capture_usage_claude( sd = session_dir(thirdeye_home, "claude", session_id) store = UsageStore(sd) state = store.read_state() - if "transcript_offset" not in state: - log_capture_error( - thirdeye_home=thirdeye_home, - phase="discover_transcript", - message=( - "first capture for session; transcript frame shape is unverified " - "against a real Claude Code transcript — remove this warning once " - "format is confirmed end-to-end" - ), - platform="claude", - session_id=session_id, - source_path=str(transcript_path), - level="warn", - ) offset = int(state.get("transcript_offset", 0)) new_rows: list[UsageRow] = [] @@ -82,42 +72,66 @@ def capture_usage_claude( def _extract_row(frame: dict, session_id: str, triggering_seq: int) -> UsageRow | None: - """Return a UsageRow if `frame` looks like an assistant turn with usage. + """Return a UsageRow for an assistant frame carrying real API usage, else None. - Handles both nested (`frame["message"]["usage"]`) and flat - (`frame["usage"]`) shapes. + Emits a row per source frame; deduplication happens on read. Cache-inclusive + input tokens follow the OTel GenAI convention: ``gen_ai.usage.input_tokens`` + includes cache reads and cache creation, which Anthropic reports separately. """ if not isinstance(frame, dict): return None + if frame.get("type") != "assistant": + return None - message = frame.get("message") if isinstance(frame.get("message"), dict) else None - if message and "usage" in message: - model = message.get("model") or frame.get("model") - usage = message.get("usage") or {} - ts = message.get("timestamp") or frame.get("timestamp") or "" - elif "usage" in frame: - model = frame.get("model") - usage = frame.get("usage") or {} - ts = frame.get("timestamp") or "" - else: + message = frame.get("message") + if not isinstance(message, dict): + return None + usage = message.get("usage") + if not isinstance(usage, dict) or not usage: return None - if not isinstance(usage, dict): + # "" is Claude's zero-token placeholder for injected messages, + # not a real API call. + model = message.get("model") + if not model or model == "": return None - input_tokens = usage.get("input_tokens") - output_tokens = usage.get("output_tokens") - if input_tokens is None or output_tokens is None or not model: + + # Claude carries the timestamp at the top level, never inside `message`. + ts = frame.get("timestamp") or "" + + # message.id is the most reliable key (no nulls over a real transcript); + # fall back to requestId, then the frame uuid. + call_id = message.get("id") or frame.get("requestId") or frame.get("uuid") + if not call_id: + return None + + # Anthropic reports input_tokens EXCLUDING cache; add both cache classes back + # in for a comparable, cache-inclusive total. + raw_input = int(usage.get("input_tokens") or 0) + cache_read = usage.get("cache_read_input_tokens") + cache_crea = usage.get("cache_creation_input_tokens") + output = usage.get("output_tokens") + + # Absent both token fields means this frame carries no usage worth recording. + if usage.get("input_tokens") is None and output is None: return None - total = int(input_tokens) + int(output_tokens) + input_tokens = raw_input + int(cache_read or 0) + int(cache_crea or 0) return UsageRow( session_id=session_id, seq=triggering_seq, - ts=str(ts) if ts else "", + call_id=str(call_id), + ts=str(ts), platform="claude", - model=str(model), - input_tokens=int(input_tokens), - output_tokens=int(output_tokens), - total_tokens=total, + provider_name="anthropic", + response_model=str(model), + input_tokens=input_tokens, + output_tokens=int(output or 0), + operation_name="chat", + # Absent vs zero: pass cache fields through only when reported. + cache_read_input_tokens=int(cache_read) if cache_read is not None else None, + cache_creation_input_tokens=(int(cache_crea) if cache_crea is not None else None), + # Anthropic does not break out thinking/reasoning tokens. + reasoning_output_tokens=None, ) diff --git a/src/thirdeye/platforms/codex/usage.py b/src/thirdeye/platforms/codex/usage.py index 3e919cb..d1c606e 100644 --- a/src/thirdeye/platforms/codex/usage.py +++ b/src/thirdeye/platforms/codex/usage.py @@ -47,9 +47,15 @@ def capture_usage_codex( rp = Path(rollout_path) last_model: str | None = state.get("last_model") new_rows: list[UsageRow] = [] + # Track each frame's absolute byte offset for a stable, unique call_id. Codex + # frames carry no natural id, and the offset is monotonic across incremental + # captures, so no two frames ever collide under the (session_id, call_id) key. + pos = offset with rp.open("rb") as f: f.seek(offset) for raw in f: + frame_offset = pos + pos += len(raw) line = raw.decode("utf-8", errors="replace").strip() if not line: continue @@ -60,7 +66,7 @@ def capture_usage_codex( inferred = _extract_model(frame) if inferred: last_model = inferred - row = _extract_usage_row(frame, session_id, triggering_seq, last_model) + row = _extract_usage_row(frame, session_id, triggering_seq, last_model, frame_offset) if row is not None: new_rows.append(row) new_offset = f.tell() @@ -100,6 +106,7 @@ def _extract_usage_row( session_id: str, triggering_seq: int, last_model: str | None, + frame_offset: int, ) -> UsageRow | None: if not isinstance(frame, dict): return None @@ -109,6 +116,8 @@ def _extract_usage_row( input_tokens = payload.get("input_tokens") output_tokens = payload.get("output_tokens") total_tokens = payload.get("total_tokens") + # A Codex usage frame reports all three counts; the presence of total_tokens + # identifies it. Totals are derived, never stored, so it is not persisted. if input_tokens is None or output_tokens is None or total_tokens is None: return None model = _extract_model(frame) or last_model or "unknown" @@ -116,10 +125,16 @@ def _extract_usage_row( return UsageRow( session_id=session_id, seq=triggering_seq, + call_id=f"{session_id}:{frame_offset}", ts=str(ts), platform="codex", - model=str(model), + provider_name="openai", + response_model=str(model), input_tokens=int(input_tokens), output_tokens=int(output_tokens), - total_tokens=int(total_tokens), + operation_name="chat", + # Codex rollout frames do not break out cache or reasoning tokens. + cache_read_input_tokens=None, + cache_creation_input_tokens=None, + reasoning_output_tokens=None, ) diff --git a/src/thirdeye/platforms/cursor/__init__.py b/src/thirdeye/platforms/cursor/__init__.py deleted file mode 100644 index e69de29..0000000 diff --git a/src/thirdeye/platforms/cursor/constants.py b/src/thirdeye/platforms/cursor/constants.py deleted file mode 100644 index aeb3e5b..0000000 --- a/src/thirdeye/platforms/cursor/constants.py +++ /dev/null @@ -1,65 +0,0 @@ -from __future__ import annotations - -from pathlib import Path - -PLATFORM_NAME = "cursor" -DISPLAY_NAME = "Cursor" - -HOOKS_FILE = Path.home() / ".cursor" / "hooks.json" - -HOOK_BIN_NAME = "thirdeye-cursor-hook" - -HOOK_TIMEOUT_S = 30 - -TRACED_EVENTS: tuple[str, ...] = ( - "sessionStart", - "sessionEnd", - "beforeSubmitPrompt", - "afterAgentResponse", - "beforeShellExecution", - "afterShellExecution", - "beforeMCPExecution", - "afterMCPExecution", - "beforeReadFile", - "afterFileEdit", - "beforeTabFileRead", - "afterTabFileEdit", - "postToolUse", - "stop", -) - -DEDICATED_TOOL_NAMES: frozenset[str] = frozenset( - { - "shell", - "terminal", - "bash", - "run_command", - "run_shell", - "read_file", - "read", - "view_file", - "view", - "edit_file", - "edit", - "write_file", - "write", - "create_file", - "delete_file", - "tab_file_read", - "tab_file_edit", - "mcp", - "mcp_execution", - } -) - -STRIP_KEYS: frozenset[str] = frozenset( - { - "conversation_id", - "conversationId", - "workspace_roots", - "transcript_path", - "user_email", - "hook_event_name", - "hookEventName", - } -) diff --git a/src/thirdeye/platforms/cursor/hook.py b/src/thirdeye/platforms/cursor/hook.py deleted file mode 100644 index 6b09300..0000000 --- a/src/thirdeye/platforms/cursor/hook.py +++ /dev/null @@ -1,239 +0,0 @@ -from __future__ import annotations - -import json -import os -import sys -from collections.abc import Callable - -from thirdeye.config import Config -from thirdeye.meta import read_meta, write_meta -from thirdeye.paths import meta_path, session_dir -from thirdeye.platforms.cursor.constants import DEDICATED_TOOL_NAMES, STRIP_KEYS -from thirdeye.platforms.cursor.usage import capture_usage_cursor -from thirdeye.store import Store -from thirdeye.tags import TagStore, extract_hashtags - -_PLATFORM = "cursor" - - -def _read_stdin() -> dict: - try: - raw = sys.stdin.read() - except OSError: - return {} - if not raw: - return {} - try: - return json.loads(raw) - except json.JSONDecodeError: - return {} - - -def _get_str(p: dict, *keys: str, default: str = "") -> str: - for k in keys: - v = p.get(k) - if v is not None and v != "": - return str(v) - return default - - -def _event_name(p: dict) -> str: - return _get_str(p, "hook_event_name", "hookEventName") - - -def _session_id(p: dict) -> str: - return _get_str(p, "conversation_id", "conversationId") - - -def _cwd(p: dict) -> str: - cwd = p.get("cwd") - if cwd: - return str(cwd) - roots = p.get("workspace_roots") - if isinstance(roots, list) and roots: - return str(roots[0]) - return os.getcwd() - - -def _strip(p: dict) -> dict: - return {k: v for k, v in p.items() if k not in STRIP_KEYS} - - -def _emit(payload: dict, t: str, extra: dict | None = None) -> int | None: - sid = _session_id(payload) - if not sid: - return None - data = _strip(payload) - if extra: - data = {**data, **extra} - return Store(Config.load()).append_event( - session_id=sid, - platform=_PLATFORM, - cwd=_cwd(payload), - t=t, - data=data, - ) - - -def _print_permissive(event: str) -> None: - out = sys.__stdout__ or sys.stdout - try: - if event.startswith("before"): - out.write('{"permission": "allow"}') - else: - out.write('{"continue": true}') - out.flush() - except Exception: - pass - - -def _h_session_start(payload: dict) -> None: - _emit(payload, "session_start") - - -def _h_session_end(payload: dict) -> None: - sid = _session_id(payload) - if not sid: - return - config = Config.load() - store = Store(config) - seq = store.append_event( - session_id=sid, - platform=_PLATFORM, - cwd=_cwd(payload), - t="session_end", - data=_strip(payload), - ) - store.close_session(sid, platform=_PLATFORM) - capture_usage_cursor( - thirdeye_home=config.root, - session_id=sid, - payload=payload, - triggering_seq=seq if seq is not None else 0, - ) - - -def _h_before_submit(payload: dict) -> None: - sid = _session_id(payload) - if not sid: - return - config = Config.load() - seq = Store(config).append_event( - session_id=sid, - platform=_PLATFORM, - cwd=_cwd(payload), - t="user_message", - data=_strip(payload), - ) - try: - prompt = _get_str(payload, "prompt", "input", "text") - tags = extract_hashtags(prompt) - if not tags: - return - sd = session_dir(config.root, _PLATFORM, sid) - ts = TagStore(sd) - for tag in tags: - ts.add(seq, tag, source="auto") - mp = meta_path(sd) - m = read_meta(mp) - if m is not None: - m.tag_count = ts.tagged_seq_count() - write_meta(mp, m) - except Exception: - pass - - -def _h_after_response(payload: dict) -> None: - _emit(payload, "assistant_message") - - -def _h_before_shell(payload: dict) -> None: - _emit(payload, "tool_call", {"tool_name": "shell"}) - - -def _h_after_shell(payload: dict) -> None: - _emit(payload, "tool_result", {"tool_name": "shell"}) - - -def _h_before_mcp(payload: dict) -> None: - name = _get_str(payload, "tool_name", "toolName", "name", default="mcp") - _emit(payload, "tool_call", {"tool_name": name}) - - -def _h_after_mcp(payload: dict) -> None: - name = _get_str(payload, "tool_name", "toolName", "name", default="mcp") - _emit(payload, "tool_result", {"tool_name": name}) - - -def _h_before_read_file(payload: dict) -> None: - _emit(payload, "tool_call", {"tool_name": "read_file"}) - - -def _h_after_file_edit(payload: dict) -> None: - _emit(payload, "tool_result", {"tool_name": "edit_file"}) - - -def _h_before_tab_read(payload: dict) -> None: - _emit(payload, "tool_call", {"tool_name": "tab_read_file"}) - - -def _h_after_tab_edit(payload: dict) -> None: - _emit(payload, "tool_result", {"tool_name": "tab_edit_file"}) - - -def _h_post_tool_use(payload: dict) -> None: - name = _get_str(payload, "tool_name", "toolName", "name", "tool") - if name and name.lower() in DEDICATED_TOOL_NAMES: - return - if not name: - name = "unknown" - _emit(payload, "tool_result", {"tool_name": name}) - - -def _h_stop(payload: dict) -> None: - sid = _session_id(payload) - if not sid: - return - config = Config.load() - # TODO: Store does not expose a max_seq(session_id) helper. Pass 0 until - # such an API is added rather than introducing a new Store method here. - triggering_seq = 0 - capture_usage_cursor( - thirdeye_home=config.root, - session_id=sid, - payload=payload, - triggering_seq=triggering_seq, - ) - - -_HANDLERS: dict[str, Callable[[dict], None]] = { - "sessionStart": _h_session_start, - "sessionEnd": _h_session_end, - "beforeSubmitPrompt": _h_before_submit, - "afterAgentResponse": _h_after_response, - "beforeShellExecution": _h_before_shell, - "afterShellExecution": _h_after_shell, - "beforeMCPExecution": _h_before_mcp, - "afterMCPExecution": _h_after_mcp, - "beforeReadFile": _h_before_read_file, - "afterFileEdit": _h_after_file_edit, - "beforeTabFileRead": _h_before_tab_read, - "afterTabFileEdit": _h_after_tab_edit, - "postToolUse": _h_post_tool_use, - "stop": _h_stop, -} - - -def main() -> int: - event = "" - try: - payload = _read_stdin() - event = _event_name(payload) - handler = _HANDLERS.get(event) - if handler is not None: - handler(payload) - except Exception: - pass - finally: - _print_permissive(event) - return 0 diff --git a/src/thirdeye/platforms/cursor/install.py b/src/thirdeye/platforms/cursor/install.py deleted file mode 100644 index 9392499..0000000 --- a/src/thirdeye/platforms/cursor/install.py +++ /dev/null @@ -1,74 +0,0 @@ -from __future__ import annotations - -import json -import shutil -from pathlib import Path - -from thirdeye.platforms.base import Platform -from thirdeye.platforms.cursor.constants import ( - DISPLAY_NAME, - HOOK_BIN_NAME, - HOOK_TIMEOUT_S, - HOOKS_FILE, - PLATFORM_NAME, - TRACED_EVENTS, -) - - -def _load(path: Path) -> dict: - if not path.exists(): - return {"version": 1, "hooks": {}} - try: - data = json.loads(path.read_text()) - if not isinstance(data, dict): - return {"version": 1, "hooks": {}} - data.setdefault("version", 1) - data.setdefault("hooks", {}) - return data - except (json.JSONDecodeError, OSError): - return {"version": 1, "hooks": {}} - - -def _save(path: Path, data: dict) -> None: - path.parent.mkdir(parents=True, exist_ok=True) - path.write_text(json.dumps(data, indent=2) + "\n") - - -def _resolve_command() -> str: - return shutil.which(HOOK_BIN_NAME) or HOOK_BIN_NAME - - -def _is_our_entry(entry: dict) -> bool: - cmd = entry.get("command", "") - return Path(cmd).name == HOOK_BIN_NAME - - -class CursorPlatform(Platform): - name = PLATFORM_NAME - display_name = DISPLAY_NAME - - def __init__(self, hooks_file: Path | None = None) -> None: - self._hooks_file = hooks_file or HOOKS_FILE - - def install(self) -> None: - data = _load(self._hooks_file) - hooks = data["hooks"] - cmd = _resolve_command() - for event in TRACED_EVENTS: - entries = hooks.setdefault(event, []) - if not any(_is_our_entry(e) for e in entries): - entries.append({"type": "command", "command": cmd, "timeout": HOOK_TIMEOUT_S}) - _save(self._hooks_file, data) - - def uninstall(self) -> None: - if not self._hooks_file.exists(): - return - data = _load(self._hooks_file) - hooks = data.get("hooks", {}) - for event in list(hooks.keys()): - entries = [e for e in hooks[event] if not _is_our_entry(e)] - if entries: - hooks[event] = entries - else: - del hooks[event] - _save(self._hooks_file, data) diff --git a/src/thirdeye/platforms/cursor/usage.py b/src/thirdeye/platforms/cursor/usage.py deleted file mode 100644 index ab12727..0000000 --- a/src/thirdeye/platforms/cursor/usage.py +++ /dev/null @@ -1,81 +0,0 @@ -from __future__ import annotations - -import time -from pathlib import Path -from typing import Any - -from thirdeye.paths import session_dir -from thirdeye.usage.errlog import safe_capture -from thirdeye.usage.store import UsageStore -from thirdeye.usage.types import UsageRow - - -def _get_str(payload: dict, *keys: str, default: str = "") -> str: - """Return the first non-empty string-coerced value among the keys.""" - for k in keys: - v = payload.get(k) - if v is not None and v != "": - return str(v) - return default - - -def _get_int(payload: dict, *keys: str) -> int | None: - """Return the first parseable int among the keys; None if none parse. - - Accepts numeric strings (including "0") but rejects "" and "--". - """ - for k in keys: - v = payload.get(k) - if v is None or v == "" or v == "--": - continue - try: - return int(v) - except (TypeError, ValueError): - continue - return None - - -@safe_capture(phase="extract_usage", platform="cursor") -def capture_usage_cursor( - *, - thirdeye_home: Path, - session_id: str, - payload: dict[str, Any], - triggering_seq: int, -) -> int: - """Append a UsageRow extracted from a Cursor stop/sessionEnd payload. - - Returns the number of rows appended (0 or 1). Wrapped in @safe_capture so - any exception is logged to usage-errors.jsonl and the function returns None - instead of raising. - """ - model = _get_str(payload, "model", "model_name") - input_tokens = _get_int(payload, "input_tokens", "inputTokens") - output_tokens = _get_int(payload, "output_tokens", "outputTokens") - - if not model or input_tokens is None or output_tokens is None: - return 0 - - sd = session_dir(thirdeye_home, "cursor", session_id) - store = UsageStore(sd) - state = store.read_state() - if state.get("last_seq", -1) >= 0: - return 0 - - row = UsageRow( - session_id=session_id, - seq=triggering_seq, - ts=_now_iso(), - platform="cursor", - model=model, - input_tokens=input_tokens, - output_tokens=output_tokens, - total_tokens=input_tokens + output_tokens, - ) - store.append([row]) - store.write_state(last_seq=triggering_seq) - return 1 - - -def _now_iso() -> str: - return time.strftime("%Y-%m-%dT%H:%M:%SZ", time.gmtime()) diff --git a/src/thirdeye/platforms/gemini/__init__.py b/src/thirdeye/platforms/gemini/__init__.py deleted file mode 100644 index e69de29..0000000 diff --git a/src/thirdeye/platforms/gemini/constants.py b/src/thirdeye/platforms/gemini/constants.py deleted file mode 100644 index 9d7f2e7..0000000 --- a/src/thirdeye/platforms/gemini/constants.py +++ /dev/null @@ -1,30 +0,0 @@ -from __future__ import annotations - -from pathlib import Path - -PLATFORM_NAME = "gemini" -DISPLAY_NAME = "Gemini CLI" -SETTINGS_DIR = Path.home() / ".gemini" -SETTINGS_FILE = SETTINGS_DIR / "settings.json" - -# Friendly hook name written into settings.json — used by uninstall to identify -# entries to remove. -HOOK_NAME = "thirdeye-tracing" - -# Per-hook timeout in milliseconds (Gemini's own default is 60000; we use 30s -# so a wedged hook can't stall the user's CLI for too long). -HOOK_TIMEOUT_MS = 30000 - -# Map of Gemini hook event name -> CLI entry-point script name registered -# in pyproject.toml [project.scripts]. Order is preserved when writing -# settings.json. -HOOK_EVENTS: dict[str, str] = { - "SessionStart": "thirdeye-gemini-session-start", - "SessionEnd": "thirdeye-gemini-session-end", - "BeforeAgent": "thirdeye-gemini-before-agent", - "AfterAgent": "thirdeye-gemini-after-agent", - "BeforeModel": "thirdeye-gemini-before-model", - "AfterModel": "thirdeye-gemini-after-model", - "BeforeTool": "thirdeye-gemini-before-tool", - "AfterTool": "thirdeye-gemini-after-tool", -} diff --git a/src/thirdeye/platforms/gemini/hooks.py b/src/thirdeye/platforms/gemini/hooks.py deleted file mode 100644 index 8564d84..0000000 --- a/src/thirdeye/platforms/gemini/hooks.py +++ /dev/null @@ -1,209 +0,0 @@ -from __future__ import annotations - -import json -import os -import sys - -from thirdeye.config import Config -from thirdeye.env_capture import capture_env, env_to_tag -from thirdeye.meta import read_meta, write_meta -from thirdeye.paths import meta_path, session_dir -from thirdeye.store import Store -from thirdeye.tags import TagStore, extract_hashtags - -_PLATFORM = "gemini" - -# Routing keys we strip from event data because they're already used as -# routing fields when calling Store.append_event, OR because they're -# variants we don't want duplicated in storage. -_STRIP_KEYS = frozenset( - { - "session_id", - "sessionId", - "cwd", - "workingDir", - "working_dir", - } -) - - -def _read_stdin() -> dict: - try: - raw = sys.stdin.read() - except OSError: - return {} - if not raw: - return {} - try: - return json.loads(raw) - except json.JSONDecodeError: - return {} - - -def _flex_get(d: dict, *keys, default=None): - """Try multiple key names, return first non-None/non-empty value.""" - for key in keys: - v = d.get(key) - if v is not None and v != "": - return v - return default - - -def _strip_payload(payload: dict) -> dict: - return {k: v for k, v in payload.items() if k not in _STRIP_KEYS} - - -def _print_response() -> None: - """Gemini hooks must print {} to stdout when they finish.""" - print(json.dumps({})) - - -def _emit(t: str, payload: dict) -> int | None: - sid = _flex_get(payload, "session_id", "sessionId") - if not sid: - return None - cwd = _flex_get(payload, "cwd", "workingDir", "working_dir") or os.getcwd() - return Store(Config.load()).append_event( - session_id=sid, - platform=_PLATFORM, - cwd=cwd, - t=t, - data=_strip_payload(payload), - ) - - -def session_start() -> None: - try: - payload = _read_stdin() - sid = _flex_get(payload, "session_id", "sessionId") - seq = _emit("session_start", payload) - if seq is None: - return - config = Config.load() - captured = capture_env(config.capture_env_patterns) - if not captured: - return - sd = session_dir(config.root, _PLATFORM, sid) - tagstore = TagStore(sd) - for name, value in captured.items(): - tag = env_to_tag(name, value) - if tag is None: - continue - tagstore.add(seq, tag, source="auto") - except Exception: - pass - finally: - _print_response() - - -def session_end() -> None: - try: - payload = _read_stdin() - if _emit("session_end", payload) is not None: - sid = _flex_get(payload, "session_id", "sessionId") - Store(Config.load()).close_session(sid, platform=_PLATFORM) - except Exception: - pass - finally: - _print_response() - - -def before_agent() -> None: - try: - payload = _read_stdin() - sid = _flex_get(payload, "session_id", "sessionId") - if sid: - cwd = _flex_get(payload, "cwd", "workingDir", "working_dir") or os.getcwd() - config = Config.load() - seq = Store(config).append_event( - session_id=sid, - platform=_PLATFORM, - cwd=cwd, - t="user_message", - data=_strip_payload(payload), - ) - try: - prompt = _flex_get(payload, "prompt", "input", "userInput", "message", "user_input") - if isinstance(prompt, str) and prompt: - tags = extract_hashtags(prompt) - if tags: - sd = session_dir(config.root, _PLATFORM, sid) - tag_store = TagStore(sd) - for tag in sorted(tags): - tag_store.add(seq, tag, source="auto") - mp = meta_path(sd) - m = read_meta(mp) - if m is not None: - m.tag_count = tag_store.tagged_seq_count() - write_meta(mp, m) - except Exception: - pass - except Exception: - pass - finally: - _print_response() - - -def after_agent() -> None: - try: - _emit("assistant_message", _read_stdin()) - except Exception: - pass - finally: - _print_response() - - -def before_model() -> None: - try: - _emit("model_request", _read_stdin()) - except Exception: - pass - finally: - _print_response() - - -def after_model() -> None: - from thirdeye.platforms.gemini.usage import capture_usage_gemini - - try: - payload = _read_stdin() - sid = _flex_get(payload, "session_id", "sessionId") - if not sid: - return - cwd = _flex_get(payload, "cwd", "workingDir", "working_dir") or os.getcwd() - config = Config.load() - seq = Store(config).append_event( - session_id=sid, - platform=_PLATFORM, - cwd=cwd, - t="model_response", - data=_strip_payload(payload), - ) - capture_usage_gemini( - thirdeye_home=config.root, - session_id=sid, - payload=payload, - triggering_seq=seq, - ) - except Exception: - pass - finally: - _print_response() - - -def before_tool() -> None: - try: - _emit("tool_call", _read_stdin()) - except Exception: - pass - finally: - _print_response() - - -def after_tool() -> None: - try: - _emit("tool_result", _read_stdin()) - except Exception: - pass - finally: - _print_response() diff --git a/src/thirdeye/platforms/gemini/install.py b/src/thirdeye/platforms/gemini/install.py deleted file mode 100644 index 66e9bae..0000000 --- a/src/thirdeye/platforms/gemini/install.py +++ /dev/null @@ -1,101 +0,0 @@ -from __future__ import annotations - -import json -import shutil -from pathlib import Path - -from thirdeye.platforms.base import Platform -from thirdeye.platforms.gemini.constants import ( - DISPLAY_NAME, - HOOK_EVENTS, - HOOK_NAME, - HOOK_TIMEOUT_MS, - PLATFORM_NAME, - SETTINGS_FILE, -) - - -def _load(path: Path) -> dict: - if not path.exists(): - return {} - try: - return json.loads(path.read_text()) - except (json.JSONDecodeError, OSError): - return {} - - -def _save(path: Path, data: dict) -> None: - path.parent.mkdir(parents=True, exist_ok=True) - path.write_text(json.dumps(data, indent=2) + "\n") - - -def _resolve_command(script_name: str) -> str: - return shutil.which(script_name) or script_name - - -class GeminiPlatform(Platform): - name = PLATFORM_NAME - display_name = DISPLAY_NAME - - def __init__(self, settings_file: Path | None = None) -> None: - self._settings_file = settings_file or SETTINGS_FILE - - def install(self) -> None: - settings = _load(self._settings_file) - hooks = settings.setdefault("hooks", {}) - - for event, script in HOOK_EVENTS.items(): - cmd = _resolve_command(script) - blocks: list = hooks.setdefault(event, []) - # Filter out any existing block that contains our hook name - blocks[:] = [ - block - for block in blocks - if not any(h.get("name") == HOOK_NAME for h in block.get("hooks", [])) - ] - # Append our block - blocks.append( - { - "matcher": "", - "hooks": [ - { - "type": "command", - "name": HOOK_NAME, - "command": cmd, - "timeout": HOOK_TIMEOUT_MS, - } - ], - } - ) - - _save(self._settings_file, settings) - - def uninstall(self) -> None: - if not self._settings_file.exists(): - return - - settings = _load(self._settings_file) - hooks = settings.get("hooks") - if not hooks: - if not settings: - self._settings_file.unlink(missing_ok=True) - return - - for event in list(hooks.keys()): - blocks = hooks[event] - blocks[:] = [ - block - for block in blocks - if not any(h.get("name") == HOOK_NAME for h in block.get("hooks", [])) - ] - if not blocks: - del hooks[event] - - if not hooks: - del settings["hooks"] - - if not settings: - self._settings_file.unlink(missing_ok=True) - return - - _save(self._settings_file, settings) diff --git a/src/thirdeye/platforms/gemini/usage.py b/src/thirdeye/platforms/gemini/usage.py deleted file mode 100644 index 3bf69df..0000000 --- a/src/thirdeye/platforms/gemini/usage.py +++ /dev/null @@ -1,56 +0,0 @@ -from __future__ import annotations - -from pathlib import Path - -from thirdeye.paths import session_dir -from thirdeye.usage.errlog import safe_capture -from thirdeye.usage.store import UsageStore -from thirdeye.usage.types import UsageRow - - -@safe_capture(phase="extract_usage", platform="gemini") -def capture_usage_gemini( - *, - thirdeye_home: Path, - session_id: str, - payload: dict, - triggering_seq: int, -) -> int: - """Build one UsageRow from a Gemini after_model payload, append it. - - Returns 1 if a row was appended, 0 if the payload had no usable usage - (e.g. an intermediate reasoning pass with empty usageMetadata). - """ - llm_response = payload.get("llm_response") if isinstance(payload, dict) else None - if not isinstance(llm_response, dict): - return 0 - usage_meta = llm_response.get("usageMetadata") - if not isinstance(usage_meta, dict): - return 0 - total = usage_meta.get("totalTokenCount") - if not total: - return 0 - - llm_request = payload.get("llm_request") if isinstance(payload, dict) else None - model = "" - if isinstance(llm_request, dict): - model = str(llm_request.get("model") or "") - if not model: - model = "unknown" - - row = UsageRow( - session_id=session_id, - seq=triggering_seq, - ts=str(payload.get("timestamp") or ""), - platform="gemini", - model=model, - input_tokens=int(usage_meta.get("promptTokenCount", 0)), - output_tokens=int(usage_meta.get("candidatesTokenCount", 0)), - total_tokens=int(total), - ) - - sd = session_dir(thirdeye_home, "gemini", session_id) - store = UsageStore(sd) - store.append([row]) - store.write_state(last_seq=triggering_seq) - return 1 diff --git a/src/thirdeye/timeparse.py b/src/thirdeye/timeparse.py index e439d82..fc4b5c9 100644 --- a/src/thirdeye/timeparse.py +++ b/src/thirdeye/timeparse.py @@ -6,7 +6,7 @@ _ROLLING_RE = re.compile(r"^(\d+)(d|h|m)$") _ISO_DATE_RE = re.compile(r"^\d{4}-\d{2}-\d{2}$") -_ERR = "could not parse {input!r}: " "expected ISO date, 'NN(d|h|m)', 'today', or 'yesterday'" +_ERR = "could not parse {input!r}: expected ISO date, 'NN(d|h|m)', 'today', or 'yesterday'" def parse_when(s: str | None, *, now: datetime | None = None) -> datetime | None: diff --git a/src/thirdeye/usage/aggregate.py b/src/thirdeye/usage/aggregate.py index eada779..683889e 100644 --- a/src/thirdeye/usage/aggregate.py +++ b/src/thirdeye/usage/aggregate.py @@ -9,7 +9,7 @@ from thirdeye.config import Config from thirdeye.paths import session_dir as _session_dir from thirdeye.store import Store -from thirdeye.usage.store import UsageStore +from thirdeye.usage.read import iter_calls @dataclass(frozen=True) @@ -56,7 +56,7 @@ def aggregate_by_day( store = Store(config) for meta in store.list_sessions(platform=platform, since=since, until=until): sd = _session_dir(config.root, meta.platform, meta.session_id) - for row in UsageStore(sd).iter_rows(): + for row in iter_calls(sd): day = _row_day(row.ts) b = sums[day] b["events"] = int(b["events"]) + 1 diff --git a/src/thirdeye/usage/index.py b/src/thirdeye/usage/index.py index 9b4c1e0..1a143a9 100644 --- a/src/thirdeye/usage/index.py +++ b/src/thirdeye/usage/index.py @@ -6,30 +6,63 @@ from thirdeye.paths import sessions_root, usage_db_path, usage_jsonl_path from thirdeye.usage.errlog import log_capture_error +from thirdeye.usage.read import iter_calls -SCHEMA_VERSION = 1 +# Bumped to 2 for the OTel GenAI schema. usage.db is purely derived from the +# sidecars, so a user_version below SCHEMA_VERSION drops and rebuilds rather +# than migrating. +SCHEMA_VERSION = 2 SCHEMA_SQL = """ CREATE TABLE IF NOT EXISTS usage ( - session_id TEXT NOT NULL, - seq INTEGER NOT NULL, - ts TEXT NOT NULL, - platform TEXT NOT NULL, - model TEXT NOT NULL, - input_tokens INTEGER NOT NULL, - output_tokens INTEGER NOT NULL, - total_tokens INTEGER NOT NULL, - PRIMARY KEY (session_id, seq) + session_id TEXT NOT NULL, + call_id TEXT NOT NULL, + seq INTEGER NOT NULL, + ts TEXT NOT NULL, + platform TEXT NOT NULL, + gen_ai_provider_name TEXT NOT NULL, + gen_ai_operation_name TEXT NOT NULL, + gen_ai_response_model TEXT NOT NULL, + gen_ai_usage_input_tokens INTEGER NOT NULL, + gen_ai_usage_output_tokens INTEGER NOT NULL, + gen_ai_usage_cache_read_input_tokens INTEGER, + gen_ai_usage_cache_creation_input_tokens INTEGER, + gen_ai_usage_reasoning_output_tokens INTEGER, + PRIMARY KEY (session_id, call_id) ); -CREATE INDEX IF NOT EXISTS idx_usage_model ON usage (model); +CREATE INDEX IF NOT EXISTS idx_usage_model ON usage (gen_ai_response_model); CREATE INDEX IF NOT EXISTS idx_usage_ts ON usage (ts); CREATE INDEX IF NOT EXISTS idx_usage_platform ON usage (platform); +CREATE INDEX IF NOT EXISTS idx_usage_seq ON usage (seq); CREATE TABLE IF NOT EXISTS usage_sync ( session_id TEXT PRIMARY KEY, last_jsonl_size INTEGER NOT NULL ); """ +# Column order for the usage table, mirrored by the upsert below. +_INSERT_SQL = """ +INSERT INTO usage ( + session_id, call_id, seq, ts, platform, + gen_ai_provider_name, gen_ai_operation_name, gen_ai_response_model, + gen_ai_usage_input_tokens, gen_ai_usage_output_tokens, + gen_ai_usage_cache_read_input_tokens, gen_ai_usage_cache_creation_input_tokens, + gen_ai_usage_reasoning_output_tokens +) VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?) +ON CONFLICT (session_id, call_id) DO UPDATE SET + seq = excluded.seq, + ts = excluded.ts, + platform = excluded.platform, + gen_ai_provider_name = excluded.gen_ai_provider_name, + gen_ai_operation_name = excluded.gen_ai_operation_name, + gen_ai_response_model = excluded.gen_ai_response_model, + gen_ai_usage_input_tokens = excluded.gen_ai_usage_input_tokens, + gen_ai_usage_output_tokens = excluded.gen_ai_usage_output_tokens, + gen_ai_usage_cache_read_input_tokens = excluded.gen_ai_usage_cache_read_input_tokens, + gen_ai_usage_cache_creation_input_tokens = excluded.gen_ai_usage_cache_creation_input_tokens, + gen_ai_usage_reasoning_output_tokens = excluded.gen_ai_usage_reasoning_output_tokens +""" + class UsageIndex: def __init__(self, thirdeye_home: Path) -> None: @@ -39,21 +72,26 @@ def __init__(self, thirdeye_home: Path) -> None: def connect(self) -> sqlite3.Connection: self.db_path.parent.mkdir(parents=True, exist_ok=True) conn = sqlite3.connect(self.db_path) - conn.executescript(SCHEMA_SQL) current = conn.execute("PRAGMA user_version").fetchone()[0] if current < SCHEMA_VERSION: + # usage.db is purely derived from the sidecars, so an out-of-date + # schema is dropped and rebuilt rather than migrated. + conn.executescript("DROP TABLE IF EXISTS usage; DROP TABLE IF EXISTS usage_sync;") + conn.executescript(SCHEMA_SQL) conn.execute(f"PRAGMA user_version = {SCHEMA_VERSION}") conn.commit() + else: + conn.executescript(SCHEMA_SQL) return conn def refresh(self, conn: sqlite3.Connection) -> int: """Pull new rows from every session's usage.jsonl into the DB. - Returns the total number of rows inserted across all sessions. - Anomalies (shrunk sidecars, malformed lines) are logged but do not - raise. + Returns the total number of rows written (inserted or updated) across + all sessions. Anomalies (shrunk sidecars, malformed lines) are logged + but do not raise. """ - inserted = 0 + written = 0 root = sessions_root(self.thirdeye_home) if not root.exists(): return 0 @@ -63,9 +101,9 @@ def refresh(self, conn: sqlite3.Connection) -> int: for session_dir_ in sorted(platform_dir_.iterdir()): if not session_dir_.is_dir(): continue - inserted += self._refresh_one(conn, session_dir_.name, session_dir_) + written += self._refresh_one(conn, session_dir_.name, session_dir_) conn.commit() - return inserted + return written def refresh_session(self, conn: sqlite3.Connection, session_id: str, session_dir_: Path) -> int: n = self._refresh_one(conn, session_id, session_dir_) @@ -99,50 +137,43 @@ def _refresh_one(self, conn: sqlite3.Connection, sid: str, session_dir_: Path) - conn.execute("DELETE FROM usage WHERE session_id = ?", (sid,)) last_size = 0 - inserted = 0 - with jsonl.open("rb") as f: - f.seek(last_size) - for raw in f: - line = raw.decode("utf-8", errors="replace").strip() - if not line: - continue - try: - row = json.loads(line) - except json.JSONDecodeError: - log_capture_error( - thirdeye_home=self.thirdeye_home, - phase="index_sync", - message="malformed jsonl line", - session_id=sid, - source_path=str(jsonl), - ) - continue - try: - cursor = conn.execute( - "INSERT OR IGNORE INTO usage " - "(session_id, seq, ts, platform, model, " - "input_tokens, output_tokens, total_tokens) " - "VALUES (?, ?, ?, ?, ?, ?, ?, ?)", - ( - row["session_id"], - int(row["seq"]), - row["ts"], - row["platform"], - row["model"], - int(row["input_tokens"]), - int(row["output_tokens"]), - int(row["total_tokens"]), - ), - ) - inserted += cursor.rowcount - except (KeyError, ValueError, sqlite3.Error) as e: - log_capture_error( - thirdeye_home=self.thirdeye_home, - phase="index_sync", - error=e, - session_id=sid, - source_path=str(jsonl), - ) + # The sidecar has grown. iter_calls is last-wins over the whole file, so + # re-read it entirely and upsert every logical call; a later append that + # corrects a call overwrites the earlier value. Malformed lines are + # silently dropped by iter_calls, so scan the grown region separately to + # log them. + self._log_malformed(sid, jsonl, last_size) + + written = 0 + for row in iter_calls(session_dir_): + try: + cursor = conn.execute( + _INSERT_SQL, + ( + row.session_id, + row.call_id, + row.seq, + row.ts, + row.platform, + row.provider_name, + row.operation_name, + row.response_model, + row.input_tokens, + row.output_tokens, + row.cache_read_input_tokens, + row.cache_creation_input_tokens, + row.reasoning_output_tokens, + ), + ) + written += cursor.rowcount + except sqlite3.Error as e: + log_capture_error( + thirdeye_home=self.thirdeye_home, + phase="index_sync", + error=e, + session_id=sid, + source_path=str(jsonl), + ) conn.execute( "INSERT INTO usage_sync (session_id, last_jsonl_size) " @@ -150,4 +181,31 @@ def _refresh_one(self, conn: sqlite3.Connection, sid: str, session_dir_: Path) - "ON CONFLICT (session_id) DO UPDATE SET last_jsonl_size = excluded.last_jsonl_size", (sid, current_size), ) - return inserted + return written + + def _log_malformed(self, sid: str, jsonl: Path, from_offset: int) -> None: + """Log any lines that are not valid JSON in the grown region. + + Purely observability: the actual upsert reads through iter_calls, which + drops malformed lines. Scans only from ``from_offset`` to avoid + re-logging lines seen in a prior refresh. + """ + try: + with jsonl.open("rb") as f: + f.seek(from_offset) + for raw in f: + line = raw.decode("utf-8", errors="replace").strip() + if not line: + continue + try: + json.loads(line) + except json.JSONDecodeError: + log_capture_error( + thirdeye_home=self.thirdeye_home, + phase="index_sync", + message="malformed jsonl line", + session_id=sid, + source_path=str(jsonl), + ) + except OSError: + return diff --git a/src/thirdeye/usage/read.py b/src/thirdeye/usage/read.py new file mode 100644 index 0000000..348afc7 --- /dev/null +++ b/src/thirdeye/usage/read.py @@ -0,0 +1,49 @@ +"""The single canonical read path over usage sidecars. + +The sidecar (``usage.jsonl``) is a faithful raw mirror: writers append one row +per source frame and never read-modify-write, so the same logical LLM call can +appear many times. Claude repeats the identical ``message.usage`` across every +content-block frame of one API call; Codex emits byte-identical repeat reports. +``iter_calls`` collapses those duplicates into one row per distinct +``(session_id, call_id)``. + +This is why capture needs no locking anywhere: two triggers racing on the same +transcript offset append the same rows, and their identical ``call_id``s +collapse here on read. Deduplication happens only in this module — every +consumer that wants logical calls (as opposed to the raw mirror) goes through +``iter_calls``. +""" + +from __future__ import annotations + +from collections.abc import Iterable, Iterator +from pathlib import Path + +from thirdeye.usage.store import UsageStore +from thirdeye.usage.types import UsageRow + + +def iter_calls(session_dir_: Path) -> Iterator[UsageRow]: + """Yield one row per distinct (session_id, call_id); last occurrence wins. + + Reads the raw sidecar via UsageStore.iter_rows() and collapses duplicates. + Row order follows first appearance of each call_id, so output is stable and + roughly chronological even though the surviving value is the last one seen. + + Last-wins is safe because duplicate rows for one call_id carry identical + token values, so which copy survives cannot change any result. + """ + latest: dict[tuple[str, str], UsageRow] = {} + for row in UsageStore(session_dir_).iter_rows(): + latest[(row.session_id, row.call_id)] = row + yield from latest.values() + + +def call_totals(rows: Iterable[UsageRow]) -> tuple[int, int]: + """Return (input_tokens, output_tokens) summed over rows.""" + input_total = 0 + output_total = 0 + for row in rows: + input_total += row.input_tokens + output_total += row.output_tokens + return input_total, output_total diff --git a/src/thirdeye/usage/store.py b/src/thirdeye/usage/store.py index 1eb5531..465fe36 100644 --- a/src/thirdeye/usage/store.py +++ b/src/thirdeye/usage/store.py @@ -39,7 +39,12 @@ def append(self, rows: list[UsageRow]) -> None: f.write(json.dumps(row.to_dict(), separators=(",", ":")) + "\n") def iter_rows(self) -> Iterator[UsageRow]: - """Yield rows from usage.jsonl. Skips empty / malformed lines silently.""" + """Yield rows from usage.jsonl. Skips empty / malformed lines silently. + + Raw and undeduplicated: this is the faithful mirror of the sidecar, so + the same logical call may appear many times. Callers wanting one row per + logical LLM call must use ``read.iter_calls``. + """ if not self.jsonl_path.exists(): return with self.jsonl_path.open("r", encoding="utf-8") as f: diff --git a/src/thirdeye/usage/types.py b/src/thirdeye/usage/types.py index 8a0136a..b514371 100644 --- a/src/thirdeye/usage/types.py +++ b/src/thirdeye/usage/types.py @@ -1,34 +1,109 @@ from __future__ import annotations -from dataclasses import asdict, dataclass +from dataclasses import dataclass from typing import Any +# Mapping from dataclass field name to its dotted OpenTelemetry GenAI spec key. +# All gen_ai.* attributes are Development status (moved to a dedicated repo in +# semconv v1.42.0). These are the exact spec names — no abbreviations, no +# invented attributes. +ATTRIBUTE_KEYS: dict[str, str] = { + "provider_name": "gen_ai.provider.name", + "operation_name": "gen_ai.operation.name", + "response_model": "gen_ai.response.model", + "input_tokens": "gen_ai.usage.input_tokens", + "output_tokens": "gen_ai.usage.output_tokens", + "cache_read_input_tokens": "gen_ai.usage.cache_read.input_tokens", + "cache_creation_input_tokens": "gen_ai.usage.cache_creation.input_tokens", + "reasoning_output_tokens": "gen_ai.usage.reasoning.output_tokens", +} + +# The thirdeye envelope fields — NOT gen_ai attributes. Serialized as bare names. +_ENVELOPE_KEYS: tuple[str, ...] = ("session_id", "seq", "call_id", "ts", "platform") + +# gen_ai.conversation.id is always emitted (mirrors session_id) but is not a +# stored dataclass field, so it lives outside ATTRIBUTE_KEYS. +_CONVERSATION_ID_KEY = "gen_ai.conversation.id" + +# Optional gen_ai.* fields, absent-vs-zero: an unreported attribute is an absent +# key (None), never 0. +_OPTIONAL_FIELDS: tuple[str, ...] = ( + "cache_read_input_tokens", + "cache_creation_input_tokens", + "reasoning_output_tokens", +) + @dataclass(frozen=True) class UsageRow: - """One row of model + token usage data, joinable back to events.alog by seq.""" + """One model call's usage, shaped as OpenTelemetry GenAI semantic conventions. + + Serializes to dotted ``gen_ai.*`` spec keys plus a small thirdeye envelope. + There is no ``total_tokens`` attribute in the spec — it is a derived property, + never stored or serialized. + """ + # --- thirdeye envelope (NOT gen_ai attributes) --- session_id: str seq: int + call_id: str ts: str platform: str - model: str + # --- gen_ai.* attributes --- + provider_name: str + response_model: str input_tokens: int output_tokens: int - total_tokens: int + operation_name: str = "chat" + cache_read_input_tokens: int | None = None + cache_creation_input_tokens: int | None = None + reasoning_output_tokens: int | None = None + + @property + def total_tokens(self) -> int: + """Derived, never stored or serialized.""" + return self.input_tokens + self.output_tokens + + def attributes(self) -> dict[str, Any]: + """Only the gen_ai.* attributes, omitting any whose value is None. + + Always includes gen_ai.conversation.id (mirrors session_id). This method + is the whole future OTLP exporter; keep it pure. + """ + out: dict[str, Any] = {_CONVERSATION_ID_KEY: self.session_id} + for field_name, key in ATTRIBUTE_KEYS.items(): + value = getattr(self, field_name) + if value is None: + continue + out[key] = value + return out def to_dict(self) -> dict[str, Any]: - return asdict(self) + """Envelope keys (bare names) merged with attributes().""" + out: dict[str, Any] = {name: getattr(self, name) for name in _ENVELOPE_KEYS} + out.update(self.attributes()) + return out @classmethod def from_dict(cls, d: dict[str, Any]) -> UsageRow: + """Inverse of to_dict. Missing optional attribute keys become None. + + Raises KeyError / ValueError on a malformed row. + """ + optional: dict[str, int | None] = {} + for field_name in _OPTIONAL_FIELDS: + key = ATTRIBUTE_KEYS[field_name] + optional[field_name] = int(d[key]) if key in d else None return cls( session_id=str(d["session_id"]), seq=int(d["seq"]), + call_id=str(d["call_id"]), ts=str(d["ts"]), platform=str(d["platform"]), - model=str(d["model"]), - input_tokens=int(d["input_tokens"]), - output_tokens=int(d["output_tokens"]), - total_tokens=int(d["total_tokens"]), + provider_name=str(d[ATTRIBUTE_KEYS["provider_name"]]), + response_model=str(d[ATTRIBUTE_KEYS["response_model"]]), + input_tokens=int(d[ATTRIBUTE_KEYS["input_tokens"]]), + output_tokens=int(d[ATTRIBUTE_KEYS["output_tokens"]]), + operation_name=str(d[ATTRIBUTE_KEYS["operation_name"]]), + **optional, ) diff --git a/src/thirdeye/web/routes/usage.py b/src/thirdeye/web/routes/usage.py index 01bafe4..54483ff 100644 --- a/src/thirdeye/web/routes/usage.py +++ b/src/thirdeye/web/routes/usage.py @@ -12,7 +12,7 @@ from thirdeye.paths import session_dir from thirdeye.timeparse import parse_when from thirdeye.usage.aggregate import aggregate_by_day -from thirdeye.usage.store import UsageStore +from thirdeye.usage.read import iter_calls async def _session_usage(request: Request) -> HTMLResponse: @@ -24,7 +24,7 @@ async def _session_usage(request: Request) -> HTMLResponse: except (KeyError, ValueError) as e: raise HTTPException(status_code=404, detail=str(e)) from e sdir = session_dir(config.root, platform, sid) - rows = list(UsageStore(sdir).iter_rows()) + rows = list(iter_calls(sdir)) aggregate = store.stats(session_id=sid) templates = request.app.state.templates return templates.TemplateResponse( diff --git a/src/thirdeye/web/templates/usage/global.html b/src/thirdeye/web/templates/usage/global.html index c1a1905..db9f3f7 100644 --- a/src/thirdeye/web/templates/usage/global.html +++ b/src/thirdeye/web/templates/usage/global.html @@ -11,7 +11,7 @@