130 lines
4.0 KiB
Python
130 lines
4.0 KiB
Python
#!/usr/bin/env python3
|
|
# /// script
|
|
# requires-python = ">=3.11"
|
|
# dependencies = ["sqlite-vec==0.1.6"]
|
|
# ///
|
|
"""SQLite storage layer for the wiki skill.
|
|
|
|
Schema, connection factory and the sqlite-vec extension load.
|
|
|
|
`chunks` is the canonical retrieval unit: `chunks_fts` gives it a BM25 rank and
|
|
`vec_chunks` a KNN rank, both keyed on `chunks.id`, so RRF merges two rankings of
|
|
the *same* set.
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
import sqlite3
|
|
from pathlib import Path
|
|
|
|
import sqlite_vec # ty: ignore[unresolved-import]
|
|
|
|
# Compiled into the vec0 table definition. `wiki_embed` guards config.dims against it —
|
|
# a mismatch means the index was built for a different model.
|
|
EMBEDDING_DIMS = 1024
|
|
|
|
SCHEMA = f"""
|
|
PRAGMA journal_mode = WAL;
|
|
PRAGMA foreign_keys = ON;
|
|
|
|
CREATE TABLE IF NOT EXISTS meta (
|
|
key TEXT PRIMARY KEY,
|
|
value TEXT NOT NULL
|
|
);
|
|
|
|
CREATE TABLE IF NOT EXISTS sources (
|
|
source_id TEXT PRIMARY KEY,
|
|
kind TEXT NOT NULL CHECK(kind IN ('git', 'workspace')),
|
|
indexed_rev TEXT,
|
|
last_sync_at TEXT
|
|
);
|
|
|
|
CREATE TABLE IF NOT EXISTS files (
|
|
source_id TEXT NOT NULL REFERENCES sources(source_id),
|
|
path TEXT NOT NULL,
|
|
title TEXT,
|
|
tags TEXT,
|
|
headings TEXT,
|
|
sha256 TEXT NOT NULL,
|
|
size INTEGER NOT NULL,
|
|
mtime REAL,
|
|
indexed_at TEXT NOT NULL,
|
|
PRIMARY KEY (source_id, path)
|
|
);
|
|
|
|
CREATE TABLE IF NOT EXISTS chunks (
|
|
id INTEGER PRIMARY KEY AUTOINCREMENT,
|
|
source_id TEXT NOT NULL,
|
|
path TEXT NOT NULL,
|
|
chunk_idx INTEGER NOT NULL,
|
|
breadcrumb TEXT NOT NULL,
|
|
text TEXT NOT NULL,
|
|
embedded_at TEXT,
|
|
UNIQUE (source_id, path, chunk_idx),
|
|
FOREIGN KEY (source_id, path) REFERENCES files(source_id, path) ON DELETE CASCADE
|
|
);
|
|
|
|
CREATE INDEX IF NOT EXISTS idx_chunks_pending ON chunks(id) WHERE embedded_at IS NULL;
|
|
|
|
CREATE VIRTUAL TABLE IF NOT EXISTS chunks_fts USING fts5(
|
|
breadcrumb, text,
|
|
content='chunks', content_rowid='id',
|
|
tokenize='unicode61 remove_diacritics 2'
|
|
);
|
|
|
|
CREATE TRIGGER IF NOT EXISTS chunks_ai AFTER INSERT ON chunks BEGIN
|
|
INSERT INTO chunks_fts(rowid, breadcrumb, text) VALUES (new.id, new.breadcrumb, new.text);
|
|
END;
|
|
|
|
CREATE TRIGGER IF NOT EXISTS chunks_ad AFTER DELETE ON chunks BEGIN
|
|
INSERT INTO chunks_fts(chunks_fts, rowid, breadcrumb, text)
|
|
VALUES('delete', old.id, old.breadcrumb, old.text);
|
|
END;
|
|
|
|
-- The WHEN guard keeps `UPDATE chunks SET embedded_at` (sync step 4b) from rewriting
|
|
-- an FTS row whose indexed text did not change.
|
|
CREATE TRIGGER IF NOT EXISTS chunks_au AFTER UPDATE ON chunks
|
|
WHEN old.text IS NOT new.text OR old.breadcrumb IS NOT new.breadcrumb
|
|
BEGIN
|
|
INSERT INTO chunks_fts(chunks_fts, rowid, breadcrumb, text)
|
|
VALUES('delete', old.id, old.breadcrumb, old.text);
|
|
INSERT INTO chunks_fts(rowid, breadcrumb, text) VALUES (new.id, new.breadcrumb, new.text);
|
|
END;
|
|
|
|
CREATE VIRTUAL TABLE IF NOT EXISTS vec_chunks USING vec0(
|
|
embedding float[{EMBEDDING_DIMS}] distance_metric=cosine
|
|
);
|
|
"""
|
|
|
|
|
|
def get_db(path: Path) -> sqlite3.Connection:
|
|
"""Return an autocommit connection with sqlite-vec loaded and foreign keys on."""
|
|
conn = sqlite3.connect(path, isolation_level=None)
|
|
conn.enable_load_extension(True)
|
|
sqlite_vec.load(conn)
|
|
conn.enable_load_extension(False)
|
|
conn.execute("PRAGMA journal_mode = WAL")
|
|
conn.execute("PRAGMA foreign_keys = ON")
|
|
conn.row_factory = sqlite3.Row
|
|
_migrate(conn)
|
|
return conn
|
|
|
|
|
|
def _migrate(conn: sqlite3.Connection) -> None:
|
|
"""Idempotently bring an existing DB up to the current schema.
|
|
|
|
init_db only runs the full SCHEMA on a missing file, so live DBs never see
|
|
later additions. Each step must be a no-op once applied.
|
|
"""
|
|
# No migrations yet — schema_version 1 is the initial shape.
|
|
|
|
|
|
def init_db(path: Path) -> None:
|
|
"""Create tables, indexes and triggers if they don't exist."""
|
|
path.parent.mkdir(parents=True, exist_ok=True)
|
|
conn = get_db(path)
|
|
try:
|
|
conn.executescript(SCHEMA)
|
|
finally:
|
|
conn.close()
|