diff --git a/scripts/test-retrieval-server.py b/scripts/test-retrieval-server.py index dc37cff..6af0a93 100644 --- a/scripts/test-retrieval-server.py +++ b/scripts/test-retrieval-server.py @@ -3,6 +3,7 @@ from array import array from concurrent.futures import ThreadPoolExecutor import os from pathlib import Path +import re from tempfile import TemporaryDirectory from types import SimpleNamespace import time @@ -16,6 +17,11 @@ MODULE = importlib.util.module_from_spec(SPEC) SPEC.loader.exec_module(MODULE) +def tokenize(text, **kwargs): + offsets = [match.span() for match in re.finditer(r'\S+', text)] + return {'offset_mapping': offsets, 'input_ids': list(range(len(offsets)))} + + def embedding(): return array('f', [1.0] + [0.0] * (MODULE.EMBEDDING_DIMENSIONS - 1)) @@ -30,6 +36,13 @@ class RetrievalServerTests(unittest.TestCase): (root / 'weather.md').write_text('The weather service fetches temperature and rainfall forecasts for a city.', encoding='utf-8') index = MODULE.Index(root, root / 'data') try: + tokenizer = index.encoder().tokenizer + for source in ['alpha beta gamma ' * 450, 'Grüße 世界 🙂 ' * 300]: + windows = list(MODULE.chunks(source, 'sample.txt', tokenizer)) + self.assertGreater(len(windows), 1) + self.assertTrue(all(len(tokenizer(value, add_special_tokens=False)['input_ids']) <= 400 + for _, value in windows)) + self.assertTrue(all('\ufffd' not in value for _, value in windows)) for query, expected in [('Where are password reset tokens created?', 'authentication.md'), ('Which service provides rainfall forecasts?', 'weather.md')]: started = time.monotonic() @@ -59,7 +72,7 @@ class RetrievalServerTests(unittest.TestCase): return [embedding() for _ in values] for index in indexes: - index.encoder = lambda: SimpleNamespace(encode=encode) + 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) @@ -92,18 +105,18 @@ class RetrievalServerTests(unittest.TestCase): root = Path(temporary) (root / 'a.md').write_text('before', encoding='utf-8') index = MODULE.Index(root, root / 'data') - index.encoder = lambda: SimpleNamespace(encode=lambda values, **kwargs: [embedding() for _ in values]) + index.encoder = lambda: SimpleNamespace(encode=lambda values, **kwargs: [embedding() for _ in values], tokenizer=tokenize) try: index.rebuild() (root / 'a.md').write_text('after', encoding='utf-8') (root / 'b.md').write_text('fail', encoding='utf-8') def encode(values, **kwargs): - if values == ['fail']: + if any(value.endswith('\n\nfail') for value in values): raise RuntimeError('encoder failed') return [embedding() for _ in values] - index.encoder = lambda: SimpleNamespace(encode=encode) + index.encoder = lambda: SimpleNamespace(encode=encode, tokenizer=tokenize) with self.assertRaisesRegex(RuntimeError, 'encoder failed'): index.rebuild() self.assertEqual(index.db.execute('select path, text from chunks').fetchall(), [('a.md', 'before')]) @@ -131,22 +144,51 @@ class RetrievalServerTests(unittest.TestCase): index = MODULE.Index.__new__(MODULE.Index) class Reranker: - def predict(self, pairs): + def predict(self, pairs, **kwargs): self.pairs = pairs return [0.1, 0.8, 0.5] reranker = Reranker() index.reranker = lambda: reranker rows = [ - {'path': 'a.md', 'score': 0.9, 'text': 'first'}, - {'path': 'b.md', 'score': 0.8, 'text': 'second'}, - {'path': 'c.md', 'score': 0.7, 'text': 'third'}, + {'path': 'a.md', 'symbol': 'A', 'score': 0.9, 'text': 'first'}, + {'path': 'b.md', 'symbol': 'B', 'score': 0.8, 'text': 'second'}, + {'path': 'c.md', 'symbol': 'C', 'score': 0.7, 'text': 'third'}, ] result = index.rerank('query', rows, 2) self.assertEqual([item['path'] for item in result], ['b.md', 'c.md']) - self.assertEqual(reranker.pairs, [('query', 'first'), ('query', 'second'), ('query', 'third')]) + self.assertEqual(reranker.pairs, [('query', 'Path: a.md\nSymbol: A\n\nfirst'), + ('query', 'Path: b.md\nSymbol: B\n\nsecond'), + ('query', 'Path: c.md\nSymbol: C\n\nthird')]) + + def test_token_windows_have_exact_overlap_and_no_redundant_tail(self): + for size, lengths in [(0, []), (400, [400]), (401, [400, 65]), (736, [400, 400]), (800, [400, 400, 128])]: + words = [f'token{number}' for number in range(size)] + result = list(MODULE.chunks(' '.join(words), 'plain.txt', tokenize)) + windows = [content.split() for _, content in result] + self.assertEqual([len(window) for window in windows], lengths) + for first, second in zip(windows, windows[1:]): + self.assertEqual(first[-64:], second[:64]) + if windows: + reconstructed = windows[0] + [word for window in windows[1:] for word in window[64:]] + self.assertEqual(reconstructed, words) + + def test_old_cache_is_preserved_but_not_reused(self): + with TemporaryDirectory() as temporary: + root = Path(temporary) + data = root / 'data' + data.mkdir() + identity = os.path.normcase(str(root.resolve())) + '\n' + MODULE.MODEL + old = data / (MODULE.hashlib.sha256(identity.encode()).hexdigest() + '.sqlite3') + old.write_bytes(b'old cache') + index = MODULE.Index(root, data) + try: + self.assertNotEqual(Path(index.db.execute('pragma database_list').fetchone()[2]), old) + self.assertEqual(old.read_bytes(), b'old cache') + finally: + index.db.close() if __name__ == '__main__': diff --git a/shared/retrieval/server.py b/shared/retrieval/server.py index e90af91..985327a 100644 --- a/shared/retrieval/server.py +++ b/shared/retrieval/server.py @@ -16,15 +16,37 @@ import threading MODEL = "Qwen/Qwen3-Embedding-0.6B" RERANKER_MODEL = "Qwen/Qwen3-Reranker-0.6B" EMBEDDING_DIMENSIONS = 1024 +CHUNK_TOKENS = 400 +OVERLAP_TOKENS = 64 +RERANK_CANDIDATES = 20 +FINAL_RESULTS = 20 +INDEX_VERSION = 0 RETRIEVAL_INSTRUCTION = "Given a codebase question, retrieve relevant code and documentation that answer the question." TEXT_EXTENSIONS = {".c", ".cpp", ".cs", ".go", ".java", ".js", ".json", ".md", ".py", ".ps1", ".rs", ".sh", ".toml", ".ts", ".tsx", ".txt", ".yaml", ".yml"} -def chunks(text: str): - for start in range(0, len(text), 1050): - value = text[start:start + 1200].strip() - if value: - yield value +def chunks(text: str, path: str, tokenizer): + for symbol, content in [("", text)]: + offsets = tokenizer(content, add_special_tokens=False, return_offsets_mapping=True)["offset_mapping"] + start = 0 + while start < len(offsets): + end = min(start + CHUNK_TOKENS, len(offsets)) + first = 0 if start == 0 else offsets[start][0] + while end > start: + last = len(content) if end == len(offsets) else offsets[end][0] + value = content[first:last].strip() + if len(tokenizer(value, add_special_tokens=False)["input_ids"]) <= CHUNK_TOKENS: + break + end -= 1 + if value: + yield symbol, value + if end == len(offsets): + break + start = max(start + 1, end - OVERLAP_TOKENS) + + +def document_text(path: str, symbol: str, content: str): + return f"Path: {path}\nSymbol: {symbol}\n\n{content}" class Index: @@ -32,10 +54,11 @@ class Index: self.root = root.resolve() self.data = data.resolve() self.data.mkdir(parents=True, exist_ok=True) - identity = os.path.normcase(str(self.root)) + "\n" + MODEL + "\n" + str(EMBEDDING_DIMENSIONS) + identity = json.dumps([os.path.normcase(str(self.root)), MODEL, EMBEDDING_DIMENSIONS, + INDEX_VERSION, CHUNK_TOKENS, OVERLAP_TOKENS]) database = hashlib.sha256(identity.encode()).hexdigest() + ".sqlite3" self.db = sqlite3.connect(self.data / database, check_same_thread=False, timeout=30) - self.db.execute("create table if not exists chunks (path text, text text, vector blob)") + 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 table if not exists files (path text primary key, digest text)") self.db.commit() @@ -74,11 +97,15 @@ class Index: return self.reranker_model def rerank(self, query: str, rows: list[dict], limit: int): - candidates = rows[:min(len(rows), max(limit, 20))] - scores = self.reranker().predict([(query, item["text"]) for item in candidates]) + candidates = rows[:RERANK_CANDIDATES] + scores = self.reranker().predict([(query, document_text(item["path"], item["symbol"], item["text"])) + for item in candidates], batch_size=8) for item, score in zip(candidates, scores): item["score"] = round(float(score), 6) - return sorted(candidates, key=lambda item: item["score"], reverse=True)[:limit] + return sorted(candidates, key=lambda item: item["score"], reverse=True)[:min(limit, FINAL_RESULTS)] + + def delete_path(self, path): + self.db.execute("delete from chunks where path = ?", (path,)) def rebuild(self): with self.lock: @@ -95,14 +122,18 @@ 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"))) - vectors = self.encoder().encode(values, normalize_embeddings=True) if values else [] - self.db.execute("delete from chunks where path = ?", (relative,)) - self.db.executemany("insert into chunks values (?, ?, ?)", - ((relative, value, vector.tobytes()) for value, vector in zip(values, vectors))) + 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), vector in zip(values, vectors): + if len(vector) != EMBEDDING_DIMENSIONS: + raise ValueError("Embedding must have 1024 dimensions") + self.db.execute("insert into chunks values (?, ?, ?, ?)", + (relative, symbol, value, vector.tobytes())) self.db.execute("insert or replace into files values (?, ?)", (relative, digest)) for relative in previous: - self.db.execute("delete from chunks where path = ?", (relative,)) + self.delete_path(relative) self.db.execute("delete from files where path = ?", (relative,)) def search(self, query: str, top_k: int): @@ -114,8 +145,10 @@ class Index: return [] import numpy as np vector = self.encoder().encode([query], prompt_name="query", normalize_embeddings=True)[0] - rows = ({"path": path, "score": float(np.dot(vector, np.frombuffer(blob, dtype=np.float32))), "text": text} - for path, text, blob in self.db.execute("select path, text, vector from chunks")) + if len(vector) != EMBEDDING_DIMENSIONS: + raise ValueError("Embedding must have 1024 dimensions") + rows = ({"path": path, "symbol": symbol, "score": float(np.dot(vector, np.frombuffer(blob, dtype=np.float32))), "text": text} + for path, symbol, text, blob in self.db.execute("select path, symbol, text, vector from chunks")) limit = max(1, min(top_k, 20)) return self.rerank(query, heapq.nlargest(20, rows, key=lambda item: item["score"]), limit)