fix(qwen): isolate and refresh indexes with code-search prompts

This commit is contained in:
Julian lechner committed 2026-09-17 15:27:04 +02:00
1 parent 5dbb1b21b4
commit 2079792a74
2 files changed
+171 -32

No files matched your search

+113
View File
@@ -1,6 +1,13 @@
import importlib.util import importlib.util
from array import array
from concurrent.futures import ThreadPoolExecutor
import os
from pathlib import Path from pathlib import Path
from tempfile import TemporaryDirectory
from types import SimpleNamespace
import time
import unittest import unittest
from unittest.mock import patch
ROOT = Path(__file__).resolve().parent.parent ROOT = Path(__file__).resolve().parent.parent
@@ -10,6 +17,112 @@ SPEC.loader.exec_module(MODULE)
class RetrievalServerTests(unittest.TestCase): 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): def test_rerank_uses_cross_encoder_scores(self):
index = MODULE.Index.__new__(MODULE.Index) index = MODULE.Index.__new__(MODULE.Index)
+58 -32
View File
@@ -4,19 +4,22 @@
from __future__ import annotations from __future__ import annotations
import argparse import argparse
import hashlib
import heapq
import json import json
import os import os
from pathlib import Path from pathlib import Path
import sqlite3 import sqlite3
import threading
MODEL = "Qwen/Qwen3-Embedding-0.6B" MODEL = "Qwen/Qwen3-Embedding-0.6B"
RERANKER_MODEL = "Qwen/Qwen3-Reranker-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"} TEXT_EXTENSIONS = {".c", ".cpp", ".cs", ".go", ".java", ".js", ".json", ".md", ".py", ".ps1", ".rs", ".sh", ".toml", ".ts", ".tsx", ".txt", ".yaml", ".yml"}
def chunks(path: Path): def chunks(text: str):
text = path.read_text(encoding="utf-8", errors="ignore")
for start in range(0, len(text), 1050): for start in range(0, len(text), 1050):
value = text[start:start + 1200].strip() value = text[start:start + 1200].strip()
if value: if value:
@@ -28,33 +31,44 @@ class Index:
self.root = root.resolve() self.root = root.resolve()
self.data = data.resolve() self.data = data.resolve()
self.data.mkdir(parents=True, exist_ok=True) 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 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.db.commit()
self.lock = threading.RLock()
self.model = None self.model = None
self.reranker_model = None self.reranker_model = None
def files(self): def files(self):
ignored = {".git", ".venv", "generated", "node_modules", ".my-ai-configuration"} ignored = {".git", ".venv", "generated", "node_modules", ".my-ai-configuration"}
result = [] result = []
for path in self.root.rglob("*"): for directory, directories, filenames in os.walk(self.root):
if len(result) >= 2000 or not path.is_file() or path.suffix.lower() not in TEXT_EXTENSIONS or path.stat().st_size > 524288: directories[:] = sorted(name for name in directories if name not in ignored
continue and not (Path(directory) / name).is_symlink()
if ignored.intersection(path.relative_to(self.root).parts): and (Path(directory) / name).resolve().is_relative_to(self.root)
continue and (Path(directory) / name).resolve() != self.data)
result.append(path) 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 return result
def encoder(self): def encoder(self):
if self.model is None: if self.model is None:
from sentence_transformers import SentenceTransformer from sentence_transformers import SentenceTransformer
self.model = SentenceTransformer(MODEL) self.model = SentenceTransformer(MODEL, prompts={"query": f"Instruct: {RETRIEVAL_INSTRUCTION}\nQuery: "})
return self.model return self.model
def reranker(self): def reranker(self):
if self.reranker_model is None: if self.reranker_model is None:
from sentence_transformers import CrossEncoder 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 return self.reranker_model
def rerank(self, query: str, rows: list[dict], limit: int): 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] return sorted(candidates, key=lambda item: item["score"], reverse=True)[:limit]
def rebuild(self): def rebuild(self):
encoder = self.encoder() with self.lock:
values = [] self.refresh(force=True)
for path in self.files(): return self.db.execute("select count(*) from chunks").fetchone()[0]
for value in chunks(path):
values.append((str(path.relative_to(self.root)), value)) def refresh(self, force=False):
self.db.execute("delete from chunks") previous = dict(self.db.execute("select path, digest from files"))
if values: with self.db:
vectors = encoder.encode([value for _, value in values], normalize_embeddings=True) for path in self.files():
self.db.executemany("insert into chunks values (?, ?, ?)", ((path, value, vector.tobytes()) for (path, value), vector in zip(values, vectors))) relative = path.relative_to(self.root).as_posix()
self.db.commit() content = path.read_bytes()
return len(values) 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): def search(self, query: str, top_k: int):
if not query.strip(): if not query.strip():
raise ValueError("query must not be empty") raise ValueError("query must not be empty")
if self.db.execute("select count(*) from chunks").fetchone()[0] == 0: with self.lock:
self.rebuild() self.refresh()
vector = self.encoder().encode([query], normalize_embeddings=True)[0] if not self.db.execute("select 1 from chunks limit 1").fetchone():
rows = [] return []
for path, text, blob in self.db.execute("select path, text, vector from chunks"):
import numpy as np import numpy as np
score = float(np.dot(vector, np.frombuffer(blob, dtype=np.float32))) vector = self.encoder().encode([query], prompt_name="query", normalize_embeddings=True)[0]
rows.append({"path": path, "score": round(score, 6), "text": text}) rows = ({"path": path, "score": float(np.dot(vector, np.frombuffer(blob, dtype=np.float32))), "text": text}
limit = max(1, min(top_k, 20)) for path, text, blob in self.db.execute("select path, text, vector from chunks"))
return self.rerank(query, sorted(rows, key=lambda item: item["score"], reverse=True), limit) limit = max(1, min(top_k, 20))
return self.rerank(query, heapq.nlargest(20, rows, key=lambda item: item["score"]), limit)
def main(): def main():
@@ -107,7 +133,7 @@ def main():
@mcp.tool() @mcp.tool()
def semantic_search(query: str, top_k: int = 5) -> list[dict]: 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) return index.search(query, top_k)
@mcp.tool() @mcp.tool()