feat(rag): hybrid FTS+vector retrieval and multi-format ingestion — name-your-tool questions find the right document

2026-08-22 01:27:02 -04:00
parent 2f738a7f19
commit 7e8d14702e
36 changed files with 2018 additions and 290 deletions
@@ -0,0 +1,129 @@
+"""Evaluate hybrid retrieval against the live knowledge base (phase 09).
+
+Embeds each question via aipi, runs the same hybrid search the chat API
+uses (cosine top-N + FTS top-N, RRF-fused), and prints the top-5 documents
+with their cosine / fts / fused scores plus the honesty-gate verdict:
+
+    uv run python -m scripts.eval_retrieval "How did I install gitlab?"
+    uv run python -m scripts.eval_retrieval --from-file questions.txt
+
+Requires ``AIPI_KEY`` in the environment (same convention as
+``scripts/llm_probe.py``) and an imported knowledge base
+(``python -m scripts.import_docs``). Exit code 0 when all questions were
+scored (deflections are a normal result — the verdict column shows them).
+"""
+from __future__ import annotations
+
+import argparse
+import asyncio
+import os
+import sys
+
+from dotenv import load_dotenv
+
+load_dotenv()
+
+
+def build_parser() -> argparse.ArgumentParser:
+    p = argparse.ArgumentParser(
+        prog="python -m scripts.eval_retrieval",
+        description="Rank hybrid retrieval results for one or more questions.",
+    )
+    p.add_argument(
+        "questions",
+        nargs="*",
+        metavar="QUESTION",
+        help="one or more questions to evaluate",
+    )
+    p.add_argument(
+        "--from-file",
+        type=str,
+        default=None,
+        metavar="PATH",
+        help="read questions from a file (one per line, blanks/# skipped)",
+    )
+    p.add_argument(
+        "--top",
+        type=int,
+        default=5,
+        help="documents to print per question (default: 5)",
+    )
+    return p
+
+
+def read_questions(args: argparse.Namespace) -> list[str]:
+    questions = list(args.questions)
+    if args.from_file:
+        with open(args.from_file, encoding="utf-8") as f:
+            questions.extend(
+                line.strip() for line in f if line.strip() and not line.lstrip().startswith("#")
+            )
+    return questions
+
+
+async def _embed_all(llm, questions: list[str]) -> list[list[float]]:
+    return await llm.embed(questions)
+
+
+def main(argv: list[str] | None = None) -> int:
+    args = build_parser().parse_args(argv)
+    api_key = os.environ.get("BOR_LLM_API_KEY") or os.environ.get("AIPI_KEY", "")
+    if not api_key or api_key == "not-needed":
+        print(
+            "eval_retrieval: AIPI_KEY is required in the environment "
+            "(same convention as scripts/llm_probe.py).",
+            file=sys.stderr,
+        )
+        return 1
+    questions = read_questions(args)
+    if not questions:
+        print("eval_retrieval: no questions given (positional or --from-file).", file=sys.stderr)
+        return 1
+
+    from app.config import get_settings
+    from app.db import SessionLocal, db_available
+    from app.rag.llm import LLMClient
+    from app.rag.retriever import RetrievedChunk, retrieve
+
+    settings = get_settings()
+    if not db_available():
+        print("eval_retrieval: Postgres is down — run `podman compose up -d db`.", file=sys.stderr)
+        return 1
+
+    llm = LLMClient(settings)
+    vectors = asyncio.run(_embed_all(llm, questions))
+
+    print(
+        f"eval: threshold={settings.relevance_threshold} "
+        f"vector_candidates={settings.hybrid_vector_candidates} "
+        f"lexical_candidates={settings.hybrid_lexical_candidates} rrf_k={settings.rrf_k}"
+    )
+    with SessionLocal() as db:
+        for question, vec in zip(questions, vectors, strict=True):
+            chunks = retrieve(db, question, vec)
+            best_cosine = max((c.cosine for c in chunks), default=0.0)
+            fts_hits = sum(1 for c in chunks if c.fts_hit)
+            verdict = (
+                "LOW (deflect)"
+                if best_cosine < settings.relevance_threshold and fts_hits == 0
+                else "HIGH (answer)"
+            )
+            print(f"\nquestion: {question!r}")
+            print(f"  gate: best_cosine={best_cosine:.4f} fts_hits={fts_hits} -> {verdict}")
+            # Best chunk per document, in fused rank order.
+            best_by_doc: dict[str, RetrievedChunk] = {}
+            for c in chunks:
+                key = f"{c.document.source}/{c.document.path}"
+                if key not in best_by_doc:
+                    best_by_doc[key] = c
+            for i, c in enumerate(list(best_by_doc.values())[: args.top], start=1):
+                print(
+                    f"  {i}. {c.document.source}/{c.document.path}  "
+                    f"cosine={c.cosine:.4f} fts={int(c.fts_hit)} fused={c.score:.5f}  "
+                    f"({c.document.title})"
+                )
+    return 0
+
+
+if __name__ == "__main__":
+    sys.exit(main())