anc_benchmark.py

import argparse
import json
import subprocess
import sys
import time
from pathlib import Path

sys.path.insert(0, str(Path(__file__).resolve().parent))

import anc
from core.codec import GlyphCodec


MODES = ['raw', 'semantic', 'glyph', 'hybrid', 'dual_glyph', 'quad_collapse_full', 'pyramid5']


def build_expected(rows):
    node_count = len(rows)
    types = [row[1] for row in rows]
    same_type = len(set(types)) == 1
    first = rows[0] if rows else None
    if first:
        first_counts = [
            len(anc.parse_slot(first[2])),
            len(anc.parse_slot(first[3])),
            len(anc.parse_slot(first[4])),
        ]
    else:
        first_counts = [0, 0, 0]
    glyph_codec = GlyphCodec(str(anc.DEFAULT_GLYPH_DB))
    keys = []
    for row in rows:
        keys.append(glyph_codec.archetype_key(row[1] or '', anc.parse_slot(row[2]), anc.parse_slot(row[3]), anc.parse_slot(row[4]), row[5] or ''))
    shared_archetype = len(set(keys)) == 1 if keys else False
    return {
        'node_count': node_count,
        'same_type': same_type,
        'first_counts': first_counts,
        'shared_archetype': shared_archetype,
    }


def build_prompt(mode, rows):
    task = None
    if mode == 'raw':
        context = anc.render_raw(rows)
    elif mode == 'semantic':
        context = anc.render_semantic(rows)
    elif mode == 'glyph':
        context = anc.render_glyph(rows)
    elif mode == 'hybrid':
        context = anc.render_hybrid(rows)
    elif mode == 'dual_glyph':
        context = anc.render_dual_glyph(rows)
    elif mode == 'quad_collapse_full':
        task = 'full'
        context = anc.render_quad_collapse(rows, task)
    else:
        context = anc.render_pyramid5(rows)
    instructions = [
        'Read the context and answer using strict JSON only.',
        'Use this schema exactly: {"node_count": int, "same_type": bool, "first_counts": [int,int,int], "shared_archetype": bool}.',
        '',
        'Definitions:',
        '- node_count = number of nodes shown',
        '- same_type = whether all nodes have the same type',
        '- first_counts = [N,M,F] counts for the first node only',
        '- shared_archetype = whether all shown nodes share the same structural archetype',
        '',
        'Rules:',
        '- If NODE_COUNT_EXACT exists, use it directly for node_count.',
        '- If TYPE_SHARED exists, use it directly for same_type.',
        '- If COUNT_VIEW exists, use FIRST_N, FIRST_M, FIRST_F directly for first_counts.',
        '- If Q_VIEW exists, extract FIRST_N, FIRST_M, FIRST_F directly for first_counts.',
        '- If ARCH_SHARED exists or ARCH_SHARED_EXACT exists, use it directly for shared_archetype.',
        '- Use S_VIEW, M_VIEW, R_VIEW, F_VIEW, ARCH_VIEW, RELATION_VIEW only as support if the explicit rails are missing.',
        '- Do not infer extra nodes.',
    ]
    parts = ['\n'.join(instructions), f'MODE\n{mode}']
    if task:
        parts.append(f'TASK\n{task}')
    parts.append(f'CONTEXT\n{context}')
    return '\n\n'.join(parts) + '\n'


def run_model(model, prompt):
    t0 = time.perf_counter()
    proc = subprocess.run(['ollama', 'run', model], input=prompt, text=True, capture_output=True, encoding='utf-8', errors='replace')
    dt = time.perf_counter() - t0
    stdout = (proc.stdout or '').strip()
    stderr = (proc.stderr or '').strip()
    return proc.returncode, stdout, stderr, dt


def parse_json(text):
    try:
        return json.loads(text)
    except Exception:
        start = text.find('{')
        end = text.rfind('}')
        if start != -1 and end != -1 and end > start:
            try:
                return json.loads(text[start:end+1])
            except Exception:
                return None
        return None


def score_answer(answer, expected):
    if not isinstance(answer, dict):
        return {'exact': False, 'field_score': 0, 'max_score': 4}
    fields = 0
    for key in ['node_count', 'same_type', 'first_counts', 'shared_archetype']:
        if answer.get(key) == expected.get(key):
            fields += 1
    return {'exact': fields == 4, 'field_score': fields, 'max_score': 4}


def main():
    parser = argparse.ArgumentParser(description='Benchmark ANC prompt modes on a local Ollama model')
    parser.add_argument('--db', default=str(anc.DEFAULT_DB))
    parser.add_argument('--model', default='qwen2.5:0.5b')
    parser.add_argument('--type', default='function')
    parser.add_argument('--search', default=None)
    parser.add_argument('--limit', type=int, default=2)
    parser.add_argument('--out', default='')
    args = parser.parse_args()

    rows = anc.fetch_rows(Path(args.db), args.limit, args.type, args.search)
    if not rows:
        raise SystemExit('No rows found for benchmark selection.')
    expected = build_expected(rows)
    results = {'model': args.model, 'expected': expected, 'cases': []}

    for mode in MODES:
        prompt = build_prompt(mode, rows)
        code, stdout, stderr, latency = run_model(args.model, prompt)
        parsed = parse_json(stdout)
        score = score_answer(parsed, expected)
        results['cases'].append({
            'mode': mode,
            'prompt_chars': len(prompt),
            'returncode': code,
            'latency_sec': round(latency, 3),
            'stdout': stdout,
            'stderr': stderr,
            'parsed': parsed,
            'score': score,
        })

    text = json.dumps(results, ensure_ascii=False, indent=2)
    if args.out:
        Path(args.out).write_text(text, encoding='utf-8')
        print(args.out)
    else:
        print(text)


if __name__ == '__main__':
    main()