fix(retrieval): honor ignore rules and scale search

scripts/tests/test-retrieval-server.py

@@ -28,6 +28,44 @@ def embedding():
     return array('f', [1.0] + [0.0] * (MODULE.EMBEDDING_DIMENSIONS - 1))
 
 
+class VectorBatch:
+    def __init__(self, blob):
+        self.values = array('f', blob)
+
+    def reshape(self, *shape):
+        return self
+
+    def __matmul__(self, vector):
+        width = MODULE.EMBEDDING_DIMENSIONS
+        return [sum(a * b for a, b in zip(self.values[offset:offset + width], vector))
+                for offset in range(0, len(self.values), width)]
+
+
+class ApproximateIndex:
+    def __init__(self, **kwargs):
+        self.vectors = {}
+
+    def add(self, keys, vectors):
+        width = MODULE.EMBEDDING_DIMENSIONS
+        for position, key in enumerate(keys):
+            start = position * width
+            self.vectors[key] = self.vectors.get(key, []) + list(vectors.values[start:start + width])
+
+    def search(self, vector, count):
+        matches = sorted(((key, sum(a * b for a, b in zip(values, vector)))
+                          for key, values in self.vectors.items()), key=lambda item: (-item[1], item[0]))[:count]
+        return [SimpleNamespace(key=key, distance=1.0 - score) for key, score in matches]
+
+
+def fake_numpy():
+    return SimpleNamespace(frombuffer=lambda blob, **kwargs: VectorBatch(blob), float32=None)
+
+
+def fake_usearch():
+    module = SimpleNamespace(Index=ApproximateIndex)
+    return SimpleNamespace(index=module), module
+
+
 def huggingface_available():
     try:
         request = urllib.request.Request('https://huggingface.co', method='HEAD')
@@ -151,19 +189,19 @@ class RetrievalServerTests(unittest.TestCase):
             for index in indexes:
                 index.encoder = lambda: SimpleNamespace(encode=encode, tokenizer=tokenize)
                 index.rerank = lambda query, rows, limit: rows[:limit]
-            numpy = SimpleNamespace(dot=lambda a, b: sum(x * y for x, y in zip(a, b)),
-                                    frombuffer=lambda blob, **kwargs: array('f', blob), float32=None)
+            numpy = fake_numpy()
+            usearch, usearch_index = fake_usearch()
             try:
-                with patch.dict('sys.modules', numpy=numpy):
+                with patch.dict('sys.modules', numpy=numpy, usearch=usearch, **{'usearch.index': usearch_index}):
                     self.assertEqual(indexes[0].search('query', 5)[0]['text'], 'alpha')
                     self.assertEqual(indexes[1].search('query', 5)[0]['text'], 'other')
                     encoded.clear()
+                    with patch.object(Path, 'read_bytes', side_effect=AssertionError('unchanged files must not be read')):
+                        self.assertEqual(indexes[0].search('query', 5)[0]['text'], 'alpha')
                     with ThreadPoolExecutor(max_workers=1) as executor:
                         self.assertEqual(executor.submit(indexes[0].search, 'query', 5).result()[0]['text'], 'alpha')
                     self.assertEqual(encoded, [])
-                    timestamp = file.stat()
                     file.write_text('bravo', encoding='utf-8')
-                    os.utime(file, ns=(timestamp.st_atime_ns, timestamp.st_mtime_ns))
                     self.assertEqual(indexes[0].search('query', 5)[0]['text'], 'bravo')
                     file.rename(first / 'renamed.md')
                     self.assertEqual(indexes[0].search('query', 5)[0]['path'], 'renamed.md')
@@ -227,6 +265,11 @@ class RetrievalServerTests(unittest.TestCase):
             root = Path(temporary)
             index = MODULE.Index(root, root / 'data')
             (root / 'visible.md').write_text('visible', encoding='utf-8')
