312 lines
14 KiB
Python
312 lines
14 KiB
Python
"""
|
|
SQLite store for the ai-agent backend.
|
|
|
|
Holds two kinds of state so we don't recompute them on every request:
|
|
|
|
* ``files`` — one row per exposed context file: content hash, size, mtime, the
|
|
*real* Claude token count and its dollar cost. Token counts are expensive
|
|
(a network round-trip to the count_tokens endpoint), so they are computed
|
|
lazily and only when a file's content hash actually changes.
|
|
* ``skills`` — aggregated skill-usage analytics mined from the Claude Code
|
|
transcripts (how many times each skill ran, and when it was last used).
|
|
* ``transcripts`` — per-transcript-file bookkeeping so skill stats can be
|
|
recomputed incrementally (only re-parse the .jsonl files that changed).
|
|
|
|
A single connection (``check_same_thread=False``) guarded by a lock is plenty
|
|
for this low-traffic, single-process service.
|
|
"""
|
|
|
|
import json
|
|
import os
|
|
import sqlite3
|
|
import threading
|
|
import time
|
|
|
|
_SCHEMA = """
|
|
CREATE TABLE IF NOT EXISTS files (
|
|
path TEXT PRIMARY KEY,
|
|
content_hash TEXT,
|
|
size INTEGER,
|
|
mtime REAL,
|
|
is_markdown INTEGER NOT NULL DEFAULT 0,
|
|
tokens INTEGER, -- real (or estimated) token count
|
|
cost REAL, -- tokens priced at the input rate
|
|
estimated INTEGER NOT NULL DEFAULT 1, -- 1 = heuristic, 0 = real API count
|
|
counted_at REAL, -- epoch secs the token count was last set
|
|
seen_at REAL, -- epoch secs the file was last seen on disk
|
|
words INTEGER, -- text stats, stored on content change so
|
|
chars INTEGER -- /api/bundle never re-reads files per request
|
|
);
|
|
|
|
CREATE TABLE IF NOT EXISTS skills (
|
|
name TEXT PRIMARY KEY,
|
|
count INTEGER NOT NULL DEFAULT 0,
|
|
last_used TEXT
|
|
);
|
|
|
|
CREATE TABLE IF NOT EXISTS transcripts (
|
|
path TEXT PRIMARY KEY,
|
|
mtime REAL,
|
|
size INTEGER,
|
|
skills_json TEXT, -- {skill: {"count": n, "last": iso}}
|
|
summary_json TEXT -- per-conversation rollup for the list view
|
|
);
|
|
|
|
CREATE TABLE IF NOT EXISTS meta (
|
|
key TEXT PRIMARY KEY,
|
|
value TEXT
|
|
);
|
|
|
|
CREATE TABLE IF NOT EXISTS conversation_diffs (
|
|
session_id TEXT PRIMARY KEY,
|
|
fingerprint TEXT, -- endedAt + repo HEADs; recompute when it moves
|
|
data_json TEXT, -- computed conversation-commits payload
|
|
computed_at REAL
|
|
);
|
|
"""
|
|
|
|
|
|
class Store:
|
|
def __init__(self, path: str):
|
|
os.makedirs(os.path.dirname(path) or ".", exist_ok=True)
|
|
self._lock = threading.Lock()
|
|
self._db = sqlite3.connect(path, check_same_thread=False)
|
|
self._db.row_factory = sqlite3.Row
|
|
# WAL keeps readers unblocked during the indexer's writes and cuts the
|
|
# fsync-per-commit cost of its frequent small updates.
|
|
self._db.execute("PRAGMA journal_mode=WAL")
|
|
self._db.execute("PRAGMA synchronous=NORMAL")
|
|
self._db.executescript(_SCHEMA)
|
|
self._migrate()
|
|
self._db.commit()
|
|
# Decoded-summary cache: json.loads over every transcript's summary is
|
|
# what every list/analytics request used to pay (~900 blobs, several MB).
|
|
# `summaries_version` bumps on any change so request-level memoizers
|
|
# (e.g. project cost rollups) can key off it.
|
|
self._summaries: dict[str, dict] | None = None
|
|
self.summaries_version = 0
|
|
|
|
def _migrate(self) -> None:
|
|
"""Add columns introduced after a DB was first created."""
|
|
cols = {r["name"] for r in self._db.execute("PRAGMA table_info(transcripts)")}
|
|
if "summary_json" not in cols:
|
|
self._db.execute("ALTER TABLE transcripts ADD COLUMN summary_json TEXT")
|
|
fcols = {r["name"] for r in self._db.execute("PRAGMA table_info(files)")}
|
|
for col in ("words", "chars"):
|
|
if col not in fcols:
|
|
self._db.execute(f"ALTER TABLE files ADD COLUMN {col} INTEGER")
|
|
|
|
# ── files ────────────────────────────────────────────────────────────
|
|
def get_file(self, path: str) -> sqlite3.Row | None:
|
|
with self._lock:
|
|
cur = self._db.execute("SELECT * FROM files WHERE path = ?", (path,))
|
|
return cur.fetchone()
|
|
|
|
def all_files(self) -> dict[str, sqlite3.Row]:
|
|
with self._lock:
|
|
cur = self._db.execute("SELECT * FROM files")
|
|
return {r["path"]: r for r in cur.fetchall()}
|
|
|
|
def upsert_file_meta(self, path: str, content_hash: str, size: int,
|
|
mtime: float, is_markdown: bool,
|
|
words: int | None = None,
|
|
chars: int | None = None) -> None:
|
|
"""Record a file's on-disk metadata. If the hash changed, the token
|
|
count is invalidated (estimated flag flipped on) so the worker recounts."""
|
|
now = time.time()
|
|
with self._lock:
|
|
row = self._db.execute(
|
|
"SELECT content_hash FROM files WHERE path = ?", (path,)
|
|
).fetchone()
|
|
if row is None:
|
|
self._db.execute(
|
|
"INSERT INTO files (path, content_hash, size, mtime, "
|
|
"is_markdown, estimated, seen_at, words, chars) "
|
|
"VALUES (?,?,?,?,?,1,?,?,?)",
|
|
(path, content_hash, size, mtime, int(is_markdown), now,
|
|
words, chars),
|
|
)
|
|
elif row["content_hash"] != content_hash:
|
|
# content changed → mark token count stale (estimated) to recount
|
|
self._db.execute(
|
|
"UPDATE files SET content_hash=?, size=?, mtime=?, "
|
|
"is_markdown=?, estimated=1, seen_at=?, words=?, chars=? "
|
|
"WHERE path=?",
|
|
(content_hash, size, mtime, int(is_markdown), now,
|
|
words, chars, path),
|
|
)
|
|
else:
|
|
self._db.execute(
|
|
"UPDATE files SET size=?, mtime=?, is_markdown=?, seen_at=?, "
|
|
"words=COALESCE(?, words), chars=COALESCE(?, chars) "
|
|
"WHERE path=?",
|
|
(size, mtime, int(is_markdown), now, words, chars, path),
|
|
)
|
|
self._db.commit()
|
|
|
|
def set_tokens(self, path: str, tokens: int, cost: float, estimated: bool) -> None:
|
|
with self._lock:
|
|
self._db.execute(
|
|
"UPDATE files SET tokens=?, cost=?, estimated=?, counted_at=? "
|
|
"WHERE path=?",
|
|
(tokens, cost, int(estimated), time.time(), path),
|
|
)
|
|
self._db.commit()
|
|
|
|
def prune_files(self, keep_paths: set[str]) -> None:
|
|
with self._lock:
|
|
existing = {r["path"] for r in self._db.execute("SELECT path FROM files")}
|
|
for gone in existing - keep_paths:
|
|
self._db.execute("DELETE FROM files WHERE path=?", (gone,))
|
|
self._db.commit()
|
|
|
|
# ── transcripts / skills ─────────────────────────────────────────────
|
|
def get_transcript(self, path: str) -> sqlite3.Row | None:
|
|
with self._lock:
|
|
return self._db.execute(
|
|
"SELECT * FROM transcripts WHERE path=?", (path,)
|
|
).fetchone()
|
|
|
|
def all_transcripts(self) -> dict[str, sqlite3.Row]:
|
|
"""Bookkeeping rows only (no JSON blobs) — this is the change-detection
|
|
index, and pulling every summary_json with it would read the whole DB."""
|
|
with self._lock:
|
|
return {r["path"]: r for r in self._db.execute(
|
|
"SELECT path, mtime, size FROM transcripts")}
|
|
|
|
def all_transcript_skills(self) -> list[str]:
|
|
"""Every transcript's raw ``skills_json`` (for the skills rollup)."""
|
|
with self._lock:
|
|
return [r["skills_json"] for r in self._db.execute(
|
|
"SELECT skills_json FROM transcripts")]
|
|
|
|
def all_transcript_skill_rows(self) -> list[tuple[str, str]]:
|
|
"""``(path, raw skills_json)`` per transcript — the same path key
|
|
:meth:`all_summaries` uses, so the skills catalog can join a skill's
|
|
calls to the spend of the conversations it ran in."""
|
|
with self._lock:
|
|
return [(r["path"], r["skills_json"]) for r in self._db.execute(
|
|
"SELECT path, skills_json FROM transcripts")]
|
|
|
|
def upsert_transcript(self, path: str, mtime: float, size: int,
|
|
skills: dict, summary: dict | None = None) -> None:
|
|
with self._lock:
|
|
self._db.execute(
|
|
"INSERT INTO transcripts (path, mtime, size, skills_json, summary_json) "
|
|
"VALUES (?,?,?,?,?) ON CONFLICT(path) DO UPDATE SET "
|
|
"mtime=excluded.mtime, size=excluded.size, "
|
|
"skills_json=excluded.skills_json, summary_json=excluded.summary_json",
|
|
(path, mtime, size, json.dumps(skills),
|
|
json.dumps(summary) if summary is not None else None),
|
|
)
|
|
self._db.commit()
|
|
if self._summaries is not None:
|
|
if summary is not None:
|
|
# round-trip through JSON so the cache holds exactly what a
|
|
# cold load would (tuples→lists etc.), decoupled from the
|
|
# parser state's live dicts
|
|
self._summaries[path] = json.loads(json.dumps(summary))
|
|
else:
|
|
self._summaries.pop(path, None)
|
|
self.summaries_version += 1
|
|
|
|
def _load_summaries(self) -> dict[str, dict]:
|
|
"""Decode-and-cache every stored summary (call with the lock held)."""
|
|
if self._summaries is None:
|
|
out: dict[str, dict] = {}
|
|
rows = self._db.execute(
|
|
"SELECT path, summary_json FROM transcripts "
|
|
"WHERE summary_json IS NOT NULL").fetchall()
|
|
for r in rows:
|
|
try:
|
|
out[r["path"]] = json.loads(r["summary_json"])
|
|
except (ValueError, TypeError):
|
|
continue
|
|
self._summaries = out
|
|
return self._summaries
|
|
|
|
def all_summaries(self) -> list[tuple[str, dict]]:
|
|
"""(abs_path, summary dict) for every transcript that has one.
|
|
|
|
Served from the decoded cache — callers must treat the dicts as
|
|
**read-only** (copy before mutating)."""
|
|
with self._lock:
|
|
return list(self._load_summaries().items())
|
|
|
|
def get_summary(self, path: str) -> dict | None:
|
|
"""The parsed summary dict for one transcript (by abs path), or None.
|
|
|
|
Read-only — same shared cache as :meth:`all_summaries`."""
|
|
with self._lock:
|
|
return self._load_summaries().get(path)
|
|
|
|
def prune_transcripts(self, keep_paths: set[str]) -> None:
|
|
with self._lock:
|
|
existing = {r["path"] for r in
|
|
self._db.execute("SELECT path FROM transcripts")}
|
|
for gone in existing - keep_paths:
|
|
self._db.execute("DELETE FROM transcripts WHERE path=?", (gone,))
|
|
if self._summaries is not None:
|
|
self._summaries.pop(gone, None)
|
|
self._db.commit()
|
|
self.summaries_version += 1
|
|
|
|
def replace_skills(self, skills: dict[str, dict]) -> None:
|
|
"""skills: {name: {"count": int, "last_used": iso|None}}"""
|
|
with self._lock:
|
|
self._db.execute("DELETE FROM skills")
|
|
self._db.executemany(
|
|
"INSERT INTO skills (name, count, last_used) VALUES (?,?,?)",
|
|
[(n, d["count"], d.get("last_used")) for n, d in skills.items()],
|
|
)
|
|
self._db.commit()
|
|
|
|
def all_skills(self) -> list[sqlite3.Row]:
|
|
with self._lock:
|
|
return list(self._db.execute(
|
|
"SELECT * FROM skills ORDER BY count DESC, name ASC"))
|
|
|
|
# ── conversation diff cache ──────────────────────────────────────────
|
|
def get_conversation_diff(self, session_id: str) -> dict | None:
|
|
"""Cached {fingerprint, data} for a conversation, or None."""
|
|
with self._lock:
|
|
row = self._db.execute(
|
|
"SELECT fingerprint, data_json FROM conversation_diffs "
|
|
"WHERE session_id=?", (session_id,)).fetchone()
|
|
if not row or not row["data_json"]:
|
|
return None
|
|
try:
|
|
return {"fingerprint": row["fingerprint"],
|
|
"data": json.loads(row["data_json"])}
|
|
except (ValueError, TypeError):
|
|
return None
|
|
|
|
def set_conversation_diff(self, session_id: str, fingerprint: str,
|
|
data: dict) -> None:
|
|
with self._lock:
|
|
self._db.execute(
|
|
"INSERT INTO conversation_diffs "
|
|
"(session_id, fingerprint, data_json, computed_at) VALUES (?,?,?,?) "
|
|
"ON CONFLICT(session_id) DO UPDATE SET "
|
|
"fingerprint=excluded.fingerprint, data_json=excluded.data_json, "
|
|
"computed_at=excluded.computed_at",
|
|
(session_id, fingerprint, json.dumps(data), time.time()),
|
|
)
|
|
self._db.commit()
|
|
|
|
# ── meta ─────────────────────────────────────────────────────────────
|
|
def get_meta(self, key: str) -> str | None:
|
|
with self._lock:
|
|
row = self._db.execute(
|
|
"SELECT value FROM meta WHERE key=?", (key,)).fetchone()
|
|
return row["value"] if row else None
|
|
|
|
def set_meta(self, key: str, value: str) -> None:
|
|
with self._lock:
|
|
self._db.execute(
|
|
"INSERT INTO meta (key, value) VALUES (?,?) "
|
|
"ON CONFLICT(key) DO UPDATE SET value=excluded.value",
|
|
(key, value),
|
|
)
|
|
self._db.commit()
|