Use LLM-as-a-judge to test QA

tests/conftest.py

@@ -3,6 +3,8 @@ from pathlib import Path
 import pytest
 from datasets import Dataset, load_dataset, load_from_disk
 
+from .llm_judge import LLMJudge
+
 
 @pytest.fixture(scope="session")
 def qa_corpus() -> Dataset:
@@ -16,3 +18,8 @@ def qa_corpus() -> Dataset:
         corpus = ds.filter(lambda doc: doc["document_topic"] == "News Stories")
         corpus.save_to_disk(ds_path)
         return corpus
+
+
+@pytest.fixture(scope="session")
+def llm_judge() -> LLMJudge:
+    return LLMJudge()

tests/generate_benchmark_db.py

@@ -1,13 +1,16 @@
+import asyncio
 from pathlib import Path
 
 from datasets import Dataset, load_dataset
+from llm_judge import LLMJudge
 from tqdm import tqdm
 
 from haiku.rag.client import HaikuRAG
+from haiku.rag.qa.ollama import QA
 
 
 async def populate_db():
-    if (Path(__file__).parent / "benchmark.sqlite").exists():
+    if (Path(__file__).parent / "data" / "benchmark.sqlite").exists():
         print("Benchmark database already exists. Skipping creation.")
         return
 
@@ -61,6 +64,7 @@ async def run_match_benchmark():
     recall_at_2 = correct_at_2 / total_queries
     recall_at_3 = correct_at_3 / total_queries
 
+    print("\n=== Retrieval Benchmark Results ===")
     print(f"Total queries: {total_queries}")
     print(f"Recall@1: {recall_at_1:.4f}")
     print(f"Recall@2: {recall_at_2:.4f}")
@@ -69,12 +73,48 @@ async def run_match_benchmark():
     return {"recall@1": recall_at_1, "recall@2": recall_at_2, "recall@3": recall_at_3}
 
 
+async def run_qa_benchmark():
+    """Run QA benchmarking on the corpus."""
+    ds: Dataset = load_dataset("ServiceNow/repliqa")["repliqa_3"]  # type: ignore
+    corpus = ds.filter(lambda doc: doc["document_topic"] == "News Stories")
+
+    judge = LLMJudge()
+    correct_answers = 0
+    total_questions = 0
+
+    async with HaikuRAG(Path(__file__).parent / "benchmark.sqlite") as rag:
+        qa = QA(rag)
+
+        for i, doc in enumerate(tqdm(corpus, desc="QA Benchmarking")):
+            question = doc["question"]  # type: ignore
+            expected_answer = doc["answer"]  # type: ignore
+
+            generated_answer = await qa.answer(question)
+            is_equivalent = await judge.judge_answers(
+                question, generated_answer, expected_answer
+            )
+
+            if is_equivalent:
+                correct_answers += 1
+            total_questions += 1
+
+    accuracy = correct_answers / total_questions if total_questions > 0 else 0
+
+    print("\n=== QA Benchmark Results ===")
+    print(f"Total questions: {total_questions}")
+    print(f"Correct answers: {correct_answers}")
+    print(f"QA Accuracy: {accuracy:.4f} ({accuracy * 100:.2f}%)")
+
+
 async def main():
     await populate_db()
+
+    print("Running retrieval benchmarks...")
     await run_match_benchmark()
 
+    print("\nRunning QA benchmarks...")
+    await run_qa_benchmark()
+
 
 if __name__ == "__main__":
-    import asyncio
-
     asyncio.run(main())

tests/llm_judge.py

@@ -0,0 +1,68 @@
+import json
+
+from ollama import AsyncClient
+from pydantic import BaseModel
+
+from haiku.rag.config import Config
+
+
+class LLMJudgeResponseSchema(BaseModel):
+    equivalent: bool
+
+
+class LLMJudge:
+    """LLM-as-judge for evaluating answer equivalence using Ollama."""
+
+    def __init__(self, model: str = "qwen3"):
+        self.model = model
+        self.client = AsyncClient(host=Config.OLLAMA_BASE_URL)
+
+    async def judge_answers(
+        self, question: str, answer: str, expected_answer: str
+    ) -> bool:
+        """
+        Judge whether two answers are equivalent for a given question.
+
+        Args:
+            question: The original question
+            answer: The generated answer to evaluate
+            expected_answer: The reference/expected answer
+
+        Returns:
+            Dictionary with judgment result:
+            - equivalent: bool indicating if answers are equivalent
+            - explanation: str explaining the reasoning
+            - score: str rating from 1-5
+        """
+
+        prompt = f"""
+        You are an expert judge evaluating the equivalence of two answers to the same question.
+
+        Question: {question}
+
+        Generated Answer: {answer}
+
+        Expected Answer: {expected_answer}
+
+        Your task is to determine if these two answers are equivalent in meaning and both correctly answer the question. Consider:
+
+        1. Do both answers provide the same answer?
+        2. Do both answers directly address the question asked?
+        3. Minor differences in wording or style are acceptable if the meaning of the answer is the same.
+
+        Be strict but fair in your evaluation. Focus on factual correctness and whether both answers would satisfy someone asking the question."""
+
+        response = await self.client.chat(
+            model=self.model,
+            messages=[{"role": "user", "content": prompt}],
+            format=LLMJudgeResponseSchema.model_json_schema(),
+            think=False,
+        )
+
+        answer = response["message"]["content"].strip()
+        try:
+            res = json.loads(answer)
+            assert "equivalent" in res, "Response must contain 'equivalent' key"
+            return res["equivalent"]
+        except json.JSONDecodeError:
+            assert False, "Response is not valid JSON"

tests/test_qa.py

@@ -0,0 +1,42 @@
+from typing import TYPE_CHECKING
+
+import pytest
+from datasets import Dataset
+
+from haiku.rag.client import HaikuRAG
+from haiku.rag.qa.ollama import QA
+
+if TYPE_CHECKING:
+    import sys
+    from pathlib import Path
+
+    sys.path.append(str(Path(__file__).parent))
+    from llm_judge import LLMJudge
+
+
+@pytest.mark.asyncio
+async def test_qa_with_dataset_question(qa_corpus: Dataset, llm_judge: "LLMJudge"):
+    """Test QA with actual question from the dataset using LLM judge."""
+    client = HaikuRAG(":memory:")
+    qa = QA(client)
+
+    # Use the first document from the corpus
+    doc = qa_corpus[1]
+
+    # Add the document to database
+    await client.create_document(
+        content=doc["document_extracted"], uri=doc["document_id"]
+    )
+
+    question = doc["question"]
+    expected_answer = doc["answer"]
+
+    answer = await qa.answer(question)
+    # Use LLM judge to evaluate answer equivalence
+    is_equivalent = await llm_judge.judge_answers(question, answer, expected_answer)
+
+    assert isinstance(answer, str)
+    assert len(answer) > 0
+    assert is_equivalent, (
+        f"Generated answer not equivalent to expected answer.\nQuestion: {question}\nGenerated: {answer}\nExpected: {expected_answer}"
+    )