refactor(retrieval): remove the Python Qwen stack and the QMD extension entry
Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
This commit is contained in:
Julian lechnerandClaude Opus 5.5 committed
1 parent
505e6ae084
commit
5bb0bdf62f
13 files changed
+10
-1785
No files matched your search
@@ -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
|
||||
|
||||
|
@@ -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):
|
||||
|
||||
@@ -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' }
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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.'
|
||||
}
|
||||
@@ -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)
|
||||
@@ -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
|
||||
}
|
||||
@@ -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'})
|
||||
|
||||
@@ -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: <module>\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', '<module>', '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', '<module>']),
|
||||
('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()
|
||||
@@ -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()
|
||||
@@ -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()
|
||||
@@ -1,2 +0,0 @@
|
||||
sentence-transformers==5.7.0
|
||||
usearch==2.26.2
|
||||
@@ -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: "<module>", len(lines): "<module>"}
|
||||
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 != "<module>" 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, "<module>")
|
||||
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, "<module>")
|
||||
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()
|
||||
Reference in new issue
Block a user