perf: persist model-isolated sqlite-vec indexes and index new notes incrementally
This commit is contained in:
@@ -2,15 +2,14 @@
|
||||
|
||||
The runtime's model_id is the authoritative space ID (including provider URL,
|
||||
endpoint, model and dimensions); equal dimensions alone never imply compatibility.
|
||||
This phase uses a lazy, rebuildable SQLite side table instead of a schema migration.
|
||||
Search scans only current blocks in one database snapshot and requires complete
|
||||
coverage. Cosine ranking costs O(blocks * dimensions) with an O(top_k) heap; this
|
||||
small-vault implementation should become a per-space ANN index at larger scale.
|
||||
Durable vectors are reused to build per-space/dimension sqlite-vec indexes lazily.
|
||||
Native exact KNN avoids Python JSON decoding and dot products on every search.
|
||||
Coverage checks and ranking share one transaction.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import heapq
|
||||
import asyncio
|
||||
import json
|
||||
import logging
|
||||
import math
|
||||
@@ -24,6 +23,7 @@ from app.operation_logs import log_event
|
||||
from app.retrieval.vectorstore import VectorHit
|
||||
from app.retrieval.provenance import record_embedding
|
||||
from app.retrieval.hybrid import rrf_fuse
|
||||
from app.retrieval import space_index
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
@@ -121,9 +121,15 @@ def _ensure_table(conn: sqlite3.Connection) -> None:
|
||||
block_id TEXT NOT NULL REFERENCES blocks(block_id) ON DELETE CASCADE,
|
||||
dimensions INTEGER NOT NULL CHECK (dimensions > 0),
|
||||
vector TEXT NOT NULL,
|
||||
PRIMARY KEY (space_id, block_id)
|
||||
PRIMARY KEY (space_id, dimensions, block_id)
|
||||
)
|
||||
""")
|
||||
primary = [row[1] for row in sorted(conn.execute('PRAGMA table_info(routed_block_vectors)'), key=lambda row: row[5]) if row[5]]
|
||||
if primary == ['space_id', 'block_id']:
|
||||
conn.execute('CREATE TABLE routed_block_vectors_upgrade (space_id TEXT NOT NULL, block_id TEXT NOT NULL REFERENCES blocks(block_id) ON DELETE CASCADE, dimensions INTEGER NOT NULL CHECK(dimensions>0), vector TEXT NOT NULL, PRIMARY KEY(space_id,dimensions,block_id))')
|
||||
conn.execute('INSERT INTO routed_block_vectors_upgrade SELECT * FROM routed_block_vectors')
|
||||
conn.execute('DROP TABLE routed_block_vectors')
|
||||
conn.execute('ALTER TABLE routed_block_vectors_upgrade RENAME TO routed_block_vectors')
|
||||
conn.execute("""
|
||||
CREATE INDEX IF NOT EXISTS routed_block_vectors_block_id
|
||||
ON routed_block_vectors(block_id)
|
||||
@@ -149,13 +155,14 @@ def store_remote(
|
||||
conn.executemany(
|
||||
"""INSERT INTO routed_block_vectors (space_id, block_id, dimensions, vector)
|
||||
VALUES (?, ?, ?, ?)
|
||||
ON CONFLICT (space_id, block_id) DO UPDATE SET
|
||||
ON CONFLICT (space_id, dimensions, block_id) DO UPDATE SET
|
||||
dimensions = excluded.dimensions, vector = excluded.vector""",
|
||||
[
|
||||
(batch.space_id, block_id, batch.dimensions, json.dumps(vector, allow_nan=False))
|
||||
for block_id, vector in zip(block_ids, batch.vectors)
|
||||
],
|
||||
)
|
||||
space_index.upsert(conn, block_ids, batch)
|
||||
except BaseException:
|
||||
conn.execute("ROLLBACK TO routed_vectors_write")
|
||||
raise
|
||||
@@ -183,6 +190,10 @@ async def search_remote(query: str, *, top_k: int, accept_local=False, strict=Fa
|
||||
if batch is None:
|
||||
return None
|
||||
|
||||
return await asyncio.to_thread(_search_space, batch, top_k, strict)
|
||||
|
||||
|
||||
def _search_space(batch, top_k, strict):
|
||||
record_embedding(attempted_space={"model_id": batch.space_id, "dimensions": batch.dimensions})
|
||||
try:
|
||||
conn = connect()
|
||||
@@ -198,29 +209,8 @@ async def search_remote(query: str, *, top_k: int, accept_local=False, strict=Fa
|
||||
if strict:
|
||||
raise ValueError("semantic index missing")
|
||||
return None
|
||||
rows = conn.execute(
|
||||
"""SELECT b.block_id, r.vector
|
||||
FROM blocks AS b
|
||||
LEFT JOIN routed_block_vectors AS r
|
||||
ON r.block_id = b.block_id AND r.space_id = ? AND r.dimensions = ?
|
||||
ORDER BY b.block_id""",
|
||||
(batch.space_id, batch.dimensions),
|
||||
)
|
||||
|
||||
def hits():
|
||||
for row in rows:
|
||||
if row["vector"] is None:
|
||||
raise ValueError("remote space has incomplete block coverage")
|
||||
vector = _unit_vector(json.loads(row["vector"]), batch.dimensions)
|
||||
score = math.fsum(a * b for a, b in zip(batch.vectors[0], vector))
|
||||
yield VectorHit(id=row["block_id"], score=max(0.0, min(1.0, score)))
|
||||
|
||||
try:
|
||||
result = heapq.nlargest(top_k, hits(), key=lambda hit: hit.score)
|
||||
finally:
|
||||
# Exceptions may retain the generator/traceback; finalize its
|
||||
# cursor now so a subsequent rebuild can acquire a write lock.
|
||||
rows.close()
|
||||
_ensure_table(conn)
|
||||
result = space_index.search(conn, batch, top_k)
|
||||
record_embedding(source=batch.source, model_id=batch.space_id,
|
||||
dimensions=batch.dimensions, fallback_reason=None)
|
||||
return result
|
||||
@@ -244,6 +234,10 @@ async def _search_partitioned(query: str, policies: set[bool], *, top_k: int, st
|
||||
if batch is None:
|
||||
return None
|
||||
batches[policy] = batch
|
||||
return await asyncio.to_thread(_search_partitions, batches, policies, top_k, strict)
|
||||
|
||||
|
||||
def _search_partitions(batches, policies, top_k, strict):
|
||||
conn = connect()
|
||||
try:
|
||||
with transaction(conn):
|
||||
@@ -253,23 +247,8 @@ async def _search_partitioned(query: str, policies: set[bool], *, top_k: int, st
|
||||
raise ValueError("embedding policies changed while querying")
|
||||
ranked = []
|
||||
for policy, batch in batches.items():
|
||||
rows = conn.execute(
|
||||
"SELECT b.block_id,r.vector FROM blocks b LEFT JOIN routed_block_vectors r "
|
||||
"ON r.block_id=b.block_id AND r.space_id=? AND r.dimensions=? "
|
||||
"WHERE b.embedding_local_only=? ORDER BY b.block_id",
|
||||
(batch.space_id, batch.dimensions, int(policy)),
|
||||
)
|
||||
def hits():
|
||||
for row in rows:
|
||||
if row['vector'] is None:
|
||||
raise ValueError("incomplete policy coverage")
|
||||
vector = _unit_vector(json.loads(row['vector']), batch.dimensions)
|
||||
score = math.fsum(a * b for a, b in zip(batch.vectors[0], vector))
|
||||
yield VectorHit(id=row['block_id'], score=max(0.0, min(1.0, score)))
|
||||
try:
|
||||
ranked.append(heapq.nlargest(top_k, hits(), key=lambda hit: hit.score))
|
||||
finally:
|
||||
rows.close()
|
||||
_ensure_table(conn)
|
||||
ranked.append(space_index.search(conn, batch, top_k, policy))
|
||||
spaces = [{"source": b.source, "model_id": b.space_id, "dimensions": b.dimensions,
|
||||
"local_only": policy} for policy, b in batches.items()]
|
||||
record_embedding(source="mixed" if len({b.source for b in batches.values()}) > 1 else batch.source,
|
||||
|
||||
Reference in New Issue
Block a user