90 lines
4.4 KiB
Python
90 lines
4.4 KiB
Python
"""Persistent vec0 indexes derived from durable routed vectors, one per space/dimension."""
|
|
import hashlib
|
|
import json
|
|
import threading
|
|
|
|
import sqlite_vec
|
|
|
|
from app.retrieval.vectorstore import VectorHit
|
|
|
|
|
|
_migration_lock = threading.Lock()
|
|
|
|
|
|
def is_ready(conn, batches):
|
|
return all(conn.execute('SELECT 1 FROM sqlite_master WHERE name=?',
|
|
(table_name(batch.space_id, batch.dimensions),)).fetchone() for batch in batches)
|
|
|
|
|
|
def prepare(conn, batches):
|
|
"""Finish lazy writes before opening a search snapshot. Warm searches do not write."""
|
|
from app.retrieval.routed_vectors import _ensure_table
|
|
batches = list(batches)
|
|
if is_ready(conn, batches):
|
|
return
|
|
# Waiting holds no read transaction, so a concurrent migration can commit.
|
|
with _migration_lock:
|
|
if is_ready(conn, batches):
|
|
return
|
|
conn.execute('BEGIN IMMEDIATE')
|
|
try:
|
|
_ensure_table(conn)
|
|
for batch in batches:
|
|
ensure(conn, batch.space_id, batch.dimensions)
|
|
conn.execute('COMMIT')
|
|
except BaseException:
|
|
conn.execute('ROLLBACK')
|
|
raise
|
|
|
|
|
|
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 = table_name(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]
|