diff --git a/src/haiku/rag/reranking/ollama.py b/src/haiku/rag/reranking/ollama.py index 727c546b..40e3ad82 100644 --- a/src/haiku/rag/reranking/ollama.py +++ b/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.""" - - messages = [ - {"role": "system", "content": system_prompt}, - {"role": "user", "content": user_prompt}, - ] +Rank these documents by relevance to the query and return the top {top_n} results as JSON.""" try: - response = await self._client.chat( - model=self._model, - messages=messages, - format=RerankResponse.model_json_schema(), - options=OLLAMA_OPTIONS, - ) + result = await self._agent.run(user_prompt) - content = response["message"]["content"] - - parsed_response = RerankResponse.model_validate(json.loads(content)) return [ - (chunks[result.index], result.relevance_score) - for result in parsed_response.results[:top_n] + (chunks[result_item.index], result_item.relevance_score) + for result_item in result.output.results[:top_n] ] except Exception: diff --git a/tests/test_reranker.py b/tests/test_reranker.py index 5ce13156..7a86209d 100644 --- a/tests/test_reranker.py +++ b/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