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_preferred_device_uses_cuda_with_cpu_fallback(self): with patch.dict('sys.modules', torch=SimpleNamespace(cuda=SimpleNamespace(is_available=lambda: True))): self.assertEqual(MODULE.preferred_device(), 'cuda') with patch.dict('sys.modules', torch=SimpleNamespace(cuda=SimpleNamespace(is_available=lambda: False))): self.assertEqual(MODULE.preferred_device(), 'cpu') def test_preload_models_warms_embedding_and_reranking_from_model_cache(self): with TemporaryDirectory() as temporary: cache = Path(temporary) / 'model-cache' calls = [] encoder = SimpleNamespace(encode=lambda values, **kwargs: calls.append(('encode', values, kwargs)) or [embedding()]) reranker = SimpleNamespace(predict=lambda pairs, **kwargs: calls.append(('predict', pairs, kwargs)) or [0.0]) with patch.object(MODULE, 'shared_encoder', return_value=encoder), \ patch.object(MODULE, 'shared_reranker', return_value=reranker), \ patch.dict(os.environ, {'HF_HOME': 'other-cache'}): getattr(MODULE, 'preload_models', lambda _: None)(cache) self.assertEqual(str(cache.resolve()), os.environ['HF_HOME']) self.assertTrue(cache.is_dir()) self.assertEqual([ ('encode', ['Verify semantic retrieval models.'], {'prompt_name': 'query'}), ('predict', [('Verify semantic retrieval models.', 'Verify semantic retrieval models.')], {'show_progress_bar': False}), ], calls) @unittest.skipUnless(os.environ.get('QWEN_TEST_MODEL_CACHE'), 'Set QWEN_TEST_MODEL_CACHE to run the real-model check.') def test_real_models(self): if not huggingface_available(): self.skipTest('huggingface.co DNS is unavailable') os.environ['HF_HOME'] = os.environ['QWEN_TEST_MODEL_CACHE'] with TemporaryDirectory() as temporary: root = Path(temporary) initialize_repository(root) (root / 'authentication.md').write_text('Password reset tokens are created by create_reset_token. The reset_password endpoint sends the user a password reset email.', encoding='utf-8') (root / 'weather.md').write_text('The weather service fetches temperature and rainfall forecasts for a city.', encoding='utf-8') index = MODULE.Index(root, root / 'data') self.addCleanup(index.db.close) try: tokenizer = index.encoder().tokenizer for source in ['alpha beta gamma ' * 450, 'Grüße 世界 🙂 ' * 300]: windows = list(MODULE.chunks(source, 'sample.txt', tokenizer)) self.assertGreater(len(windows), 1) self.assertTrue(all(len(tokenizer(value, add_special_tokens=False)['input_ids']) <= 400 for _, value in windows)) self.assertTrue(all('\ufffd' not in value for _, value in windows)) for query, expected in [('Where are password reset tokens created?', 'authentication.md'), ('Which service provides rainfall forecasts?', 'weather.md')]: started = time.monotonic() results = index.search(query, 2) self.assertEqual(results[0]['path'], expected) self.assertEqual(len(results), 2) self.assertGreater(results[0]['score'], results[1]['score']) print(f'Qwen real search: {expected}, {time.monotonic() - started:.2f}s, embedding device {index.encoder().device}, reranker device {index.reranker().device}', flush=True) except OSError as error: if huggingface_network_unavailable(error): self.skipTest('Hugging Face network is unavailable') raise def test_isolation_freshness_and_incremental_encoding(self): with TemporaryDirectory() as temporary: root = Path(temporary) first, second = root / 'first', root / 'second' first.mkdir() second.mkdir() initialize_repository(first) initialize_repository(second) file = first / 'a.md' file.write_text('alpha', encoding='utf-8') (second / 'a.md').write_text('other', encoding='utf-8') indexes = [MODULE.Index(repository, root / 'data') for repository in (first, second)] encoded = [] def encode(values, **kwargs): self.assertEqual(kwargs.get('prompt_name'), 'query' if values == ['query'] else None) encoded.extend(values) return [embedding() for _ in values] for index in indexes: index.encoder = lambda: SimpleNamespace(encode=encode, tokenizer=tokenize) index.rerank = lambda query, rows, limit: rows[:limit] numpy = fake_numpy() usearch, usearch_index = fake_usearch() try: with patch.dict('sys.modules', numpy=numpy, usearch=usearch, **{'usearch.index': usearch_index}): self.assertEqual(indexes[0].search('query', 5)[0]['text'], 'alpha') self.assertEqual(indexes[1].search('query', 5)[0]['text'], 'other') encoded.clear() with patch.object(Path, 'read_bytes', side_effect=AssertionError('unchanged files must not be read')): self.assertEqual(indexes[0].search('query', 5)[0]['text'], 'alpha') with ThreadPoolExecutor(max_workers=1) as executor: self.assertEqual(executor.submit(indexes[0].search, 'query', 5).result()[0]['text'], 'alpha') self.assertEqual(encoded, []) file.write_text('bravo', encoding='utf-8') self.assertEqual(indexes[0].search('query', 5)[0]['text'], 'bravo') file.rename(first / 'renamed.md') self.assertEqual(indexes[0].search('query', 5)[0]['path'], 'renamed.md') (first / 'renamed.md').unlink() encoded.clear() self.assertEqual(indexes[0].search('query', 5), []) self.assertEqual(indexes[0].search('query', 5), []) self.assertEqual(encoded, []) self.assertEqual(indexes[1].search('query', 5)[0]['text'], 'other') finally: for index in indexes: index.db.close() def test_failed_refresh_preserves_previous_index(self): with TemporaryDirectory() as temporary: root = Path(temporary) initialize_repository(root) (root / 'a.md').write_text('before', encoding='utf-8') index = MODULE.Index(root, root / 'data') index.encoder = lambda: SimpleNamespace(encode=lambda values, **kwargs: [embedding() for _ in values], tokenizer=tokenize) try: index.rebuild() (root / 'a.md').write_text('after', encoding='utf-8') (root / 'b.md').write_text('fail', encoding='utf-8') def encode(values, **kwargs): if any(value.endswith('\n\nfail') for value in values): raise RuntimeError('encoder failed') return [embedding() for _ in values] index.encoder = lambda: SimpleNamespace(encode=encode, tokenizer=tokenize) with self.assertRaisesRegex(RuntimeError, 'encoder failed'): index.rebuild() self.assertEqual(index.db.execute('select path, text from chunks').fetchall(), [('a.md', 'before')]) self.assertEqual(index.db.execute('select document from lexical').fetchall(), [('Path: a.md\nSymbol: \n\nbefore',)]) finally: index.db.close() def test_refresh_batches_embeddings_for_changed_files(self): with TemporaryDirectory() as temporary: root = Path(temporary) initialize_repository(root) (root / 'first.md').write_text('first', encoding='utf-8') (root / 'second.md').write_text('second', encoding='utf-8') index = MODULE.Index(root, root / 'data') encoded = [] def encode(values, **kwargs): encoded.append(list(values)) return [embedding() for _ in values] index.encoder = lambda: SimpleNamespace(encode=encode, tokenizer=tokenize) try: index.rebuild() self.assertEqual(len(encoded), 1) self.assertEqual(len(encoded[0]), 2) finally: index.db.close() def test_scan_prunes_ignored_directories_and_warns_at_limit(self): with TemporaryDirectory() as temporary: root = Path(temporary) initialize_repository(root) index = MODULE.Index(root, root / 'data') (root / 'visible.md').write_text('visible', encoding='utf-8') (root / '.gitignore').write_text('secrets.json\n**/private/\n', encoding='utf-8') (root / 'secrets.json').write_text('{"token":"hidden"}', encoding='utf-8') private = root / 'nested' / 'private' private.mkdir(parents=True) (private / 'hidden.md').write_text('hidden', encoding='utf-8') for name in ('node_modules', '.git', 'generated', 'data'): directory = root / name directory.mkdir(exist_ok=True) (directory / 'hidden.md').write_text('hidden', encoding='utf-8') try: self.assertEqual(index.files(), [(root / 'visible.md').resolve()]) filenames = ['visible.md'] * 2001 with patch.object(MODULE.subprocess, 'run', return_value=SimpleNamespace(stdout=b'visible.md\0' * 2001)), \ patch.object(MODULE.sys, 'stderr') as stderr: self.assertEqual(len(index.files()), 2000) self.assertIn('stopped scanning after 2000', stderr.write.call_args_list[0].args[0]) finally: index.db.close() def test_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'): sentinel = SimpleNamespace() with patch.dict('sys.modules', sentence_transformers=SimpleNamespace( SentenceTransformer=lambda *args, **kwargs: sentinel, CrossEncoder=lambda *args, **kwargs: sentinel)): self.assertIs(MODULE.shared_encoder(), sentinel) self.assertIs(MODULE.shared_encoder(), MODULE.shared_encoder()) self.assertIs(MODULE.shared_reranker(), sentinel) self.assertIs(MODULE.shared_reranker(), MODULE.shared_reranker()) def test_rerank_uses_cross_encoder_scores(self): index = MODULE.Index.__new__(MODULE.Index) class Reranker: def predict(self, pairs, **kwargs): self.pairs = pairs return [0.1, 0.8, 0.5] reranker = Reranker() index.reranker = lambda: reranker rows = [ {'path': 'a.md', 'symbol': 'A', 'score': 0.9, 'text': 'first'}, {'path': 'b.md', 'symbol': 'B', 'score': 0.8, 'text': 'second'}, {'path': 'c.md', 'symbol': 'C', 'score': 0.7, 'text': 'third'}, ] result = index.rerank('query', rows, 2) self.assertEqual([item['path'] for item in result], ['b.md', 'c.md']) self.assertEqual(reranker.pairs, [('query', 'Path: a.md\nSymbol: A\n\nfirst'), ('query', 'Path: b.md\nSymbol: B\n\nsecond'), ('query', 'Path: c.md\nSymbol: C\n\nthird')]) def test_token_windows_have_exact_overlap_and_no_redundant_tail(self): for size, lengths in [(0, []), (400, [400]), (401, [400, 65]), (736, [400, 400]), (800, [400, 400, 128])]: words = [f'token{number}' for number in range(size)] result = list(MODULE.chunks(' '.join(words), 'plain.txt', tokenize)) windows = [content.split() for _, content in result] self.assertEqual([len(window) for window in windows], lengths) for first, second in zip(windows, windows[1:]): self.assertEqual(first[-64:], second[:64]) if windows: reconstructed = windows[0] + [word for window in windows[1:] for word in window[64:]] self.assertEqual(reconstructed, words) def test_python_symbols_include_decorators_and_restore_parent(self): source = 'class Service:\n @cached\n def load(self):\n return 1\n\nvalue = 2\n\nasync def refresh():\n return value\n' result = list(MODULE.sections(source, 'service.py')) self.assertEqual([symbol for symbol, _ in result], ['Service', 'Service.load', '', 'refresh']) self.assertIn('@cached', result[1][1]) self.assertEqual(''.join(content for _, content in result), source) broken = 'def incomplete(\n' + 'word ' * 450 self.assertEqual(len(list(MODULE.chunks(broken, 'broken.py', tokenize))), 2) def test_markdown_sections_ignore_fenced_headings(self): source = '# Setup\nIntro\n## Install\n```sh\n# not a heading\n```\nSteps\n## Run\nStart\n' result = list(MODULE.chunks(source, 'guide.md', tokenize)) self.assertEqual([symbol for symbol, _ in result], ['Setup', 'Setup / Install', 'Setup / Run']) self.assertIn('# not a heading', result[1][1]) def test_common_declarations_and_config_sections(self): for path, source, symbols in [ ('service.ts', 'export function load() {\n return 1;\n}\nexport const save = () => 2;\n', ['load', 'save']), ('service.cs', 'public class Service\n{\n public void Load()\n {\n }\n}\n', ['Service', 'Load']), ('setup.ps1', 'function Get-Result {\n return 1\n}\n', ['Get-Result']), ('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()