jhonny-editor/rag/bench.py

"""Benchmark da RAG: mede recall@k e custo de tokens da busca.

Roda um conjunto dourado de consultas contra o índice e reporta, para cada
uma, se o arquivo esperado apareceu no top-k e em que posição. No fim mostra
recall@k agregado, latência média e o tamanho médio da saída — as três coisas
que o RAG precisa otimizar (achar o arquivo certo, rápido, gastando pouco
token).

Uso:
    rag/bench.sh                 # roda tudo
    rag/bench.sh --mode snippet  # mede o custo de outro modo de saída
    rag/bench.sh --baseline      # usa a busca densa pura (pré-otimização)

O resultado é impresso e também salvo em rag/bench-results/<rotulo>.json,
para comparar antes/depois.
"""

import argparse
import json
import os
import statistics
import sys
import time

RAG_DIR = os.path.dirname(os.path.abspath(__file__))
if RAG_DIR not in sys.path:
    sys.path.insert(0, RAG_DIR)

# Carrega o env ANTES de escolher o conjunto dourado. Sem isso o schema seria
# lido do ambiente do processo, onde não está — quem o define é o arquivo
# .env, e o search.py só o carrega quando é importado, tarde demais.
from dotenv import load_dotenv  # noqa: E402

load_dotenv(os.environ.get("RAG_ENV_FILE", os.path.join(RAG_DIR, ".env")))

# Conjuntos dourados, um por sistema indexado: (consulta, arquivo esperado,
# categoria).
# - "exato": o usuário sabe o nome do tipo/função e quer o arquivo dele.
# - "conceitual": descreve comportamento, sem citar identificador.
# - "modulo": pergunta ampla sobre uma área do sistema.
GOLDEN_TIGRE = [
    # --- nome exato de tipo (onde a busca densa pura falha) ---
    ("FCPXMLValidator", "TigreVideoEditing/FCPXMLValidator.swift", "exato"),
    ("SilenceCutPipeline", "TigreVideoEditing/SilenceCutPipeline.swift", "exato"),
    ("TranscriptIndexBuilder", "TigreTranscription/Pipeline/TranscriptIndexBuilder.swift", "exato"),
    ("SQLiteProjectStore", "TigreFoundation/Persistence/SQLiteProjectStore.swift", "exato"),
    ("EditPromptBuilder", "TigreAI/EditPromptBuilder.swift", "exato"),
    ("UniformFrameScheduler", "TigreVideoAnalysis/Pipeline/UniformFrameScheduler.swift", "exato"),
    ("TemplateRenderer", "TigreFoundation/Web/TemplateRenderer.swift", "exato"),

    # --- conceitual: descreve o comportamento ---
    ("como o wizard decide qual e o proximo passo", "TigreAppUI/WizardStep.swift", "conceitual"),
    ("extrair o audio de um video com ffmpeg", "TigreTranscription/Pipeline/FFmpegAudioExtractor.swift", "conceitual"),
    ("detectar trechos de silencio no audio", "TigreVideoEditing/SilenceAnalyzer.swift", "conceitual"),
    ("baixar o modelo do whisper", "TigreTranscription/Pipeline/WhisperModelDownloader.swift", "conceitual"),
    ("escrever o arquivo final de FCPXML", "TigreVideoEditing/FCPXMLWriter.swift", "conceitual"),
    ("detectar rostos nos quadros do video", "TigreVideoAnalysis/Pipeline/FacePerceiver.swift", "conceitual"),
    ("servidor HTTP que atende a webapp", "TigreApp/Web/HTTPServer.swift", "conceitual"),
    ("gravar arquivo em disco de forma atomica", "TigreFoundation/Files/AtomicFileIO.swift", "conceitual"),
    ("cliente que fala com o Ollama", "TigreAI/OllamaClient.swift", "conceitual"),
    ("impedir caminho de projeto fora da pasta permitida", "TigreFoundation/Persistence/ProjectPathSafety.swift", "conceitual"),
    ("mapear tempo da timeline para o tempo da fonte", "TigreVideoEditing/TimelineMapper.swift", "conceitual"),
]

