fix(qwen): isolate and refresh indexes with code-search prompts
This commit is contained in:
1 parent
5dbb1b21b4
commit
2079792a74
2 files changed
+167
-28
No files matched your search
@@ -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)
|
||||||
|
|
||||||
|
|||||||
+54
-28
@@ -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)
|
||||||
|
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
|
continue
|
||||||
result.append(path)
|
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)
|
||||||
|
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():
|
for path in self.files():
|
||||||
for value in chunks(path):
|
relative = path.relative_to(self.root).as_posix()
|
||||||
values.append((str(path.relative_to(self.root)), value))
|
content = path.read_bytes()
|
||||||
self.db.execute("delete from chunks")
|
digest = hashlib.sha256(content).hexdigest()
|
||||||
if values:
|
old_digest = previous.pop(relative, None)
|
||||||
vectors = encoder.encode([value for _, value in values], normalize_embeddings=True)
|
if not force and digest == old_digest:
|
||||||
self.db.executemany("insert into chunks values (?, ?, ?)", ((path, value, vector.tobytes()) for (path, value), vector in zip(values, vectors)))
|
continue
|
||||||
self.db.commit()
|
values = list(chunks(content.decode("utf-8", errors="ignore")))
|
||||||
return len(values)
|
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}
|
||||||
|
for path, text, blob in self.db.execute("select path, text, vector from chunks"))
|
||||||
limit = max(1, min(top_k, 20))
|
limit = max(1, min(top_k, 20))
|
||||||
return self.rerank(query, sorted(rows, key=lambda item: item["score"], reverse=True), limit)
|
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()
|
||||||
|
|||||||
Reference in new issue
Block a user