scripts/test_retrieval_cli.py

#!/usr/bin/env python3
"""CLI para probar retrieval del backend (sin frontend).

Uso:
  python scripts/test_retrieval_cli.py
  python scripts/test_retrieval_cli.py --question "O que é NORM?"
  python scripts/test_retrieval_cli.py --question "césio-137" --top-k 20
"""

from __future__ import annotations

import argparse
import sys
from typing import Any, Dict, List

import app.api_server as srv


def _print_header(top_k: int, normalize: bool) -> None:
    print("=" * 90)
    print("Teste de retrieval (modo chatbot)")
    print(f"top_k efetivo: {top_k}")
    print(f"index.type: {srv.CONFIG.get('index', {}).get('type', 'faiss')}")
    print(f"normalize L2: {normalize}")
    print("=" * 90)


def _print_chunk(i: int, chunk: Dict[str, Any], max_chars: int) -> None:
    doc_id = chunk.get("document_id", "N/A")
    title = chunk.get("document_title", "N/A")
    topic = chunk.get("topic", "N/A")
    frag_id = chunk.get("fragment_id", chunk.get("idx", "N/A"))
    score = chunk.get("score")
    cit_id = chunk.get("citation_id")
    text = srv._normalizar_texto(chunk.get("content", ""))
    preview = text[:max_chars] + ("..." if len(text) > max_chars else "")

    print(f"\n[{i}] citation_id={cit_id} | score={score} | doc_id={doc_id} | frag={frag_id}")
    print(f"    title: {title}")
    print(f"    topic: {topic}")
    print(f"    chunk: {preview}")


def _run_retrieval(question: str, top_k_override: int | None, max_chars: int) -> None:
    question = (question or "").strip()
    if not question:
        print("Pergunta vazia. Tente novamente.")
        return

    pregunta_lower = question.lower().strip()

    help_triggers = [str(t).lower() for t in srv.TRIGGERS.get("help", [])]
    smalltalk_triggers = [str(t).lower() for t in srv.TRIGGERS.get("smalltalk", [])]
    greeting_triggers = [str(t).lower() for t in srv.TRIGGERS.get("greeting", [])]

    # Replica o comportamento do endpoint: alguns casos não passam pelo retrieval.
    if srv._contains_any_trigger(pregunta_lower, help_triggers):
        print("[SKIP] A pergunta bateu em trigger de help; no endpoint isso retorna sem retrieval.")
        print(f"Resposta esperada: {srv.HELP_SCOPE_TEXT}")
        return
    if srv._contains_any_trigger(pregunta_lower, smalltalk_triggers):
        print("[SKIP] A pergunta bateu em trigger de smalltalk; no endpoint isso retorna sem retrieval.")
        print(f"Resposta esperada: {srv.ABOUT_BOT_TEXT}")
        return
    if srv._contains_any_trigger(pregunta_lower, greeting_triggers):
        print("[SKIP] A pergunta bateu em trigger de greeting; no endpoint isso retorna sem retrieval.")
        print(f"Resposta esperada: {srv.GREETING_TEXT}")
        return

    top_k = top_k_override or srv.CONFIG.get("retrieve", {}).get("top_k", 4)
    chatbot_min_top_k = int(srv.CONFIG.get("retrieve", {}).get("chatbot_top_k", 16))
    top_k = max(int(top_k), chatbot_min_top_k)

    index_type = srv.CONFIG.get("index", {}).get("type", "faiss").lower()
    normalize = index_type == "faiss"

    _print_header(top_k=top_k, normalize=normalize)

    q_vec = srv.embed_query(srv.EMBED_MODEL, question, normalize=normalize)
    retrieved: List[Dict[str, Any]] = srv.RETRIEVER.retrieve(q_vec, top_k)

    citation_ids: Dict[str, int] = {}
    next_id = 1
    for m in retrieved:
        doc_id = m.get("document_id")
        if not doc_id:
            continue
        if doc_id not in citation_ids:
            citation_ids[doc_id] = next_id
            next_id += 1
        m["citation_id"] = citation_ids[doc_id]

    print(f"Pergunta: {question}")
    print(f"Chunks recuperados: {len(retrieved)}")

    if not retrieved:
        print("Nenhum chunk recuperado.")
        return

    for i, chunk in enumerate(retrieved, start=1):
        _print_chunk(i=i, chunk=chunk, max_chars=max_chars)

    print("\nFim.\n")


def main() -> int:
    parser = argparse.ArgumentParser(description="Testa retrieval do backend no modo chatbot")
    parser.add_argument("--question", "-q", type=str, help="Pergunta para execução única")
    parser.add_argument("--top-k", type=int, default=None, help="Override de top_k")
    parser.add_argument(
        "--max-chars",
        type=int,
        default=320,
        help="Máximo de caracteres exibidos por chunk",
    )
    args = parser.parse_args()

    if args.question:
        _run_retrieval(args.question, args.top_k, args.max_chars)
        return 0

    print("Modo interativo. Digite sua pergunta (ou 'sair' para encerrar).")
    while True:
        try:
            q = input("\nPergunta> ").strip()
        except (EOFError, KeyboardInterrupt):
            print("\nEncerrado.")
            return 0

        if q.lower() in {"sair", "exit", "quit"}:
            print("Encerrado.")
            return 0

        _run_retrieval(q, args.top_k, args.max_chars)


if __name__ == "__main__":
    sys.exit(main())