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
@@ -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()
|
||||
Reference in new issue
Block a user