perf(retrieval): batch changed embeddings

scripts/test-retrieval-server.py

@@ -125,6 +125,26 @@ class RetrievalServerTests(unittest.TestCase):
             finally:
                 index.db.close()
 
+    def test_refresh_batches_embeddings_for_changed_files(self):
+        with TemporaryDirectory() as temporary:
+            root = Path(temporary)
+            (root / 'first.md').write_text('first', encoding='utf-8')
+            (root / 'second.md').write_text('second', encoding='utf-8')
+            index = MODULE.Index(root, root / 'data')
+            encoded = []
+
+            def encode(values, **kwargs):
+                encoded.append(list(values))
+                return [embedding() for _ in values]
+
+            index.encoder = lambda: SimpleNamespace(encode=encode, tokenizer=tokenize)
+            try:
+                index.rebuild()
+                self.assertEqual(len(encoded), 1)
+                self.assertEqual(len(encoded[0]), 2)
+            finally:
+                index.db.close()
+
     def test_scan_prunes_ignored_directories_and_stops_at_limit(self):
         with TemporaryDirectory() as temporary:
             root = Path(temporary)

shared/retrieval/server.py

@@ -193,7 +193,10 @@ class Index:
 
     def refresh(self, force=False):
         previous = dict(self.db.execute("select path, digest from files"))
-        with self.db:
+        encoder = None
+        changed = []
+        documents = []
+        document_offsets = []
         for path in self.files():
             relative = path.relative_to(self.root).as_posix()
             content = path.read_bytes()
@@ -201,11 +204,20 @@ class Index:
             old_digest = previous.pop(relative, None)
             if not force and digest == old_digest:
                 continue
-                values = list(chunks(content.decode("utf-8", errors="ignore"), relative, self.encoder().tokenizer))
-                documents = [document_text(relative, symbol, value) for symbol, value in values]
-                vectors = self.encoder().encode(documents, normalize_embeddings=True) if documents else []
+            if encoder is None:
+                encoder = self.encoder()
+            values = list(chunks(content.decode("utf-8", errors="ignore"), relative, encoder.tokenizer))
+            file_documents = [document_text(relative, symbol, value) for symbol, value in values]
+            document_offsets.append(len(documents))
+            documents.extend(file_documents)
+            changed.append((relative, digest, values, file_documents))
+        vectors = encoder.encode(documents, normalize_embeddings=True) if documents else []
+        with self.db:
+            for number, (relative, digest, values, file_documents) in enumerate(changed):
+                start = document_offsets[number]
+                file_vectors = vectors[start:start + len(file_documents)]
                 self.delete_path(relative)
-                for (symbol, value), document, vector in zip(values, documents, vectors):
+                for (symbol, value), document, vector in zip(values, file_documents, file_vectors):
                     if len(vector) != EMBEDDING_DIMENSIONS:
                         raise ValueError("Embedding must have 1024 dimensions")
                     cursor = self.db.execute("insert into chunks values (?, ?, ?, ?)",