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,19 +193,31 @@ class Index:
 
     def refresh(self, force=False):
         previous = dict(self.db.execute("select path, digest from files"))
+        encoder = None
+        changed = []
+        documents = []
+        document_offsets = []
+        for path in self.files():
+            relative = path.relative_to(self.root).as_posix()
+            content = path.read_bytes()
+            digest = hashlib.sha256(content).hexdigest()
+            old_digest = previous.pop(relative, None)
+            if not force and digest == old_digest:
+                continue
+            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 path in self.files():
+            for number, (relative, digest, values, file_documents) in enumerate(changed):
-                relative = path.relative_to(self.root).as_posix()
+                start = document_offsets[number]
-                content = path.read_bytes()
+                file_vectors = vectors[start:start + len(file_documents)]
-                digest = hashlib.sha256(content).hexdigest()
-                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 []
                 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 (?, ?, ?, ?)",