MyAiConfiguration/scripts/tests/test-retrieval-server.py

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_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_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', '<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']),
            ('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()