perf: persist model-isolated sqlite-vec indexes and index new notes incrementally

backend/app/retrieval/space_index.py

@@ -0,0 +1,59 @@
+"""Persistent vec0 indexes derived from durable routed vectors, one per space/dimension."""
+import hashlib
+import json
+
+import sqlite_vec
+
+from app.retrieval.vectorstore import VectorHit
+
+
+def table_name(space, dimensions):
+    return 'routed_vec_' + hashlib.sha256(json.dumps([space, dimensions]).encode()).hexdigest()
+
+
+def ensure(conn, space, dimensions):
+    from app.retrieval.routed_vectors import _unit_vector
+    table = table_name(space, dimensions)
+    if conn.execute('SELECT 1 FROM sqlite_master WHERE name=?', (table,)).fetchone():
+        return table
+    if type(dimensions) is not int or not 0 < dimensions <= 8192:
+        raise ValueError('unsupported vector dimensions')
+    conn.execute(f'CREATE VIRTUAL TABLE {table} USING vec0(block_id TEXT PRIMARY KEY, embedding float[{dimensions}], local_only INTEGER)')
+    for row in conn.execute('SELECT r.block_id,r.vector,b.embedding_local_only FROM routed_block_vectors r JOIN blocks b USING(block_id) WHERE r.space_id=? AND r.dimensions=?', (space, dimensions)):
+        conn.execute(f'INSERT INTO {table}(block_id,embedding,local_only) VALUES (?,?,?)',
+                     (row[0], sqlite_vec.serialize_float32(_unit_vector(json.loads(row[1]), dimensions)), row[2]))
+    literal = conn.execute('SELECT quote(?)', (space,)).fetchone()[0]
+    for event in ('DELETE', 'UPDATE'):
+        conn.execute(f'''CREATE TRIGGER {table}_{event.lower()} AFTER {event} ON routed_block_vectors
+            WHEN old.space_id={literal} AND old.dimensions={dimensions}
+            BEGIN DELETE FROM {table} WHERE block_id=old.block_id; END''')
+    return table
+
+
+def upsert(conn, block_ids, batch):
+    from app.retrieval.routed_vectors import _unit_vector
+    table = ensure(conn, batch.space_id, batch.dimensions)
+    for block_id, vector in zip(block_ids, batch.vectors):
+        conn.execute(f'DELETE FROM {table} WHERE block_id=?', (block_id,))
+        conn.execute(f'INSERT INTO {table}(block_id,embedding,local_only) SELECT block_id,?,embedding_local_only FROM blocks WHERE block_id=?',
+                     (sqlite_vec.serialize_float32(_unit_vector(vector, batch.dimensions)), block_id))
+
+
+def search(conn, batch, top_k, policy=None):
+    table = ensure(conn, batch.space_id, batch.dimensions)
+    # Coverage checks stay relational; no JSON decoding or Python dot products on the hot path.
+    where = '' if policy is None else ' AND b.embedding_local_only=?'
+    params = () if policy is None else (int(policy),)
+    missing = conn.execute(f'''SELECT 1 FROM blocks b LEFT JOIN routed_block_vectors r
+        ON r.block_id=b.block_id AND r.space_id=? AND r.dimensions=?
+        WHERE r.block_id IS NULL{where} LIMIT 1''', (batch.space_id, batch.dimensions, *params)).fetchone()
+    expected = conn.execute('SELECT COUNT(*) FROM blocks' + ('' if policy is None else ' WHERE embedding_local_only=?'), params).fetchone()[0]
+    actual = conn.execute(f'SELECT COUNT(*) FROM {table}' + ('' if policy is None else ' WHERE local_only=?'), params).fetchone()[0]
+    if missing or actual != expected:
+        raise ValueError('incomplete vector space coverage')
+    if top_k <= 0:
+        return []
+    rows = conn.execute(f'SELECT block_id,distance FROM {table} WHERE embedding MATCH ? AND k=?'
+                        + ('' if policy is None else ' AND local_only=?'),
+                        (sqlite_vec.serialize_float32(batch.vectors[0]), top_k, *params)).fetchall()
+    return [VectorHit(id=row[0], score=max(0.0, min(1.0, 1 - row[1] ** 2 / 2))) for row in rows]