# O legado é Python e NÃO segue "uma classe por arquivo": há módulos de 5000
# linhas com dezenas de funções soltas. O conjunto reflete isso — as
# consultas exatas citam função, não só classe.
GOLDEN_DOZA = [
    ("SpineSegment", "doza_assist/fcpxml/parser.py", "exato"),
    ("OllamaProvider", "ai_providers/ollama_provider.py", "exato"),
    ("detect_hardware_tier", "model_config.py", "exato"),
    ("snap_framerate", "exporters/media_probe.py", "exato"),
    ("whisper_catalog", "whisper_catalog.py", "exato"),
    ("premiere_xml", "exporters/premiere_xml.py", "exato"),
    ("parakeet worker", "parakeet_worker.py", "exato"),

    ("converter segundos em timecode", "exporters/edl.py", "conceitual"),
    ("descobrir a resolucao do video com ffprobe", "exporters/media_probe.py", "conceitual"),
    ("detectar quanta memoria RAM a maquina tem", "model_config.py", "conceitual"),
    ("escrever o arquivo FCPXML de saida", "doza_assist/fcpxml/writer.py", "conceitual"),
    ("ler e interpretar um arquivo FCPXML", "doza_assist/fcpxml/parser.py", "conceitual"),
    ("perfis de DNA editorial salvos em disco", "editorial_dna/profiles.py", "conceitual"),
    ("instalar e baixar modelos do whisper", "whisper_catalog.py", "conceitual"),
    ("exportar lista de decisao EDL", "exporters/edl.py", "conceitual"),
    ("transcrever o audio para texto", "transcribe.py", "conceitual"),
    ("analisar video usando o framework Vision", "video_analysis.py", "conceitual"),
    ("provedor de IA da Anthropic", "ai_providers/anthropic_provider.py", "conceitual"),
]

# Projeto Jhonny: MCP em TypeScript para controlar o Premiere Pro (code/src,
# code/tools). Consultas exatas citam a classe/funcao exportada; conceituais
# descrevem o comportamento sem citar identificador.
GOLDEN_JHONNY = [
    ("OAuthResourceServer", "src/oauth-resource-server.ts", "exato"),
    ("UxpWebSocketBridge", "src/bridge/uxp-websocket-bridge.ts", "exato"),
    ("ProjectContextRepository", "src/context/project-context-store.ts", "exato"),
    ("applyDoctorRepairPlan", "src/doctor-repairs.ts", "exato"),
    ("buildContentSecurityPolicy", "src/http-security.ts", "exato"),
    ("planDerivedSilenceRemoval", "src/tools/silence-removal.ts", "exato"),
    ("buildCaptionTimingPlan", "src/ai/caption-timing.ts", "exato"),
    ("emitAudit", "src/security/audit.ts", "exato"),

    ("cabecalhos de seguranca HTTP e CSP", "src/http-security.ts", "conceitual"),
    ("registrar evento de telemetria de ativacao", "src/telemetry.ts", "conceitual"),
    ("ponte websocket com o plugin UXP", "src/bridge/uxp-websocket-bridge.ts", "conceitual"),
    ("interpretar arquivo de legenda SRT ou VTT", "src/ai/caption-timing.ts", "conceitual"),
    ("plano de remocao de silencio derivado da transcricao", "src/tools/silence-removal.ts", "conceitual"),
    ("armazenar contexto do projeto em disco", "src/context/project-context-store.ts", "conceitual"),
    ("checagem e reparo automatico de instalacao (doctor)", "src/doctor-repairs.ts", "conceitual"),
    ("relatorio de intake do projeto", "src/intake/project-intake.ts", "conceitual"),
    ("servidor OAuth para autenticacao de recursos", "src/oauth-resource-server.ts", "conceitual"),
    ("registrar operacao de auditoria de seguranca", "src/security/audit.ts", "conceitual"),
]