+            (root / '.gitignore').write_text('secrets.json\n**/private/\n', encoding='utf-8')
+            (root / 'secrets.json').write_text('{"token":"hidden"}', encoding='utf-8')
+            private = root / 'nested' / 'private'
+            private.mkdir(parents=True)
+            (private / 'hidden.md').write_text('hidden', encoding='utf-8')
             for name in ('node_modules', '.git', 'generated', 'data'):
                 directory = root / name
                 directory.mkdir(exist_ok=True)
@@ -234,7 +277,7 @@ class RetrievalServerTests(unittest.TestCase):
             try:
                 self.assertEqual(index.files(), [(root / 'visible.md').resolve()])
                 filenames = ['visible.md'] * 2001
-                with patch.object(MODULE.os, 'walk', return_value=iter([(str(root), [], filenames)])), \
+                with patch.object(MODULE.subprocess, 'run', return_value=SimpleNamespace(stdout=b'visible.md\0' * 2001)), \
                      patch.object(MODULE.sys, 'stderr') as stderr:
                     self.assertEqual(len(index.files()), 2000)
                 self.assertIn('stopped scanning after 2000', stderr.write.call_args_list[0].args[0])
@@ -255,9 +298,9 @@ class RetrievalServerTests(unittest.TestCase):
                     self.assertEqual(MODULE.read_max_files(), MODULE.DEFAULT_MAX_FILES)
                 (root / 'visible.md').write_text('visible', encoding='utf-8')
                 index.max_files = 3
-                filenames = ['visible.md'] * 5
-                with patch.object(MODULE.os, 'walk', return_value=iter([(str(root), [], filenames)])), \
-                     patch.object(MODULE.sys, 'stderr'):
+                for number in range(5):
+                    (root / f'{number}.md').write_text('visible', encoding='utf-8')
+                with patch.object(MODULE.sys, 'stderr'):
                     self.assertEqual(len(index.files()), 3)
             finally:
                 index.db.close()
@@ -365,8 +408,7 @@ class RetrievalServerTests(unittest.TestCase):
             index.encoder = lambda: SimpleNamespace(encode=encode, tokenizer=tokenize)
             index.reranker = lambda: SimpleNamespace(predict=lambda values, **kwargs:
                 pairs.extend(values) or [float('109.md' in document) for _, document in values])
-            numpy = SimpleNamespace(dot=lambda a, b: 1.0,
-                                    frombuffer=lambda blob, **kwargs: array('f', blob), float32=None)
+            numpy = fake_numpy()
             original_fusion = MODULE.reciprocal_rank_fusion
             rankings = []
 
@@ -375,7 +417,9 @@ class RetrievalServerTests(unittest.TestCase):
                 return original_fusion(*values)
 
             try:
-                with patch.dict('sys.modules', numpy=numpy), patch.object(MODULE, 'reciprocal_rank_fusion', fuse):
+                usearch, usearch_index = fake_usearch()
+                with patch.dict('sys.modules', numpy=numpy, usearch=usearch, **{'usearch.index': usearch_index}), \
+                     patch.object(MODULE, 'reciprocal_rank_fusion', fuse):
                     results = index.search('109.md Symbol109 content', 20)
                 self.assertEqual([len(ranking) for ranking in rankings], [50, 50])
                 self.assertEqual(len(pairs), 50)
@@ -384,7 +428,7 @@ class RetrievalServerTests(unittest.TestCase):
                 self.assertEqual(results[0]['symbol'], 'Symbol109')
                 self.assertIn('Path: 109.md\nSymbol: Symbol109\n\n# Symbol109\ncontent 109', documents)
                 self.assertTrue(all(document in documents for _, document in pairs))
-                with patch.dict('sys.modules', numpy=numpy):
+                with patch.dict('sys.modules', numpy=numpy, usearch=usearch, **{'usearch.index': usearch_index}):
                     self.assertEqual(len(index.search('" OR * () : -')), 5)
                     self.assertEqual(len(index.search('***')), 5)
                     (root / '109.md').unlink()

