backend/scripts/evalRetrieval.mjs

import { config } from "#config/index.js";
import { searchDocs } from "#services/searchService.js";
import { rewriteQuery } from "#chat/chatChain.js";
import { generateHydeDocument } from "#services/hydeService.js";
import { HYDE_SYSTEM_PROMPT } from "#chat/prompts.js";

function parseArgs(argv) {
  const flags = { hyde: false, rewrite: false, minScore: 0, topK: 5, label: "" };
  for (const arg of argv) {
    if (arg === "--hyde") flags.hyde = true;
    else if (arg === "--no-hyde") flags.hyde = false;
    else if (arg === "--rewrite") flags.rewrite = true;
    else if (arg === "--no-rewrite") flags.rewrite = false;
    else if (arg.startsWith("--min-score=")) flags.minScore = Number(arg.split("=")[1]);
    else if (arg.startsWith("--top-k=")) flags.topK = Number(arg.split("=")[1]);
    else if (arg.startsWith("--label=")) flags.label = arg.split("=").slice(1).join("=");
  }
  return flags;
}


const CASES = [
  { query: "como ver os usuários conectados no equipamento", expectedSourceIncludes: "Huawei", type: "positive" },
  { query: "comando para exibir usuários de acesso", expectedSourceIncludes: "Huawei", type: "positive" },
  { query: "como configurar Eth-Trunk", expectedSourceIncludes: "Huawei", type: "positive" },
  { query: "como configurar OSPF no equipamento", expectedSourceIncludes: "Huawei", type: "positive" },
  { query: "comando display access-user", expectedSourceIncludes: "Huawei", type: "positive" },
  { query: "como resetar o roteador TP-Link", expectedSourceIncludes: "TL-WR940N", type: "positive" },
  { query: "como instalar o roteador TP-Link pela primeira vez", expectedSourceIncludes: "TL-WR940N", type: "positive" },
  { query: "qual o endereço padrão de acesso ao roteador", expectedSourceIncludes: "TL-WR940N", type: "positive" },
  { query: "qual o valor do vale refeição da empresa", expectedSourceIncludes: null, type: "negative" },
  { query: "quantos dias de férias tenho direito por ano", expectedSourceIncludes: null, type: "negative" }
];

function sourceMatches(hit, expectedIncludes) {
  if (!expectedIncludes) return false;
  return String(hit?.source ?? "").toLowerCase().includes(expectedIncludes.toLowerCase());
}

async function resolveEmbeddingQuery(originalQuery, flags) {
  const searchQuery = flags.rewrite ? await rewriteQuery(originalQuery, []) : originalQuery;
  if (!flags.hyde) return { searchQuery, embeddingQuery: searchQuery, embedRole: "query" };
  const hydeText = await generateHydeDocument(searchQuery, HYDE_SYSTEM_PROMPT);
  return hydeText
    ? { searchQuery, embeddingQuery: hydeText, embedRole: "passage" }
    : { searchQuery, embeddingQuery: searchQuery, embedRole: "query" };
}

async function runCase(testCase, flags) {
  const { searchQuery, embeddingQuery, embedRole } = await resolveEmbeddingQuery(testCase.query, flags);
  const hits = await searchDocs({
    query: embeddingQuery,
    topK: flags.topK,
    embedRole,
    minScore: flags.minScore
  });

  const top1 = hits[0] ?? null;
  const top3 = hits.slice(0, 3);
  const top1Match = testCase.type === "positive" ? sourceMatches(top1, testCase.expectedSourceIncludes) : null;
  const top3Match = testCase.type === "positive" ? top3.some((h) => sourceMatches(h, testCase.expectedSourceIncludes)) : null;

  return {
    query: testCase.query,
    searchQuery: searchQuery !== testCase.query ? searchQuery : "",
    type: testCase.type,
    top1_score: top1 ? top1.score.toFixed(4) : "-",
    top1_source: top1?.source ?? "-",
    top1_match: testCase.type === "positive" ? (top1Match ? "OK" : "MISS") : "-",
    top3_match: testCase.type === "positive" ? (top3Match ? "OK" : "MISS") : "-"
  };
}

async function main() {
  const flags = parseArgs(process.argv.slice(2));

  console.log(`\n=== evalRetrieval ${flags.label ? `[${flags.label}] ` : ""}===`);
  console.log(
    `modelo=${config.ollama.embeddingsModel} minScore=${flags.minScore} topK=${flags.topK} ` +
    `rewrite=${flags.rewrite} hyde=${flags.hyde} colecao=${config.qdrant.collection}\n`
  );

  const rows = [];
  for (const testCase of CASES) {
    
    rows.push(await runCase(testCase, flags));
  }

  console.table(rows);

  const positives = rows.filter((r) => r.type === "positive");
  const negatives = rows.filter((r) => r.type === "negative");

  const positiveScores = positives.map((r) => Number(r.top1_score)).filter((n) => !Number.isNaN(n));
  const negativeScores = negatives.map((r) => Number(r.top1_score)).filter((n) => !Number.isNaN(n));

  const avg = (arr) => (arr.length ? arr.reduce((a, b) => a + b, 0) / arr.length : NaN);
  const hitRate = (arr, key) => (arr.length ? arr.filter((r) => r[key] === "OK").length / arr.length : NaN);

  console.log("\n--- resumo ---");
  console.log(`positivos: score médio top-1 = ${avg(positiveScores).toFixed(4)} | hit-rate top-1 = ${(hitRate(positives, "top1_match") * 100).toFixed(0)}% | hit-rate top-3 = ${(hitRate(positives, "top3_match") * 100).toFixed(0)}%`);
  console.log(`positivos: score mínimo top-1 (piso de segurança) = ${positiveScores.length ? Math.min(...positiveScores).toFixed(4) : "-"}`);
  console.log(`negativos: score máximo top-1 (piso de ruído) = ${negativeScores.length ? Math.max(...negativeScores).toFixed(4) : "-"}`);
  console.log("");
}

main()
  .then(() => process.exit(0))
  .catch((err) => {
    console.error("[evalRetrieval] erro:", err);
    process.exit(1);
  });