Allow any provider/model supported by pydantic AI

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

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.

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

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,
+    )

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."""

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(

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]