feat(retrieval): offer device benchmark

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 }
+        }
     }
 }
 

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() {

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.'
+}

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 ''))

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
+}

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()