286 lines
10 KiB
JavaScript
286 lines
10 KiB
JavaScript
/**
|
|
* The SQLite index.
|
|
*
|
|
* `node:sqlite` rather than a dependency: the plugin ships with no build step
|
|
* and no install-time code, so every package it imports is one the user has to
|
|
* trust and keep current. Node has carried a synchronous SQLite binding since
|
|
* 22.5, and a synchronous API is genuinely the right shape here — the work is
|
|
* local, short, and on the critical path of a tool call, so an async wrapper
|
|
* would only add interleaving to reason about.
|
|
*
|
|
* **Why vectors are stored and not merely ranked at write time.** A search must
|
|
* be able to answer a query it has never seen, so the index has to keep the
|
|
* vectors themselves. They are written as unit-length `Float32Array` blobs, so
|
|
* scoring is a dot product and nothing else.
|
|
*
|
|
* **Why the model and width are recorded.** Vectors from two different models,
|
|
* or two widths of the same model, are not comparable, and mixing them produces
|
|
* plausible-looking but meaningless ranking. A change to either is detected on
|
|
* open and the index is discarded rather than silently degraded.
|
|
*/
|
|
|
|
import { mkdirSync } from 'node:fs'
|
|
import { dirname } from 'node:path'
|
|
import { DatabaseSync } from 'node:sqlite'
|
|
|
|
/** Schema version. A bump discards existing indexes. */
|
|
const SCHEMA_VERSION = '1'
|
|
|
|
const SCHEMA = `
|
|
CREATE TABLE IF NOT EXISTS meta (
|
|
key TEXT PRIMARY KEY,
|
|
value TEXT NOT NULL
|
|
);
|
|
CREATE TABLE IF NOT EXISTS files (
|
|
path TEXT PRIMARY KEY,
|
|
hash TEXT NOT NULL,
|
|
bytes INTEGER NOT NULL,
|
|
chunks INTEGER NOT NULL,
|
|
indexed_at INTEGER NOT NULL
|
|
);
|
|
CREATE TABLE IF NOT EXISTS chunks (
|
|
id INTEGER PRIMARY KEY AUTOINCREMENT,
|
|
path TEXT NOT NULL,
|
|
start_line INTEGER NOT NULL,
|
|
end_line INTEGER NOT NULL,
|
|
symbol TEXT,
|
|
text TEXT NOT NULL,
|
|
dim INTEGER NOT NULL,
|
|
vec BLOB NOT NULL
|
|
);
|
|
CREATE INDEX IF NOT EXISTS chunks_path_idx ON chunks(path);
|
|
`
|
|
|
|
/** One indexed chunk, metadata only. */
|
|
/** @typedef {{id: number, path: string, startLine: number, endLine: number, symbol?: string, text: string}} ChunkRow */
|
|
|
|
export class RecallIndex {
|
|
/** @param {string} file absolute path to the index database */
|
|
constructor(file) {
|
|
this.file = file
|
|
/** @type {DatabaseSync | undefined} */
|
|
this.db = undefined
|
|
/** Cached scoring matrix, dropped on every write. */
|
|
this.cached = undefined
|
|
}
|
|
|
|
/**
|
|
* Open the database, creating it when absent.
|
|
*
|
|
* @param {{model: string, dimensions: number, apiBase: string}} identity
|
|
* the embedding identity this index is only valid under
|
|
* @returns {{discarded?: string}} what was thrown away, when the identity moved
|
|
*/
|
|
open(identity) {
|
|
mkdirSync(dirname(this.file), { recursive: true })
|
|
this.db = new DatabaseSync(this.file)
|
|
// WAL keeps a search from blocking a concurrent sync; the index is a cache
|
|
// and losing the tail of a write on power loss is acceptable.
|
|
this.db.exec('PRAGMA journal_mode = WAL')
|
|
this.db.exec('PRAGMA synchronous = NORMAL')
|
|
this.db.exec(SCHEMA)
|
|
return this.recheck(identity)
|
|
}
|
|
|
|
/**
|
|
* Re-validate the stored embedding identity against a live settings snapshot.
|
|
*
|
|
* Called on every tool call, not just on open, because the settings namespace
|
|
* is editable at runtime: pointing the plugin at a different model must
|
|
* invalidate the existing vectors rather than silently mix two incompatible
|
|
* spaces into one ranking.
|
|
*
|
|
* @param {{model: string, dimensions: number, apiBase: string}} identity
|
|
* @returns {{discarded?: string}} what was thrown away, when the identity moved
|
|
*/
|
|
recheck(identity) {
|
|
const previous = {
|
|
version: this.#meta('schema_version'),
|
|
model: this.#meta('model'),
|
|
dimensions: this.#meta('dimensions'),
|
|
}
|
|
|
|
let discarded
|
|
if (previous.version !== undefined && previous.version !== SCHEMA_VERSION) {
|
|
discarded = `schema version changed from ${previous.version} to ${SCHEMA_VERSION}`
|
|
} else if (previous.model !== undefined && previous.model !== identity.model) {
|
|
discarded = `embedding model changed from "${previous.model}" to "${identity.model}"`
|
|
} else if (previous.dimensions !== undefined && previous.dimensions !== String(identity.dimensions)) {
|
|
discarded = `embedding width changed from ${previous.dimensions} to ${identity.dimensions}`
|
|
}
|
|
|
|
if (discarded !== undefined) this.#wipe()
|
|
|
|
this.#setMeta('schema_version', SCHEMA_VERSION)
|
|
this.#setMeta('model', identity.model)
|
|
this.#setMeta('dimensions', String(identity.dimensions))
|
|
this.#setMeta('api_base', identity.apiBase)
|
|
return discarded === undefined ? {} : { discarded }
|
|
}
|
|
|
|
/** Close the database, if open. Safe to call twice. */
|
|
close() {
|
|
try {
|
|
this.db?.close()
|
|
} catch {
|
|
// Already closed, or the handle was lost with the process. Nothing to do.
|
|
}
|
|
this.db = undefined
|
|
this.cached = undefined
|
|
}
|
|
|
|
/** Drop every row and reset the identity. */
|
|
#wipe() {
|
|
this.db.exec('DELETE FROM chunks; DELETE FROM files; DELETE FROM meta;')
|
|
this.cached = undefined
|
|
}
|
|
|
|
/** One meta value, or `undefined`. */
|
|
#meta(key) {
|
|
const row = this.db.prepare('SELECT value FROM meta WHERE key = ?').get(key)
|
|
return row === undefined ? undefined : row.value
|
|
}
|
|
|
|
/** Write one meta value. */
|
|
#setMeta(key, value) {
|
|
this.db
|
|
.prepare('INSERT INTO meta (key, value) VALUES (?, ?) ON CONFLICT(key) DO UPDATE SET value = excluded.value')
|
|
.run(key, value)
|
|
}
|
|
|
|
/** `path -> content hash` for every indexed file. */
|
|
fileHashes() {
|
|
const out = new Map()
|
|
for (const row of this.db.prepare('SELECT path, hash FROM files').all()) out.set(row.path, row.hash)
|
|
return out
|
|
}
|
|
|
|
/**
|
|
* Byte size of each named file, for the files that are indexed.
|
|
*
|
|
* This is what makes an honest "returned less than it would have read" figure
|
|
* possible: the index already recorded every file's size when it embedded it,
|
|
* so the alternative to an excerpt can be stated without reading a single
|
|
* source file at query time.
|
|
*
|
|
* @param {readonly string[]} paths workspace-relative paths
|
|
* @returns {Map<string, number>} only the paths that are indexed
|
|
*/
|
|
bytesFor(paths) {
|
|
const out = new Map()
|
|
if (paths.length === 0) return out
|
|
const statement = this.db.prepare('SELECT bytes FROM files WHERE path = ?')
|
|
for (const path of paths) {
|
|
const row = statement.get(path)
|
|
if (row !== undefined) out.set(path, row.bytes)
|
|
}
|
|
return out
|
|
}
|
|
|
|
/** Counts for the status surface. */
|
|
stats() {
|
|
const files = this.db.prepare('SELECT COUNT(*) AS n FROM files').get()?.n ?? 0
|
|
const chunks = this.db.prepare('SELECT COUNT(*) AS n FROM chunks').get()?.n ?? 0
|
|
const bytes = this.db.prepare('SELECT COALESCE(SUM(bytes), 0) AS n FROM files').get()?.n ?? 0
|
|
return { files, chunks, bytes }
|
|
}
|
|
|
|
/**
|
|
* Replace one file's chunks atomically.
|
|
*
|
|
* @param {string} path workspace-relative path
|
|
* @param {string} hash content hash
|
|
* @param {number} bytes file size
|
|
* @param {{startLine: number, endLine: number, symbol?: string, text: string, vec: Float32Array}[]} chunks
|
|
*/
|
|
replaceFile(path, hash, bytes, chunks) {
|
|
this.db.exec('BEGIN IMMEDIATE')
|
|
try {
|
|
this.db.prepare('DELETE FROM chunks WHERE path = ?').run(path)
|
|
const insert = this.db.prepare(
|
|
'INSERT INTO chunks (path, start_line, end_line, symbol, text, dim, vec) VALUES (?, ?, ?, ?, ?, ?, ?)',
|
|
)
|
|
for (const chunk of chunks) {
|
|
insert.run(
|
|
path,
|
|
chunk.startLine,
|
|
chunk.endLine,
|
|
chunk.symbol ?? null,
|
|
chunk.text,
|
|
chunk.vec.length,
|
|
new Uint8Array(chunk.vec.buffer, chunk.vec.byteOffset, chunk.vec.byteLength),
|
|
)
|
|
}
|
|
this.db
|
|
.prepare(
|
|
`INSERT INTO files (path, hash, bytes, chunks, indexed_at) VALUES (?, ?, ?, ?, ?)
|
|
ON CONFLICT(path) DO UPDATE SET
|
|
hash = excluded.hash, bytes = excluded.bytes,
|
|
chunks = excluded.chunks, indexed_at = excluded.indexed_at`,
|
|
)
|
|
.run(path, hash, bytes, chunks.length, Date.now())
|
|
this.db.exec('COMMIT')
|
|
} catch (error) {
|
|
this.db.exec('ROLLBACK')
|
|
throw error
|
|
}
|
|
this.cached = undefined
|
|
}
|
|
|
|
/** Forget files that no longer exist on disk. Returns how many were dropped. */
|
|
prune(toRemove) {
|
|
if (toRemove.length === 0) return 0
|
|
this.db.exec('BEGIN IMMEDIATE')
|
|
try {
|
|
const dropChunks = this.db.prepare('DELETE FROM chunks WHERE path = ?')
|
|
const dropFile = this.db.prepare('DELETE FROM files WHERE path = ?')
|
|
for (const path of toRemove) {
|
|
dropChunks.run(path)
|
|
dropFile.run(path)
|
|
}
|
|
this.db.exec('COMMIT')
|
|
} catch (error) {
|
|
this.db.exec('ROLLBACK')
|
|
throw error
|
|
}
|
|
this.cached = undefined
|
|
return toRemove.length
|
|
}
|
|
|
|
/**
|
|
* Load every chunk for scoring.
|
|
*
|
|
* One contiguous `Float32Array` and a parallel row list, built once and reused
|
|
* until the next write. At 300k lines of source this is a few tens of
|
|
* megabytes — the alternative, a blob read per query, would pay SQLite's row
|
|
* overhead on every search for no benefit, since a search visits every row
|
|
* anyway.
|
|
*
|
|
* @returns {{vectors: Float32Array, rows: ChunkRow[], dimensions: number}}
|
|
*/
|
|
matrix() {
|
|
if (this.cached !== undefined) return this.cached
|
|
const raw = this.db
|
|
.prepare('SELECT id, path, start_line, end_line, symbol, text, dim, vec FROM chunks ORDER BY id')
|
|
.all()
|
|
const dimensions = raw.length === 0 ? 0 : raw[0].dim
|
|
const vectors = new Float32Array(raw.length * dimensions)
|
|
const rows = new Array(raw.length)
|
|
for (let at = 0; at < raw.length; at += 1) {
|
|
const row = raw[at]
|
|
const view = new Float32Array(row.vec.buffer, row.vec.byteOffset, row.vec.byteLength / 4)
|
|
vectors.set(view, at * dimensions)
|
|
rows[at] = {
|
|
id: row.id,
|
|
path: row.path,
|
|
startLine: row.start_line,
|
|
endLine: row.end_line,
|
|
...(row.symbol === null ? {} : { symbol: row.symbol }),
|
|
text: row.text,
|
|
}
|
|
}
|
|
this.cached = { vectors, rows, dimensions }
|
|
return this.cached
|
|
}
|
|
}
|