shared/retrieval/requirements.txt

@@ -1 +1,2 @@
 sentence-transformers==5.7.0
+usearch==2.26.2

shared/retrieval/server.py

@@ -6,7 +6,6 @@ from __future__ import annotations
 import argparse
 import ast
 import hashlib
-import heapq
 import json
 import multiprocessing
 from multiprocessing import AuthenticationError
@@ -379,23 +378,32 @@ class Index:
         self.db.execute("create table if not exists chunks (path text, symbol text, text text, vector blob)")
         self.db.execute("create index if not exists chunks_path on chunks (path)")
         self.db.execute("create virtual table if not exists lexical using fts5(document)")
-        self.db.execute("create table if not exists files (path text primary key, digest text)")
+        self.db.execute("create table if not exists files (path text primary key, digest text, size integer, mtime integer, ctime integer)")
+        columns = {row[1] for row in self.db.execute("pragma table_info(files)")}
+        for name in ("size", "mtime", "ctime"):
+            if name not in columns:
+                self.db.execute(f"alter table files add column {name} integer")
         self.db.execute("create table if not exists query_cache (cache_key text primary key, result text not null)")
         self.db.commit()
         self.lock = threading.RLock()
         self.max_files = read_max_files()
+        self.vector_index = None
 
     def files(self):
-        ignored = {".git", ".venv", "generated", "node_modules", ".my-ai-configuration"}
+        ignored = {".git", ".venv", "generated", "node_modules", ".my-ai-configuration", "bin", "obj"}
+        output = subprocess.run(["git", "-C", str(self.root), "ls-files", "--cached", "--others",
+                                 "--exclude-standard", "-z", "--", "."],
+                                check=True, capture_output=True, timeout=30).stdout
+        candidates = [self.root / name for name in
+                      output.decode("utf-8", errors="surrogateescape").split("\0") if name
+                      and not any(part in ignored for part in Path(name).parts)]
         result = []
-        for directory, directories, filenames in os.walk(self.root):
-            directories[:] = sorted(name for name in directories if name not in ignored
-                                    and not (Path(directory) / name).is_symlink()
-                                    and (Path(directory) / name).resolve().is_relative_to(self.root)
-                                    and (Path(directory) / name).resolve() != self.data)
-            for name in sorted(filenames):
-                path = Path(directory) / name
-                if path.is_symlink() or path.suffix.lower() not in TEXT_EXTENSIONS or not path.is_file() or path.stat().st_size > 524288:
+        for candidate in sorted(candidates):
+            if candidate.is_symlink():
+                continue
+            path = candidate.resolve()
+            if (not path.is_relative_to(self.root) or path.is_relative_to(self.data)
+                    or path.suffix.lower() not in TEXT_EXTENSIONS or not path.is_file() or path.stat().st_size > 524288):
                 continue
             result.append(path)
             if len(result) == self.max_files:
@@ -429,17 +437,25 @@ class Index:
             return self.db.execute("select count(*) from chunks").fetchone()[0]
 
     def refresh(self, force=False):
-        previous = dict(self.db.execute("select path, digest from files"))
+        previous = {path: (digest, size, mtime, ctime)
+                    for path, digest, size, mtime, ctime in self.db.execute("select path, digest, size, mtime, ctime from files")}
         encoder = None
         changed = []
         documents = []
         document_offsets = []
         for path in self.files():
             relative = path.relative_to(self.root).as_posix()
+            metadata = path.stat()
+            size, mtime, ctime = metadata.st_size, metadata.st_mtime_ns, metadata.st_ctime_ns
+            old_digest, old_size, old_mtime, old_ctime = previous.pop(relative, (None, None, None, None))
+            if not force and old_digest and (size, mtime, ctime) == (old_size, old_mtime, old_ctime):
+                continue
             content = path.read_bytes()
             digest = hashlib.sha256(content).hexdigest()
