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

2026-09-06 16:34:26 +08:00
parent 3b9490e3fb
commit b03b168920
12 changed files with 295 additions and 48 deletions
@@ -0,0 +1,64 @@
+"""Synthetic, isolated exact-search comparison; does not access the user Vault."""
+import heapq
+import json
+import math
+import random
+import sqlite3
+import statistics
+import sys
+import tempfile
+import time
+from pathlib import Path
+
+import sqlite_vec
+
+sys.path.insert(0, str(Path(__file__).resolve().parents[1]))
+from app.retrieval import space_index
+from app.retrieval.routed_vectors import RemoteEmbeddings, _unit_vector
+
+
+def main():
+    rng = random.Random(42)
+    count, dimensions = 4000, 384
+    with tempfile.TemporaryDirectory(prefix='notes-vec-bench-') as temporary:
+        conn = sqlite3.connect(Path(temporary) / 'vectors.db')
+        conn.row_factory = sqlite3.Row
+        conn.enable_load_extension(True)
+        sqlite_vec.load(conn)
+        conn.enable_load_extension(False)
+        conn.execute('CREATE TABLE blocks(block_id TEXT PRIMARY KEY, embedding_local_only INTEGER)')
+        conn.execute('CREATE TABLE routed_block_vectors(space_id TEXT,block_id TEXT,dimensions INTEGER,vector TEXT,PRIMARY KEY(space_id,dimensions,block_id))')
+        vectors = [_unit_vector([rng.uniform(-1, 1) for _ in range(dimensions)], dimensions) for _ in range(count)]
+        conn.executemany('INSERT INTO blocks VALUES (?,0)', [(str(i),) for i in range(count)])
+        conn.executemany('INSERT INTO routed_block_vectors VALUES (?,?,?,?)', [('benchmark', str(i), dimensions, json.dumps(v)) for i, v in enumerate(vectors)])
+        start = time.perf_counter()
+        space_index.ensure(conn, 'benchmark', dimensions)
+        migration_ms = (time.perf_counter() - start) * 1000
+        conn.commit()
+        query = vectors[0]
+        batch = RemoteEmbeddings('benchmark', dimensions, [query])
+        def legacy():
+            def hits():
+                for row in conn.execute('SELECT block_id,vector FROM routed_block_vectors'):
+                    vector = _unit_vector(json.loads(row[1]), dimensions)
+                    yield row[0], max(0., min(1., math.fsum(a*b for a,b in zip(query,vector))))
+            return heapq.nlargest(20, hits(), key=lambda hit:hit[1])
+        def native():
+            return [(hit.id,hit.score) for hit in space_index.search(conn,batch,20)]
+        measurements = {}
+        results = {}
+        for name, operation in [('python_json_scan', legacy), ('sqlite_vec',native)]:
+            elapsed = []
+            for _ in range(5):
+                start = time.perf_counter()
+                results[name] = operation()
+                elapsed.append((time.perf_counter()-start)*1000)
+            measurements[name] = {'median_ms':statistics.median(elapsed), 'samples_ms':elapsed}
+        assert [hit[0] for hit in results['python_json_scan']] == [hit[0] for hit in results['sqlite_vec']]
+        print(json.dumps({'blocks':count,'dimensions':dimensions,'top_k':20,'migration_ms':migration_ms,
+                          'same_top_k':True,'measurements':measurements},indent=2))
+        conn.close()
+
+
+if __name__ == '__main__':
+    main()