Julian lechnerandClaude Sonnet 5
9ad793289d
Each server.py search/rebuild invocation was a fresh process that reloaded the embedding and reranker models from disk before answering a single query. Add an optional daemon (server.py daemon) that keeps both models resident in memory and serves repeated queries, across repository roots, over a local socket/pipe. search and rebuild now try the daemon first, spawn one on a cache miss (guarded by an exclusive lock file to avoid duplicate spawns), and fall back to the original one-shot in-process worker if the daemon cannot be reached or started, so behavior is unchanged when the daemon is unavailable. The daemon exits on its own after 30 minutes of inactivity. Co-Authored-By: Claude Sonnet 5 <noreply@anthropic.com>
518 lines
22 KiB
Python
518 lines
22 KiB
Python
#!/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 heapq
|
|
import json
|
|
import multiprocessing
|
|
from multiprocessing.connection import Client, Listener
|
|
import os
|
|
from pathlib import Path
|
|
import re
|
|
import sqlite3
|
|
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 = 2
|
|
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"}
|
|
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_AUTHKEY = b"ai-config-semantic-retrieval"
|
|
|
|
|
|
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():
|
|
runtime = runtime_python()
|
|
if runtime is not None:
|
|
os.execv(str(runtime), [str(runtime), str(Path(__file__).resolve()), *sys.argv[1:]])
|
|
|
|
|
|
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 in {".js", ".ts", ".tsx", ".cs", ".java", ".go", ".rs", ".c", ".cpp", ".ps1", ".sh"}:
|
|
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|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 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": f"Instruct: {RETRIEVAL_INSTRUCTION}\nQuery: "},
|
|
device=preferred_device())
|
|
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())
|
|
return _shared_reranker
|
|
|
|
|
|
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 daemon_lock_path(data_dir: Path) -> Path:
|
|
return data_dir / "daemon.starting"
|
|
|
|
|
|
def daemon_address(data_dir: Path) -> str:
|
|
digest = hashlib.sha256(os.path.normcase(str(data_dir.resolve())).encode()).hexdigest()[:16]
|
|
if os.name == "nt":
|
|
return r"\\.\pipe\ai-config-retrieval-" + digest
|
|
return str(data_dir / f"daemon-{digest}.sock")
|
|
|
|
|
|
def request_daemon(address: str, payload: dict):
|
|
with Client(address, authkey=DAEMON_AUTHKEY) as connection:
|
|
connection.send(payload)
|
|
if not connection.poll(CLI_TIMEOUT_SECONDS):
|
|
raise TimeoutError(f"Semantic retrieval daemon exceeded the {CLI_TIMEOUT_SECONDS} second timeout.")
|
|
success, result = connection.recv()
|
|
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)
|
|
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)
|
|
except (OSError, EOFError):
|
|
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)
|
|
daemon_lock_path(data_dir).unlink(missing_ok=True)
|
|
try:
|
|
while True:
|
|
connection = listener.accept()
|
|
try:
|
|
response = (True, None)
|
|
try:
|
|
request = connection.recv()
|
|
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(response)
|
|
except (OSError, EOFError):
|
|
pass
|
|
finally:
|
|
connection.close()
|
|
with activity_lock:
|
|
last_activity = time.monotonic()
|
|
finally:
|
|
listener.close()
|
|
|
|
|
|
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)
|
|
identity = json.dumps([os.path.normcase(str(self.root)), MODEL, EMBEDDING_DIMENSIONS,
|
|
INDEX_VERSION, CHUNK_TOKENS, OVERLAP_TOKENS])
|
|
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, 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)")
|
|
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()
|
|
|
|
def files(self):
|
|
ignored = {".git", ".venv", "generated", "node_modules", ".my-ai-configuration"}
|
|
result = []
|
|
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) == self.max_files:
|
|
print(f"Semantic retrieval: stopped scanning after {self.max_files} indexable files; "
|
|
f"some files under {self.root} were not indexed. Set {MAX_FILES_ENV} to raise this limit.",
|
|
file=sys.stderr)
|
|
return result
|
|
return result
|
|
|
|
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 = dict(self.db.execute("select path, digest from files"))
|
|
encoder = None
|
|
changed = []
|
|
documents = []
|
|
document_offsets = []
|
|
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
|
|
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, values, file_documents))
|
|
vectors = encoder.encode(documents, normalize_embeddings=True) if documents else []
|
|
with self.db:
|
|
for number, (relative, digest, 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))
|
|
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")
|
|
|
|
def cache_key(self, query: str, top_k: int):
|
|
files = self.db.execute("select path, digest from files order by path").fetchall()
|
|
state = json.dumps(files, separators=(",", ":"))
|
|
return hashlib.sha256(f"{top_k}\0{query}\0{state}".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,))
|
|
import numpy as np
|
|
vector = self.encoder().encode([query], prompt_name="query", normalize_embeddings=True)[0]
|
|
if len(vector) != EMBEDDING_DIMENSIONS:
|
|
raise ValueError("Embedding must have 1024 dimensions")
|
|
scores = ((identifier, float(np.dot(vector, np.frombuffer(blob, dtype=np.float32))))
|
|
for identifier, blob in self.db.execute("select rowid, vector from chunks"))
|
|
dense = [identifier for identifier, _ in heapq.nlargest(DENSE_CANDIDATES, scores, key=lambda item: item[1])]
|
|
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 run_direct(action: str, root: Path, data_dir: Path, model_cache: Path, query: str | None, top_k: int):
|
|
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", "daemon"))
|
|
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", required=True, type=Path)
|
|
parser.add_argument("--model-cache", required=True, type=Path)
|
|
args = parser.parse_args()
|
|
data_dir = args.data_dir.resolve()
|
|
model_cache = args.model_cache.resolve()
|
|
if args.action == "daemon":
|
|
data_dir.mkdir(parents=True, exist_ok=True)
|
|
run_daemon(data_dir, model_cache)
|
|
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 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()
|