diff --git a/scripts/lib/install-options.ps1 b/scripts/lib/install-options.ps1 index 2e8ba9d..c4ceaa4 100644 --- a/scripts/lib/install-options.ps1 +++ b/scripts/lib/install-options.ps1 @@ -24,7 +24,23 @@ function Read-InstallOptions { return @{ update_agents = Read-InstallBoolean -Prompt "Update selected agent CLIs ($Client)?" -Default $defaults.update_agents flashbang = Read-InstallBoolean -Prompt 'Enable the Flashbang notification hook?' -Default $defaults.flashbang - semantic_retrieval = Read-InstallBoolean -Prompt 'Enable Qwen3 embedding and reranking semantic retrieval?' -Default $defaults.semantic_retrieval + semantic_retrieval = Read-SemanticRetrievalOption -Default $defaults.semantic_retrieval -RepositoryRoot $RepositoryRoot -HomePath $HomePath + } +} + +function Read-SemanticRetrievalOption { + param([bool]$Default,[string]$RepositoryRoot,[string]$HomePath) + $hint = if ($Default) { 'Y/n/a' } else { 'y/N/a' } + while ($true) { + Write-Host "Enable Qwen3 embedding and reranking semantic retrieval? [$hint] " -NoNewline + $answer = [Console]::ReadLine() + if ($null -eq $answer) { throw 'Input ended before semantic retrieval was confirmed.' } + switch ($answer.Trim().ToLowerInvariant()) { + '' { return $Default } + { $_ -in @('y', 'yes') } { return $true } + { $_ -in @('n', 'no') } { return $false } + { $_ -in @('a', 'auto') } { return Test-SemanticRetrievalDevice -RepositoryRoot $RepositoryRoot -HomePath $HomePath } + } } } diff --git a/scripts/lib/install-options.sh b/scripts/lib/install-options.sh index 845a73a..e5e0a61 100644 --- a/scripts/lib/install-options.sh +++ b/scripts/lib/install-options.sh @@ -34,7 +34,28 @@ read_install_options() { [ "$dry_run" = false ] || return 0 update_agents=$(read_install_boolean "Update selected agent CLIs ($client)?" "$update_agents") || return flashbang=$(read_install_boolean 'Enable the Flashbang notification hook?' "$flashbang") || return - semantic_retrieval=$(read_install_boolean 'Enable Qwen3 embedding and reranking semantic retrieval?' "$semantic_retrieval") || return + semantic_retrieval=$(read_semantic_retrieval_option "$semantic_retrieval") || return +} + +read_semantic_retrieval_option() { + local default=$1 hint=y/N/a answer + [ "$default" = false ] || hint=Y/n/a + while :; do + printf 'Enable Qwen3 embedding and reranking semantic retrieval? [%s] ' "$hint" >&2 + IFS= read -r answer || { printf 'Input ended before semantic retrieval was confirmed.\n' >&2; return 1; } + answer=${answer%$'\r'} + case "$answer" in + '') printf '%s\n' "$default"; return ;; + y|Y|yes|Yes|YES) printf 'true\n'; return ;; + n|N|no|No|NO) printf 'false\n'; return ;; + a|A|auto|Auto|AUTO) + if test_semantic_retrieval_device "$root" "$home_path"; then printf 'true\n'; return; fi + status=$? + [ "$status" -eq 1 ] && { printf 'false\n'; return; } + return "$status" + ;; + esac + done } global_npm_package() { diff --git a/scripts/lib/semantic-retrieval.ps1 b/scripts/lib/semantic-retrieval.ps1 index cb26307..39a7e3f 100644 --- a/scripts/lib/semantic-retrieval.ps1 +++ b/scripts/lib/semantic-retrieval.ps1 @@ -8,3 +8,11 @@ function Sync-SemanticRetrieval { & python @arguments if ($LASTEXITCODE -ne 0) { throw 'Semantic retrieval reconciliation failed.' } } + +function Test-SemanticRetrievalDevice { + param([Parameter(Mandatory=$true)][string]$RepositoryRoot,[Parameter(Mandatory=$true)][string]$HomePath,[switch]$Summary) + & python (Join-Path $PSScriptRoot 'semantic-retrieval.py') benchmark --root $RepositoryRoot --home $HomePath $(if ($Summary) { '--summary' }) | Write-Host + if ($LASTEXITCODE -eq 0) { return $true } + if ($LASTEXITCODE -eq 2) { return $false } + throw 'Semantic retrieval benchmark failed.' +} diff --git a/scripts/lib/semantic-retrieval.py b/scripts/lib/semantic-retrieval.py index f66a03f..f0ca25e 100644 --- a/scripts/lib/semantic-retrieval.py +++ b/scripts/lib/semantic-retrieval.py @@ -8,7 +8,7 @@ import subprocess import sys -FILES = ('server.py', 'requirements.txt') +FILES = ('server.py', 'benchmark.py', 'requirements.txt') CUDA_INDEX = 'https://download.pytorch.org/whl/cu126' @@ -106,17 +106,55 @@ def sync(root, home, clients, enabled, dry_run, update, summary=False): save_state(state_path, state, False) +def benchmark(root, home, dry_run, summary=False): + base = home / '.my-ai-configuration' + target = base / 'semantic-retrieval' + state_path = base / 'semantic-retrieval.json' + if dry_run: + if not summary: + print('DRYRUN benchmark semantic retrieval: Qwen3-Embedding-0.6B + Qwen3-Reranker-0.6B') + return True + state = load_state(state_path) + source = root / 'shared' / 'retrieval' + if target.exists() and not state['files']: + if not all((target / name).is_file() and digest(target / name) == digest(source / name) for name in FILES): + raise ValueError('Semantic retrieval directory exists without setup ownership; refusing to overwrite it.') + target.mkdir(parents=True, exist_ok=True) + for name in FILES: + shutil.copy2(source / name, target / name) + state['files'][name] = digest(target / name) + python = ensure_runtime(target, False) + result = subprocess.run([str(python), str(target / 'benchmark.py'), '--model-cache', str(target / 'model-cache')], check=True, + stdout=subprocess.PIPE, text=True, timeout=300) + outcome = json.loads(result.stdout) + if not isinstance(outcome, dict) or not isinstance(outcome.get('recommended'), bool): + raise ValueError('Semantic retrieval benchmark returned an invalid result.') + if outcome['recommended']: + save_state(state_path, state, False) + else: + shutil.rmtree(target) + state_path.unlink(missing_ok=True) + if not summary: + print('Semantic retrieval benchmark: ' + ('recommended' if outcome['recommended'] else 'not recommended') + f" (embedding {outcome['embeddingSeconds']}s, reranking {outcome['rerankingSeconds']}s)") + return outcome['recommended'] + + def main(): parser = argparse.ArgumentParser() - parser.add_argument('action', choices=('sync',)) + parser.add_argument('action', choices=('sync', 'benchmark')) 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('--client', choices=('codex', 'claude', 'both')) + parser.add_argument('--enabled', choices=('true', 'false')) parser.add_argument('--dry-run', action='store_true') parser.add_argument('--update', action='store_true') parser.add_argument('--summary', action='store_true') args = parser.parse_args() + if args.action == 'benchmark': + recommended = benchmark(args.root.resolve(), args.home.resolve(), args.dry_run, args.summary) + raise SystemExit(0 if recommended else 2) + if args.client is None or args.enabled is None: + parser.error('sync requires --client and --enabled') 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, args.summary) print('SEMANTIC RETRIEVAL: reconciliation complete' + (' (dry run)' if args.dry_run else '')) diff --git a/scripts/lib/semantic-retrieval.sh b/scripts/lib/semantic-retrieval.sh index 822d362..ad607ba 100644 --- a/scripts/lib/semantic-retrieval.sh +++ b/scripts/lib/semantic-retrieval.sh @@ -9,3 +9,12 @@ sync_semantic_retrieval() { [ "$summary" != true ] || arguments+=(--summary) python3 "${arguments[@]}" } + +test_semantic_retrieval_device() { + local root=$1 home_path=$2 summary=${3:-false} + local arguments=("$root/scripts/lib/semantic-retrieval.py" benchmark --root "$root" --home "$home_path") + [ "$summary" != true ] || arguments+=(--summary) + if python3 "${arguments[@]}" >&2; then return 0; fi + [ "$?" -eq 2 ] && return 1 + return 2 +} diff --git a/shared/retrieval/benchmark.py b/shared/retrieval/benchmark.py new file mode 100644 index 0000000..a30e076 --- /dev/null +++ b/shared/retrieval/benchmark.py @@ -0,0 +1,56 @@ +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()