nanobot: 2026-09-10 12:33:37
This commit is contained in:
lachtan
280
skills/wiki/scripts/wiki_store.py
Normal file
280
skills/wiki/scripts/wiki_store.py
Normal file
@@ -0,0 +1,280 @@
|
||||
#!/usr/bin/env python3
|
||||
# /// script
|
||||
# requires-python = ">=3.11"
|
||||
# dependencies = ["sqlite-vec==0.1.6"]
|
||||
# ///
|
||||
"""Data-access layer for the wiki skill.
|
||||
|
||||
Pure SQL plus lifecycle helpers. No printing, no argparse, no sys.exit.
|
||||
|
||||
`vec_chunks` is a vec0 virtual table and therefore NOT reachable by the foreign key
|
||||
cascade that cleans up `chunks` and `chunks_fts`. Every path that removes chunks must
|
||||
delete their vectors explicitly — that is why the deletes here go through
|
||||
`_delete_vectors_for_file`.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import json
|
||||
import sqlite3
|
||||
from collections.abc import Iterator, Sequence
|
||||
from contextlib import contextmanager
|
||||
from pathlib import Path
|
||||
|
||||
from sqlite_vec import serialize_float32 # ty: ignore[unresolved-import]
|
||||
from wiki_db import get_db, init_db
|
||||
|
||||
|
||||
@contextmanager
|
||||
def connection(db_path: Path) -> Iterator[sqlite3.Connection]:
|
||||
"""Open a connection, initialising the DB if missing."""
|
||||
if not db_path.exists():
|
||||
init_db(db_path)
|
||||
conn = get_db(db_path)
|
||||
try:
|
||||
yield conn
|
||||
finally:
|
||||
conn.close()
|
||||
|
||||
|
||||
@contextmanager
|
||||
def tx(conn: sqlite3.Connection) -> Iterator[sqlite3.Connection]:
|
||||
"""Wrap an already-open connection in an explicit transaction.
|
||||
|
||||
The connection is in autocommit mode (`isolation_level=None`), so transactions are
|
||||
ours to open — sqlite3's implicit handling would otherwise fail a nested BEGIN.
|
||||
"""
|
||||
conn.execute("BEGIN")
|
||||
try:
|
||||
yield conn
|
||||
conn.execute("COMMIT")
|
||||
except Exception:
|
||||
conn.execute("ROLLBACK")
|
||||
raise
|
||||
|
||||
|
||||
@contextmanager
|
||||
def transaction(db_path: Path) -> Iterator[sqlite3.Connection]:
|
||||
"""Open a connection wrapped in an explicit transaction."""
|
||||
with connection(db_path) as conn, tx(conn):
|
||||
yield conn
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# meta — identity of the embedding space
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
def read_meta(conn: sqlite3.Connection) -> dict[str, str]:
|
||||
return {row["key"]: row["value"] for row in conn.execute("SELECT key, value FROM meta")}
|
||||
|
||||
|
||||
def write_meta(conn: sqlite3.Connection, values: dict[str, str]) -> None:
|
||||
conn.executemany(
|
||||
"INSERT INTO meta (key, value) VALUES (?, ?) ON CONFLICT(key) DO UPDATE SET value = excluded.value",
|
||||
sorted(values.items()),
|
||||
)
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# sources
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
def upsert_source(conn: sqlite3.Connection, source_id: str, kind: str) -> None:
|
||||
conn.execute(
|
||||
"INSERT INTO sources (source_id, kind) VALUES (?, ?) ON CONFLICT(source_id) DO UPDATE SET kind = excluded.kind",
|
||||
(source_id, kind),
|
||||
)
|
||||
|
||||
|
||||
def get_source(conn: sqlite3.Connection, source_id: str) -> sqlite3.Row | None:
|
||||
return conn.execute("SELECT * FROM sources WHERE source_id = ?", (source_id,)).fetchone()
|
||||
|
||||
|
||||
def mark_synced(conn: sqlite3.Connection, source_id: str, indexed_rev: str | None, now: str) -> None:
|
||||
conn.execute(
|
||||
"UPDATE sources SET indexed_rev = ?, last_sync_at = ? WHERE source_id = ?",
|
||||
(indexed_rev, now, source_id),
|
||||
)
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# files
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
def get_file(conn: sqlite3.Connection, source_id: str, path: str) -> sqlite3.Row | None:
|
||||
return conn.execute("SELECT * FROM files WHERE source_id = ? AND path = ?", (source_id, path)).fetchone()
|
||||
|
||||
|
||||
def list_source_files(conn: sqlite3.Connection, source_id: str) -> dict[str, sqlite3.Row]:
|
||||
"""Indexed files of one source, keyed by path — the basis for deletion detection."""
|
||||
rows = conn.execute("SELECT * FROM files WHERE source_id = ?", (source_id,))
|
||||
return {row["path"]: row for row in rows}
|
||||
|
||||
|
||||
def upsert_file(
|
||||
conn: sqlite3.Connection,
|
||||
source_id: str,
|
||||
path: str,
|
||||
title: str | None,
|
||||
tags: list[str],
|
||||
headings: list[str],
|
||||
sha256: str,
|
||||
size: int,
|
||||
mtime: float | None,
|
||||
now: str,
|
||||
) -> None:
|
||||
conn.execute(
|
||||
"INSERT INTO files (source_id, path, title, tags, headings, sha256, size, mtime, indexed_at) "
|
||||
"VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?) "
|
||||
"ON CONFLICT(source_id, path) DO UPDATE SET "
|
||||
"title = excluded.title, tags = excluded.tags, headings = excluded.headings, "
|
||||
"sha256 = excluded.sha256, size = excluded.size, mtime = excluded.mtime, "
|
||||
"indexed_at = excluded.indexed_at",
|
||||
(
|
||||
source_id,
|
||||
path,
|
||||
title,
|
||||
json.dumps(tags, ensure_ascii=False),
|
||||
json.dumps(headings, ensure_ascii=False),
|
||||
sha256,
|
||||
size,
|
||||
mtime,
|
||||
now,
|
||||
),
|
||||
)
|
||||
|
||||
|
||||
def touch_file_stat(conn: sqlite3.Connection, source_id: str, path: str, size: int, mtime: float | None) -> None:
|
||||
"""Refresh the cheap change-detection stats after a sha256 match (content unchanged)."""
|
||||
conn.execute(
|
||||
"UPDATE files SET size = ?, mtime = ? WHERE source_id = ? AND path = ?",
|
||||
(size, mtime, source_id, path),
|
||||
)
|
||||
|
||||
|
||||
def delete_file(conn: sqlite3.Connection, source_id: str, path: str) -> None:
|
||||
"""Drop a file and everything derived from it, vec0 rows included."""
|
||||
_delete_vectors_for_file(conn, source_id, path)
|
||||
conn.execute("DELETE FROM files WHERE source_id = ? AND path = ?", (source_id, path))
|
||||
|
||||
|
||||
def list_toc_files(conn: sqlite3.Connection, source_id: str | None = None, tag: str | None = None) -> list[sqlite3.Row]:
|
||||
sql = "SELECT source_id, path, title, tags FROM files"
|
||||
clauses: list[str] = []
|
||||
params: list[str] = []
|
||||
if source_id:
|
||||
clauses.append("source_id = ?")
|
||||
params.append(source_id)
|
||||
if tag:
|
||||
clauses.append("EXISTS (SELECT 1 FROM json_each(files.tags) WHERE value = ?)")
|
||||
params.append(tag)
|
||||
if clauses:
|
||||
sql += " WHERE " + " AND ".join(clauses)
|
||||
sql += " ORDER BY source_id, path"
|
||||
return list(conn.execute(sql, params))
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# chunks
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
def _delete_vectors_for_file(conn: sqlite3.Connection, source_id: str, path: str) -> None:
|
||||
ids = [
|
||||
row["id"] for row in conn.execute("SELECT id FROM chunks WHERE source_id = ? AND path = ?", (source_id, path))
|
||||
]
|
||||
conn.executemany("DELETE FROM vec_chunks WHERE rowid = ?", [(i,) for i in ids])
|
||||
|
||||
|
||||
def replace_chunks(conn: sqlite3.Connection, source_id: str, path: str, chunks: Sequence[tuple[str, str]]) -> None:
|
||||
"""Swap a file's chunks for a freshly built set, as (breadcrumb, text) in order.
|
||||
|
||||
New chunks land with `embedded_at IS NULL`; the embed pass picks them up, so an
|
||||
unreachable Ollama degrades to FTS-only instead of failing the sync.
|
||||
"""
|
||||
_delete_vectors_for_file(conn, source_id, path)
|
||||
conn.execute("DELETE FROM chunks WHERE source_id = ? AND path = ?", (source_id, path))
|
||||
conn.executemany(
|
||||
"INSERT INTO chunks (source_id, path, chunk_idx, breadcrumb, text) VALUES (?, ?, ?, ?, ?)",
|
||||
[(source_id, path, idx, breadcrumb, text) for idx, (breadcrumb, text) in enumerate(chunks)],
|
||||
)
|
||||
|
||||
|
||||
def pending_chunks(conn: sqlite3.Connection, limit: int) -> list[sqlite3.Row]:
|
||||
return list(
|
||||
conn.execute(
|
||||
"SELECT id, breadcrumb, text FROM chunks WHERE embedded_at IS NULL ORDER BY id LIMIT ?",
|
||||
(limit,),
|
||||
)
|
||||
)
|
||||
|
||||
|
||||
def has_pending(conn: sqlite3.Connection) -> bool:
|
||||
return conn.execute("SELECT 1 FROM chunks WHERE embedded_at IS NULL LIMIT 1").fetchone() is not None
|
||||
|
||||
|
||||
def store_embedding(conn: sqlite3.Connection, chunk_id: int, vector: Sequence[float], now: str) -> None:
|
||||
conn.execute("DELETE FROM vec_chunks WHERE rowid = ?", (chunk_id,))
|
||||
conn.execute(
|
||||
"INSERT INTO vec_chunks (rowid, embedding) VALUES (?, ?)",
|
||||
(chunk_id, serialize_float32(list(vector))),
|
||||
)
|
||||
conn.execute("UPDATE chunks SET embedded_at = ? WHERE id = ?", (now, chunk_id))
|
||||
|
||||
|
||||
def reset_index(conn: sqlite3.Connection) -> None:
|
||||
"""Drop every indexed artifact, keeping the source rows. Used by `--full`."""
|
||||
ids = [row["id"] for row in conn.execute("SELECT id FROM chunks")]
|
||||
conn.executemany("DELETE FROM vec_chunks WHERE rowid = ?", [(i,) for i in ids])
|
||||
conn.execute("DELETE FROM files")
|
||||
conn.execute("UPDATE sources SET indexed_rev = NULL")
|
||||
|
||||
|
||||
def index_stats(conn: sqlite3.Connection) -> dict[str, int]:
|
||||
return {
|
||||
"files": conn.execute("SELECT count(*) AS n FROM files").fetchone()["n"],
|
||||
"chunks": conn.execute("SELECT count(*) AS n FROM chunks").fetchone()["n"],
|
||||
"vectors": conn.execute("SELECT count(*) AS n FROM vec_chunks").fetchone()["n"],
|
||||
"pending": conn.execute("SELECT count(*) AS n FROM chunks WHERE embedded_at IS NULL").fetchone()["n"],
|
||||
}
|
||||
|
||||
|
||||
def orphan_vector_ids(conn: sqlite3.Connection) -> list[int]:
|
||||
"""Vector rowids with no surviving chunk — must always be empty (regression guard)."""
|
||||
rows = conn.execute("SELECT rowid AS rid FROM vec_chunks WHERE rowid NOT IN (SELECT id FROM chunks)")
|
||||
return [row["rid"] for row in rows]
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# retrieval
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
def bm25_ranked_ids(conn: sqlite3.Connection, match_expr: str, limit: int) -> list[int]:
|
||||
rows = conn.execute(
|
||||
"SELECT rowid AS rid FROM chunks_fts WHERE chunks_fts MATCH ? ORDER BY rank LIMIT ?",
|
||||
(match_expr, limit),
|
||||
)
|
||||
return [row["rid"] for row in rows]
|
||||
|
||||
|
||||
def knn_ranked_ids(conn: sqlite3.Connection, vector: Sequence[float], k: int) -> list[int]:
|
||||
rows = conn.execute(
|
||||
"SELECT rowid AS rid FROM vec_chunks WHERE embedding MATCH ? AND k = ? ORDER BY distance",
|
||||
(serialize_float32(list(vector)), k),
|
||||
)
|
||||
return [row["rid"] for row in rows]
|
||||
|
||||
|
||||
def fetch_chunks(conn: sqlite3.Connection, ids: Sequence[int]) -> dict[int, sqlite3.Row]:
|
||||
if not ids:
|
||||
return {}
|
||||
placeholders = ",".join("?" * len(ids))
|
||||
rows = conn.execute(
|
||||
f"SELECT id, source_id, path, chunk_idx, breadcrumb, text FROM chunks WHERE id IN ({placeholders})",
|
||||
list(ids),
|
||||
)
|
||||
return {row["id"]: row for row in rows}
|
||||
Reference in New Issue
Block a user