searchi/tools/mixed-eval/evaluate.py

"""Frozen mixed diagnostic via the installed CLI's shared engine and MCP API."""
import argparse
from datetime import datetime, timezone
import hashlib
import importlib.util
import json
import os
from pathlib import Path
import shutil
import subprocess
import tempfile
import time
import platform

ROOT = Path(__file__).resolve().parents[2]
FIXTURE = ROOT / 'fixtures/mixed'


def digest(path):
    with path.open('rb') as stream:
        return hashlib.file_digest(stream, 'sha256').hexdigest()


def load_fixture(directory=FIXTURE):
    raw = (directory / 'fixture.json').read_bytes()
    lock = json.loads((directory / 'labels.lock.json').read_bytes())
    if hashlib.sha256(raw).hexdigest() != lock['fixture_sha256']:
        raise ValueError('Frozen labels changed; create a new version, never refreeze v1')
    fixture = json.loads(raw)
    allowed = directory.parent.resolve()
    ids, names = set(), set()
    for asset in fixture['assets']:
        path = (directory / asset['path']).resolve()
        if not path.is_relative_to(allowed) or not path.is_file():
            raise ValueError('Asset must be an existing file below fixtures')
        if asset['id'] in ids or path.name in names:
            raise ValueError('Asset IDs and materialized filenames must be unique')
        ids.add(asset['id']); names.add(path.name)
        if path.stat().st_size != asset['bytes'] or digest(path) != asset['sha256']:
            raise ValueError('Frozen asset changed: ' + asset['id'])
    for query in fixture['queries']:
        # Temporal/media labels can include structured entries; normalize below.
        if any(identifier not in ids for identifier in expected_ids(query)):
            raise ValueError('Unknown expected asset')
    return fixture, lock


def expected_ids(query):
    return {item if isinstance(item, str) else item['asset'] for item in query['expected']}


def score(query, hits):
    expected = expected_ids(query)
    unique = list(dict.fromkeys(hit['asset'] for hit in hits))
    ranks = [i + 1 for i, asset in enumerate(unique) if asset in expected]
    first = min(ranks, default=None)
    return {'first_relevant_rank': first,
            'top1': first == 1 if expected else None,
            'recall_at_5': len(expected.intersection(unique[:5])) / len(expected) if expected else None,
            'reciprocal_rank': 1 / first if first else 0 if expected else None,
            'no_answer_returned_any': bool(hits) if not expected else None,
            'hard_negative_in_top5': any(asset in query.get('hard_negatives', []) for asset in unique[:5])}


def summarize(rows):
    groups = {}
    for row in rows:
        if 'unsupported' in row:
            continue
        key = (row['mode'], row['scope'], row['query']['category'], row['query']['language'])
        groups.setdefault(key, []).append(row['score'])
    result = []
    for key, values in sorted(groups.items()):
        group = dict(zip(('mode', 'scope', 'category', 'language'), key))
        group['queries'] = len(values)
        for metric in ('top1', 'recall_at_5', 'reciprocal_rank', 'no_answer_returned_any', 'hard_negative_in_top5'):
            numbers = [v[metric] for v in values if v[metric] is not None]
            group[metric] = sum(numbers) / len(numbers) if numbers else None
        result.append(group)
    return result


