From a67c34f78d198cdfeca9c96af4ad003c0438617b Mon Sep 17 00:00:00 2001 From: Julian lechner Date: Wed, 30 Sep 2026 09:56:57 +0200 Subject: [PATCH] fix(retrieval): honor ignore rules and scale search --- scripts/tests/test-retrieval-server.py | 70 +++++++++++++++++---- shared/retrieval/requirements.txt | 1 + shared/retrieval/server.py | 86 +++++++++++++++++--------- 3 files changed, 115 insertions(+), 42 deletions(-) diff --git a/scripts/tests/test-retrieval-server.py b/scripts/tests/test-retrieval-server.py index 627a63e..a6e4cf6 100644 --- a/scripts/tests/test-retrieval-server.py +++ b/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() diff --git a/shared/retrieval/requirements.txt b/shared/retrieval/requirements.txt index 28cee65..9e68d87 100644 --- a/shared/retrieval/requirements.txt +++ b/shared/retrieval/requirements.txt @@ -1 +1,2 @@ sentence-transformers==5.7.0 +usearch==2.26.2 diff --git a/shared/retrieval/server.py b/shared/retrieval/server.py index 9a0d89a..525aefb 100644 --- a/shared/retrieval/server.py +++ b/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,30 +378,39 @@ 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: - continue - result.append(path) - if len(result) == self.max_files: - print(f"Semantic retrieval: stopped scanning after {self.max_files} indexable files; " - f"some files under {self.root} were not indexed. Set {MAX_FILES_ENV} to raise this limit.", - file=sys.stderr) - return result + 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: + print(f"Semantic retrieval: stopped scanning after {self.max_files} indexable files; " + f"some files under {self.root} were not indexed. Set {MAX_FILES_ENV} to raise this limit.", + file=sys.stderr) + return result return result def encoder(self): @@ -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(