# O conjunto e o prefixo de caminho seguem o schema apontado pelo env, para o
# mesmo bench servir os tres bancos sem flag extra.
_SCHEMA = os.environ.get("RAG_DB_SCHEMA", "tigre")
if _SCHEMA == "doza":
    GOLDEN, PREFIX = GOLDEN_DOZA, "code/"
elif _SCHEMA == "jhonny-rag":
    GOLDEN, PREFIX = GOLDEN_JHONNY, "code/"
else:
    GOLDEN, PREFIX = GOLDEN_TIGRE, "codeclass/Sources/"


def _run(top_k, use_baseline, mode):
    from search import rag_search

    rows = []
    for query, expected, kind in GOLDEN:
        expected_path = PREFIX + expected
        t0 = time.perf_counter()
        if use_baseline:
            results = rag_search(query, top_k=top_k, dense_only=True)
        else:
            results = rag_search(query, top_k=top_k)
        elapsed = time.perf_counter() - t0

        paths = [r["file_path"] for r in results]
        rank = paths.index(expected_path) + 1 if expected_path in paths else None
        rows.append({
            "query": query,
            "kind": kind,
            "expected": expected,
            "rank": rank,
            "elapsed": elapsed,
            "bytes": _render_bytes(results, mode),
            "got": [p[len(PREFIX):] if p.startswith(PREFIX) else p for p in paths],
        })
    return rows


def _render_bytes(results, mode):
    """Tamanho da saída que o agente realmente receberia neste modo."""
    from search import format_results
    return len(format_results(results, mode=mode))


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--top-k", type=int, default=5)
    ap.add_argument("--mode", default="map", help="map | snippet | full")
    ap.add_argument("--baseline", action="store_true",
                    help="busca densa pura, sem fusao lexica (estado pre-otimizacao)")
    ap.add_argument("--label", default=None, help="nome do arquivo de resultado")
    args = ap.parse_args()

    rows = _run(args.top_k, args.baseline, args.mode)

    hits = [r for r in rows if r["rank"] is not None]
    print(f"\n{'ok':>3} {'#':>2}  {'cat':<11} consulta")
    print("-" * 72)
    for r in rows:
        mark = "OK " if r["rank"] else "-- "
        pos = str(r["rank"]) if r["rank"] else "x"
        print(f"{mark:>3} {pos:>2}  {r['kind']:<11} {r['query'][:48]}")
        if not r["rank"]:
            print(f"{'':>18} esperado: {r['expected']}")
            print(f"{'':>18} veio:     {', '.join(r['got'][:3])}")

    recall = len(hits) / len(rows)
    top1 = sum(1 for r in hits if r["rank"] == 1) / len(rows)
    summary = {
        "label": args.label or ("baseline" if args.baseline else "otimizado"),
        "mode": args.mode,
        "top_k": args.top_k,
        "n": len(rows),
        f"recall@{args.top_k}": round(recall, 3),
        "top1": round(top1, 3),
        "latencia_media_s": round(statistics.mean(r["elapsed"] for r in rows), 3),
        "bytes_medios": round(statistics.mean(r["bytes"] for r in rows)),
    }
    print("-" * 72)
    for k, v in summary.items():
        print(f"  {k:<18} {v}")

    out_dir = os.path.join(RAG_DIR, "bench-results")
    os.makedirs(out_dir, exist_ok=True)
    out = os.path.join(out_dir, f"{summary['label']}-{args.mode}.json")
    with open(out, "w", encoding="utf-8") as f:
        json.dump({"summary": summary, "rows": rows}, f, indent=2, ensure_ascii=False)
    print(f"\n  -> {os.path.relpath(out, os.path.dirname(RAG_DIR))}")


if __name__ == "__main__":
    main()