refactor(retrieval): remove the Python Qwen stack and the QMD extension entry
Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
This commit is contained in:
Julian lechnerandClaude Opus 5.5 committed
1 parent
505e6ae084
commit
5bb0bdf62f
13 files changed
+10
-1785
No files matched your search
@@ -1,56 +0,0 @@
|
||||
import argparse
|
||||
import json
|
||||
import os
|
||||
from pathlib import Path
|
||||
import statistics
|
||||
import sys
|
||||
import time
|
||||
|
||||
|
||||
def elapsed(call):
|
||||
started = time.monotonic()
|
||||
call()
|
||||
return time.monotonic() - started
|
||||
|
||||
|
||||
def main():
|
||||
parser = argparse.ArgumentParser()
|
||||
parser.add_argument('--model-cache', required=True, type=Path)
|
||||
parser.add_argument('--max-seconds', type=float, default=float(os.environ.get('QWEN_BENCHMARK_MAX_SECONDS', '8')))
|
||||
args = parser.parse_args()
|
||||
if args.max_seconds <= 0:
|
||||
raise ValueError('--max-seconds must be greater than zero.')
|
||||
os.environ['HF_HOME'] = str(args.model_cache)
|
||||
from server import shared_encoder, shared_reranker
|
||||
|
||||
encoder = shared_encoder()
|
||||
reranker = shared_reranker()
|
||||
query = 'Find the implementation that updates the selected client configuration.'
|
||||
documents = ['Path: scripts/commands/install.py\n\n' + query] * 50
|
||||
print('Benchmark: warming up embedding model...', file=sys.stderr, flush=True)
|
||||
encoder.encode([query], prompt_name='query')
|
||||
print('Benchmark: probing reranker speed...', file=sys.stderr, flush=True)
|
||||
probe = elapsed(lambda: reranker.predict([(query, documents[0])], show_progress_bar=False))
|
||||
if probe * len(documents) > args.max_seconds:
|
||||
print(json.dumps({
|
||||
'embeddingSeconds': None,
|
||||
'rerankingSeconds': round(probe, 3),
|
||||
'maxSeconds': args.max_seconds,
|
||||
'recommended': False,
|
||||
}))
|
||||
return
|
||||
print('Benchmark: measuring full retrieval batch...', file=sys.stderr, flush=True)
|
||||
reranker.predict([(query, document) for document in documents], show_progress_bar=False)
|
||||
embedding = statistics.median(elapsed(lambda: encoder.encode([query], prompt_name='query')) for _ in range(3))
|
||||
reranking = statistics.median(elapsed(lambda: reranker.predict([(query, document) for document in documents], show_progress_bar=False)) for _ in range(3))
|
||||
result = {
|
||||
'embeddingSeconds': round(embedding, 3),
|
||||
'rerankingSeconds': round(reranking, 3),
|
||||
'maxSeconds': args.max_seconds,
|
||||
'recommended': embedding <= args.max_seconds and reranking <= args.max_seconds,
|
||||
}
|
||||
print(json.dumps(result))
|
||||
|
||||
|
||||
if __name__ == '__main__':
|
||||
main()
|
||||
@@ -1,2 +0,0 @@
|
||||
sentence-transformers==5.7.0
|
||||
usearch==2.26.2
|
||||
@@ -1,695 +0,0 @@
|
||||
#!/usr/bin/env python3
|
||||
"""Local CLI semantic search backed by Qwen3 embedding and reranking."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import argparse
|
||||
import ast
|
||||
import hashlib
|
||||
import json
|
||||
import multiprocessing
|
||||
from multiprocessing import AuthenticationError
|
||||
from multiprocessing.connection import Client, Listener
|
||||
import os
|
||||
import queue
|
||||
from pathlib import Path
|
||||
import re
|
||||
import sqlite3
|
||||
import stat
|
||||
import subprocess
|
||||
import sys
|
||||
import threading
|
||||
import time
|
||||
|
||||
|
||||
MODEL = "Qwen/Qwen3-Embedding-0.6B"
|
||||
MODEL_REVISION = "97b0c614be4d77ee51c0cef4e5f07c00f9eb65b3"
|
||||
RERANKER_MODEL = "Qwen/Qwen3-Reranker-0.6B"
|
||||
RERANKER_MODEL_REVISION = "e61197ed45024b0ed8a2d74b80b4d909f1255473"
|
||||
EMBEDDING_DIMENSIONS = 1024
|
||||
CHUNK_TOKENS = 400
|
||||
OVERLAP_TOKENS = 64
|
||||
DENSE_CANDIDATES = 50
|
||||
LEXICAL_CANDIDATES = 50
|
||||
RERANK_CANDIDATES = 50
|
||||
FINAL_RESULTS = 5
|
||||
CLI_TIMEOUT_SECONDS = 300
|
||||
RRF_K = 60
|
||||
INDEX_VERSION = 3
|
||||
RETRIEVAL_INSTRUCTION = "Given a codebase question, retrieve relevant code and documentation that answer the question."
|
||||
# Qwen3-Embedding query format from the model card: no space after "Query:".
|
||||
QUERY_PROMPT = f"Instruct: {RETRIEVAL_INSTRUCTION}\nQuery:"
|
||||
TEXT_EXTENSIONS = {".c", ".cjs", ".cpp", ".cs", ".csproj", ".css", ".go", ".gradle", ".html", ".java", ".js", ".json", ".jsx", ".kt", ".kts",
|
||||
".md", ".mjs", ".props", ".ps1", ".psm1", ".py", ".razor", ".rs", ".scss", ".sh", ".sql", ".targets", ".toml", ".ts",
|
||||
".tsx", ".txt", ".xaml", ".xml", ".yaml", ".yml"}
|
||||
DEFAULT_MAX_FILES = 2000
|
||||
MAX_FILES_ENV = "SEMANTIC_RETRIEVAL_MAX_FILES"
|
||||
|
||||
# A daemon keeps the embedding/reranker models resident in memory across CLI
|
||||
# invocations, since each `server.py search` call is otherwise a fresh process
|
||||
# that would reload both models from disk before it can answer a single query.
|
||||
DAEMON_IDLE_SECONDS = 1800
|
||||
DAEMON_READY_TIMEOUT_SECONDS = 120
|
||||
DAEMON_MESSAGE_LIMIT = 16 * 1024 * 1024
|
||||
|
||||
|
||||
def runtime_python(script=__file__, executable=sys.executable, system=os.name):
|
||||
relative = Path('.venv') / ('Scripts/python.exe' if system == 'nt' else 'bin/python')
|
||||
candidate = (Path(script).resolve().parent / relative)
|
||||
if candidate.is_file() and candidate != Path(executable).resolve():
|
||||
return candidate
|
||||
return None
|
||||
|
||||
|
||||
def use_runtime_python(system=os.name):
|
||||
runtime = runtime_python()
|
||||
if runtime is not None:
|
||||
command = [str(runtime), str(Path(__file__).resolve()), *sys.argv[1:]]
|
||||
if system == 'nt':
|
||||
# os.execv does not quote arguments on Windows, so a multi-word --query would be split.
|
||||
sys.exit(subprocess.run(command).returncode)
|
||||
os.execv(command[0], command)
|
||||
|
||||
|
||||
def sections(text: str, path: str):
|
||||
lines = text.splitlines(keepends=True)
|
||||
boundaries = {0: "<module>", len(lines): "<module>"}
|
||||
suffix = Path(path).suffix.lower()
|
||||
if suffix == ".py":
|
||||
try:
|
||||
tree = ast.parse(text)
|
||||
except (SyntaxError, ValueError):
|
||||
tree = None
|
||||
|
||||
def visit(node, parent):
|
||||
symbol = parent
|
||||
if isinstance(node, (ast.FunctionDef, ast.AsyncFunctionDef, ast.ClassDef)):
|
||||
symbol = f"{parent}.{node.name}" if parent != "<module>" else node.name
|
||||
start = min([node.lineno, *(item.lineno for item in node.decorator_list)]) - 1
|
||||
boundaries[start] = symbol
|
||||
boundaries.setdefault(node.end_lineno, parent)
|
||||
for child in ast.iter_child_nodes(node):
|
||||
visit(child, symbol)
|
||||
|
||||
if tree is not None:
|
||||
visit(tree, "<module>")
|
||||
else:
|
||||
fence = None
|
||||
headings = []
|
||||
for number, line in enumerate(lines):
|
||||
if suffix == ".md":
|
||||
marker = re.match(r"^\s{0,3}(`{3,}|~{3,})", line)
|
||||
if marker:
|
||||
value = marker[1]
|
||||
if fence is None:
|
||||
fence = value
|
||||
elif value[0] == fence[0] and len(value) >= len(fence):
|
||||
fence = None
|
||||
continue
|
||||
heading = re.match(r"^\s{0,3}(#{1,6})\s+(.+?)\s*#*\s*$", line)
|
||||
if heading and fence is None:
|
||||
level = len(heading[1])
|
||||
headings = [(depth, title) for depth, title in headings if depth < level]
|
||||
headings.append((level, heading[2]))
|
||||
boundaries[number] = " / ".join(title for _, title in headings)
|
||||
elif suffix == ".toml":
|
||||
heading = re.match(r"^\s*\[\[?([^\]]+)\]\]?\s*(?:#.*)?$", line)
|
||||
if heading:
|
||||
boundaries[number] = heading[1]
|
||||
elif suffix == ".sql":
|
||||
heading = re.match(r"^\s*create\s+(?:or\s+(?:alter|replace)\s+)?(?:procedure|proc|function|view|table|trigger|index)\s+([\w.\[\]\"]+)",
|
||||
line, re.IGNORECASE)
|
||||
if heading:
|
||||
boundaries[number] = heading[1]
|
||||
elif suffix in {".js", ".jsx", ".mjs", ".cjs", ".ts", ".tsx", ".cs", ".java", ".kt", ".kts", ".go", ".rs", ".c", ".cpp", ".ps1", ".psm1", ".sh"}:
|
||||
# A closing brace in column 0 ends a top-level declaration; code after it belongs to the module again.
|
||||
if line.startswith("}"):
|
||||
boundaries.setdefault(number + 1, "<module>")
|
||||
symbol = re.match(
|
||||
r"^\s*(?:(?:export|default|public|private|protected|internal|static|abstract|async|sealed|partial|pub)\s+)*"
|
||||
r"(?:(?:class|interface|struct|enum|record|function|fun|fn|def)\s+([\w$-]+)"
|
||||
r"|func\s+(?:\([^)]*\)\s*)?([\w]+)"
|
||||
r"|(?:const|let|var)\s+([\w$]+)\s*=.*=>"
|
||||
r"|([\w$-]+)\s*\([^;]*\)\s*\{"
|
||||
r"|[\w<>\[\],?]+\s+([\w]+)\s*\([^;]*\)\s*(?:\{|=>|$))", line)
|
||||
if symbol:
|
||||
name = next(value for value in symbol.groups() if value)
|
||||
if name not in {"if", "for", "foreach", "while", "switch", "catch", "using", "lock"}:
|
||||
boundaries[number] = name
|
||||
positions = sorted(boundaries)
|
||||
for start, end in zip(positions, positions[1:]):
|
||||
content = "".join(lines[start:end])
|
||||
if content.strip():
|
||||
yield boundaries[start], content
|
||||
|
||||
|
||||
def chunks(text: str, path: str, tokenizer):
|
||||
for symbol, content in sections(text, path):
|
||||
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}"
|
||||
|
||||
|
||||
def run_action(action: str, root: Path, data_dir: Path, model_cache: Path, query: str | None, top_k: int, output):
|
||||
model_cache.mkdir(parents=True, exist_ok=True)
|
||||
os.environ.setdefault("HF_HOME", str(model_cache.resolve()))
|
||||
index = Index(root, data_dir)
|
||||
result = index.search(query, top_k) if action == "search" else {"chunks": index.rebuild()}
|
||||
output.put((True, result))
|
||||
|
||||
|
||||
def preferred_device():
|
||||
import torch
|
||||
return "cuda" if torch.cuda.is_available() else "cpu"
|
||||
|
||||
|
||||
def model_kwargs():
|
||||
"""Use bfloat16 on GPUs that support it; CPUs keep float32, where bfloat16 is often slower."""
|
||||
import torch
|
||||
return {"dtype": torch.bfloat16} if torch.cuda.is_available() and torch.cuda.is_bf16_supported() else {}
|
||||
|
||||
|
||||
def read_max_files():
|
||||
value = os.environ.get(MAX_FILES_ENV)
|
||||
if value is None:
|
||||
return DEFAULT_MAX_FILES
|
||||
try:
|
||||
parsed = int(value)
|
||||
except ValueError:
|
||||
parsed = 0
|
||||
if parsed <= 0:
|
||||
raise ValueError(f"{MAX_FILES_ENV} must be a positive integer, got {value!r}")
|
||||
return parsed
|
||||
|
||||
|
||||
_shared_model_lock = threading.Lock()
|
||||
_shared_encoder = None
|
||||
_shared_reranker = None
|
||||
|
||||
|
||||
def shared_encoder():
|
||||
global _shared_encoder
|
||||
with _shared_model_lock:
|
||||
if _shared_encoder is None:
|
||||
from sentence_transformers import SentenceTransformer
|
||||
_shared_encoder = SentenceTransformer(MODEL, revision=MODEL_REVISION, truncate_dim=EMBEDDING_DIMENSIONS,
|
||||
prompts={"query": QUERY_PROMPT}, device=preferred_device(),
|
||||
model_kwargs=model_kwargs())
|
||||
return _shared_encoder
|
||||
|
||||
|
||||
def shared_reranker():
|
||||
global _shared_reranker
|
||||
with _shared_model_lock:
|
||||
if _shared_reranker is None:
|
||||
from sentence_transformers import CrossEncoder
|
||||
_shared_reranker = CrossEncoder(RERANKER_MODEL, revision=RERANKER_MODEL_REVISION, prompts={"query": RETRIEVAL_INSTRUCTION},
|
||||
default_prompt_name="query", device=preferred_device(),
|
||||
model_kwargs=model_kwargs())
|
||||
return _shared_reranker
|
||||
|
||||
|
||||
def preload_models(model_cache: Path):
|
||||
model_cache.mkdir(parents=True, exist_ok=True)
|
||||
os.environ["HF_HOME"] = str(model_cache.resolve())
|
||||
query = "Verify semantic retrieval models."
|
||||
embeddings = shared_encoder().encode([query], prompt_name="query")
|
||||
if len(embeddings) != 1 or len(embeddings[0]) != EMBEDDING_DIMENSIONS:
|
||||
raise ValueError("Embedding model returned unexpected dimensions.")
|
||||
scores = shared_reranker().predict([(query, query)], show_progress_bar=False)
|
||||
if len(scores) != 1:
|
||||
raise ValueError("Reranking model returned unexpected scores.")
|
||||
|
||||
|
||||
def reciprocal_rank_fusion(*rankings):
|
||||
scores = {}
|
||||
for ranking in rankings:
|
||||
for rank, identifier in enumerate(ranking, 1):
|
||||
scores[identifier] = scores.get(identifier, 0.0) + 1.0 / (RRF_K + rank)
|
||||
return sorted(scores, key=lambda identifier: (-scores[identifier], identifier))
|
||||
|
||||
|
||||
def code_version() -> str:
|
||||
"""Hash of this file, so a daemon started by older code is never reused after an update."""
|
||||
return hashlib.sha256(Path(__file__).read_bytes()).hexdigest()[:12]
|
||||
|
||||
|
||||
def daemon_lock_path(data_dir: Path) -> Path:
|
||||
return data_dir / f"daemon-{code_version()}.starting"
|
||||
|
||||
|
||||
def daemon_address(data_dir: Path) -> str:
|
||||
identity = os.path.normcase(str(data_dir.resolve())) + "|" + code_version()
|
||||
digest = hashlib.sha256(identity.encode()).hexdigest()[:16]
|
||||
if os.name == "nt":
|
||||
return r"\\.\pipe\ai-config-retrieval-" + digest
|
||||
return str(data_dir / f"daemon-{digest}.sock")
|
||||
|
||||
|
||||
def daemon_authkey(data_dir: Path) -> bytes:
|
||||
if data_dir.is_symlink():
|
||||
raise OSError("Semantic retrieval data directory must not be a link.")
|
||||
data_dir.mkdir(parents=True, exist_ok=True, mode=0o700)
|
||||
path = data_dir / "daemon.key"
|
||||
try:
|
||||
descriptor = os.open(path, os.O_CREAT | os.O_EXCL | os.O_WRONLY, 0o600)
|
||||
except FileExistsError:
|
||||
pass
|
||||
else:
|
||||
with os.fdopen(descriptor, "wb") as stream:
|
||||
stream.write(os.urandom(32))
|
||||
stream.flush()
|
||||
os.fsync(stream.fileno())
|
||||
metadata = path.lstat()
|
||||
if not stat.S_ISREG(metadata.st_mode) or (os.name != "nt" and metadata.st_mode & 0o077):
|
||||
raise OSError("Semantic retrieval daemon credential is not private.")
|
||||
key = path.read_bytes()
|
||||
if len(key) != 32:
|
||||
raise OSError("Semantic retrieval daemon credential is incomplete.")
|
||||
return key
|
||||
|
||||
|
||||
def request_daemon(address: str, payload: dict, data_dir: Path):
|
||||
with Client(address, authkey=daemon_authkey(data_dir)) as connection:
|
||||
connection.send_bytes(json.dumps(payload).encode("utf-8"))
|
||||
if not connection.poll(CLI_TIMEOUT_SECONDS):
|
||||
raise TimeoutError(f"Semantic retrieval daemon exceeded the {CLI_TIMEOUT_SECONDS} second timeout.")
|
||||
success, result = json.loads(connection.recv_bytes(DAEMON_MESSAGE_LIMIT))
|
||||
if not success:
|
||||
raise RuntimeError(str(result))
|
||||
return result
|
||||
|
||||
|
||||
def spawn_daemon(data_dir: Path, model_cache: Path) -> bool:
|
||||
lock_path = daemon_lock_path(data_dir)
|
||||
try:
|
||||
os.close(os.open(lock_path, os.O_CREAT | os.O_EXCL | os.O_WRONLY))
|
||||
except FileExistsError:
|
||||
try:
|
||||
if time.time() - lock_path.stat().st_mtime > DAEMON_READY_TIMEOUT_SECONDS:
|
||||
lock_path.unlink(missing_ok=True)
|
||||
except OSError:
|
||||
pass
|
||||
return True
|
||||
try:
|
||||
subprocess.Popen(
|
||||
[sys.executable, str(Path(__file__).resolve()), "daemon", "--data-dir", str(data_dir), "--model-cache", str(model_cache)],
|
||||
stdin=subprocess.DEVNULL, stdout=subprocess.DEVNULL, stderr=subprocess.DEVNULL,
|
||||
start_new_session=(os.name != "nt"),
|
||||
creationflags=(subprocess.DETACHED_PROCESS | subprocess.CREATE_NEW_PROCESS_GROUP) if os.name == "nt" else 0)
|
||||
return True
|
||||
except OSError:
|
||||
lock_path.unlink(missing_ok=True)
|
||||
return False
|
||||
|
||||
|
||||
def try_daemon(root: Path, data_dir: Path, model_cache: Path, action: str, query: str | None, top_k: int):
|
||||
address = daemon_address(data_dir)
|
||||
payload = {"action": action, "root": str(root), "query": query, "top_k": top_k}
|
||||
try:
|
||||
return request_daemon(address, payload, data_dir)
|
||||
except AuthenticationError:
|
||||
return None
|
||||
except (OSError, EOFError):
|
||||
pass
|
||||
if not spawn_daemon(data_dir, model_cache):
|
||||
return None
|
||||
deadline = time.monotonic() + DAEMON_READY_TIMEOUT_SECONDS
|
||||
while time.monotonic() < deadline:
|
||||
try:
|
||||
return request_daemon(address, payload, data_dir)
|
||||
except (OSError, EOFError, AuthenticationError):
|
||||
time.sleep(0.5)
|
||||
return None
|
||||
|
||||
|
||||
def run_daemon(data_dir: Path, model_cache: Path):
|
||||
address = daemon_address(data_dir)
|
||||
if os.name != "nt":
|
||||
Path(address).unlink(missing_ok=True)
|
||||
model_cache.mkdir(parents=True, exist_ok=True)
|
||||
os.environ.setdefault("HF_HOME", str(model_cache.resolve()))
|
||||
indexes: dict[str, Index] = {}
|
||||
activity_lock = threading.Lock()
|
||||
last_activity = time.monotonic()
|
||||
|
||||
def watchdog():
|
||||
while True:
|
||||
time.sleep(5)
|
||||
with activity_lock:
|
||||
idle_for = time.monotonic() - last_activity
|
||||
if idle_for > DAEMON_IDLE_SECONDS:
|
||||
os._exit(0)
|
||||
|
||||
threading.Thread(target=watchdog, daemon=True).start()
|
||||
listener = Listener(address, family="AF_PIPE" if os.name == "nt" else "AF_UNIX", authkey=daemon_authkey(data_dir))
|
||||
daemon_lock_path(data_dir).unlink(missing_ok=True)
|
||||
try:
|
||||
while True:
|
||||
connection = listener.accept()
|
||||
try:
|
||||
response = (True, None)
|
||||
try:
|
||||
request = json.loads(connection.recv_bytes(DAEMON_MESSAGE_LIMIT))
|
||||
if isinstance(request, dict) and request.get("action") == "stop":
|
||||
connection.send_bytes(json.dumps((True, None)).encode("utf-8"))
|
||||
os._exit(0)
|
||||
if not isinstance(request, dict) or request.get("action") not in ("search", "rebuild") or not isinstance(request.get("root"), str):
|
||||
raise ValueError("Invalid semantic retrieval daemon request.")
|
||||
index = indexes.get(request["root"])
|
||||
if index is None:
|
||||
if len(indexes) >= 8:
|
||||
indexes.pop(next(iter(indexes))).db.close()
|
||||
index = Index(Path(request["root"]), data_dir)
|
||||
indexes[request["root"]] = index
|
||||
if request["action"] == "search":
|
||||
response = (True, index.search(request["query"], request["top_k"]))
|
||||
else:
|
||||
response = (True, {"chunks": index.rebuild()})
|
||||
except EOFError:
|
||||
continue
|
||||
except Exception as error:
|
||||
response = (False, str(error))
|
||||
connection.send_bytes(json.dumps(response).encode("utf-8"))
|
||||
except (OSError, EOFError):
|
||||
pass
|
||||
finally:
|
||||
connection.close()
|
||||
with activity_lock:
|
||||
last_activity = time.monotonic()
|
||||
finally:
|
||||
listener.close()
|
||||
|
||||
|
||||
def index_database(root: Path, data: Path) -> Path:
|
||||
identity = json.dumps([os.path.normcase(str(root.resolve())), MODEL, MODEL_REVISION,
|
||||
EMBEDDING_DIMENSIONS, RETRIEVAL_INSTRUCTION, INDEX_VERSION,
|
||||
CHUNK_TOKENS, OVERLAP_TOKENS])
|
||||
return data.resolve() / (hashlib.sha256(identity.encode()).hexdigest() + ".sqlite3")
|
||||
|
||||
|
||||
def daemon_running(data_dir: Path) -> bool:
|
||||
if not pipes_available():
|
||||
return False
|
||||
try:
|
||||
Client(daemon_address(data_dir), authkey=daemon_authkey(data_dir)).close()
|
||||
return True
|
||||
except (OSError, EOFError, AuthenticationError):
|
||||
return False
|
||||
|
||||
|
||||
def index_status(root: Path, data_dir: Path) -> dict:
|
||||
"""Report what the next search costs without loading models or starting the daemon."""
|
||||
database = index_database(root, data_dir)
|
||||
previous = {}
|
||||
if database.is_file():
|
||||
connection = sqlite3.connect(database.as_uri() + "?mode=ro", uri=True)
|
||||
try:
|
||||
previous = {path: (size, mtime, ctime) for path, size, mtime, ctime
|
||||
in connection.execute("select path, size, mtime, ctime from files")}
|
||||
except sqlite3.OperationalError:
|
||||
pass
|
||||
finally:
|
||||
connection.close()
|
||||
files = indexable_files(root.resolve(), data_dir.resolve(), read_max_files())
|
||||
stale = 0
|
||||
for path in files:
|
||||
metadata = path.stat()
|
||||
stale += previous.get(path.relative_to(root.resolve()).as_posix()) != (metadata.st_size, metadata.st_mtime_ns, metadata.st_ctime_ns)
|
||||
return {"indexable_files": len(files), "indexed_files": len(previous), "stale_files": stale,
|
||||
"models_loaded": daemon_running(data_dir), "in_process": not pipes_available()}
|
||||
|
||||
|
||||
def indexable_files(root: Path, data: Path, max_files: int):
|
||||
ignored = {".git", ".venv", "generated", "node_modules", ".my-ai-configuration", "bin", "obj"}
|
||||
output = subprocess.run(["git", "-C", str(root), "ls-files", "--cached", "--others",
|
||||
"--exclude-standard", "-z", "--", "."],
|
||||
check=True, capture_output=True, timeout=30).stdout
|
||||
candidates = [root / name for name in
|
||||
output.decode("utf-8", errors="surrogateescape").split("\0") if name
|
||||
and not any(part in ignored for part in Path(name).parts)]
|
||||
result = []
|
||||
for candidate in sorted(candidates):
|
||||
if candidate.is_symlink():
|
||||
continue
|
||||
path = candidate.resolve()
|
||||
if (not path.is_relative_to(root) or path.is_relative_to(data)
|
||||
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) == max_files:
|
||||
print(f"Semantic retrieval: stopped scanning after {max_files} indexable files; "
|
||||
f"some files under {root} were not indexed. Set {MAX_FILES_ENV} to raise this limit.",
|
||||
file=sys.stderr)
|
||||
return result
|
||||
return result
|
||||
|
||||
|
||||
class Index:
|
||||
def __init__(self, root: Path, data: Path):
|
||||
self.root = root.resolve()
|
||||
self.data = data.resolve()
|
||||
self.data.mkdir(parents=True, exist_ok=True)
|
||||
self.db = sqlite3.connect(index_database(self.root, self.data), check_same_thread=False, timeout=30)
|
||||
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 virtual table if not exists lexical using fts5(document)")
|
||||
self.db.execute("create table if not exists files (path text primary key, digest text, size integer, mtime integer, ctime integer)")
|
||||
columns = {row[1] for row in self.db.execute("pragma table_info(files)")}
|
||||
for name in ("size", "mtime", "ctime"):
|
||||
if name not in columns:
|
||||
self.db.execute(f"alter table files add column {name} integer")
|
||||
self.db.execute("create table if not exists query_cache (cache_key text primary key, result text not null)")
|
||||
self.db.commit()
|
||||
self.lock = threading.RLock()
|
||||
self.max_files = read_max_files()
|
||||
self.vector_index = None
|
||||
|
||||
def files(self):
|
||||
return indexable_files(self.root, self.data, self.max_files)
|
||||
|
||||
def encoder(self):
|
||||
return shared_encoder()
|
||||
|
||||
def reranker(self):
|
||||
return shared_reranker()
|
||||
|
||||
def rerank(self, query: str, rows: list[dict], limit: int):
|
||||
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)[:min(limit, FINAL_RESULTS)]
|
||||
|
||||
def delete_path(self, path):
|
||||
self.db.execute("delete from lexical where rowid in (select rowid from chunks where path = ?)", (path,))
|
||||
self.db.execute("delete from chunks where path = ?", (path,))
|
||||
|
||||
def rebuild(self):
|
||||
with self.lock:
|
||||
self.refresh(force=True)
|
||||
return self.db.execute("select count(*) from chunks").fetchone()[0]
|
||||
|
||||
def refresh(self, force=False):
|
||||
previous = {path: (digest, size, mtime, ctime)
|
||||
for path, digest, size, mtime, ctime in self.db.execute("select path, digest, size, mtime, ctime from files")}
|
||||
encoder = None
|
||||
changed = []
|
||||
documents = []
|
||||
document_offsets = []
|
||||
for path in self.files():
|
||||
relative = path.relative_to(self.root).as_posix()
|
||||
metadata = path.stat()
|
||||
size, mtime, ctime = metadata.st_size, metadata.st_mtime_ns, metadata.st_ctime_ns
|
||||
old_digest, old_size, old_mtime, old_ctime = previous.pop(relative, (None, None, None, None))
|
||||
if not force and old_digest and (size, mtime, ctime) == (old_size, old_mtime, old_ctime):
|
||||
continue
|
||||
content = path.read_bytes()
|
||||
digest = hashlib.sha256(content).hexdigest()
|
||||
if not force and digest == old_digest:
|
||||
with self.db:
|
||||
self.db.execute("update files set size = ?, mtime = ?, ctime = ? where path = ?",
|
||||
(size, mtime, ctime, relative))
|
||||
continue
|
||||
if encoder is None:
|
||||
encoder = self.encoder()
|
||||
values = list(chunks(content.decode("utf-8", errors="ignore"), relative, encoder.tokenizer))
|
||||
file_documents = [document_text(relative, symbol, value) for symbol, value in values]
|
||||
document_offsets.append(len(documents))
|
||||
documents.extend(file_documents)
|
||||
changed.append((relative, digest, size, mtime, ctime, values, file_documents))
|
||||
vectors = encoder.encode(documents, normalize_embeddings=True) if documents else []
|
||||
with self.db:
|
||||
for number, (relative, digest, size, mtime, ctime, values, file_documents) in enumerate(changed):
|
||||
start = document_offsets[number]
|
||||
file_vectors = vectors[start:start + len(file_documents)]
|
||||
self.delete_path(relative)
|
||||
for (symbol, value), document, vector in zip(values, file_documents, file_vectors):
|
||||
if len(vector) != EMBEDDING_DIMENSIONS:
|
||||
raise ValueError("Embedding must have 1024 dimensions")
|
||||
cursor = self.db.execute("insert into chunks values (?, ?, ?, ?)",
|
||||
(relative, symbol, value, vector.tobytes()))
|
||||
self.db.execute("insert into lexical(rowid, document) values (?, ?)", (cursor.lastrowid, document))
|
||||
self.db.execute("insert or replace into files values (?, ?, ?, ?, ?)",
|
||||
(relative, digest, size, mtime, ctime))
|
||||
for relative in previous:
|
||||
self.delete_path(relative)
|
||||
self.db.execute("delete from files where path = ?", (relative,))
|
||||
if changed or previous:
|
||||
self.db.execute("delete from query_cache")
|
||||
self.vector_index = None
|
||||
|
||||
def dense_candidates(self, vector):
|
||||
import numpy as np
|
||||
from usearch.index import Index as VectorIndex
|
||||
|
||||
if self.vector_index is None:
|
||||
index = VectorIndex(ndim=EMBEDDING_DIMENSIONS, metric="cos", dtype="f32",
|
||||
connectivity=16, expansion_add=128, expansion_search=64)
|
||||
cursor = self.db.execute("select rowid, vector from chunks")
|
||||
while rows := cursor.fetchmany(4096):
|
||||
vectors = np.frombuffer(b"".join(row[1] for row in rows), dtype=np.float32).reshape(-1, EMBEDDING_DIMENSIONS)
|
||||
index.add([row[0] for row in rows], vectors)
|
||||
self.vector_index = index
|
||||
return [int(match.key) for match in self.vector_index.search(vector, DENSE_CANDIDATES)]
|
||||
|
||||
def cache_key(self, query: str, top_k: int):
|
||||
identity = json.dumps([
|
||||
query, top_k,
|
||||
RERANKER_MODEL, RERANKER_MODEL_REVISION, RETRIEVAL_INSTRUCTION, QUERY_PROMPT,
|
||||
DENSE_CANDIDATES, LEXICAL_CANDIDATES, RERANK_CANDIDATES, FINAL_RESULTS, RRF_K,
|
||||
], ensure_ascii=False, separators=(",", ":"))
|
||||
return hashlib.sha256(identity.encode("utf-8")).hexdigest()
|
||||
|
||||
def search(self, query: str, top_k: int = FINAL_RESULTS):
|
||||
if not query.strip():
|
||||
raise ValueError("query must not be empty")
|
||||
with self.lock:
|
||||
self.refresh()
|
||||
if not self.db.execute("select 1 from chunks limit 1").fetchone():
|
||||
return []
|
||||
key = self.cache_key(query, top_k)
|
||||
cached = self.db.execute("select result from query_cache where cache_key = ?", (key,)).fetchone()
|
||||
if cached:
|
||||
try:
|
||||
return json.loads(cached[0])
|
||||
except (TypeError, json.JSONDecodeError):
|
||||
with self.db:
|
||||
self.db.execute("delete from query_cache where cache_key = ?", (key,))
|
||||
vector = self.encoder().encode([query], prompt_name="query", normalize_embeddings=True)[0]
|
||||
if len(vector) != EMBEDDING_DIMENSIONS:
|
||||
raise ValueError("Embedding must have 1024 dimensions")
|
||||
dense = self.dense_candidates(vector)
|
||||
terms = list(dict.fromkeys(re.findall(r"[^\W_]+", query)))[:128]
|
||||
expression = " OR ".join(f'"{term}"' for term in terms)
|
||||
lexical = [row[0] for row in self.db.execute(
|
||||
"select rowid from lexical where lexical match ? order by bm25(lexical), rowid limit ?",
|
||||
(expression, LEXICAL_CANDIDATES))] if terms else []
|
||||
identifiers = reciprocal_rank_fusion(dense, lexical)[:RERANK_CANDIDATES]
|
||||
placeholders = ",".join("?" for _ in identifiers)
|
||||
rows = {identifier: {"path": path, "symbol": symbol, "text": text}
|
||||
for identifier, path, symbol, text in self.db.execute(
|
||||
f"select rowid, path, symbol, text from chunks where rowid in ({placeholders})", identifiers)}
|
||||
candidates = [rows[identifier] for identifier in identifiers]
|
||||
result = self.rerank(query, candidates, max(1, min(top_k, FINAL_RESULTS)))
|
||||
with self.db:
|
||||
self.db.execute("insert or replace into query_cache values (?, ?)",
|
||||
(key, json.dumps(result, ensure_ascii=False)))
|
||||
return result
|
||||
|
||||
|
||||
def pipes_available() -> bool:
|
||||
"""Sandboxes such as Codex on Windows forbid the pipes that the daemon and worker process need."""
|
||||
try:
|
||||
first, second = multiprocessing.Pipe()
|
||||
except PermissionError:
|
||||
return False
|
||||
first.close()
|
||||
second.close()
|
||||
return True
|
||||
|
||||
|
||||
def run_direct(action: str, root: Path, data_dir: Path, model_cache: Path, query: str | None, top_k: int):
|
||||
if not pipes_available():
|
||||
# ponytail: in-process fallback has no worker timeout; the client's own command timeout still applies.
|
||||
output = queue.Queue(1)
|
||||
run_action(action, root, data_dir, model_cache, query, top_k, output)
|
||||
return output.get()[1]
|
||||
output = multiprocessing.Queue(1)
|
||||
worker = multiprocessing.Process(target=run_action, args=(action, root, data_dir, model_cache, query, top_k, output))
|
||||
worker.start()
|
||||
worker.join(CLI_TIMEOUT_SECONDS)
|
||||
if worker.is_alive():
|
||||
worker.terminate()
|
||||
worker.join()
|
||||
raise TimeoutError(f"Qwen retrieval exceeded the {CLI_TIMEOUT_SECONDS} second timeout.")
|
||||
if worker.exitcode != 0 or output.empty():
|
||||
raise RuntimeError("Qwen retrieval worker failed.")
|
||||
success, result = output.get()
|
||||
if not success:
|
||||
raise RuntimeError(str(result))
|
||||
return result
|
||||
|
||||
|
||||
def main():
|
||||
parser = argparse.ArgumentParser()
|
||||
parser.add_argument("action", choices=("search", "rebuild", "status", "stop", "daemon", "preload"))
|
||||
parser.add_argument("--query")
|
||||
parser.add_argument("--top-k", type=int, default=FINAL_RESULTS)
|
||||
parser.add_argument("--root", default=os.getcwd(), type=Path)
|
||||
parser.add_argument("--data-dir", type=Path)
|
||||
parser.add_argument("--model-cache", required=True, type=Path)
|
||||
args = parser.parse_args()
|
||||
model_cache = args.model_cache.resolve()
|
||||
if args.action == "preload":
|
||||
preload_models(model_cache)
|
||||
return
|
||||
if args.data_dir is None:
|
||||
parser.error(f"{args.action} requires --data-dir")
|
||||
data_dir = args.data_dir.resolve()
|
||||
if args.action == "daemon":
|
||||
data_dir.mkdir(parents=True, exist_ok=True)
|
||||
run_daemon(data_dir, model_cache)
|
||||
return
|
||||
if args.action == "status":
|
||||
print(json.dumps(index_status(args.root, data_dir)))
|
||||
return
|
||||
if args.action == "stop":
|
||||
try:
|
||||
request_daemon(daemon_address(data_dir), {"action": "stop"}, data_dir)
|
||||
stopped = True
|
||||
except (OSError, EOFError, AuthenticationError):
|
||||
stopped = False
|
||||
print(json.dumps({"stopped": stopped}))
|
||||
return
|
||||
if args.action == "search" and not args.query:
|
||||
parser.error("search requires --query")
|
||||
root = args.root.resolve()
|
||||
data_dir.mkdir(parents=True, exist_ok=True)
|
||||
result = try_daemon(root, data_dir, model_cache, args.action, args.query, args.top_k) if pipes_available() else None
|
||||
if result is None:
|
||||
result = run_direct(args.action, root, data_dir, model_cache, args.query, args.top_k)
|
||||
print(json.dumps(result, ensure_ascii=False))
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
use_runtime_python()
|
||||
main()
|
||||
Reference in new issue
Block a user