diff --git a/adapters/plugins.tsv b/adapters/plugins.tsv index bac1a20..18d301c 100644 --- a/adapters/plugins.tsv +++ b/adapters/plugins.tsv @@ -4,7 +4,6 @@ i-have-adhd ayghri/i-have-adhd i-have-adhd@i-have-adhd ayghri/i-have-adhd i-have Superpowers anthropics/claude-plugins-official superpowers@claude-plugins-official openai-curated-remote superpowers@openai-curated-remote plugin - - Context7 upstash/context7 context7@context7-marketplace upstash/context7 context7@context7-marketplace plugin - - context7 https://mcp.context7.com/mcp Caveman JuliusBrussee/caveman caveman@caveman thinkhome-org/caveman-codex caveman@thinkhome-caveman plugin - - -QMD tobi/qmd qmd@qmd - - qmd @tobilu/qmd qmd MCPorter - mcporter - mcporter cli mcporter - Humanizer blader/humanizer humanizer@humanizer - - skill blader/humanizer humanizer Impeccable pbakaus/impeccable impeccable@impeccable pbakaus/impeccable impeccable skill pbakaus/impeccable impeccable diff --git a/scripts/lib/managed-extensions.py b/scripts/lib/managed-extensions.py index e5f3b86..d33966e 100644 --- a/scripts/lib/managed-extensions.py +++ b/scripts/lib/managed-extensions.py @@ -304,9 +304,6 @@ class Manager: if entry['codex_method'] == 'cli': self.ensure_cli(client, entry) return - if entry['codex_method'] == 'qmd': - self.ensure_qmd(client, entry) - return skill = client == 'codex' and entry['codex_method'] != 'plugin' kind = 'skill' if skill else 'plugin' record = next((r for r in self.state['resources'] if r['client'] == client and r['name'] == entry['name']), None) @@ -417,28 +414,6 @@ class Manager: atomic_write(path, (json.dumps(content, indent=2) + '\n').encode('utf-8')) print('Repaired Caveman Windows hook commands; review the changed hooks in Codex /hooks.') - def ensure_qmd(self, client, entry): - package = entry['codex_source'] - if package != '@tobilu/qmd': - raise ValueError('Invalid QMD package.') - record = next((r for r in self.state['resources'] if r['client'] == 'shared' and r['name'] == entry['name']), None) - installed = package in self.global_packages() - if record is not None: - record['clients'] = sorted(set(record.get('clients', ['claude', 'codex'])) | {client}) - self.save() - if record is None and installed and not self.summary: - print(f'PASS Preserving pre-existing global package: {package}') - if (record is not None or not installed) and (not installed or (self.update and package not in self.updated_packages)): - self.command(['npm', 'update' if installed else 'install', '--global', package]) - self.global_packages().add(package) - self.updated_packages.add(package) - if record is None: - self.state['resources'].append(dict(client='shared', name=entry['name'], kind='qmd', package=package, clients=[client])) - self.save() - if client == 'claude': - entry = dict(entry, codex_method='plugin') - self.ensure(client, entry) - def ensure_mcporter(self, entry, clients): server, url = entry['mcporter_name'], entry['mcporter_url'] if not re.fullmatch(r'[a-z0-9-]+', server) or not re.fullmatch(r'https://[^\s]+', url): diff --git a/scripts/lib/plugins.ps1 b/scripts/lib/plugins.ps1 index 1bfffdd..0920d7a 100644 --- a/scripts/lib/plugins.ps1 +++ b/scripts/lib/plugins.ps1 @@ -125,9 +125,7 @@ function Select-ConfiguredPlugins { } else { $false } foreach ($referenceClient in $referenceClients) { $selector = if ($referenceClient -eq 'Claude') { $entry.claude_plugin } else { $entry.codex_plugin } - if ($referenceClient -eq 'Codex' -and $entry.codex_method -eq 'qmd') { - $found = $found -or ($null -ne (Get-Command qmd -ErrorAction SilentlyContinue)) - } elseif ($referenceClient -eq 'Codex' -and $entry.codex_method -ne 'plugin') { + if ($referenceClient -eq 'Codex' -and $entry.codex_method -ne 'plugin') { $found = $found -or (Test-Path -LiteralPath (Join-Path $HomePath ('.agents/skills/' + $entry.codex_skill))) -or (Test-Path -LiteralPath (Join-Path $HomePath ('.codex/skills/' + $entry.codex_skill))) } else { $idField = if ($referenceClient -eq 'Claude') { 'id' } else { 'pluginId' } diff --git a/scripts/lib/plugins.sh b/scripts/lib/plugins.sh index 813eebc..d46fa2c 100755 --- a/scripts/lib/plugins.sh +++ b/scripts/lib/plugins.sh @@ -161,9 +161,7 @@ select_configured_plugins() { fi for reference_client in "${reference_clients[@]}"; do if [ "$reference_client" = claude ]; then selector=$(plugin_field "$root" "$name" claude_plugin); else selector=$(plugin_field "$root" "$name" codex_plugin); fi - if [ "$reference_client" = codex ] && [ "$(plugin_field "$root" "$name" codex_method)" = qmd ]; then - command -v qmd >/dev/null 2>&1 && found=true - elif [ "$reference_client" = codex ] && [ "$(plugin_field "$root" "$name" codex_method)" != plugin ]; then + if [ "$reference_client" = codex ] && [ "$(plugin_field "$root" "$name" codex_method)" != plugin ]; then local skill_name skill_name=$(plugin_field "$root" "$name" codex_skill) { [ -d "$HOME/.agents/skills/$skill_name" ] || [ -d "$HOME/.codex/skills/$skill_name" ]; } && found=true diff --git a/scripts/lib/semantic-retrieval.ps1 b/scripts/lib/semantic-retrieval.ps1 deleted file mode 100644 index 39a7e3f..0000000 --- a/scripts/lib/semantic-retrieval.ps1 +++ /dev/null @@ -1,18 +0,0 @@ -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' } - if ($Summary) { $arguments += '--summary' } - & 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 deleted file mode 100644 index fc5f8f7..0000000 --- a/scripts/lib/semantic-retrieval.py +++ /dev/null @@ -1,230 +0,0 @@ -import argparse -import hashlib -import json -import os -from pathlib import Path -import re -import shutil -import subprocess -import sys -import tempfile - - -FILES = ('server.py', 'benchmark.py', 'requirements.txt') -CUDA_INDEX = 'https://download.pytorch.org/whl/cu126' -MODEL_PRELOAD_TIMEOUT = 900 - - -def digest(path): - return hashlib.sha256(path.read_bytes()).hexdigest() - - -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.') - arguments = [str(python), '-m', 'pip', 'install', '--disable-pip-version-check', '-r', str(target / 'requirements.txt')] - if shutil.which('nvidia-smi'): - try: - subprocess.run([*arguments, '--extra-index-url', CUDA_INDEX], check=True, timeout=900) - except subprocess.CalledProcessError: - print('Semantic retrieval: CUDA runtime installation failed; falling back to CPU.', file=sys.stderr) - subprocess.run(arguments, check=True, timeout=900) - else: - subprocess.run(arguments, check=True, timeout=900) - return python - - -def ensure_models(python, target): - subprocess.run([str(python), str(target / 'server.py'), 'preload', '--model-cache', str(target / 'model-cache')], - check=True, timeout=MODEL_PRELOAD_TIMEOUT) - - -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) - temporary = None - try: - with tempfile.NamedTemporaryFile(mode='w', encoding='utf-8', dir=path.parent, delete=False) as stream: - temporary = Path(stream.name) - json.dump(state, stream, indent=2) - stream.write('\n') - stream.flush() - os.fsync(stream.fileno()) - os.replace(temporary, path) - finally: - if temporary is not None: - temporary.unlink(missing_ok=True) - - -def check_runtime(target, state): - if target.is_symlink(): - raise ValueError('Semantic retrieval directory must not be a symbolic link.') - if not target.exists(): - return - if not target.is_dir(): - raise ValueError('Semantic retrieval target is not a directory.') - for name, expected in state['files'].items(): - path = target / name - if path.is_symlink() or not path.is_file() or digest(path) != expected: - raise ValueError(f'Owned semantic retrieval file changed: {name}') - allowed = set(state['files']) | {'.venv', 'model-cache', '__pycache__', 'data'} - for path in target.iterdir(): - if path.is_symlink() or path.name not in allowed: - raise ValueError(f'Foreign or linked semantic retrieval path: {path.name}') - if path.name == 'data' and (not path.is_dir() or any( - child.is_symlink() or not ( - child.is_file() and (child.name in ('daemon.key', 'daemon.starting') or re.fullmatch( - r'[0-9a-f]{64}\.sqlite3(?:-(?:wal|shm|journal))?|daemon-[0-9a-f]{12}\.starting', child.name)) - or child.is_socket() and re.fullmatch(r'daemon-[0-9a-f]{16}\.sock', child.name)) - for child in path.iterdir())): - raise ValueError('Foreign or linked semantic retrieval path: data') - if path.name == '__pycache__' and (not path.is_dir() or any( - child.is_symlink() or not child.is_file() or not re.fullmatch( - r'(?:server|benchmark)\.[A-Za-z0-9_-]+(?:\.opt-\d+)?\.pyc', child.name) - for child in path.iterdir())): - raise ValueError('Foreign or linked semantic retrieval path: __pycache__') - - -def sync(root, home, clients, enabled, dry_run, update, summary=False): - 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: - if not summary: - print('DRYRUN disable semantic retrieval: ' + (', '.join(owned) if owned else 'nothing owned')) - return - state['clients'] = [client for client in state['clients'] if client not in clients] - if state['clients']: - save_state(state_path, state, False) - return - if not state['files']: - state_path.unlink(missing_ok=True) - return - check_runtime(target, state) - if target.exists() and target.is_dir(): - shutil.rmtree(target) - state_path.unlink(missing_ok=True) - return - - check_runtime(target, state) - if target.exists() and not state['files']: - raise ValueError('Semantic retrieval directory exists without setup ownership; refusing to overwrite it.') - if dry_run: - if not summary: - print('DRYRUN enable semantic retrieval: Qwen3-Embedding-0.6B + Qwen3-Reranker-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) - save_state(state_path, state, False) - python = ensure_runtime(target, False) - ensure_models(python, target) - # Sandboxed clients may write only to data/; it cannot be created from inside the sandbox. - (target / 'data').mkdir(exist_ok=True) - state['clients'] = list(dict.fromkeys(state['clients'])) - save_state(state_path, state, False) - for client in clients: - if client not in state['clients']: - state['clients'].append(client) - save_state(state_path, state, 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 state['files']: - check_runtime(target, state) - elif target.is_symlink(): - raise ValueError('Semantic retrieval directory must not be a symbolic link.') - 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.') - state['files'] = {name: digest(target / name) for name in FILES} - check_runtime(target, state) - 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'] or state['clients']: - save_state(state_path, state, False) - else: - check_runtime(target, state) - shutil.rmtree(target) - state_path.unlink(missing_ok=True) - if not summary: - print('Semantic retrieval warm-model check: ' + ('passed' if outcome['recommended'] else 'failed') + f" (embedding {outcome['embeddingSeconds']}s, reranking {outcome['rerankingSeconds']}s; excludes model startup and indexing)") - return outcome['recommended'] - - -def main(): - parser = argparse.ArgumentParser() - 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')) - 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 '')) - - -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 deleted file mode 100644 index 1ba36aa..0000000 --- a/scripts/lib/semantic-retrieval.sh +++ /dev/null @@ -1,21 +0,0 @@ -#!/usr/bin/env bash - -sync_semantic_retrieval() { - local root=$1 client=$2 home_path=$3 enabled=$4 dry_run=$5 update=$6 summary=${7:-false} - [ "$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) - [ "$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) - local status - if python3 "${arguments[@]}" >&2; then return 0; else status=$?; fi - [ "$status" -eq 2 ] && return 1 - return 2 -} diff --git a/scripts/tests/test-managed-extensions.py b/scripts/tests/test-managed-extensions.py index 31c1e92..fe29d2d 100644 --- a/scripts/tests/test-managed-extensions.py +++ b/scripts/tests/test-managed-extensions.py @@ -260,27 +260,15 @@ class ManagedExtensionTests(unittest.TestCase): self.manager(dry_run=True, summary=True).sync([plugin(), skill()], ['codex'], {'Example', 'Humanizer'}) self.assertEqual('EXTENSIONS: reconciliation complete (dry run)\n', output.getvalue()) - def test_qmd_installs_once_as_a_cli_without_registering_mcp(self): - entry = dict(name='QMD', codex_method='qmd', codex_source='@tobilu/qmd', codex_skill='qmd', codex_plugin='-', codex_marketplace='-', claude_plugin='qmd@qmd', claude_marketplace='tobi/qmd') - self.manager().sync([entry], ['codex'], {'QMD'}) - self.assertEqual([['npm', 'install', '--global', '@tobilu/qmd']], self.commands) - self.manager().sync([], ['codex'], set()) - self.assertEqual(['npm', 'uninstall', '--global', '@tobilu/qmd'], self.commands[-1]) + def test_manifest_has_no_qmd_extension(self): + rows = (root / 'adapters/plugins.tsv').read_text(encoding='utf-8').splitlines() + self.assertFalse(any(row.split(' ')[0] == 'QMD' for row in rows[1:])) - def test_qmd_installs_once_for_both_clients(self): - entry = dict(name='QMD', codex_method='qmd', codex_source='@tobilu/qmd', codex_skill='qmd', codex_plugin='-', codex_marketplace='-', claude_plugin='qmd@qmd', claude_marketplace='tobi/qmd') - self.manager().sync([entry], ['claude', 'codex'], {'QMD'}) - self.assertEqual(1, self.commands.count(['npm', 'install', '--global', '@tobilu/qmd'])) - self.assertIn(['claude', 'plugin', 'install', 'qmd@qmd', '--scope', 'user'], self.commands) - self.assertNotIn(['codex', 'mcp', 'add', 'qmd', '--', 'qmd', 'mcp'], self.commands) - - def test_owned_qmd_uses_npm_update_instead_of_reinstalling(self): - entry = dict(name='QMD', codex_method='qmd', codex_source='@tobilu/qmd', codex_skill='qmd', codex_plugin='-', codex_marketplace='-', claude_plugin='qmd@qmd', claude_marketplace='tobi/qmd') - self.manager().sync([entry], ['codex'], {'QMD'}) - self.commands.clear() - with patch.object(extensions.Manager, 'global_packages', return_value={'@tobilu/qmd'}): - self.manager(update=True).sync([entry], ['codex'], {'QMD'}) - self.assertEqual([['npm', 'update', '--global', '@tobilu/qmd']], self.commands) + def test_legacy_qmd_record_still_validates(self): + (self.home / '.my-ai-configuration').mkdir() + (self.home / '.my-ai-configuration/extensions.json').write_text(json.dumps({'version': 1, 'resources': [ + {'client': 'shared', 'name': 'QMD', 'kind': 'qmd', 'package': '@tobilu/qmd', 'clients': ['codex']}]})) + self.manager() # must not raise def test_managed_mcporter_updates_when_requested(self): self.manager().sync([mcporter()], ['codex'], {'MCPorter'}) diff --git a/scripts/tests/test-retrieval-server.py b/scripts/tests/test-retrieval-server.py deleted file mode 100644 index be3f7b8..0000000 --- a/scripts/tests/test-retrieval-server.py +++ /dev/null @@ -1,548 +0,0 @@ -import importlib.util -from array import array -from concurrent.futures import ThreadPoolExecutor -import os -from pathlib import Path -import re -import socket -import subprocess -from tempfile import TemporaryDirectory -from types import SimpleNamespace -import time -import unittest -import urllib.request -from unittest.mock import patch - - -ROOT = Path(__file__).resolve().parent.parent.parent -SPEC = importlib.util.spec_from_file_location('retrieval_server', ROOT / 'shared/retrieval/server.py') -MODULE = importlib.util.module_from_spec(SPEC) -SPEC.loader.exec_module(MODULE) - - -def tokenize(text, **kwargs): - offsets = [match.span() for match in re.finditer(r'\S+', text)] - return {'offset_mapping': offsets, 'input_ids': list(range(len(offsets)))} - - -def embedding(): - return array('f', [1.0] + [0.0] * (MODULE.EMBEDDING_DIMENSIONS - 1)) - - -class VectorBatch: - def __init__(self, blob): - self.values = array('f', blob) - - def reshape(self, *shape): - return self - - def __matmul__(self, vector): - width = MODULE.EMBEDDING_DIMENSIONS - return [sum(a * b for a, b in zip(self.values[offset:offset + width], vector)) - for offset in range(0, len(self.values), width)] - - -class ApproximateIndex: - def __init__(self, **kwargs): - self.vectors = {} - - def add(self, keys, vectors): - width = MODULE.EMBEDDING_DIMENSIONS - for position, key in enumerate(keys): - start = position * width - self.vectors[key] = self.vectors.get(key, []) + list(vectors.values[start:start + width]) - - def search(self, vector, count): - matches = sorted(((key, sum(a * b for a, b in zip(values, vector))) - for key, values in self.vectors.items()), key=lambda item: (-item[1], item[0]))[:count] - return [SimpleNamespace(key=key, distance=1.0 - score) for key, score in matches] - - -def fake_numpy(): - return SimpleNamespace(frombuffer=lambda blob, **kwargs: VectorBatch(blob), float32=None) - - -def fake_usearch(): - module = SimpleNamespace(Index=ApproximateIndex) - return SimpleNamespace(index=module), module - - -def initialize_repository(root): - subprocess.run(['git', '-C', str(root), 'init', '--quiet'], check=True) - - -def huggingface_available(): - try: - request = urllib.request.Request('https://huggingface.co', method='HEAD') - with urllib.request.urlopen(request, timeout=5): - pass - except OSError: - return False - return True - - -def huggingface_network_unavailable(error): - while error: - if isinstance(error, socket.gaierror): - return True - error = error.__cause__ or error.__context__ - return False - - -class RetrievalServerTests(unittest.TestCase): - def test_huggingface_availability_handles_connection_failure(self): - with patch('urllib.request.urlopen', side_effect=OSError): - self.assertFalse(huggingface_available()) - - def test_huggingface_network_failure_is_detected_from_cause_chain(self): - error = OSError('model download failed') - error.__cause__ = socket.gaierror('name resolution failed') - self.assertTrue(huggingface_network_unavailable(error)) - self.assertFalse(huggingface_network_unavailable(OSError('model is missing'))) - - def test_ci_disables_xet_for_model_downloads(self): - content = (ROOT / '.gitea/workflows/ci.yml').read_text(encoding='utf-8') - self.assertIn("HF_HUB_DISABLE_XET: '1'", content) - - def test_client_smoke_builds_packages(self): - content = (ROOT / '.gitea/workflows/ci.yml').read_text(encoding='utf-8') - self.assertIn('name: Build packages\n run: bash scripts/commands/build.sh --summary', content) - - def test_runtime_python_uses_adjacent_virtualenv(self): - with TemporaryDirectory() as temporary: - script = Path(temporary) / 'semantic-retrieval' / 'server.py' - runtime = script.parent / '.venv' / 'Scripts' / 'python.exe' - runtime.parent.mkdir(parents=True) - runtime.touch() - - self.assertEqual(MODULE.runtime_python(script, 'C:/Python/python.exe', 'nt'), runtime.resolve()) - - def test_windows_runtime_handoff_keeps_multi_word_arguments(self): - calls = [] - with patch.object(MODULE, 'runtime_python', return_value=Path('C:/runtime/python.exe')), patch.object(MODULE.sys, 'argv', ['server.py', 'search', '--query', 'two words']), patch.object(MODULE.subprocess, 'run', side_effect=lambda command: calls.append(command) or SimpleNamespace(returncode=3)), self.assertRaises(SystemExit) as stopped: - MODULE.use_runtime_python('nt') - self.assertEqual(stopped.exception.code, 3) - self.assertEqual(calls[0][2:], ['search', '--query', 'two words']) - - def test_daemon_identity_changes_with_code_version(self): - with TemporaryDirectory() as temporary: - data_dir = Path(temporary) - with patch.object(MODULE, 'code_version', return_value='aaaaaaaaaaaa'): - old_address, old_lock = MODULE.daemon_address(data_dir), MODULE.daemon_lock_path(data_dir) - with patch.object(MODULE, 'code_version', return_value='bbbbbbbbbbbb'): - self.assertNotEqual(MODULE.daemon_address(data_dir), old_address) - self.assertNotEqual(MODULE.daemon_lock_path(data_dir), old_lock) - - def test_run_direct_runs_in_process_when_pipes_are_forbidden(self): - def action(action, root, data_dir, model_cache, query, top_k, output): - output.put((True, [query, top_k])) - - with patch.object(MODULE, 'pipes_available', return_value=False), patch.object(MODULE, 'run_action', action): - self.assertEqual(MODULE.run_direct('search', Path('.'), Path('.'), Path('.'), 'two words', 3), ['two words', 3]) - - def test_model_kwargs_use_bfloat16_only_on_supporting_gpus(self): - bfloat16 = object() - for available, supported, expected in [(True, True, {'dtype': bfloat16}), (True, False, {}), (False, True, {})]: - cuda = SimpleNamespace(is_available=lambda available=available: available, is_bf16_supported=lambda supported=supported: supported) - with patch.dict('sys.modules', torch=SimpleNamespace(cuda=cuda, bfloat16=bfloat16)): - self.assertEqual(MODULE.model_kwargs(), expected) - - def test_query_prompt_matches_qwen_format(self): - self.assertTrue(MODULE.QUERY_PROMPT.startswith('Instruct: ')) - self.assertTrue(MODULE.QUERY_PROMPT.endswith('\nQuery:')) - - def test_preferred_device_uses_cuda_with_cpu_fallback(self): - with patch.dict('sys.modules', torch=SimpleNamespace(cuda=SimpleNamespace(is_available=lambda: True))): - self.assertEqual(MODULE.preferred_device(), 'cuda') - with patch.dict('sys.modules', torch=SimpleNamespace(cuda=SimpleNamespace(is_available=lambda: False))): - self.assertEqual(MODULE.preferred_device(), 'cpu') - - def test_preload_models_warms_embedding_and_reranking_from_model_cache(self): - with TemporaryDirectory() as temporary: - cache = Path(temporary) / 'model-cache' - calls = [] - encoder = SimpleNamespace(encode=lambda values, **kwargs: calls.append(('encode', values, kwargs)) or [embedding()]) - reranker = SimpleNamespace(predict=lambda pairs, **kwargs: calls.append(('predict', pairs, kwargs)) or [0.0]) - - with patch.object(MODULE, 'shared_encoder', return_value=encoder), \ - patch.object(MODULE, 'shared_reranker', return_value=reranker), \ - patch.dict(os.environ, {'HF_HOME': 'other-cache'}): - getattr(MODULE, 'preload_models', lambda _: None)(cache) - self.assertEqual(str(cache.resolve()), os.environ['HF_HOME']) - - self.assertTrue(cache.is_dir()) - self.assertEqual([ - ('encode', ['Verify semantic retrieval models.'], {'prompt_name': 'query'}), - ('predict', [('Verify semantic retrieval models.', 'Verify semantic retrieval models.')], {'show_progress_bar': False}), - ], calls) - - @unittest.skipUnless(os.environ.get('QWEN_TEST_MODEL_CACHE'), 'Set QWEN_TEST_MODEL_CACHE to run the real-model check.') - def test_real_models(self): - if not huggingface_available(): - self.skipTest('huggingface.co DNS is unavailable') - os.environ['HF_HOME'] = os.environ['QWEN_TEST_MODEL_CACHE'] - with TemporaryDirectory() as temporary: - root = Path(temporary) - initialize_repository(root) - (root / 'authentication.md').write_text('Password reset tokens are created by create_reset_token. The reset_password endpoint sends the user a password reset email.', encoding='utf-8') - (root / 'weather.md').write_text('The weather service fetches temperature and rainfall forecasts for a city.', encoding='utf-8') - index = MODULE.Index(root, root / 'data') - self.addCleanup(index.db.close) - try: - tokenizer = index.encoder().tokenizer - for source in ['alpha beta gamma ' * 450, 'Grüße 世界 🙂 ' * 300]: - windows = list(MODULE.chunks(source, 'sample.txt', tokenizer)) - self.assertGreater(len(windows), 1) - self.assertTrue(all(len(tokenizer(value, add_special_tokens=False)['input_ids']) <= 400 - for _, value in windows)) - self.assertTrue(all('\ufffd' not in value for _, value in windows)) - for query, expected in [('Where are password reset tokens created?', 'authentication.md'), - ('Which service provides rainfall forecasts?', 'weather.md')]: - started = time.monotonic() - results = index.search(query, 2) - self.assertEqual(results[0]['path'], expected) - self.assertEqual(len(results), 2) - self.assertGreater(results[0]['score'], results[1]['score']) - print(f'Qwen real search: {expected}, {time.monotonic() - started:.2f}s, embedding device {index.encoder().device}, reranker device {index.reranker().device}', flush=True) - except OSError as error: - if huggingface_network_unavailable(error): - self.skipTest('Hugging Face network is unavailable') - raise - - def test_isolation_freshness_and_incremental_encoding(self): - with TemporaryDirectory() as temporary: - root = Path(temporary) - first, second = root / 'first', root / 'second' - first.mkdir() - second.mkdir() - initialize_repository(first) - initialize_repository(second) - file = first / 'a.md' - file.write_text('alpha', encoding='utf-8') - (second / 'a.md').write_text('other', encoding='utf-8') - indexes = [MODULE.Index(repository, root / 'data') for repository in (first, second)] - encoded = [] - - def encode(values, **kwargs): - self.assertEqual(kwargs.get('prompt_name'), 'query' if values == ['query'] else None) - encoded.extend(values) - return [embedding() for _ in values] - - for index in indexes: - index.encoder = lambda: SimpleNamespace(encode=encode, tokenizer=tokenize) - index.rerank = lambda query, rows, limit: rows[:limit] - numpy = fake_numpy() - usearch, usearch_index = fake_usearch() - try: - with patch.dict('sys.modules', numpy=numpy, usearch=usearch, **{'usearch.index': usearch_index}): - self.assertEqual(indexes[0].search('query', 5)[0]['text'], 'alpha') - self.assertEqual(indexes[1].search('query', 5)[0]['text'], 'other') - encoded.clear() - with patch.object(Path, 'read_bytes', side_effect=AssertionError('unchanged files must not be read')): - self.assertEqual(indexes[0].search('query', 5)[0]['text'], 'alpha') - with ThreadPoolExecutor(max_workers=1) as executor: - self.assertEqual(executor.submit(indexes[0].search, 'query', 5).result()[0]['text'], 'alpha') - self.assertEqual(encoded, []) - file.write_text('bravo', encoding='utf-8') - self.assertEqual(indexes[0].search('query', 5)[0]['text'], 'bravo') - file.rename(first / 'renamed.md') - self.assertEqual(indexes[0].search('query', 5)[0]['path'], 'renamed.md') - (first / 'renamed.md').unlink() - encoded.clear() - self.assertEqual(indexes[0].search('query', 5), []) - self.assertEqual(indexes[0].search('query', 5), []) - self.assertEqual(encoded, []) - self.assertEqual(indexes[1].search('query', 5)[0]['text'], 'other') - finally: - for index in indexes: - index.db.close() - - def test_failed_refresh_preserves_previous_index(self): - with TemporaryDirectory() as temporary: - root = Path(temporary) - initialize_repository(root) - (root / 'a.md').write_text('before', encoding='utf-8') - index = MODULE.Index(root, root / 'data') - index.encoder = lambda: SimpleNamespace(encode=lambda values, **kwargs: [embedding() for _ in values], tokenizer=tokenize) - try: - index.rebuild() - (root / 'a.md').write_text('after', encoding='utf-8') - (root / 'b.md').write_text('fail', encoding='utf-8') - - def encode(values, **kwargs): - if any(value.endswith('\n\nfail') for value in values): - raise RuntimeError('encoder failed') - return [embedding() for _ in values] - - index.encoder = lambda: SimpleNamespace(encode=encode, tokenizer=tokenize) - with self.assertRaisesRegex(RuntimeError, 'encoder failed'): - index.rebuild() - self.assertEqual(index.db.execute('select path, text from chunks').fetchall(), [('a.md', 'before')]) - self.assertEqual(index.db.execute('select document from lexical').fetchall(), - [('Path: a.md\nSymbol: \n\nbefore',)]) - finally: - index.db.close() - - def test_refresh_batches_embeddings_for_changed_files(self): - with TemporaryDirectory() as temporary: - root = Path(temporary) - initialize_repository(root) - (root / 'first.md').write_text('first', encoding='utf-8') - (root / 'second.md').write_text('second', encoding='utf-8') - index = MODULE.Index(root, root / 'data') - encoded = [] - - def encode(values, **kwargs): - encoded.append(list(values)) - return [embedding() for _ in values] - - index.encoder = lambda: SimpleNamespace(encode=encode, tokenizer=tokenize) - try: - index.rebuild() - self.assertEqual(len(encoded), 1) - self.assertEqual(len(encoded[0]), 2) - finally: - index.db.close() - - def test_scan_prunes_ignored_directories_and_warns_at_limit(self): - with TemporaryDirectory() as temporary: - root = Path(temporary) - initialize_repository(root) - index = MODULE.Index(root, root / 'data') - (root / 'visible.md').write_text('visible', encoding='utf-8') - (root / '.gitignore').write_text('secrets.json\n**/private/\n', encoding='utf-8') - (root / 'secrets.json').write_text('{"token":"hidden"}', encoding='utf-8') - private = root / 'nested' / 'private' - private.mkdir(parents=True) - (private / 'hidden.md').write_text('hidden', encoding='utf-8') - for name in ('node_modules', '.git', 'generated', 'data'): - directory = root / name - directory.mkdir(exist_ok=True) - (directory / 'hidden.md').write_text('hidden', encoding='utf-8') - try: - self.assertEqual(index.files(), [(root / 'visible.md').resolve()]) - filenames = ['visible.md'] * 2001 - with patch.object(MODULE.subprocess, 'run', return_value=SimpleNamespace(stdout=b'visible.md\0' * 2001)), \ - patch.object(MODULE.sys, 'stderr') as stderr: - self.assertEqual(len(index.files()), 2000) - self.assertIn('stopped scanning after 2000', stderr.write.call_args_list[0].args[0]) - finally: - index.db.close() - - def test_status_reports_index_freshness_without_creating_an_index(self): - with TemporaryDirectory() as temporary: - root, data = Path(temporary), Path(temporary) / 'data' - initialize_repository(root) - source = root / 'module.py' - source.write_text('value = 1\n', encoding='utf-8') - status = MODULE.index_status(root, data) - self.assertEqual((status['indexable_files'], status['indexed_files'], status['stale_files']), (1, 0, 1)) - self.assertFalse(status['models_loaded']) - self.assertFalse(MODULE.index_database(root, data).exists()) - index = MODULE.Index(root, data) - try: - metadata = source.stat() - with index.db: - index.db.execute('insert into files values (?, ?, ?, ?, ?)', - ('module.py', 'digest', metadata.st_size, metadata.st_mtime_ns, metadata.st_ctime_ns)) - finally: - index.db.close() - self.assertEqual(MODULE.index_status(root, data)['stale_files'], 0) - source.write_text('value = 22\n', encoding='utf-8') - self.assertEqual(MODULE.index_status(root, data)['stale_files'], 1) - - def test_max_files_is_configurable_via_environment(self): - with TemporaryDirectory() as temporary: - root = Path(temporary) - initialize_repository(root) - index = MODULE.Index(root, root / 'data') - try: - with patch.dict(os.environ, {MODULE.MAX_FILES_ENV: '3'}): - self.assertEqual(MODULE.read_max_files(), 3) - with patch.dict(os.environ, {MODULE.MAX_FILES_ENV: 'not-a-number'}): - with self.assertRaises(ValueError): - MODULE.read_max_files() - with patch.dict(os.environ, {}, clear=True): - self.assertEqual(MODULE.read_max_files(), MODULE.DEFAULT_MAX_FILES) - (root / 'visible.md').write_text('visible', encoding='utf-8') - index.max_files = 3 - for number in range(5): - (root / f'{number}.md').write_text('visible', encoding='utf-8') - with patch.object(MODULE.sys, 'stderr'): - self.assertEqual(len(index.files()), 3) - finally: - index.db.close() - - def test_daemon_address_is_deterministic_and_platform_specific(self): - with TemporaryDirectory() as temporary: - data_dir = Path(temporary) - first = MODULE.daemon_address(data_dir) - second = MODULE.daemon_address(data_dir) - self.assertEqual(first, second) - if os.name == 'nt': - self.assertTrue(first.startswith(r'\\.\pipe\ai-config-retrieval-')) - else: - self.assertTrue(first.startswith(str(data_dir))) - - def test_shared_encoder_and_reranker_are_cached_module_singletons(self): - with patch.object(MODULE, '_shared_encoder', None), patch.object(MODULE, '_shared_reranker', None), \ - patch.object(MODULE, 'preferred_device', lambda: 'cpu'), patch.object(MODULE, 'model_kwargs', dict): - sentinel = SimpleNamespace() - with patch.dict('sys.modules', sentence_transformers=SimpleNamespace( - SentenceTransformer=lambda *args, **kwargs: sentinel, - CrossEncoder=lambda *args, **kwargs: sentinel)): - self.assertIs(MODULE.shared_encoder(), sentinel) - self.assertIs(MODULE.shared_encoder(), MODULE.shared_encoder()) - self.assertIs(MODULE.shared_reranker(), sentinel) - self.assertIs(MODULE.shared_reranker(), MODULE.shared_reranker()) - - def test_rerank_uses_cross_encoder_scores(self): - index = MODULE.Index.__new__(MODULE.Index) - - class Reranker: - def predict(self, pairs, **kwargs): - self.pairs = pairs - return [0.1, 0.8, 0.5] - - reranker = Reranker() - index.reranker = lambda: reranker - rows = [ - {'path': 'a.md', 'symbol': 'A', 'score': 0.9, 'text': 'first'}, - {'path': 'b.md', 'symbol': 'B', 'score': 0.8, 'text': 'second'}, - {'path': 'c.md', 'symbol': 'C', 'score': 0.7, 'text': 'third'}, - ] - - result = index.rerank('query', rows, 2) - - self.assertEqual([item['path'] for item in result], ['b.md', 'c.md']) - self.assertEqual(reranker.pairs, [('query', 'Path: a.md\nSymbol: A\n\nfirst'), - ('query', 'Path: b.md\nSymbol: B\n\nsecond'), - ('query', 'Path: c.md\nSymbol: C\n\nthird')]) - - def test_token_windows_have_exact_overlap_and_no_redundant_tail(self): - for size, lengths in [(0, []), (400, [400]), (401, [400, 65]), (736, [400, 400]), (800, [400, 400, 128])]: - words = [f'token{number}' for number in range(size)] - result = list(MODULE.chunks(' '.join(words), 'plain.txt', tokenize)) - windows = [content.split() for _, content in result] - self.assertEqual([len(window) for window in windows], lengths) - for first, second in zip(windows, windows[1:]): - self.assertEqual(first[-64:], second[:64]) - if windows: - reconstructed = windows[0] + [word for window in windows[1:] for word in window[64:]] - self.assertEqual(reconstructed, words) - - def test_python_symbols_include_decorators_and_restore_parent(self): - source = 'class Service:\n @cached\n def load(self):\n return 1\n\nvalue = 2\n\nasync def refresh():\n return value\n' - result = list(MODULE.sections(source, 'service.py')) - self.assertEqual([symbol for symbol, _ in result], ['Service', 'Service.load', '', 'refresh']) - self.assertIn('@cached', result[1][1]) - self.assertEqual(''.join(content for _, content in result), source) - broken = 'def incomplete(\n' + 'word ' * 450 - self.assertEqual(len(list(MODULE.chunks(broken, 'broken.py', tokenize))), 2) - - def test_markdown_sections_ignore_fenced_headings(self): - source = '# Setup\nIntro\n## Install\n```sh\n# not a heading\n```\nSteps\n## Run\nStart\n' - result = list(MODULE.chunks(source, 'guide.md', tokenize)) - self.assertEqual([symbol for symbol, _ in result], ['Setup', 'Setup / Install', 'Setup / Run']) - self.assertIn('# not a heading', result[1][1]) - - def test_common_declarations_and_config_sections(self): - for path, source, symbols in [ - ('service.ts', 'export function load() {\n return 1;\n}\nexport const save = () => 2;\n', ['load', 'save']), - ('service.cs', 'public class Service\n{\n public void Load()\n {\n }\n}\n', ['Service', 'Load']), - ('setup.ps1', 'function Get-Result {\n return 1\n}\n', ['Get-Result']), - ('build.sh', 'quote() {\n echo "$1"\n}\nfor shell in a b; do\n quote "$shell"\ndone\n', ['quote', '']), - ('schema.sql', 'CREATE TABLE users (id int);\nCREATE OR ALTER PROCEDURE dbo.GetUser AS SELECT 1;\n', ['users', 'dbo.GetUser']), - ('Main.kt', 'fun main() {\n println(1)\n}\n', ['main']), - ('config.toml', '[server]\nport = 1\n[client]\nport = 2\n', ['server', 'client']), - ]: - self.assertEqual([symbol for symbol, _ in MODULE.sections(source, path)], symbols) - - def test_rrf_rewards_agreement_and_breaks_ties_deterministically(self): - self.assertEqual(MODULE.reciprocal_rank_fusion([1, 2, 3], [4, 2, 3]), [2, 3, 1, 4]) - self.assertEqual(MODULE.reciprocal_rank_fusion([], [4, 2]), [4, 2]) - - def test_dense_candidates_converts_vector_keys_for_sqlite(self): - class VectorKey: - def __int__(self): - return 9 - - with TemporaryDirectory() as temporary: - root = Path(temporary) - index = MODULE.Index(root, root / 'data') - index.vector_index = SimpleNamespace(search=lambda vector, count: [SimpleNamespace(key=VectorKey())]) - usearch, usearch_index = fake_usearch() - try: - with patch.dict('sys.modules', numpy=fake_numpy(), usearch=usearch, - **{'usearch.index': usearch_index}): - candidates = index.dense_candidates(embedding()) - self.assertEqual(candidates, [9]) - self.assertIs(type(candidates[0]), int) - finally: - index.db.close() - - def test_hybrid_candidates_and_document_metadata(self): - with TemporaryDirectory() as temporary: - root = Path(temporary) - initialize_repository(root) - for number in range(110): - (root / f'{number:03}.md').write_text(f'# Symbol{number}\ncontent {number}', encoding='utf-8', newline='\n') - index = MODULE.Index(root, root / 'data') - documents = [] - - def encode(values, **kwargs): - if kwargs.get('prompt_name') != 'query': - documents.extend(values) - return [embedding() for _ in values] - - pairs = [] - index.encoder = lambda: SimpleNamespace(encode=encode, tokenizer=tokenize) - index.reranker = lambda: SimpleNamespace(predict=lambda values, **kwargs: - pairs.extend(values) or [float('109.md' in document) for _, document in values]) - numpy = fake_numpy() - original_fusion = MODULE.reciprocal_rank_fusion - rankings = [] - - def fuse(*values): - rankings.extend(values) - return original_fusion(*values) - - try: - usearch, usearch_index = fake_usearch() - with patch.dict('sys.modules', numpy=numpy, usearch=usearch, **{'usearch.index': usearch_index}), \ - patch.object(MODULE, 'reciprocal_rank_fusion', fuse): - results = index.search('109.md Symbol109 content', 20) - self.assertEqual([len(ranking) for ranking in rankings], [50, 50]) - self.assertEqual(len(pairs), 50) - self.assertEqual(len(results), 5) - self.assertEqual(results[0]['path'], '109.md') - self.assertEqual(results[0]['symbol'], 'Symbol109') - self.assertIn('Path: 109.md\nSymbol: Symbol109\n\n# Symbol109\ncontent 109', documents) - self.assertTrue(all(document in documents for _, document in pairs)) - with patch.dict('sys.modules', numpy=numpy, usearch=usearch, **{'usearch.index': usearch_index}): - self.assertEqual(len(index.search('" OR * () : -')), 5) - self.assertEqual(len(index.search('***')), 5) - (root / '109.md').unlink() - index.search('Symbol109') - self.assertEqual(index.db.execute('select count(*) from lexical where lexical match ?', ('"Symbol109"',)).fetchone()[0], 0) - finally: - index.db.close() - - def test_old_cache_is_preserved_but_not_reused(self): - with TemporaryDirectory() as temporary: - root = Path(temporary) - data = root / 'data' - data.mkdir() - identity = os.path.normcase(str(root.resolve())) + '\n' + MODULE.MODEL - old = data / (MODULE.hashlib.sha256(identity.encode()).hexdigest() + '.sqlite3') - old.write_bytes(b'old cache') - index = MODULE.Index(root, data) - try: - self.assertNotEqual(Path(index.db.execute('pragma database_list').fetchone()[2]), old) - self.assertEqual(old.read_bytes(), b'old cache') - finally: - index.db.close() - - -if __name__ == '__main__': - unittest.main() diff --git a/scripts/tests/test-semantic-retrieval.py b/scripts/tests/test-semantic-retrieval.py deleted file mode 100644 index a394ad4..0000000 --- a/scripts/tests/test-semantic-retrieval.py +++ /dev/null @@ -1,163 +0,0 @@ -import hashlib -import importlib.util -import io -import json -from contextlib import redirect_stdout -from pathlib import Path -import py_compile -import subprocess -import tempfile -import unittest -from unittest.mock import patch - - -ROOT = Path(__file__).resolve().parent.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_sync_does_not_register_mcp(self): - with tempfile.TemporaryDirectory() as directory: - home = Path(directory) / 'home' - output = io.StringIO() - with redirect_stdout(output): - MODULE.sync(ROOT, home, ['codex', 'claude'], True, True, False) - self.assertNotIn('mcp', output.getvalue().lower()) - - def test_runtime_prefers_cuda_packages_for_nvidia_gpu(self): - with tempfile.TemporaryDirectory() as directory: - target = Path(directory) - python = MODULE.python_path(target / '.venv') - python.parent.mkdir(parents=True) - python.touch() - with patch.object(MODULE.shutil, 'which', return_value='nvidia-smi'), patch.object(MODULE.subprocess, 'run') as run: - MODULE.ensure_runtime(target, False) - - self.assertEqual(run.call_args.args[0][-2:], ['--extra-index-url', MODULE.CUDA_INDEX]) - - def test_runtime_falls_back_to_cpu_packages_when_cuda_install_fails(self): - with tempfile.TemporaryDirectory() as directory: - target = Path(directory) - python = MODULE.python_path(target / '.venv') - python.parent.mkdir(parents=True) - python.touch() - failure = MODULE.subprocess.CalledProcessError(1, 'pip') - with patch.object(MODULE.shutil, 'which', return_value='nvidia-smi'), patch.object(MODULE.subprocess, 'run', side_effect=[failure, None]) as run: - MODULE.ensure_runtime(target, False) - - self.assertEqual(run.call_count, 2) - self.assertNotIn('--extra-index-url', run.call_args.args[0]) - - def test_sync_preloads_models_after_installing_runtime(self): - with tempfile.TemporaryDirectory() as directory: - home = Path(directory) / 'home' - target = home / '.my-ai-configuration/semantic-retrieval' - python = MODULE.python_path(target / '.venv') - with patch.object(MODULE, 'ensure_runtime', return_value=python), \ - patch.object(MODULE.subprocess, 'run') as run: - MODULE.sync(ROOT, home, ['codex'], True, False, False) - - self.assertEqual([[ - str(python), str(target / 'server.py'), 'preload', '--model-cache', str(target / 'model-cache') - ]], [call.args[0] for call in run.call_args_list]) - - def test_summary_dry_run_omits_details_without_writes(self): - with tempfile.TemporaryDirectory() as directory: - home = Path(directory) / 'home' - for enabled in (True, False): - output = io.StringIO() - with redirect_stdout(output): - MODULE.sync(ROOT, home, ['codex'], enabled, True, False, summary=True) - self.assertEqual(output.getvalue(), '') - self.assertFalse(home.exists()) - - def test_disable_removes_owned_runtime_without_cli_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}})) - MODULE.sync(ROOT, home, ['codex'], False, False, False) - self.assertFalse(target.exists()) - self.assertFalse(state_path.exists()) - - def test_update_accepts_python_cache_but_rejects_foreign_cache_content(self): - with tempfile.TemporaryDirectory() as directory: - home = Path(directory) / 'home' - target = home / '.my-ai-configuration/semantic-retrieval' - target.mkdir(parents=True) - files = {} - for name in MODULE.FILES: - content = (ROOT / 'shared/retrieval' / name).read_bytes() - (target / name).write_bytes(content) - files[name] = hashlib.sha256(content).hexdigest() - MODULE.save_state(target.parent / 'semantic-retrieval.json', {'version': 1, 'clients': ['codex'], 'files': files}, False) - py_compile.compile(str(target / 'server.py'), doraise=True) - cache = target / '__pycache__' - with patch.object(MODULE, 'ensure_runtime', return_value=Path('python')), patch.object(MODULE, 'ensure_models'): - MODULE.sync(ROOT, home, ['codex'], True, False, True) - (cache / 'foreign.txt').write_text('keep') - with self.assertRaisesRegex(ValueError, '__pycache__'): - MODULE.sync(ROOT, home, ['codex'], True, False, True) - - def test_install_search_update_and_remove_owned_data(self): - with tempfile.TemporaryDirectory() as directory: - home = Path(directory) / 'home' - repository = Path(directory) / 'repository' - repository.mkdir() - subprocess.run(['git', '-C', str(repository), 'init', '--quiet'], check=True) - target = home / '.my-ai-configuration/semantic-retrieval' - state_path = target.parent / 'semantic-retrieval.json' - with patch.object(MODULE, 'ensure_runtime', return_value=Path('python')), \ - patch.object(MODULE, 'ensure_models'): - MODULE.sync(ROOT, home, ['codex'], True, False, False) - installed = importlib.util.spec_from_file_location('installed_retrieval', target / 'server.py') - server = importlib.util.module_from_spec(installed) - installed.loader.exec_module(server) - index = server.Index(repository, target / 'data') - try: - self.assertEqual(index.search('query'), []) - finally: - index.db.close() - self.assertEqual(len(list((target / 'data').glob('*.sqlite3'))), 1) - server.daemon_authkey(target / 'data') - MODULE.sync(ROOT, home, ['codex'], True, False, True) - self.assertTrue((target / 'data').is_dir()) - foreign = target / 'data' / 'foreign.txt' - foreign.write_text('keep') - with self.assertRaisesRegex(ValueError, 'data'): - MODULE.sync(ROOT, home, ['codex'], False, False, False) - self.assertTrue(foreign.exists()) - foreign.unlink() - MODULE.sync(ROOT, home, ['codex'], False, False, False) - self.assertFalse(target.exists()) - self.assertFalse(state_path.exists()) - - def test_negative_benchmark_keeps_active_runtime(self): - with tempfile.TemporaryDirectory() as directory: - home = Path(directory) / 'home' - target = home / '.my-ai-configuration/semantic-retrieval' - target.mkdir(parents=True) - state_path = target.parent / 'semantic-retrieval.json' - files = {} - for name in MODULE.FILES: - content = (ROOT / 'shared/retrieval' / name).read_bytes() - (target / name).write_bytes(content) - files[name] = hashlib.sha256(content).hexdigest() - MODULE.save_state(state_path, {'version': 1, 'clients': ['codex'], 'files': files}, False) - result = type('Result', (), {'stdout': '{"recommended": false, "embeddingSeconds": 1, "rerankingSeconds": 1}'})() - with patch.object(MODULE, 'ensure_runtime', return_value=Path('python')), \ - patch.object(MODULE.subprocess, 'run', return_value=result): - self.assertFalse(MODULE.benchmark(ROOT, home, False, True)) - self.assertTrue(target.is_dir()) - self.assertEqual(MODULE.load_state(state_path)['clients'], ['codex']) - - -if __name__ == '__main__': - unittest.main() diff --git a/shared/retrieval/benchmark.py b/shared/retrieval/benchmark.py deleted file mode 100644 index a30e076..0000000 --- a/shared/retrieval/benchmark.py +++ /dev/null @@ -1,56 +0,0 @@ -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() diff --git a/shared/retrieval/requirements.txt b/shared/retrieval/requirements.txt deleted file mode 100644 index 9e68d87..0000000 --- a/shared/retrieval/requirements.txt +++ /dev/null @@ -1,2 +0,0 @@ -sentence-transformers==5.7.0 -usearch==2.26.2 diff --git a/shared/retrieval/server.py b/shared/retrieval/server.py deleted file mode 100644 index 194bc45..0000000 --- a/shared/retrieval/server.py +++ /dev/null @@ -1,695 +0,0 @@ -#!/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 json -import multiprocessing -from multiprocessing import AuthenticationError -from multiprocessing.connection import Client, Listener -import os -import queue -from pathlib import Path -import re -import sqlite3 -import stat -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 = 3 -RETRIEVAL_INSTRUCTION = "Given a codebase question, retrieve relevant code and documentation that answer the question." -# Qwen3-Embedding query format from the model card: no space after "Query:". -QUERY_PROMPT = f"Instruct: {RETRIEVAL_INSTRUCTION}\nQuery:" -TEXT_EXTENSIONS = {".c", ".cjs", ".cpp", ".cs", ".csproj", ".css", ".go", ".gradle", ".html", ".java", ".js", ".json", ".jsx", ".kt", ".kts", - ".md", ".mjs", ".props", ".ps1", ".psm1", ".py", ".razor", ".rs", ".scss", ".sh", ".sql", ".targets", ".toml", ".ts", - ".tsx", ".txt", ".xaml", ".xml", ".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_MESSAGE_LIMIT = 16 * 1024 * 1024 - - -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(system=os.name): - runtime = runtime_python() - if runtime is not None: - command = [str(runtime), str(Path(__file__).resolve()), *sys.argv[1:]] - if system == 'nt': - # os.execv does not quote arguments on Windows, so a multi-word --query would be split. - sys.exit(subprocess.run(command).returncode) - os.execv(command[0], command) - - -def sections(text: str, path: str): - lines = text.splitlines(keepends=True) - boundaries = {0: "", len(lines): ""} - 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 != "" 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, "") - 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 == ".sql": - heading = re.match(r"^\s*create\s+(?:or\s+(?:alter|replace)\s+)?(?:procedure|proc|function|view|table|trigger|index)\s+([\w.\[\]\"]+)", - line, re.IGNORECASE) - if heading: - boundaries[number] = heading[1] - elif suffix in {".js", ".jsx", ".mjs", ".cjs", ".ts", ".tsx", ".cs", ".java", ".kt", ".kts", ".go", ".rs", ".c", ".cpp", ".ps1", ".psm1", ".sh"}: - # A closing brace in column 0 ends a top-level declaration; code after it belongs to the module again. - if line.startswith("}"): - boundaries.setdefault(number + 1, "") - 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|fun|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 model_kwargs(): - """Use bfloat16 on GPUs that support it; CPUs keep float32, where bfloat16 is often slower.""" - import torch - return {"dtype": torch.bfloat16} if torch.cuda.is_available() and torch.cuda.is_bf16_supported() else {} - - -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": QUERY_PROMPT}, device=preferred_device(), - model_kwargs=model_kwargs()) - 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(), - model_kwargs=model_kwargs()) - return _shared_reranker - - -def preload_models(model_cache: Path): - model_cache.mkdir(parents=True, exist_ok=True) - os.environ["HF_HOME"] = str(model_cache.resolve()) - query = "Verify semantic retrieval models." - embeddings = shared_encoder().encode([query], prompt_name="query") - if len(embeddings) != 1 or len(embeddings[0]) != EMBEDDING_DIMENSIONS: - raise ValueError("Embedding model returned unexpected dimensions.") - scores = shared_reranker().predict([(query, query)], show_progress_bar=False) - if len(scores) != 1: - raise ValueError("Reranking model returned unexpected scores.") - - -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 code_version() -> str: - """Hash of this file, so a daemon started by older code is never reused after an update.""" - return hashlib.sha256(Path(__file__).read_bytes()).hexdigest()[:12] - - -def daemon_lock_path(data_dir: Path) -> Path: - return data_dir / f"daemon-{code_version()}.starting" - - -def daemon_address(data_dir: Path) -> str: - identity = os.path.normcase(str(data_dir.resolve())) + "|" + code_version() - digest = hashlib.sha256(identity.encode()).hexdigest()[:16] - if os.name == "nt": - return r"\\.\pipe\ai-config-retrieval-" + digest - return str(data_dir / f"daemon-{digest}.sock") - - -def daemon_authkey(data_dir: Path) -> bytes: - if data_dir.is_symlink(): - raise OSError("Semantic retrieval data directory must not be a link.") - data_dir.mkdir(parents=True, exist_ok=True, mode=0o700) - path = data_dir / "daemon.key" - try: - descriptor = os.open(path, os.O_CREAT | os.O_EXCL | os.O_WRONLY, 0o600) - except FileExistsError: - pass - else: - with os.fdopen(descriptor, "wb") as stream: - stream.write(os.urandom(32)) - stream.flush() - os.fsync(stream.fileno()) - metadata = path.lstat() - if not stat.S_ISREG(metadata.st_mode) or (os.name != "nt" and metadata.st_mode & 0o077): - raise OSError("Semantic retrieval daemon credential is not private.") - key = path.read_bytes() - if len(key) != 32: - raise OSError("Semantic retrieval daemon credential is incomplete.") - return key - - -def request_daemon(address: str, payload: dict, data_dir: Path): - with Client(address, authkey=daemon_authkey(data_dir)) as connection: - connection.send_bytes(json.dumps(payload).encode("utf-8")) - if not connection.poll(CLI_TIMEOUT_SECONDS): - raise TimeoutError(f"Semantic retrieval daemon exceeded the {CLI_TIMEOUT_SECONDS} second timeout.") - success, result = json.loads(connection.recv_bytes(DAEMON_MESSAGE_LIMIT)) - 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, data_dir) - except AuthenticationError: - return None - 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, data_dir) - except (OSError, EOFError, AuthenticationError): - 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(data_dir)) - daemon_lock_path(data_dir).unlink(missing_ok=True) - try: - while True: - connection = listener.accept() - try: - response = (True, None) - try: - request = json.loads(connection.recv_bytes(DAEMON_MESSAGE_LIMIT)) - if isinstance(request, dict) and request.get("action") == "stop": - connection.send_bytes(json.dumps((True, None)).encode("utf-8")) - os._exit(0) - if not isinstance(request, dict) or request.get("action") not in ("search", "rebuild") or not isinstance(request.get("root"), str): - raise ValueError("Invalid semantic retrieval daemon request.") - 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_bytes(json.dumps(response).encode("utf-8")) - except (OSError, EOFError): - pass - finally: - connection.close() - with activity_lock: - last_activity = time.monotonic() - finally: - listener.close() - - -def index_database(root: Path, data: Path) -> Path: - identity = json.dumps([os.path.normcase(str(root.resolve())), MODEL, MODEL_REVISION, - EMBEDDING_DIMENSIONS, RETRIEVAL_INSTRUCTION, INDEX_VERSION, - CHUNK_TOKENS, OVERLAP_TOKENS]) - return data.resolve() / (hashlib.sha256(identity.encode()).hexdigest() + ".sqlite3") - - -def daemon_running(data_dir: Path) -> bool: - if not pipes_available(): - return False - try: - Client(daemon_address(data_dir), authkey=daemon_authkey(data_dir)).close() - return True - except (OSError, EOFError, AuthenticationError): - return False - - -def index_status(root: Path, data_dir: Path) -> dict: - """Report what the next search costs without loading models or starting the daemon.""" - database = index_database(root, data_dir) - previous = {} - if database.is_file(): - connection = sqlite3.connect(database.as_uri() + "?mode=ro", uri=True) - try: - previous = {path: (size, mtime, ctime) for path, size, mtime, ctime - in connection.execute("select path, size, mtime, ctime from files")} - except sqlite3.OperationalError: - pass - finally: - connection.close() - files = indexable_files(root.resolve(), data_dir.resolve(), read_max_files()) - stale = 0 - for path in files: - metadata = path.stat() - stale += previous.get(path.relative_to(root.resolve()).as_posix()) != (metadata.st_size, metadata.st_mtime_ns, metadata.st_ctime_ns) - return {"indexable_files": len(files), "indexed_files": len(previous), "stale_files": stale, - "models_loaded": daemon_running(data_dir), "in_process": not pipes_available()} - - -def indexable_files(root: Path, data: Path, max_files: int): - ignored = {".git", ".venv", "generated", "node_modules", ".my-ai-configuration", "bin", "obj"} - output = subprocess.run(["git", "-C", str(root), "ls-files", "--cached", "--others", - "--exclude-standard", "-z", "--", "."], - check=True, capture_output=True, timeout=30).stdout - candidates = [root / name for name in - output.decode("utf-8", errors="surrogateescape").split("\0") if name - and not any(part in ignored for part in Path(name).parts)] - result = [] - for candidate in sorted(candidates): - if candidate.is_symlink(): - continue - path = candidate.resolve() - if (not path.is_relative_to(root) or path.is_relative_to(data) - 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) == max_files: - print(f"Semantic retrieval: stopped scanning after {max_files} indexable files; " - f"some files under {root} were not indexed. Set {MAX_FILES_ENV} to raise this limit.", - file=sys.stderr) - return result - return result - - -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(index_database(self.root, self.data), 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, size integer, mtime integer, ctime integer)") - columns = {row[1] for row in self.db.execute("pragma table_info(files)")} - for name in ("size", "mtime", "ctime"): - if name not in columns: - self.db.execute(f"alter table files add column {name} integer") - 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() - self.vector_index = None - - def files(self): - return indexable_files(self.root, self.data, self.max_files) - - 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 = {path: (digest, size, mtime, ctime) - for path, digest, size, mtime, ctime in self.db.execute("select path, digest, size, mtime, ctime from files")} - encoder = None - changed = [] - documents = [] - document_offsets = [] - for path in self.files(): - relative = path.relative_to(self.root).as_posix() - metadata = path.stat() - size, mtime, ctime = metadata.st_size, metadata.st_mtime_ns, metadata.st_ctime_ns - old_digest, old_size, old_mtime, old_ctime = previous.pop(relative, (None, None, None, None)) - if not force and old_digest and (size, mtime, ctime) == (old_size, old_mtime, old_ctime): - continue - content = path.read_bytes() - digest = hashlib.sha256(content).hexdigest() - if not force and digest == old_digest: - with self.db: - self.db.execute("update files set size = ?, mtime = ?, ctime = ? where path = ?", - (size, mtime, ctime, relative)) - 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, size, mtime, ctime, values, file_documents)) - vectors = encoder.encode(documents, normalize_embeddings=True) if documents else [] - with self.db: - for number, (relative, digest, size, mtime, ctime, 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, size, mtime, ctime)) - 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") - self.vector_index = None - - def dense_candidates(self, vector): - import numpy as np - from usearch.index import Index as VectorIndex - - if self.vector_index is None: - index = VectorIndex(ndim=EMBEDDING_DIMENSIONS, metric="cos", dtype="f32", - connectivity=16, expansion_add=128, expansion_search=64) - cursor = self.db.execute("select rowid, vector from chunks") - while rows := cursor.fetchmany(4096): - vectors = np.frombuffer(b"".join(row[1] for row in rows), dtype=np.float32).reshape(-1, EMBEDDING_DIMENSIONS) - index.add([row[0] for row in rows], vectors) - self.vector_index = index - return [int(match.key) for match in self.vector_index.search(vector, DENSE_CANDIDATES)] - - def cache_key(self, query: str, top_k: int): - identity = json.dumps([ - query, top_k, - RERANKER_MODEL, RERANKER_MODEL_REVISION, RETRIEVAL_INSTRUCTION, QUERY_PROMPT, - DENSE_CANDIDATES, LEXICAL_CANDIDATES, RERANK_CANDIDATES, FINAL_RESULTS, RRF_K, - ], ensure_ascii=False, separators=(",", ":")) - return hashlib.sha256(identity.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,)) - vector = self.encoder().encode([query], prompt_name="query", normalize_embeddings=True)[0] - if len(vector) != EMBEDDING_DIMENSIONS: - raise ValueError("Embedding must have 1024 dimensions") - dense = self.dense_candidates(vector) - 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 pipes_available() -> bool: - """Sandboxes such as Codex on Windows forbid the pipes that the daemon and worker process need.""" - try: - first, second = multiprocessing.Pipe() - except PermissionError: - return False - first.close() - second.close() - return True - - -def run_direct(action: str, root: Path, data_dir: Path, model_cache: Path, query: str | None, top_k: int): - if not pipes_available(): - # ponytail: in-process fallback has no worker timeout; the client's own command timeout still applies. - output = queue.Queue(1) - run_action(action, root, data_dir, model_cache, query, top_k, output) - return output.get()[1] - 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", "status", "stop", "daemon", "preload")) - 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", type=Path) - parser.add_argument("--model-cache", required=True, type=Path) - args = parser.parse_args() - model_cache = args.model_cache.resolve() - if args.action == "preload": - preload_models(model_cache) - return - if args.data_dir is None: - parser.error(f"{args.action} requires --data-dir") - data_dir = args.data_dir.resolve() - if args.action == "daemon": - data_dir.mkdir(parents=True, exist_ok=True) - run_daemon(data_dir, model_cache) - return - if args.action == "status": - print(json.dumps(index_status(args.root, data_dir))) - return - if args.action == "stop": - try: - request_daemon(daemon_address(data_dir), {"action": "stop"}, data_dir) - stopped = True - except (OSError, EOFError, AuthenticationError): - stopped = False - print(json.dumps({"stopped": stopped})) - 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 pipes_available() else None - 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()