281 lines
10 KiB
Python
281 lines
10 KiB
Python
#!/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}
|