kg-agent/scripts/benchmark_qa_planning.py

"""Compare planning with both optimizations off/on; real LLM calls require authorization.

--execute additionally runs schema-validated read-only Neo4j queries. No raw query
results are saved. Planning times exclude confirmation, query execution and answer.
"""
from __future__ import annotations
import argparse
import hashlib
import json
import os
from pathlib import Path
from time import perf_counter
from unittest.mock import patch
from step3_qa_agent.agent import nodes, production, llm
from step3_qa_agent.retrieval.production_query import compile_query
from step2_graph_building.runtime import release_snapshot

def query_fingerprint(plan, context):
    if release_snapshot()['data_version'] != context['build_id']:
        raise RuntimeError('数据版本变化，取消对比')
    driver = production.get_driver()
    results = []
    for step in plan['steps']:
        spec = step['params']
        compiled = compile_query(spec, context['schema'], context['build_id'])
        records = driver.execute_query(compiled.query, **compiled.params).records
        aliases = {n['alias']: n['type'] for n in spec['nodes']}
        columns = [(str((aliases[c['alias']], c.get('op', 'field'), c.get('field'))), c['as'])
                   for c in spec.get('select', []) + spec.get('aggregates', [])]
        rows = [sorted((identity, row.get(label)) for identity, label in columns) for row in records]
        encoded = sorted(json.dumps(row, ensure_ascii=False, sort_keys=True, default=str) for row in rows)
        results.append({'rows': len(records), 'truncated': len(records) > compiled.limit,
                        'sha256': hashlib.sha256(json.dumps(encoded, ensure_ascii=False).encode()).hexdigest()})
    if release_snapshot()['data_version'] != context['build_id']:
        raise RuntimeError('查询期间数据版本变化，取消对比')
    return results

def main():
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument('--question', action='append')
    parser.add_argument('--repeat', type=int, default=1)
    parser.add_argument('--execute', action='store_true', help='实际执行受控只读查询，仅保存结果摘要')
    parser.add_argument('--output', default='.runtime/qa/planning-benchmark.json')
    args = parser.parse_args()
    if args.repeat < 1:
        parser.error('--repeat 必须大于0')
    questions = args.question or ['共有多少名员工？', '岗位名称包含保安的员工有多少名？',
                                   '张三的岗位名称是什么？', '按岗位名称分别统计员工人数']
    output = Path(args.output)
    output.parent.mkdir(parents=True, exist_ok=True)
    report = {'scope': 'understand + production_plan only; both switches off vs on',
              'executed_read_only_queries': args.execute, 'runs': []}
    original_chat = llm.get_chat
    def bounded_chat(*a, **kw):
        model = original_chat(*a, **kw)
        model.request_timeout = 90
        model.max_retries = 0
        return model
    for repeat in range(args.repeat):
        for index, question in enumerate(questions):
            for enabled in ([False, True] if (repeat + index) % 2 == 0 else [True, False]):
                row = {'question': question, 'repeat': repeat, 'optimizations_enabled': enabled}
                started = perf_counter()
                try:
                    with patch.dict(os.environ, {
                            'QA_FAST_PLAN_ENABLED': str(enabled).lower(),
                            'QA_COMPACT_INTERMEDIATE_ENABLED': str(enabled).lower()}), \
                         patch.object(llm, 'get_chat', side_effect=bounded_chat):
                        state = {'question': question}
                        state.update(nodes.understand(state))
                        state.update(production.plan(state))
                    row.update(total_ms=round((perf_counter() - started) * 1000, 2),
                               build_id=state['qa_context']['build_id'], trace=state['trace'],
                               plan=state['plan'], candidate=state.get('simple_query'))
                    if args.execute:
                        row['query_results'] = query_fingerprint(state['plan'], state['qa_context'])
                    row['status'] = 'ok'
                except Exception as exc:
                    row.update(status='error', error_type=type(exc).__name__,
                               total_ms=round((perf_counter() - started) * 1000, 2))
                report['runs'].append(row)
                output.write_text(json.dumps(report, ensure_ascii=False, indent=2), encoding='utf-8')
                trace = row.get('trace', [])
                print(json.dumps({'question': question, 'optimized': enabled, 'status': row['status'],
                                  'ms': row['total_ms'], 'mode': trace[-1].get('mode') if trace else None,
                                  'reason': trace[-1].get('reason') if trace else None,
                                  'llm_streams': sum(len(t.get('llm_calls', [])) for t in trace)}, ensure_ascii=False), flush=True)
    print('Saved: ' + str(output), flush=True)
    return 1 if any(r['status'] != 'ok' for r in report['runs']) else 0

if __name__ == '__main__':
    raise SystemExit(main())