def run(args):
    if os.environ.get('SEARCHI_FORCE_CPU') != '1':
        raise ValueError('Set SEARCHI_FORCE_CPU=1; this evaluator must never run GPU inference')
    fixture, lock = load_fixture()
    manifest_path = ROOT / 'tools/embeddinggemma2/model-manifest.json'
    manifest = json.loads(manifest_path.read_bytes())
    for model in manifest['files']:
        path = args.models / model['path']
        if path.stat().st_size != model['bytes'] or digest(path) != model['sha256']:
            raise ValueError('Model bytes differ from the pinned Google manifest')
    spec = importlib.util.spec_from_file_location('searchi_benchmark', ROOT / 'tools/embeddinggemma2/benchmark.py')
    benchmark = importlib.util.module_from_spec(spec); spec.loader.exec_module(benchmark)
    args.output.parent.mkdir(parents=True, exist_ok=True)
    rows, stages = [], []
    provenance = {'fixture_sha256': lock['fixture_sha256'], 'model_manifest_sha256': digest(manifest_path),
                  'model': manifest, 'device': 'cpu', 'cpu_threads': 1,
                  'cli_sha256': digest(args.cli), 'worker_sha256': digest(args.worker),
                  'platform': platform.platform(), 'python': platform.python_version()}
    def checkpoint():
        args.output.write_text(json.dumps({**provenance, 'complete': False, 'stages': stages, 'rows': rows}, indent=2, ensure_ascii=False), encoding='utf-8')
    def stage(name, command, timeout):
        started = time.perf_counter()
        try:
            result = subprocess.run(command, check=False, timeout=timeout, capture_output=True, text=True, encoding='utf-8')
        except subprocess.TimeoutExpired as error:
            def text(value):
                return value.decode('utf-8', errors='replace') if isinstance(value, bytes) else value or ''
            stages.append({'name': name, 'elapsed_ms': (time.perf_counter() - started) * 1000,
                           'timed_out': True, 'stdout': text(error.stdout), 'stderr': text(error.stderr)})
            checkpoint()
            raise
        entry = {'name': name, 'elapsed_ms': (time.perf_counter() - started) * 1000,
                 'returncode': result.returncode, 'stdout': result.stdout, 'stderr': result.stderr}
        try:
            entry['report'] = json.loads(result.stdout)
        except json.JSONDecodeError:
            pass
        stages.append(entry); checkpoint()
        result.check_returncode()
    checkpoint()
    with tempfile.TemporaryDirectory(prefix='searchi-mixed-') as temporary:
        base = Path(temporary); corpus = base / 'corpus'; corpus.mkdir(); database = base / 'index.sqlite'
        names = {}
        for asset in fixture['assets']:
            source = (FIXTURE / asset['path']).resolve()
            shutil.copyfile(source, corpus / source.name)
            names[source.name] = asset['id']
        stage('lexical_index', [str(args.cli), 'index', '--db', str(database), '--root', str(corpus)], 300)
        stage('semantic_index', [str(args.cli), 'embed', '--db', str(database), '--model-dir', str(args.models),
                        '--embedding-worker', str(args.worker), '--cpu-threads', '1',
                        '--ffmpeg', str(args.ffmpeg), '--ffprobe', str(args.ffprobe)], 3600)
        for mode in ('lexical_only', 'hybrid'):
            command = [str(args.cli), 'mcp', '--db', str(database)]
            if mode == 'hybrid':
                command += ['--model-dir', str(args.models), '--embedding-worker', str(args.worker), '--cpu-threads', '1']
            session = benchmark.Session(command)
            try:
                session.mcp(1, 'initialize', {'protocolVersion': '2025-03-26', 'capabilities': {}, 'clientInfo': {'name': 'mixed-eval', 'version': '1'}})
                session.send({'jsonrpc': '2.0', 'method': 'notifications/initialized'})
                request = 2
                for query in fixture['queries']:
                    scopes = ('all_files', 'code') if query['category'] == 'code' else ('all_files',)
                    for scope in scopes:
                        options = {'query': query['text'], 'limit': 20, 'code_only': scope == 'code'}
                        filters = query.get('filters', {})
                        if 'time_range_s' in filters:
                            rows.append({'mode': mode, 'scope': scope, 'query': query, 'unsupported': 'Temporal-range search filter is not supported'})
                            continue
                        if 'extension' in filters:
                            options['extension'] = filters['extension']
                        response, elapsed = session.mcp(request, 'tools/call', {'name': 'search', 'arguments': options}); request += 1
                        if response.get('isError'):
                            raise RuntimeError('Shared search failed: ' + json.dumps(response))
                        content = response['structuredContent']
                        if content['mode'] != mode:
                            raise RuntimeError('Unexpected semantic fallback; not valid quality evidence')
                        hits = [{'asset': names[hit['filename']], 'score': hit['score'], 'match_kind': hit['match_kind'], 'source': hit.get('source')} for hit in content['results']]
                        rows.append({'mode': mode, 'scope': scope, 'query': query, 'hits': hits, 'score': score(query, hits), 'roundtrip_ms': elapsed, 'metrics': content.get('metrics')})
                        # Retain partial actual evidence if a later request fails.
                        checkpoint()
            finally:
                session.close()
    report = {**provenance, 'stages': stages, 'complete': True, 'recorded_at': datetime.now(timezone.utc).isoformat(), 'fixture_sha256': lock['fixture_sha256'],
              'model_manifest_sha256': digest(manifest_path), 'model': manifest,
              'device': 'cpu', 'cpu_threads': 1, 'cli_sha256': digest(args.cli), 'worker_sha256': digest(args.worker),
              'scope_note': 'Small mixed diagnostic, reused previously tested media labels; file-level relevance only. No temporal-localization or production-scale claims. MCP roundtrip excludes initialization and is not UI latency.',
              'rows': rows, 'summary': summarize(rows)}
    args.output.write_text(json.dumps(report, indent=2, ensure_ascii=False), encoding='utf-8')


if __name__ == '__main__':
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument('--verify-only', action='store_true')
    for flag in ('cli', 'models', 'worker', 'ffmpeg', 'ffprobe', 'output'):
        parser.add_argument('--' + flag, type=Path)
    args = parser.parse_args()
    if args.verify_only:
        fixture, lock = load_fixture(); print('Verified', len(fixture['assets']), 'assets and', len(fixture['queries']), 'frozen queries', lock['fixture_sha256'])
    else:
        if any(getattr(args, flag) is None for flag in ('cli', 'models', 'worker', 'ffmpeg', 'ffprobe', 'output')):
            parser.error('Supply all explicit runtime paths and output, or use --verify-only')
        run(args)