feat(qwen): add token windows and document metadata

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__':

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 [("<module>", 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)