From 2079792a741fde9a82821cccf10f00652dad4219 Mon Sep 17 00:00:00 2001 From: Julian lechner Date: Thu, 17 Sep 2026 15:27:04 +0200 Subject: [PATCH] fix(qwen): isolate and refresh indexes with code-search prompts --- scripts/test-retrieval-server.py | 113 +++++++++++++++++++++++++++++++ shared/retrieval/server.py | 90 +++++++++++++++--------- 2 files changed, 171 insertions(+), 32 deletions(-) diff --git a/scripts/test-retrieval-server.py b/scripts/test-retrieval-server.py index 49a1784..2990e98 100644 --- a/scripts/test-retrieval-server.py +++ b/scripts/test-retrieval-server.py @@ -1,6 +1,13 @@ import importlib.util +from array import array +from concurrent.futures import ThreadPoolExecutor +import os from pathlib import Path +from tempfile import TemporaryDirectory +from types import SimpleNamespace +import time import unittest +from unittest.mock import patch ROOT = Path(__file__).resolve().parent.parent @@ -10,6 +17,112 @@ SPEC.loader.exec_module(MODULE) class RetrievalServerTests(unittest.TestCase): + @unittest.skipUnless(os.environ.get('QWEN_TEST_MODEL_CACHE'), 'Set QWEN_TEST_MODEL_CACHE to run the real-model check.') + def test_real_models(self): + os.environ['HF_HOME'] = os.environ['QWEN_TEST_MODEL_CACHE'] + with TemporaryDirectory() as temporary: + root = Path(temporary) + (root / 'authentication.md').write_text('Password reset tokens are created by create_reset_token. The reset_password endpoint sends the user a password reset email.', encoding='utf-8') + (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: + for query, expected in [('Where are password reset tokens created?', 'authentication.md'), + ('Which service provides rainfall forecasts?', 'weather.md')]: + started = time.monotonic() + results = index.search(query, 2) + self.assertEqual(results[0]['path'], expected) + self.assertEqual(len(results), 2) + self.assertGreater(results[0]['score'], results[1]['score']) + print(f'Qwen real search: {expected}, {time.monotonic() - started:.2f}s, embedding device {index.encoder().device}, reranker device {index.reranker().device}', flush=True) + finally: + index.db.close() + + def test_isolation_freshness_and_incremental_encoding(self): + with TemporaryDirectory() as temporary: + root = Path(temporary) + first, second = root / 'first', root / 'second' + first.mkdir() + second.mkdir() + file = first / 'a.md' + file.write_text('alpha', encoding='utf-8') + (second / 'a.md').write_text('other', encoding='utf-8') + indexes = [MODULE.Index(repository, root / 'data') for repository in (first, second)] + encoded = [] + + def encode(values, **kwargs): + self.assertEqual(kwargs.get('prompt_name'), 'query' if values == ['query'] else None) + encoded.extend(values) + return [array('f', [1.0]) for _ in values] + + for index in indexes: + index.encoder = lambda: SimpleNamespace(encode=encode) + 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) + try: + with patch.dict('sys.modules', numpy=numpy): + self.assertEqual(indexes[0].search('query', 5)[0]['text'], 'alpha') + self.assertEqual(indexes[1].search('query', 5)[0]['text'], 'other') + encoded.clear() + with ThreadPoolExecutor(max_workers=1) as executor: + self.assertEqual(executor.submit(indexes[0].search, 'query', 5).result()[0]['text'], 'alpha') + self.assertEqual(encoded, ['query']) + 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') + (first / 'renamed.md').unlink() + encoded.clear() + self.assertEqual(indexes[0].search('query', 5), []) + self.assertEqual(indexes[0].search('query', 5), []) + self.assertEqual(encoded, []) + self.assertEqual(indexes[1].search('query', 5)[0]['text'], 'other') + finally: + for index in indexes: + index.db.close() + + def test_failed_refresh_preserves_previous_index(self): + with TemporaryDirectory() as temporary: + 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: [array('f', [1.0]) for _ in values]) + 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']: + raise RuntimeError('encoder failed') + return [array('f', [1.0]) for _ in values] + + index.encoder = lambda: SimpleNamespace(encode=encode) + with self.assertRaisesRegex(RuntimeError, 'encoder failed'): + index.rebuild() + self.assertEqual(index.db.execute('select path, text from chunks').fetchall(), [('a.md', 'before')]) + finally: + index.db.close() + + def test_scan_prunes_ignored_directories_and_stops_at_limit(self): + with TemporaryDirectory() as temporary: + root = Path(temporary) + index = MODULE.Index(root, root / 'data') + (root / 'visible.md').write_text('visible', encoding='utf-8') + for name in ('node_modules', '.git', 'generated', 'data'): + directory = root / name + directory.mkdir(exist_ok=True) + (directory / 'hidden.md').write_text('hidden', encoding='utf-8') + 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)])): + self.assertEqual(len(index.files()), 2000) + finally: + index.db.close() + def test_rerank_uses_cross_encoder_scores(self): index = MODULE.Index.__new__(MODULE.Index) diff --git a/shared/retrieval/server.py b/shared/retrieval/server.py index 3e21da8..34cf30d 100644 --- a/shared/retrieval/server.py +++ b/shared/retrieval/server.py @@ -4,19 +4,22 @@ from __future__ import annotations import argparse +import hashlib +import heapq import json import os from pathlib import Path import sqlite3 +import threading MODEL = "Qwen/Qwen3-Embedding-0.6B" RERANKER_MODEL = "Qwen/Qwen3-Reranker-0.6B" +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(path: Path): - text = path.read_text(encoding="utf-8", errors="ignore") +def chunks(text: str): for start in range(0, len(text), 1050): value = text[start:start + 1200].strip() if value: @@ -28,33 +31,44 @@ class Index: self.root = root.resolve() self.data = data.resolve() self.data.mkdir(parents=True, exist_ok=True) - self.db = sqlite3.connect(self.data / "index.sqlite3") + identity = os.path.normcase(str(self.root)) + "\n" + MODEL + 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 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() + self.lock = threading.RLock() self.model = None self.reranker_model = None def files(self): ignored = {".git", ".venv", "generated", "node_modules", ".my-ai-configuration"} result = [] - for path in self.root.rglob("*"): - if len(result) >= 2000 or not path.is_file() or path.suffix.lower() not in TEXT_EXTENSIONS or path.stat().st_size > 524288: - continue - if ignored.intersection(path.relative_to(self.root).parts): - continue - result.append(path) + 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) == 2000: + return result return result def encoder(self): if self.model is None: from sentence_transformers import SentenceTransformer - self.model = SentenceTransformer(MODEL) + self.model = SentenceTransformer(MODEL, prompts={"query": f"Instruct: {RETRIEVAL_INSTRUCTION}\nQuery: "}) return self.model def reranker(self): if self.reranker_model is None: from sentence_transformers import CrossEncoder - self.reranker_model = CrossEncoder(RERANKER_MODEL) + self.reranker_model = CrossEncoder(RERANKER_MODEL, prompts={"query": RETRIEVAL_INSTRUCTION}, default_prompt_name="query") return self.reranker_model def rerank(self, query: str, rows: list[dict], limit: int): @@ -65,31 +79,43 @@ class Index: return sorted(candidates, key=lambda item: item["score"], reverse=True)[:limit] def rebuild(self): - encoder = self.encoder() - values = [] - for path in self.files(): - for value in chunks(path): - values.append((str(path.relative_to(self.root)), value)) - self.db.execute("delete from chunks") - if values: - vectors = encoder.encode([value for _, value in values], normalize_embeddings=True) - self.db.executemany("insert into chunks values (?, ?, ?)", ((path, value, vector.tobytes()) for (path, value), vector in zip(values, vectors))) - self.db.commit() - return len(values) + with self.lock: + self.refresh(force=True) + 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")) + with self.db: + 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 + 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))) + 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.db.execute("delete from files where path = ?", (relative,)) def search(self, query: str, top_k: int): if not query.strip(): raise ValueError("query must not be empty") - if self.db.execute("select count(*) from chunks").fetchone()[0] == 0: - self.rebuild() - vector = self.encoder().encode([query], normalize_embeddings=True)[0] - rows = [] - for path, text, blob in self.db.execute("select path, text, vector from chunks"): + with self.lock: + self.refresh() + if not self.db.execute("select 1 from chunks limit 1").fetchone(): + return [] import numpy as np - score = float(np.dot(vector, np.frombuffer(blob, dtype=np.float32))) - rows.append({"path": path, "score": round(score, 6), "text": text}) - limit = max(1, min(top_k, 20)) - return self.rerank(query, sorted(rows, key=lambda item: item["score"], reverse=True), limit) + 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")) + limit = max(1, min(top_k, 20)) + return self.rerank(query, heapq.nlargest(20, rows, key=lambda item: item["score"]), limit) def main(): @@ -107,7 +133,7 @@ def main(): @mcp.tool() def semantic_search(query: str, top_k: int = 5) -> list[dict]: - """Find repository passages with Qwen3 embedding and reranking.""" + """Search repository concepts, behavior, and implementation patterns first with Qwen3 embedding and reranking. Verify returned passages against current files; use direct reads for known paths.""" return index.search(query, top_k) @mcp.tool()