feat(retrieval): add optional Qwen semantic layer
This commit is contained in:
1 parent
bfbe001248
commit
27f1c1ec79
6 files changed
+311
No files matched your search
@@ -0,0 +1,9 @@
|
|||||||
|
function Sync-SemanticRetrieval {
|
||||||
|
param([Parameter(Mandatory=$true)][string]$RepositoryRoot,[Parameter(Mandatory=$true)][string]$HomePath,[ValidateSet('Codex','Claude','Both')][Parameter(Mandatory=$true)][string]$Client,[Parameter(Mandatory=$true)][bool]$Enabled,[switch]$DryRun,[switch]$Update,[switch]$Summary)
|
||||||
|
$statePath = Join-Path $HomePath '.my-ai-configuration/semantic-retrieval.json'
|
||||||
|
if (-not $Enabled -and -not (Test-Path -LiteralPath $statePath)) { return }
|
||||||
|
$arguments = @((Join-Path $PSScriptRoot 'semantic-retrieval.py'), 'sync', '--root', $RepositoryRoot, '--home', $HomePath, '--client', $Client.ToLowerInvariant(), '--enabled', $Enabled.ToString().ToLowerInvariant())
|
||||||
|
if ($DryRun) { $arguments += '--dry-run' }; if ($Update) { $arguments += '--update' }
|
||||||
|
& python @arguments
|
||||||
|
if ($LASTEXITCODE -ne 0) { throw 'Semantic retrieval reconciliation failed.' }
|
||||||
|
}
|
||||||
@@ -0,0 +1,147 @@
|
|||||||
|
import argparse
|
||||||
|
import hashlib
|
||||||
|
import json
|
||||||
|
import os
|
||||||
|
from pathlib import Path
|
||||||
|
import shutil
|
||||||
|
import subprocess
|
||||||
|
import sys
|
||||||
|
|
||||||
|
|
||||||
|
NAME = 'my-ai-qwen3-retrieval'
|
||||||
|
FILES = ('server.py', 'requirements.txt')
|
||||||
|
|
||||||
|
|
||||||
|
def digest(path):
|
||||||
|
return hashlib.sha256(path.read_bytes()).hexdigest()
|
||||||
|
|
||||||
|
|
||||||
|
def command(client, arguments, home, dry_run):
|
||||||
|
if dry_run:
|
||||||
|
print('DRYRUN ' + ' '.join(arguments))
|
||||||
|
return True
|
||||||
|
executable = shutil.which(client)
|
||||||
|
if executable is None:
|
||||||
|
raise ValueError(f'{client} CLI is required for semantic retrieval.')
|
||||||
|
result = subprocess.run([executable, *arguments[1:]], stdin=subprocess.DEVNULL, capture_output=True, text=True,
|
||||||
|
timeout=180, env=dict(os.environ, HOME=str(home), USERPROFILE=str(home), CODEX_HOME=str(home / '.codex'), CLAUDE_CONFIG_DIR=str(home / '.claude')))
|
||||||
|
if result.returncode:
|
||||||
|
raise ValueError(f'Semantic retrieval command failed: {client} {" ".join(arguments[1:])}\n{result.stdout}{result.stderr}')
|
||||||
|
return True
|
||||||
|
|
||||||
|
|
||||||
|
def python_path(target):
|
||||||
|
return target / ('Scripts/python.exe' if os.name == 'nt' else 'bin/python')
|
||||||
|
|
||||||
|
|
||||||
|
def ensure_runtime(target, dry_run):
|
||||||
|
if dry_run:
|
||||||
|
return python_path(target / '.venv')
|
||||||
|
venv = target / '.venv'
|
||||||
|
if not venv.exists():
|
||||||
|
subprocess.run([sys.executable, '-m', 'venv', str(venv)], check=True, timeout=180)
|
||||||
|
python = python_path(venv)
|
||||||
|
if not python.is_file():
|
||||||
|
raise ValueError('Semantic retrieval virtual environment is incomplete.')
|
||||||
|
subprocess.run([str(python), '-m', 'pip', 'install', '--disable-pip-version-check', '-r', str(target / 'requirements.txt')],
|
||||||
|
check=True, timeout=900)
|
||||||
|
return python
|
||||||
|
|
||||||
|
|
||||||
|
def load_state(path):
|
||||||
|
if not path.exists():
|
||||||
|
return {'version': 1, 'clients': [], 'files': {}}
|
||||||
|
state = json.loads(path.read_text(encoding='utf-8'))
|
||||||
|
if state.get('version') != 1 or not isinstance(state.get('clients'), list) or not isinstance(state.get('files'), dict):
|
||||||
|
raise ValueError('Invalid semantic retrieval ownership state.')
|
||||||
|
if any(client not in ('codex', 'claude') for client in state['clients']):
|
||||||
|
raise ValueError('Invalid semantic retrieval client state.')
|
||||||
|
if any(name not in FILES or not isinstance(value, str) or len(value) != 64 or any(character not in '0123456789abcdef' for character in value) for name, value in state['files'].items()):
|
||||||
|
raise ValueError('Invalid semantic retrieval file state.')
|
||||||
|
return state
|
||||||
|
|
||||||
|
|
||||||
|
def save_state(path, state, dry_run):
|
||||||
|
if not dry_run:
|
||||||
|
path.parent.mkdir(parents=True, exist_ok=True)
|
||||||
|
path.write_text(json.dumps(state, indent=2) + '\n', encoding='utf-8')
|
||||||
|
|
||||||
|
|
||||||
|
def sync(root, home, clients, enabled, dry_run, update):
|
||||||
|
base = home / '.my-ai-configuration'
|
||||||
|
target = base / 'semantic-retrieval'
|
||||||
|
state_path = base / 'semantic-retrieval.json'
|
||||||
|
state = load_state(state_path)
|
||||||
|
clients = list(dict.fromkeys(clients))
|
||||||
|
if base.exists() and base.is_symlink():
|
||||||
|
raise ValueError('Semantic retrieval state directory must not be a symbolic link.')
|
||||||
|
if not enabled:
|
||||||
|
owned = [client for client in state['clients'] if client in clients]
|
||||||
|
if dry_run:
|
||||||
|
print('DRYRUN disable semantic retrieval: ' + (', '.join(owned) if owned else 'nothing owned'))
|
||||||
|
return
|
||||||
|
for client in owned:
|
||||||
|
command(client, [client, 'mcp', 'remove', NAME] + (['--scope', 'user'] if client == 'claude' else []), home, False)
|
||||||
|
state['clients'] = [client for client in state['clients'] if client not in clients]
|
||||||
|
if state['clients']:
|
||||||
|
save_state(state_path, state, False)
|
||||||
|
return
|
||||||
|
if target.is_symlink():
|
||||||
|
raise ValueError('Semantic retrieval directory must not be a symbolic link.')
|
||||||
|
if target.exists() and target.is_dir():
|
||||||
|
shutil.rmtree(target)
|
||||||
|
state_path.unlink(missing_ok=True)
|
||||||
|
return
|
||||||
|
|
||||||
|
if target.exists() and not state['files']:
|
||||||
|
raise ValueError('Semantic retrieval directory exists without setup ownership; refusing to overwrite it.')
|
||||||
|
if dry_run:
|
||||||
|
print('DRYRUN enable semantic retrieval: Qwen3-Embedding-0.6B')
|
||||||
|
return
|
||||||
|
target.mkdir(parents=True, exist_ok=True)
|
||||||
|
source = root / 'shared' / 'retrieval'
|
||||||
|
for name in FILES:
|
||||||
|
shutil.copy2(source / name, target / name)
|
||||||
|
state['files'][name] = digest(target / name)
|
||||||
|
python = ensure_runtime(target, False)
|
||||||
|
existing_clients = set(state['clients'])
|
||||||
|
state['clients'] = list(dict.fromkeys(state['clients']))
|
||||||
|
save_state(state_path, state, False)
|
||||||
|
arguments = ['--', str(python), str(target / 'server.py'), '--root', '.', '--data-dir', str(target / 'data'), '--model-cache', str(target / 'model-cache')]
|
||||||
|
for client in clients:
|
||||||
|
command_args = [client, 'mcp', 'add']
|
||||||
|
if client == 'claude':
|
||||||
|
command_args += ['--scope', 'user', '-e', f'HF_HOME={target / "model-cache"}']
|
||||||
|
else:
|
||||||
|
command_args += ['--env', f'HF_HOME={target / "model-cache"}']
|
||||||
|
command_args += [NAME] + arguments
|
||||||
|
if update and client in existing_clients:
|
||||||
|
command(client, [client, 'mcp', 'remove', NAME] + (['--scope', 'user'] if client == 'claude' else []), home, False)
|
||||||
|
command(client, command_args, home, False)
|
||||||
|
if client not in state['clients']:
|
||||||
|
state['clients'].append(client)
|
||||||
|
save_state(state_path, state, False)
|
||||||
|
save_state(state_path, state, False)
|
||||||
|
|
||||||
|
|
||||||
|
def main():
|
||||||
|
parser = argparse.ArgumentParser()
|
||||||
|
parser.add_argument('action', choices=('sync',))
|
||||||
|
parser.add_argument('--root', required=True, type=Path)
|
||||||
|
parser.add_argument('--home', required=True, type=Path)
|
||||||
|
parser.add_argument('--client', choices=('codex', 'claude', 'both'), required=True)
|
||||||
|
parser.add_argument('--enabled', choices=('true', 'false'), required=True)
|
||||||
|
parser.add_argument('--dry-run', action='store_true')
|
||||||
|
parser.add_argument('--update', action='store_true')
|
||||||
|
args = parser.parse_args()
|
||||||
|
clients = ('codex', 'claude') if args.client == 'both' else (args.client,)
|
||||||
|
sync(args.root.resolve(), args.home.resolve(), clients, args.enabled == 'true', args.dry_run, args.update)
|
||||||
|
print('SEMANTIC RETRIEVAL: reconciliation complete' + (' (dry run)' if args.dry_run else ''))
|
||||||
|
|
||||||
|
|
||||||
|
if __name__ == '__main__':
|
||||||
|
try:
|
||||||
|
main()
|
||||||
|
except (OSError, ValueError, json.JSONDecodeError, subprocess.SubprocessError) as error:
|
||||||
|
print(f'Semantic retrieval: {error}', file=sys.stderr)
|
||||||
|
raise SystemExit(1)
|
||||||
@@ -0,0 +1,8 @@
|
|||||||
|
sync_semantic_retrieval() {
|
||||||
|
local root=$1 client=$2 home_path=$3 enabled=$4 dry_run=$5 update=$6
|
||||||
|
[ "$enabled" = true ] || [ -f "$home_path/.my-ai-configuration/semantic-retrieval.json" ] || return 0
|
||||||
|
local arguments=("$root/scripts/lib/semantic-retrieval.py" sync --root "$root" --home "$home_path" --client "$client" --enabled "$enabled")
|
||||||
|
[ "$dry_run" != true ] || arguments+=(--dry-run)
|
||||||
|
[ "$update" != true ] || arguments+=(--update)
|
||||||
|
python3 "${arguments[@]}"
|
||||||
|
}
|
||||||
@@ -0,0 +1,39 @@
|
|||||||
|
import hashlib
|
||||||
|
import importlib.util
|
||||||
|
import json
|
||||||
|
from pathlib import Path
|
||||||
|
import tempfile
|
||||||
|
import unittest
|
||||||
|
|
||||||
|
|
||||||
|
ROOT = Path(__file__).resolve().parent.parent
|
||||||
|
SPEC = importlib.util.spec_from_file_location('semantic_retrieval', ROOT / 'scripts/lib/semantic-retrieval.py')
|
||||||
|
MODULE = importlib.util.module_from_spec(SPEC)
|
||||||
|
SPEC.loader.exec_module(MODULE)
|
||||||
|
|
||||||
|
|
||||||
|
class SemanticRetrievalTests(unittest.TestCase):
|
||||||
|
def test_disable_removes_owned_runtime_and_registration(self):
|
||||||
|
with tempfile.TemporaryDirectory() as directory:
|
||||||
|
home = Path(directory) / 'home'
|
||||||
|
target = home / '.my-ai-configuration/semantic-retrieval'
|
||||||
|
target.mkdir(parents=True)
|
||||||
|
for name in MODULE.FILES:
|
||||||
|
(target / name).write_text(name)
|
||||||
|
state_path = target.parent / 'semantic-retrieval.json'
|
||||||
|
state_path.write_text(json.dumps({'version': 1, 'clients': ['codex'], 'files': {
|
||||||
|
name: hashlib.sha256((target / name).read_bytes()).hexdigest() for name in MODULE.FILES}}))
|
||||||
|
calls = []
|
||||||
|
original = MODULE.command
|
||||||
|
MODULE.command = lambda client, arguments, home, dry_run: calls.append((client, arguments))
|
||||||
|
try:
|
||||||
|
MODULE.sync(ROOT, home, ['codex'], False, False, False)
|
||||||
|
finally:
|
||||||
|
MODULE.command = original
|
||||||
|
self.assertEqual(calls[0][1], ['codex', 'mcp', 'remove', MODULE.NAME])
|
||||||
|
self.assertFalse(target.exists())
|
||||||
|
self.assertFalse(state_path.exists())
|
||||||
|
|
||||||
|
|
||||||
|
if __name__ == '__main__':
|
||||||
|
unittest.main()
|
||||||
@@ -0,0 +1,2 @@
|
|||||||
|
mcp>=1.9,<2
|
||||||
|
sentence-transformers>=4.1,<5
|
||||||
@@ -0,0 +1,106 @@
|
|||||||
|
#!/usr/bin/env python3
|
||||||
|
"""Local MCP semantic search backed by Qwen3-Embedding-0.6B."""
|
||||||
|
|
||||||
|
from __future__ import annotations
|
||||||
|
|
||||||
|
import argparse
|
||||||
|
import json
|
||||||
|
import os
|
||||||
|
from pathlib import Path
|
||||||
|
import sqlite3
|
||||||
|
|
||||||
|
|
||||||
|
MODEL = "Qwen/Qwen3-Embedding-0.6B"
|
||||||
|
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")
|
||||||
|
for start in range(0, len(text), 1050):
|
||||||
|
value = text[start:start + 1200].strip()
|
||||||
|
if value:
|
||||||
|
yield value
|
||||||
|
|
||||||
|
|
||||||
|
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(self.data / "index.sqlite3")
|
||||||
|
self.db.execute("create table if not exists chunks (path text, text text, vector blob)")
|
||||||
|
self.db.commit()
|
||||||
|
self.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)
|
||||||
|
return result
|
||||||
|
|
||||||
|
def encoder(self):
|
||||||
|
if self.model is None:
|
||||||
|
from sentence_transformers import SentenceTransformer
|
||||||
|
self.model = SentenceTransformer(MODEL)
|
||||||
|
return self.model
|
||||||
|
|
||||||
|
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)
|
||||||
|
|
||||||
|
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"):
|
||||||
|
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})
|
||||||
|
return sorted(rows, key=lambda item: item["score"], reverse=True)[:max(1, min(top_k, 20))]
|
||||||
|
|
||||||
|
|
||||||
|
def main():
|
||||||
|
parser = argparse.ArgumentParser()
|
||||||
|
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()
|
||||||
|
args.model_cache.mkdir(parents=True, exist_ok=True)
|
||||||
|
os.environ.setdefault("HF_HOME", str(args.model_cache.resolve()))
|
||||||
|
from mcp.server.fastmcp import FastMCP
|
||||||
|
|
||||||
|
index = Index(args.root, args.data_dir)
|
||||||
|
mcp = FastMCP("my-ai-qwen3-retrieval")
|
||||||
|
|
||||||
|
@mcp.tool()
|
||||||
|
def semantic_search(query: str, top_k: int = 5) -> list[dict]:
|
||||||
|
"""Find repository passages by meaning with Qwen3-Embedding-0.6B."""
|
||||||
|
return index.search(query, top_k)
|
||||||
|
|
||||||
|
@mcp.tool()
|
||||||
|
def rebuild_index() -> str:
|
||||||
|
"""Rebuild the local repository embedding index."""
|
||||||
|
return json.dumps({"chunks": index.rebuild()})
|
||||||
|
|
||||||
|
mcp.run()
|
||||||
|
|
||||||
|
|
||||||
|
if __name__ == "__main__":
|
||||||
|
main()
|
||||||
Reference in new issue
Block a user