Convert LLM Judge

tests/llm_judge.py

@@ -1,21 +1,57 @@
-import json
-
-from ollama import AsyncClient
 from pydantic import BaseModel
+from pydantic_ai import Agent
+from pydantic_ai.models.openai import OpenAIModel
+from pydantic_ai.providers.ollama import OllamaProvider
 
 from haiku.rag.config import Config
 
+# Shared rubric/prompt for answer equivalence evaluation
+ANSWER_EQUIVALENCE_RUBRIC = """You are evaluating whether two answers to the same question are semantically equivalent.
+
+EVALUATION CRITERIA:
+Rate as EQUIVALENT if:
+✓ Both answers contain the same core factual information
+✓ Both directly address the question asked
+✓ The key claims and conclusions are consistent
+✓ Any additional detail in one answer doesn't contradict the other
+
+Rate as NOT EQUIVALENT if:
+✗ Factual contradictions exist between the answers
+✗ One answer fails to address the core question
+✗ Key information is missing that changes the meaning
+✗ The answers lead to different conclusions or implications
+
+GUIDELINES:
+- Ignore minor differences in phrasing, style, or formatting
+- Focus on semantic meaning rather than exact wording
+- Consider both answers correct if they convey the same essential information
+- Be tolerant of different levels of detail if the core answer is preserved
+- Evaluate based on what a person asking this question would need to know
+/no_think"""
+
 
 class LLMJudgeResponseSchema(BaseModel):
     equivalent: bool
 
 
 class LLMJudge:
-    """LLM-as-judge for evaluating answer equivalence using Ollama."""
+    """LLM-as-judge for evaluating answer equivalence using Pydantic AI."""
 
     def __init__(self, model: str = Config.QA_MODEL):
         self.model = model
-        self.client = AsyncClient(host=Config.OLLAMA_BASE_URL)
+
+        # Create Ollama model
+        ollama_model = OpenAIModel(
+            model_name=model,
+            provider=OllamaProvider(base_url=f"{Config.OLLAMA_BASE_URL}/v1"),
+        )
+
+        # Create Pydantic AI agent
+        self._agent = Agent(
+            model=ollama_model,
+            output_type=LLMJudgeResponseSchema,
+            system_prompt=ANSWER_EQUIVALENCE_RUBRIC,
+        )
 
     async def judge_answers(
         self, question: str, answer: str, expected_answer: str
@@ -29,53 +65,14 @@ class LLMJudge:
             expected_answer: The reference/expected answer
 
         Returns:
-            Dictionary with judgment result:
+            bool indicating if answers are equivalent
-            - equivalent: bool indicating if answers are equivalent
-            - explanation: str explaining the reasoning
-            - score: str rating from 1-5
         """
 
-        prompt = f"""You are an expert evaluator determining whether two answers to the same question are semantically equivalent.
+        prompt = f"""QUESTION: {question}
-
-QUESTION: {question}
 
 GENERATED ANSWER: {answer}
 
-EXPECTED ANSWER: {expected_answer}
+EXPECTED ANSWER: {expected_answer}"""
 
-EVALUATION CRITERIA:
+        result = await self._agent.run(prompt)
-Rate as EQUIVALENT (true) if:
+        return result.output.equivalent
-✓ Both answers contain the same core factual information
-✓ Both directly address the question asked
-✓ The key claims and conclusions are consistent
-✓ Any additional detail in one answer doesn't contradict the other
-
-Rate as NOT EQUIVALENT (false) if:
-✗ Factual contradictions exist between the answers
-✗ One answer fails to address the core question
-✗ Key information is missing from one answer that changes the meaning
-✗ The answers lead to different conclusions or implications
-
-GUIDELINES:
-- Ignore minor differences in phrasing, style, or formatting
-- Focus on semantic meaning rather than exact wording
-- Consider both answers correct if they convey the same essential information
-- Be tolerant of different levels of detail if the core answer is preserved
-- Evaluate based on what a person asking this question would need to know
-
-Respond with JSON containing only: {{"equivalent": true}} or {{"equivalent": false}}"""
-
-        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"