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

+58 -32
View File
@@ -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()