-            old_digest = previous.pop(relative, None)
             if not force and digest == old_digest:
+                with self.db:
+                    self.db.execute("update files set size = ?, mtime = ?, ctime = ? where path = ?",
+                                    (size, mtime, ctime, relative))
                 continue
             if encoder is None:
                 encoder = self.encoder()
@@ -447,10 +463,10 @@ class Index:
             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))
+            changed.append((relative, digest, size, mtime, ctime, 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):
+            for number, (relative, digest, size, mtime, ctime, values, file_documents) in enumerate(changed):
                 start = document_offsets[number]
                 file_vectors = vectors[start:start + len(file_documents)]
                 self.delete_path(relative)
@@ -460,17 +476,32 @@ class Index:
                     cursor = self.db.execute("insert into chunks values (?, ?, ?, ?)",
                                              (relative, symbol, value, vector.tobytes()))
                     self.db.execute("insert into lexical(rowid, document) values (?, ?)", (cursor.lastrowid, document))
-                self.db.execute("insert or replace into files values (?, ?)", (relative, digest))
+                self.db.execute("insert or replace into files values (?, ?, ?, ?, ?)",
+                                (relative, digest, size, mtime, ctime))
             for relative in previous:
                 self.delete_path(relative)
                 self.db.execute("delete from files where path = ?", (relative,))
             if changed or previous:
                 self.db.execute("delete from query_cache")
+                self.vector_index = None
+
+    def dense_candidates(self, vector):
+        import numpy as np
+        from usearch.index import Index as VectorIndex
+
+        if self.vector_index is None:
+            index = VectorIndex(ndim=EMBEDDING_DIMENSIONS, metric="cos", dtype="f32",
+                                connectivity=16, expansion_add=128, expansion_search=64)
+            cursor = self.db.execute("select rowid, vector from chunks")
+            while rows := cursor.fetchmany(4096):
+                vectors = np.frombuffer(b"".join(row[1] for row in rows), dtype=np.float32).reshape(-1, EMBEDDING_DIMENSIONS)
+                index.add([row[0] for row in rows], vectors)
+            self.vector_index = index
+        return [match.key for match in self.vector_index.search(vector, DENSE_CANDIDATES)]
 
     def cache_key(self, query: str, top_k: int):
-        files = self.db.execute("select path, digest from files order by path").fetchall()
         identity = json.dumps([
-            query, top_k, files,
+            query, top_k,
             RERANKER_MODEL, RERANKER_MODEL_REVISION, RETRIEVAL_INSTRUCTION,
             DENSE_CANDIDATES, LEXICAL_CANDIDATES, RERANK_CANDIDATES, FINAL_RESULTS, RRF_K,
         ], ensure_ascii=False, separators=(",", ":"))
@@ -491,13 +522,10 @@ class Index:
                 except (TypeError, json.JSONDecodeError):
                     with self.db:
                         self.db.execute("delete from query_cache where cache_key = ?", (key,))
-            import numpy as np
             vector = self.encoder().encode([query], prompt_name="query", normalize_embeddings=True)[0]
             if len(vector) != EMBEDDING_DIMENSIONS:
                 raise ValueError("Embedding must have 1024 dimensions")
-            scores = ((identifier, float(np.dot(vector, np.frombuffer(blob, dtype=np.float32))))
-                      for identifier, blob in self.db.execute("select rowid, vector from chunks"))
-            dense = [identifier for identifier, _ in heapq.nlargest(DENSE_CANDIDATES, scores, key=lambda item: item[1])]
+            dense = self.dense_candidates(vector)
             terms = list(dict.fromkeys(re.findall(r"[^\W_]+", query)))[:128]
             expression = " OR ".join(f'"{term}"' for term in terms)
             lexical = [row[0] for row in self.db.execute(