From d6d79d5e9401890f7f2047736c9214b358be18f3 Mon Sep 17 00:00:00 2001 From: Yiorgis Gozadinos Date: Sun, 17 Aug 2025 14:55:50 +0200 Subject: [PATCH] Allow any provider/model supported by pydantic AI --- README.md | 2 +- docs/configuration.md | 28 +++++++++++++++++++++------- src/haiku/rag/config.py | 3 +-- src/haiku/rag/qa/__init__.py | 17 ++++++----------- src/haiku/rag/qa/agent.py | 23 ++++++++++++----------- tests/llm_judge.py | 24 +++++++++++++++++------- tests/test_qa.py | 8 +++++--- 7 files changed, 63 insertions(+), 42 deletions(-) diff --git a/README.md b/README.md index be1a33b2..b7dcc337 100644 --- a/README.md +++ b/README.md @@ -8,7 +8,7 @@ Retrieval-Augmented Generation (RAG) library on SQLite. - **Local SQLite**: No external servers required - **Multiple embedding providers**: Ollama, VoyageAI, OpenAI -- **Multiple QA providers**: Ollama, OpenAI, Anthropic +- **Multiple QA providers**: Any provider/model supported by Pydantic AI - **Hybrid search**: Vector + full-text search with Reciprocal Rank Fusion - **Reranking**: Default search result reranking with MixedBread AI or Cohere - **Question answering**: Built-in QA agents on your documents diff --git a/docs/configuration.md b/docs/configuration.md index ae73b177..b76d7ed7 100644 --- a/docs/configuration.md +++ b/docs/configuration.md @@ -55,13 +55,12 @@ OPENAI_API_KEY="your-api-key" ## Question Answering Providers -Configure which LLM provider to use for question answering. +Configure which LLM provider to use for question answering using the `provider:model` format. Any provider and model supported by [Pydantic AI](https://ai.pydantic.dev/models/) can be used. ### Ollama (Default) ```bash -QA_PROVIDER="ollama" -QA_MODEL="qwen3" +QA_PROVIDER="ollama:qwen3" OLLAMA_BASE_URL="http://localhost:11434" ``` @@ -70,8 +69,7 @@ OLLAMA_BASE_URL="http://localhost:11434" OpenAI QA is included in the default installation. Simply configure: ```bash -QA_PROVIDER="openai" -QA_MODEL="gpt-4o-mini" # or gpt-4, gpt-3.5-turbo, etc. +QA_PROVIDER="openai:gpt-4o-mini" # or openai:gpt-4, openai:gpt-3.5-turbo, etc. OPENAI_API_KEY="your-api-key" ``` @@ -80,11 +78,27 @@ OPENAI_API_KEY="your-api-key" Anthropic QA is included in the default installation. Simply configure: ```bash -QA_PROVIDER="anthropic" -QA_MODEL="claude-3-5-haiku-20241022" # or claude-3-5-sonnet-20241022, etc. +QA_PROVIDER="anthropic:claude-3-5-haiku-20241022" # or anthropic:claude-3-5-sonnet-20241022, etc. ANTHROPIC_API_KEY="your-api-key" ``` +### Other Providers + +Any provider supported by Pydantic AI can be used. Examples include: + +```bash +# Google Gemini +QA_PROVIDER="gemini:gemini-1.5-flash" + +# Groq +QA_PROVIDER="groq:llama-3.3-70b-versatile" + +# Mistral +QA_PROVIDER="mistral:mistral-small-latest" +``` + +See the [Pydantic AI documentation](https://ai.pydantic.dev/models/) for the complete list of supported providers and models. + ## Reranking Reranking improves search quality by re-ordering the initial search results using specialized models. When enabled, the system retrieves more candidates (3x the requested limit) and then reranks them to return the most relevant results. diff --git a/src/haiku/rag/config.py b/src/haiku/rag/config.py index 78328fe6..478cf6e2 100644 --- a/src/haiku/rag/config.py +++ b/src/haiku/rag/config.py @@ -22,8 +22,7 @@ class AppConfig(BaseModel): RERANK_PROVIDER: str = "ollama" RERANK_MODEL: str = "qwen3" - QA_PROVIDER: str = "ollama" - QA_MODEL: str = "qwen3" + QA_PROVIDER: str = "ollama:qwen3" CHUNK_SIZE: int = 256 CONTEXT_CHUNK_RADIUS: int = 0 diff --git a/src/haiku/rag/qa/__init__.py b/src/haiku/rag/qa/__init__.py index 3d5a056b..d401597b 100644 --- a/src/haiku/rag/qa/__init__.py +++ b/src/haiku/rag/qa/__init__.py @@ -4,15 +4,10 @@ from haiku.rag.qa.agent import QuestionAnswerAgent def get_qa_agent(client: HaikuRAG, use_citations: bool = False) -> QuestionAnswerAgent: - provider = Config.QA_PROVIDER - model_name = Config.QA_MODEL + provider_model = Config.QA_PROVIDER - if provider in ("openai", "anthropic", "ollama"): - return QuestionAnswerAgent( - client=client, - provider=provider, - model=model_name, - use_citations=use_citations, - ) - - raise ValueError(f"Unsupported QA provider: {provider}") + return QuestionAnswerAgent( + client=client, + provider_model=provider_model, + use_citations=use_citations, + ) diff --git a/src/haiku/rag/qa/agent.py b/src/haiku/rag/qa/agent.py index 8ad14a59..bb0028fa 100644 --- a/src/haiku/rag/qa/agent.py +++ b/src/haiku/rag/qa/agent.py @@ -23,15 +23,14 @@ class QuestionAnswerAgent: def __init__( self, client: HaikuRAG, - provider: str, - model: str, + provider_model: str, use_citations: bool = False, q: float = 0.0, ): self._client = client system_prompt = SYSTEM_PROMPT_WITH_CITATIONS if use_citations else SYSTEM_PROMPT - model_obj = self._get_model(provider, model) + model_obj = self._get_model(provider_model) self._agent = Agent( model=model_obj, @@ -58,19 +57,21 @@ class QuestionAnswerAgent: for chunk, score in expanded_results ] - def _get_model(self, provider: str, model: str): - """Get the appropriate model object for the provider.""" - if provider == "openai": - return f"openai:{model}" - elif provider == "anthropic": - return f"anthropic:{model}" - elif provider == "ollama": + def _get_model(self, provider_model: str): + """Get the appropriate model object for the provider:model format.""" + if ":" not in provider_model: + raise ValueError(f"Invalid provider:model format: {provider_model}") + + provider, model = provider_model.split(":", 1) + + if provider == "ollama": return OpenAIModel( model_name=model, provider=OllamaProvider(base_url=f"{Config.OLLAMA_BASE_URL}/v1"), ) else: - raise ValueError(f"Unsupported provider: {provider}") + # For other providers, use the provider:model string directly + return provider_model async def answer(self, question: str) -> str: """Answer a question using the RAG system.""" diff --git a/tests/llm_judge.py b/tests/llm_judge.py index a738cd41..f9233e6b 100644 --- a/tests/llm_judge.py +++ b/tests/llm_judge.py @@ -37,14 +37,24 @@ class LLMJudgeResponseSchema(BaseModel): class LLMJudge: """LLM-as-judge for evaluating answer equivalence using Pydantic AI.""" - def __init__(self, model: str = Config.QA_MODEL): - self.model = model + def __init__(self, provider_model: str = Config.QA_PROVIDER): + self.provider_model = provider_model - # Create Ollama model - ollama_model = OpenAIModel( - model_name=model, - provider=OllamaProvider(base_url=f"{Config.OLLAMA_BASE_URL}/v1"), - ) + # Parse provider:model format + if ":" not in provider_model: + raise ValueError(f"Invalid provider:model format: {provider_model}") + + provider, model = provider_model.split(":", 1) + + if provider == "ollama": + # Create Ollama model + ollama_model = OpenAIModel( + model_name=model, + provider=OllamaProvider(base_url=f"{Config.OLLAMA_BASE_URL}/v1"), + ) + else: + # For other providers, use the provider:model string directly + ollama_model = provider_model # Create Pydantic AI agent self._agent = Agent( diff --git a/tests/test_qa.py b/tests/test_qa.py index af8addd7..49348fdf 100644 --- a/tests/test_qa.py +++ b/tests/test_qa.py @@ -15,7 +15,7 @@ ANTHROPIC_AVAILABLE = bool(Config.ANTHROPIC_API_KEY) async def test_qa_ollama(qa_corpus: Dataset): """Test Ollama QA with LLM judge.""" client = HaikuRAG(":memory:") - qa = QuestionAnswerAgent(client, "ollama", "qwen3") + qa = QuestionAnswerAgent(client, provider_model="ollama:qwen3") llm_judge = LLMJudge() doc = qa_corpus[1] @@ -39,7 +39,7 @@ async def test_qa_ollama(qa_corpus: Dataset): async def test_qa_openai(qa_corpus: Dataset): """Test OpenAI QA with LLM judge.""" client = HaikuRAG(":memory:") - qa = QuestionAnswerAgent(client, "openai", "gpt-4o-mini") + qa = QuestionAnswerAgent(client, provider_model="openai:gpt-4o-mini") llm_judge = LLMJudge() doc = qa_corpus[1] @@ -63,7 +63,9 @@ async def test_qa_openai(qa_corpus: Dataset): async def test_qa_anthropic(qa_corpus: Dataset): """Test Anthropic QA with LLM judge.""" client = HaikuRAG(":memory:") - qa = QuestionAnswerAgent(client, "anthropic", "claude-3-5-haiku-20241022") + qa = QuestionAnswerAgent( + client, provider_model="anthropic:claude-3-5-haiku-20241022" + ) llm_judge = LLMJudge() doc = qa_corpus[1]