"""Research session tracking + cost attribution for the Teleo pipeline. This module adds three capabilities: 1. research_sessions table — tracks WHY agents researched, what they found interesting, session cost, and links to generated sources 2. Extraction cost attribution — writes per-source cost to sources.cost_usd after extraction 3. Source → claim linkage — ensures prs.source_path is always populated Designed for Epimetheus to integrate into the pipeline. Argus built the spec; Ganymede reviews; Epimetheus wires it in. Data flow: Agent research session → research_sessions row (with reasoning + summary) → sources created (with session_id FK) → extraction runs (cost written to sources.cost_usd + costs table) → PRs created (source_path populated) → claims merged (traceable back to session) """ import json import logging import sqlite3 from datetime import datetime from typing import Optional logger = logging.getLogger("research_tracking") # --------------------------------------------------------------------------- # Migration v11: research_sessions table + sources.session_id FK # (v9 is current; v10 is Epimetheus's eval pipeline migration) # --------------------------------------------------------------------------- MIGRATION_V11_SQL = """ -- Research session tracking table CREATE TABLE IF NOT EXISTS research_sessions ( id INTEGER PRIMARY KEY AUTOINCREMENT, agent TEXT NOT NULL, -- Which agent ran the research (leo, rio, astra, etc.) domain TEXT, -- Primary domain of the research topic TEXT NOT NULL, -- What they researched (short description) reasoning TEXT, -- WHY they chose this topic (agent's own explanation) summary TEXT, -- What they found most interesting/relevant sources_planned INTEGER DEFAULT 0, -- How many sources they intended to produce sources_produced INTEGER DEFAULT 0, -- How many actually materialized model TEXT, -- Model used for research (e.g. claude-opus-4-6) input_tokens INTEGER DEFAULT 0, output_tokens INTEGER DEFAULT 0, cost_usd REAL DEFAULT 0, -- Total research session cost (LLM calls for discovery + writing) status TEXT DEFAULT 'running', -- running, completed, failed, partial started_at TEXT DEFAULT (datetime('now')), completed_at TEXT, metadata TEXT DEFAULT '{}' -- JSON: any extra context (prompt version, search queries used, etc.) ); CREATE INDEX IF NOT EXISTS idx_rs_agent ON research_sessions(agent); CREATE INDEX IF NOT EXISTS idx_rs_domain ON research_sessions(domain); CREATE INDEX IF NOT EXISTS idx_rs_started ON research_sessions(started_at); -- Add session_id FK to sources table ALTER TABLE sources ADD COLUMN session_id INTEGER REFERENCES research_sessions(id); CREATE INDEX IF NOT EXISTS idx_sources_session ON sources(session_id); -- Record migration INSERT INTO schema_version (version) VALUES (11); """ # --------------------------------------------------------------------------- # Cost attribution: write extraction cost to sources.cost_usd # --------------------------------------------------------------------------- # Pricing per million tokens (as of March 2026) MODEL_PRICING = { "anthropic/claude-sonnet-4.5": {"input": 3.00, "output": 15.00}, "anthropic/claude-sonnet-4-5": {"input": 3.00, "output": 15.00}, "anthropic/claude-haiku-4.5": {"input": 0.80, "output": 4.00}, "anthropic/claude-haiku-4-5-20251001": {"input": 0.80, "output": 4.00}, "minimax/minimax-m2.5": {"input": 0.14, "output": 0.56}, } def calculate_cost(model: str, input_tokens: int, output_tokens: int) -> float: """Calculate USD cost from model name and token counts.""" pricing = MODEL_PRICING.get(model) if not pricing: # Default to Sonnet 4.5 pricing as conservative estimate logger.warning("Unknown model %s — using Sonnet 4.5 pricing", model) pricing = {"input": 3.00, "output": 15.00} return (input_tokens * pricing["input"] + output_tokens * pricing["output"]) / 1_000_000 def record_extraction_cost( conn: sqlite3.Connection, source_path: str, model: str, input_tokens: int, output_tokens: int, ): """Write extraction cost to both sources.cost_usd and costs table. Call this after each successful extraction call in openrouter-extract-v2.py. This is the missing link — the CSV logger records tokens but never writes cost back to the DB. """ cost = calculate_cost(model, input_tokens, output_tokens) # Update source row conn.execute( "UPDATE sources SET cost_usd = cost_usd + ?, extraction_model = ? WHERE path = ?", (cost, model, source_path), ) # Also record in costs table for dashboard aggregation date = datetime.utcnow().strftime("%Y-%m-%d") conn.execute( """INSERT INTO costs (date, model, stage, calls, input_tokens, output_tokens, cost_usd) VALUES (?, ?, 'extraction', 1, ?, ?, ?) ON CONFLICT(date, model, stage) DO UPDATE SET calls = calls + 1, input_tokens = input_tokens + excluded.input_tokens, output_tokens = output_tokens + excluded.output_tokens, cost_usd = cost_usd + excluded.cost_usd""", (date, model, input_tokens, output_tokens, cost), ) conn.commit() logger.info( "Recorded extraction cost for %s: $%.4f (%d in, %d out, %s)", source_path, cost, input_tokens, output_tokens, model, ) return cost # --------------------------------------------------------------------------- # Research session lifecycle # --------------------------------------------------------------------------- def start_session( conn: sqlite3.Connection, agent: str, topic: str, domain: Optional[str] = None, reasoning: Optional[str] = None, sources_planned: int = 0, model: Optional[str] = None, metadata: Optional[dict] = None, ) -> int: """Call at the START of a research session. Returns session_id. The agent should call this before it begins producing sources, explaining what it plans to research and why. """ cur = conn.execute( """INSERT INTO research_sessions (agent, domain, topic, reasoning, sources_planned, model, metadata) VALUES (?, ?, ?, ?, ?, ?, ?)""", ( agent, domain, topic, reasoning, sources_planned, model, json.dumps(metadata or {}), ), ) conn.commit() session_id = cur.lastrowid logger.info("Started research session #%d: %s / %s", session_id, agent, topic) return session_id def link_source_to_session( conn: sqlite3.Connection, source_path: str, session_id: int, ): """Link a source file to its research session. Call this when a source is written to inbox/ during a research session. """ conn.execute( "UPDATE sources SET session_id = ? WHERE path = ?", (session_id, source_path), ) conn.execute( """UPDATE research_sessions SET sources_produced = sources_produced + 1 WHERE id = ?""", (session_id,), ) conn.commit() def complete_session( conn: sqlite3.Connection, session_id: int, summary: str, input_tokens: int = 0, output_tokens: int = 0, cost_usd: float = 0, status: str = "completed", ): """Call at the END of a research session. The agent should summarize what it found most interesting/relevant. Cost should include ALL LLM calls made during the session (web search, analysis, source writing — everything). """ conn.execute( """UPDATE research_sessions SET summary = ?, input_tokens = ?, output_tokens = ?, cost_usd = ?, status = ?, completed_at = datetime('now') WHERE id = ?""", (summary, input_tokens, output_tokens, cost_usd, status, session_id), ) conn.commit() logger.info("Completed research session #%d: %s", session_id, status) # --------------------------------------------------------------------------- # Source → PR linkage fix # --------------------------------------------------------------------------- def ensure_source_path_on_pr( conn: sqlite3.Connection, pr_number: int, source_path: str, ): """Ensure prs.source_path is populated. Call during PR creation. Currently 0/1451 PRs have source_path set. This is the fix. """ conn.execute( "UPDATE prs SET source_path = ? WHERE number = ? AND (source_path IS NULL OR source_path = '')", (source_path, pr_number), ) conn.commit() # --------------------------------------------------------------------------- # Backfill: attribute extraction costs from existing CSV log # --------------------------------------------------------------------------- def backfill_extraction_costs(conn: sqlite3.Connection, csv_path: str): """One-time backfill: read openrouter-usage.csv and write costs to sources + costs tables. Run once to fill in the ~$338 of extraction costs that were logged to CSV but never written to the database. Safe to re-run — only updates sources where cost_usd = 0, so partial runs can be resumed without double-counting. """ import csv count = 0 total_cost = 0.0 with open(csv_path) as f: reader = csv.DictReader(f) for row in reader: source_file = row.get("source_file", "") model = row.get("model", "") try: in_tok = int(row.get("input_tokens", 0) or 0) out_tok = int(row.get("output_tokens", 0) or 0) except (ValueError, TypeError): continue cost = calculate_cost(model, in_tok, out_tok) if cost <= 0: continue # Try to match source_file to sources.path # CSV has filename, DB has full path — match on exact suffix # Use ORDER BY length(path) to prefer shortest (most specific) match matched = conn.execute( "SELECT path FROM sources WHERE path LIKE ? AND cost_usd = 0 ORDER BY length(path) LIMIT 1", (f"%/{source_file}" if "/" not in source_file else f"%{source_file}",), ).fetchone() if matched: conn.execute( "UPDATE sources SET cost_usd = ?, extraction_model = ? WHERE path = ?", (cost, model, matched[0]), ) # Always record in costs table date = row.get("date", "unknown") conn.execute( """INSERT INTO costs (date, model, stage, calls, input_tokens, output_tokens, cost_usd) VALUES (?, ?, 'extraction', 1, ?, ?, ?) ON CONFLICT(date, model, stage) DO UPDATE SET calls = calls + 1, input_tokens = input_tokens + excluded.input_tokens, output_tokens = output_tokens + excluded.output_tokens, cost_usd = cost_usd + excluded.cost_usd""", (date, model, in_tok, out_tok, cost), ) count += 1 total_cost += cost conn.commit() logger.info("Backfilled %d extraction cost records, total $%.2f", count, total_cost) return count, total_cost # --------------------------------------------------------------------------- # Backfill: populate prs.source_path from branch naming convention # --------------------------------------------------------------------------- def backfill_source_paths(conn: sqlite3.Connection): """One-time backfill: derive source_path for existing PRs from branch names. Branch format: extract/YYYY-MM-DD-source-name or similar patterns. Source path format: inbox/queue/YYYY-MM-DD-source-name.md """ rows = conn.execute( "SELECT number, branch FROM prs WHERE source_path IS NULL AND branch IS NOT NULL" ).fetchall() count = 0 for number, branch in rows: # Try to extract source name from branch # Common patterns: extract/source-name, claims/source-name parts = branch.split("/", 1) if len(parts) < 2: continue source_stem = parts[1] # Try to find matching source in DB — exact suffix match, shortest path wins matched = conn.execute( "SELECT path FROM sources WHERE path LIKE ? ORDER BY length(path) LIMIT 1", (f"%/{source_stem}%" if source_stem else "",), ).fetchone() if matched: conn.execute( "UPDATE prs SET source_path = ? WHERE number = ?", (matched[0], number), ) count += 1 conn.commit() logger.info("Backfilled source_path for %d PRs", count) return count # --------------------------------------------------------------------------- # Integration points (for Epimetheus to wire in) # --------------------------------------------------------------------------- INTEGRATION_GUIDE = """ ## Where to wire this in ### 1. openrouter-extract-v2.py — after successful extraction call from research_tracking import record_extraction_cost # After line 430 (content, usage = call_openrouter(...)) # After line 672 (log_usage(...)) record_extraction_cost( conn, args.source_file, args.model, usage.get("prompt_tokens", 0), usage.get("completion_tokens", 0), ) ### 2. Agent research scripts — wrap research sessions from research_tracking import start_session, link_source_to_session, complete_session # At start of research: session_id = start_session(conn, agent="leo", topic="weapons stigmatization campaigns", domain="grand-strategy", reasoning="Following up on EU AI Act national security exclusion — exploring how stigmatization campaigns have historically driven arms control policy", sources_planned=6, model="claude-opus-4-6") # As each source is written: link_source_to_session(conn, source_path, session_id) # At end of research: complete_session(conn, session_id, summary="Ottawa Treaty mine ban model is the strongest parallel to AI weapons — same 3-condition framework (humanitarian harm + low military utility + civil society coalition). Ukraine Shahed case is a near-miss triggering event.", input_tokens=total_in, output_tokens=total_out, cost_usd=total_cost) ### 3. PR creation in lib/merge.py or lib/validate.py — ensure source_path from research_tracking import ensure_source_path_on_pr # When creating a PR, pass the source: ensure_source_path_on_pr(conn, pr_number, source_path) ### 4. One-time backfills (run manually after migration) from research_tracking import backfill_extraction_costs, backfill_source_paths backfill_extraction_costs(conn, "/opt/teleo-eval/logs/openrouter-usage.csv") backfill_source_paths(conn) ### 5. Migration Run MIGRATION_V11_SQL against pipeline.db after backing up. """