diff --git a/scripts/lib/semantic-retrieval.ps1 b/scripts/lib/semantic-retrieval.ps1 new file mode 100644 index 0000000..f10216e --- /dev/null +++ b/scripts/lib/semantic-retrieval.ps1 @@ -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.' } +} diff --git a/scripts/lib/semantic-retrieval.py b/scripts/lib/semantic-retrieval.py new file mode 100644 index 0000000..67a5521 --- /dev/null +++ b/scripts/lib/semantic-retrieval.py @@ -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) diff --git a/scripts/lib/semantic-retrieval.sh b/scripts/lib/semantic-retrieval.sh new file mode 100644 index 0000000..6c6e05d --- /dev/null +++ b/scripts/lib/semantic-retrieval.sh @@ -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[@]}" +} diff --git a/scripts/test-semantic-retrieval.py b/scripts/test-semantic-retrieval.py new file mode 100644 index 0000000..4853d96 --- /dev/null +++ b/scripts/test-semantic-retrieval.py @@ -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() diff --git a/shared/retrieval/requirements.txt b/shared/retrieval/requirements.txt new file mode 100644 index 0000000..95325d9 --- /dev/null +++ b/shared/retrieval/requirements.txt @@ -0,0 +1,2 @@ +mcp>=1.9,<2 +sentence-transformers>=4.1,<5 diff --git a/shared/retrieval/server.py b/shared/retrieval/server.py new file mode 100644 index 0000000..7824e6d --- /dev/null +++ b/shared/retrieval/server.py @@ -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()