Refactor Ollama reranker around pydantic AI

src/haiku/rag/reranking/ollama.py

@@ -1,14 +1,12 @@
-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
 from haiku.rag.reranking.base import RerankerBase
 from haiku.rag.store.models.chunk import Chunk
 
-OLLAMA_OPTIONS = {"temperature": 0.0, "seed": 42, "num_ctx": 16384}
-
 
 class RerankResult(BaseModel):
     """Individual rerank result with index and relevance score."""
@@ -26,7 +24,26 @@ class RerankResponse(BaseModel):
 class OllamaReranker(RerankerBase):
     def __init__(self, model: str = Config.RERANK_MODEL):
         self._model = model
-        self._client = AsyncClient(host=Config.OLLAMA_BASE_URL)
+
+        # Create the reranking prompt
+        system_prompt = """You are a document reranking assistant. Given a query and a list of document chunks, you must rank them by relevance to the query.
+
+Return your response as a JSON object with a "results" array. Each result should have:
+- "index": the original index of the document (integer)
+- "relevance_score": a score between 0.0 and 1.0 indicating relevance (float, where 1.0 is most relevant)
+
+Only return the top documents up to the requested limit, ordered by decreasing relevance score."""
+
+        model_obj = OpenAIModel(
+            model_name=model,
+            provider=OllamaProvider(base_url=f"{Config.OLLAMA_BASE_URL}/v1"),
+        )
+
+        self._agent = Agent(
+            model=model_obj,
+            output_type=RerankResponse,
+            system_prompt=system_prompt,
+        )
 
     async def rerank(
         self, query: str, chunks: list[Chunk], top_n: int = 10
@@ -38,15 +55,6 @@ class OllamaReranker(RerankerBase):
         for i, chunk in enumerate(chunks):
             documents.append({"index": i, "content": chunk.content})
 
-        # Create the prompt for reranking
-        system_prompt = """You are a document reranking assistant. Given a query and a list of document chunks, you must rank them by relevance to the query.
-
-Return your response as a JSON object with a "results" array. Each result should have:
-- "index": the original index of the document (integer)
-- "relevance_score": a score between 0.0 and 1.0 indicating relevance (float, where 1.0 is most relevant)
-
-Only return the top documents up to the requested limit, ordered by decreasing relevance score."""
-
         documents_text = ""
         for doc in documents:
             documents_text += f"Index {doc['index']}: {doc['content']}\n\n"
@@ -56,27 +64,14 @@ Only return the top documents up to the requested limit, ordered by decreasing r
 Documents to rerank:
 {documents_text.strip()}
 
-Please rank these documents by relevance to the query and return the top {top_n} results as JSON."""
+Rank these documents by relevance to the query and return the top {top_n} results as JSON."""
-
-        messages = [
-            {"role": "system", "content": system_prompt},
-            {"role": "user", "content": user_prompt},
-        ]
 
         try:
-            response = await self._client.chat(
+            result = await self._agent.run(user_prompt)
-                model=self._model,
-                messages=messages,
-                format=RerankResponse.model_json_schema(),
-                options=OLLAMA_OPTIONS,
-            )
 
-            content = response["message"]["content"]
-
-            parsed_response = RerankResponse.model_validate(json.loads(content))
             return [
-                (chunks[result.index], result.relevance_score)
+                (chunks[result_item.index], result_item.relevance_score)
-                for result in parsed_response.results[:top_n]
+                for result_item in result.output.results[:top_n]
             ]
 
         except Exception:

tests/test_reranker.py

@@ -1,8 +1,11 @@
 import pytest
 
+from haiku.rag.config import Config
 from haiku.rag.reranking.base import RerankerBase
 from haiku.rag.store.models.chunk import Chunk
 
+COHERE_AVAILABLE = bool(Config.COHERE_API_KEY)
+
 chunks = [
     Chunk(content=content, document_id=i)
     for i, content in enumerate(
@@ -43,6 +46,7 @@ async def test_mxbai_reranker():
 
 
 @pytest.mark.asyncio
+@pytest.mark.skipif(not COHERE_AVAILABLE, reason="Cohere API key not available")
 async def test_cohere_reranker():
     try:
         from haiku.rag.reranking.cohere import CohereReranker