#!/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())