Spaces:
Paused
Paused
| """Every generation request/response pair, stored in one SQLite file for Stage 2 evaluation. | |
| One row per request: what the tester asked, what was retrieved, every model | |
| attempt with its validation result, the final output, plus two review | |
| channels: the tester's own feedback (did it work when they ran it?) and the | |
| Orcanion reviewer's verdict (``python -m api.review``). | |
| """ | |
| from __future__ import annotations | |
| import datetime as dt | |
| import json | |
| import sqlite3 | |
| import uuid | |
| from contextlib import contextmanager | |
| from pathlib import Path | |
| SCHEMA = """ | |
| CREATE TABLE IF NOT EXISTS generations ( | |
| id TEXT PRIMARY KEY, | |
| created_at TEXT NOT NULL, | |
| tester TEXT NOT NULL, | |
| kind TEXT NOT NULL, -- 'simulate' | 'code' | |
| request TEXT NOT NULL, -- JSON of the request body | |
| model TEXT, | |
| sources TEXT, -- JSON list of retrieved chunk sources | |
| attempts TEXT, -- JSON list of {raw_output, validation, ...} | |
| output TEXT, -- final MJCF, or JSON {path: content} for code | |
| valid INTEGER, -- automatic validity check of the final output | |
| validation TEXT, -- JSON of the final validation result | |
| error TEXT, -- server-side failure (LLM down, etc.) | |
| duration_ms INTEGER, | |
| tester_rating TEXT, -- 'works' | 'partly' | 'broken' | |
| tester_note TEXT, | |
| review_verdict TEXT, -- 'correct' | 'partial' | 'incorrect' | |
| review_note TEXT, | |
| reviewed_at TEXT | |
| ); | |
| CREATE INDEX IF NOT EXISTS generations_created ON generations(created_at); | |
| -- Every backend job: simulation runs, sweeps, ROS 2 checks, video reviews, dataset exports, Isaac renders | |
| -- (usage visibility: who ran what, when, and what it cost). | |
| CREATE TABLE IF NOT EXISTS runs ( | |
| id TEXT PRIMARY KEY, | |
| created_at TEXT NOT NULL, | |
| tester TEXT NOT NULL, | |
| kind TEXT NOT NULL DEFAULT 'sim', -- sim | sweep | ros_check | review | export | isaac | |
| parent_id TEXT, -- the run a sweep/review/render was made from | |
| hardware TEXT, -- price key: t4 | l4 | cpu | runpod | |
| params TEXT, -- JSON: kind-specific request parameters | |
| cost_usd REAL, -- estimated, set when the job ends | |
| generation_id TEXT, -- source generation, if the scene came from /simulate/generate | |
| robot TEXT, | |
| duration_s REAL, | |
| controller TEXT, | |
| backend TEXT, -- modal | runpod | |
| job_id TEXT, -- the backend's id for the job | |
| status TEXT NOT NULL, -- SUBMITTING | IN_QUEUE | IN_PROGRESS | COMPLETED | FAILED | ... | |
| error TEXT, | |
| output TEXT, -- JSON job output: summary + R2 keys | |
| execution_ms INTEGER, -- GPU execution time (RunPod's figure, else the worker's own) | |
| delay_ms INTEGER, -- queue + cold start (RunPod only) | |
| updated_at TEXT | |
| ); | |
| CREATE INDEX IF NOT EXISTS runs_tester_created ON runs(tester, created_at); | |
| -- Every AI model call, for the rupee budget (api/budget.py). | |
| CREATE TABLE IF NOT EXISTS usage ( | |
| id INTEGER PRIMARY KEY AUTOINCREMENT, | |
| created_at TEXT NOT NULL, | |
| tester TEXT NOT NULL, | |
| generation_id TEXT, | |
| category TEXT NOT NULL, -- coding | simulation_ai | debugging | testing | |
| model TEXT, | |
| tier TEXT, | |
| provider TEXT, -- who billed it: openrouter, ... | |
| input_tokens INTEGER, | |
| output_tokens INTEGER, | |
| cost_usd REAL, | |
| cost_inr REAL, | |
| billable INTEGER NOT NULL, -- 0: OpenRouter-exempt, recorded but not deducted | |
| exempt_reason TEXT, | |
| estimated INTEGER NOT NULL DEFAULT 0, -- 1: cost computed from list prices, not reported by the provider | |
| estimate_inr REAL -- what the plan expected before running | |
| ); | |
| CREATE INDEX IF NOT EXISTS usage_tester ON usage(tester, created_at); | |
| -- Lightweight feedback on anything: generations, runs, checks, routing choices. | |
| CREATE TABLE IF NOT EXISTS feedback ( | |
| id INTEGER PRIMARY KEY AUTOINCREMENT, | |
| created_at TEXT NOT NULL, | |
| tester TEXT NOT NULL, | |
| target_type TEXT NOT NULL, -- generation | job | |
| target_id TEXT NOT NULL, | |
| rating TEXT NOT NULL, -- helpful | not_helpful | |
| note TEXT, | |
| context TEXT -- JSON: task, routing tier/model, simulator, result state, ... | |
| ); | |
| """ | |
| JSON_FIELDS = ("request", "sources", "attempts", "validation") | |
| def _decode_run(d: dict) -> dict: | |
| d["output"] = json.loads(d["output"]) if d.get("output") else None | |
| d["params"] = json.loads(d["params"]) if d.get("params") else {} | |
| return d | |
| def _now() -> str: | |
| return dt.datetime.now(dt.timezone.utc).isoformat(timespec="seconds") | |
| class LogStore: | |
| def __init__(self, path: Path): | |
| self.path = Path(path) | |
| self.path.parent.mkdir(parents=True, exist_ok=True) | |
| with self._conn() as c: | |
| c.executescript(SCHEMA) | |
| self._migrate(c) | |
| def _migrate(c) -> None: | |
| """Bring run tables created by earlier versions up to the current schema.""" | |
| cols = {r[1] for r in c.execute("PRAGMA table_info(runs)")} | |
| if "runpod_job_id" in cols: # before Modal support | |
| c.execute("ALTER TABLE runs RENAME COLUMN runpod_job_id TO job_id") | |
| added = { | |
| "backend": "TEXT", | |
| "kind": "TEXT NOT NULL DEFAULT 'sim'", | |
| "parent_id": "TEXT", | |
| "hardware": "TEXT", | |
| "params": "TEXT", | |
| "cost_usd": "REAL", | |
| } | |
| for name, decl in added.items(): | |
| if name not in cols: | |
| c.execute(f"ALTER TABLE runs ADD COLUMN {name} {decl}") | |
| if "runpod_job_id" in cols and "backend" not in cols: | |
| c.execute("UPDATE runs SET backend = 'runpod' WHERE job_id IS NOT NULL") | |
| c.execute("CREATE INDEX IF NOT EXISTS runs_created ON runs(created_at)") | |
| def _conn(self): | |
| conn = sqlite3.connect(self.path, timeout=30) | |
| conn.row_factory = sqlite3.Row | |
| try: | |
| conn.execute("PRAGMA journal_mode=WAL") | |
| yield conn | |
| conn.commit() | |
| finally: | |
| conn.close() | |
| def record( | |
| self, | |
| *, | |
| tester: str, | |
| kind: str, | |
| request: dict, | |
| model: str | None = None, | |
| sources: list | None = None, | |
| attempts: list | None = None, | |
| output: str | None = None, | |
| valid: bool | None = None, | |
| validation: dict | None = None, | |
| error: str | None = None, | |
| duration_ms: int | None = None, | |
| ) -> str: | |
| gen_id = uuid.uuid4().hex[:12] | |
| with self._conn() as c: | |
| c.execute( | |
| "INSERT INTO generations (id, created_at, tester, kind, request, model, sources, attempts, output, " | |
| "valid, validation, error, duration_ms) VALUES (?,?,?,?,?,?,?,?,?,?,?,?,?)", | |
| ( | |
| gen_id, | |
| _now(), | |
| tester, | |
| kind, | |
| json.dumps(request), | |
| model, | |
| json.dumps(sources or []), | |
| json.dumps(attempts or []), | |
| output, | |
| None if valid is None else int(valid), | |
| json.dumps(validation or {}), | |
| error, | |
| duration_ms, | |
| ), | |
| ) | |
| return gen_id | |
| def get(self, gen_id: str) -> dict | None: | |
| with self._conn() as c: | |
| row = c.execute("SELECT * FROM generations WHERE id = ?", (gen_id,)).fetchone() | |
| return self._decode(row) if row else None | |
| def list( | |
| self, kind: str | None = None, tester: str | None = None, unreviewed: bool = False, limit: int = 50 | |
| ) -> list[dict]: | |
| sql, args = "SELECT * FROM generations WHERE 1=1", [] | |
| if kind: | |
| sql += " AND kind = ?" | |
| args.append(kind) | |
| if tester: | |
| sql += " AND tester = ?" | |
| args.append(tester) | |
| if unreviewed: | |
| sql += " AND review_verdict IS NULL" | |
| sql += " ORDER BY created_at DESC LIMIT ?" | |
| args.append(limit) | |
| with self._conn() as c: | |
| return [self._decode(r) for r in c.execute(sql, args).fetchall()] | |
| def all(self) -> list[dict]: | |
| with self._conn() as c: | |
| return [self._decode(r) for r in c.execute("SELECT * FROM generations ORDER BY created_at").fetchall()] | |
| def set_feedback(self, gen_id: str, tester: str, rating: str, note: str) -> bool: | |
| with self._conn() as c: | |
| cur = c.execute( | |
| "UPDATE generations SET tester_rating = ?, tester_note = ? WHERE id = ? AND tester = ?", | |
| (rating, note, gen_id, tester), | |
| ) | |
| return cur.rowcount == 1 | |
| def set_review(self, gen_id: str, verdict: str, note: str) -> bool: | |
| with self._conn() as c: | |
| cur = c.execute( | |
| "UPDATE generations SET review_verdict = ?, review_note = ?, reviewed_at = ? WHERE id = ?", | |
| (verdict, note, _now(), gen_id), | |
| ) | |
| return cur.rowcount == 1 | |
| def summary(self) -> list[dict]: | |
| sql = """SELECT kind, COUNT(*) AS total, SUM(valid = 1) AS auto_valid, | |
| SUM(tester_rating = 'works') AS tester_works, SUM(tester_rating IS NOT NULL) AS tester_rated, | |
| SUM(review_verdict = 'correct') AS reviewed_correct, | |
| SUM(review_verdict IS NOT NULL) AS reviewed | |
| FROM generations GROUP BY kind ORDER BY kind""" | |
| with self._conn() as c: | |
| return [dict(r) for r in c.execute(sql).fetchall()] | |
| # ------------------------------------------------------------------ AI usage and feedback | |
| def record_usage(self, **row) -> None: | |
| allowed = {"created_at", "tester", "generation_id", "category", "model", "tier", | |
| "provider", "input_tokens", "output_tokens", "cost_usd", "cost_inr", | |
| "billable", "exempt_reason", "estimated", "estimate_inr"} | |
| row = {k: v for k, v in row.items() if k in allowed} | |
| row = {"created_at": _now(), **row} | |
| cols = ", ".join(row) | |
| with self._conn() as c: | |
| c.execute(f"INSERT INTO usage ({cols}) VALUES ({', '.join('?' * len(row))})", tuple(row.values())) | |
| def usage_rows(self, tester: str | None = None, since_iso: str | None = None) -> list[dict]: | |
| sql, args = "SELECT * FROM usage WHERE 1=1", [] | |
| if tester: | |
| sql += " AND tester = ?" | |
| args.append(tester) | |
| if since_iso: | |
| sql += " AND created_at >= ?" | |
| args.append(since_iso) | |
| with self._conn() as c: | |
| return [dict(r) for r in c.execute(sql + " ORDER BY id", args).fetchall()] | |
| def record_feedback(self, tester: str, target_type: str, target_id: str, rating: str, note: str, context: dict): | |
| with self._conn() as c: | |
| c.execute( | |
| "INSERT INTO feedback (created_at, tester, target_type, target_id, rating, note, context) " | |
| "VALUES (?,?,?,?,?,?,?)", | |
| (_now(), tester, target_type, target_id, rating, note, json.dumps(context)), | |
| ) | |
| def feedback_rows(self) -> list[dict]: | |
| with self._conn() as c: | |
| rows = [dict(r) for r in c.execute("SELECT * FROM feedback ORDER BY id").fetchall()] | |
| for r in rows: | |
| r["context"] = json.loads(r["context"]) if r.get("context") else {} | |
| return rows | |
| # ------------------------------------------------------------------ runs | |
| def record_run( | |
| self, | |
| *, | |
| run_id: str, | |
| tester: str, | |
| generation_id: str | None, | |
| robot: str | None, | |
| duration_s: float | None, | |
| controller: str | None, | |
| kind: str = "sim", | |
| parent_id: str | None = None, | |
| hardware: str | None = None, | |
| params: dict | None = None, | |
| ) -> None: | |
| now = _now() | |
| with self._conn() as c: | |
| c.execute( | |
| "INSERT INTO runs (id, created_at, tester, kind, parent_id, hardware, params, generation_id, robot, " | |
| "duration_s, controller, status, updated_at) VALUES (?,?,?,?,?,?,?,?,?,?,?,?,?)", | |
| ( | |
| run_id, | |
| now, | |
| tester, | |
| kind, | |
| parent_id, | |
| hardware, | |
| json.dumps(params) if params is not None else None, | |
| generation_id, | |
| robot, | |
| duration_s, | |
| controller, | |
| "SUBMITTING", | |
| now, | |
| ), | |
| ) | |
| def update_run(self, run_id: str, **fields) -> None: | |
| allowed = {"backend", "job_id", "status", "error", "output", "execution_ms", "delay_ms", "cost_usd"} | |
| if set(fields) - allowed: | |
| raise ValueError(f"unknown run fields {set(fields) - allowed}") | |
| if "output" in fields and not isinstance(fields["output"], (str, type(None))): | |
| fields["output"] = json.dumps(fields["output"]) | |
| fields["updated_at"] = _now() | |
| cols = ", ".join(f"{k} = ?" for k in fields) | |
| with self._conn() as c: | |
| c.execute(f"UPDATE runs SET {cols} WHERE id = ?", (*fields.values(), run_id)) | |
| def get_run(self, run_id: str) -> dict | None: | |
| with self._conn() as c: | |
| row = c.execute("SELECT * FROM runs WHERE id = ?", (run_id,)).fetchone() | |
| if not row: | |
| return None | |
| return _decode_run(dict(row)) | |
| def list_runs(self, tester: str | None = None, limit: int = 50) -> list[dict]: | |
| sql, args = "SELECT * FROM runs", [] | |
| if tester: | |
| sql += " WHERE tester = ?" | |
| args.append(tester) | |
| sql += " ORDER BY created_at DESC LIMIT ?" | |
| args.append(limit) | |
| with self._conn() as c: | |
| rows = [dict(r) for r in c.execute(sql, args).fetchall()] | |
| return [_decode_run(r) for r in rows] | |
| def jobs_since(self, since_iso: str, kind: str | None = None) -> list[dict]: | |
| """Every job created at or after `since_iso` (for spend and daily-limit checks).""" | |
| sql, args = "SELECT * FROM runs WHERE created_at >= ?", [since_iso] | |
| if kind: | |
| sql += " AND kind = ?" | |
| args.append(kind) | |
| with self._conn() as c: | |
| return [_decode_run(dict(r)) for r in c.execute(sql, args).fetchall()] | |
| def children(self, parent_id: str, kind: str | None = None) -> list[dict]: | |
| sql, args = "SELECT * FROM runs WHERE parent_id = ?", [parent_id] | |
| if kind: | |
| sql += " AND kind = ?" | |
| args.append(kind) | |
| sql += " ORDER BY created_at" | |
| with self._conn() as c: | |
| return [_decode_run(dict(r)) for r in c.execute(sql, args).fetchall()] | |
| def runs_since(self, tester: str, since_iso: str) -> int: | |
| with self._conn() as c: | |
| return c.execute( | |
| "SELECT COUNT(*) FROM runs WHERE tester = ? AND created_at >= ?", (tester, since_iso) | |
| ).fetchone()[0] | |
| def usage(self, window_minutes: int = 10) -> list[dict]: | |
| """Per tester: total jobs, jobs in the last 24 h, busiest window, compute seconds, estimated cost.""" | |
| with self._conn() as c: | |
| rows = c.execute( | |
| "SELECT tester, created_at, execution_ms, cost_usd FROM runs ORDER BY tester, created_at" | |
| ).fetchall() | |
| day_ago = (dt.datetime.now(dt.timezone.utc) - dt.timedelta(days=1)).isoformat(timespec="seconds") | |
| out: dict[str, dict] = {} | |
| times: dict[str, list[dt.datetime]] = {} | |
| for tester, created, exec_ms, cost in rows: | |
| u = out.setdefault( | |
| tester, | |
| { | |
| "tester": tester, | |
| "runs": 0, | |
| "last_24h": 0, | |
| "gpu_seconds": 0.0, | |
| "cost_usd": 0.0, | |
| "first": created, | |
| "last": created, | |
| }, | |
| ) | |
| u["runs"] += 1 | |
| u["last"] = created | |
| u["last_24h"] += created >= day_ago | |
| u["gpu_seconds"] += (exec_ms or 0) / 1000 | |
| u["cost_usd"] += cost or 0.0 | |
| times.setdefault(tester, []).append(dt.datetime.fromisoformat(created)) | |
| window = dt.timedelta(minutes=window_minutes) | |
| for tester, ts in times.items(): | |
| best, j = 0, 0 | |
| for i, t in enumerate(ts): | |
| while t - ts[j] > window: | |
| j += 1 | |
| best = max(best, i - j + 1) | |
| out[tester][f"max_in_{window_minutes}min"] = best | |
| out[tester]["gpu_seconds"] = round(out[tester]["gpu_seconds"], 1) | |
| out[tester]["cost_usd"] = round(out[tester]["cost_usd"], 3) | |
| return sorted(out.values(), key=lambda u: -u["runs"]) | |
| def _decode(row: sqlite3.Row) -> dict: | |
| d = dict(row) | |
| for f in JSON_FIELDS: | |
| if d.get(f): | |
| d[f] = json.loads(d[f]) | |
| if d.get("valid") is not None: | |
| d["valid"] = bool(d["valid"]) | |
| return d | |