130 lines
4.5 KiB
Python
130 lines
4.5 KiB
Python
"""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())
|