From 1f61d781656a352a257f20b917842dd6054fcf71 Mon Sep 17 00:00:00 2001 From: Yiorgis Gozadinos Date: Sun, 17 Aug 2025 15:27:03 +0200 Subject: [PATCH] Bring back QA_MODEL and keep compatibility with previous versions --- docs/configuration.md | 20 +++++++++++++------- src/haiku/rag/config.py | 3 ++- src/haiku/rag/qa/__init__.py | 6 ++++-- src/haiku/rag/qa/agent.py | 18 +++++++----------- tests/llm_judge.py | 24 ++++++------------------ tests/test_qa.py | 8 +++----- 6 files changed, 35 insertions(+), 44 deletions(-) diff --git a/docs/configuration.md b/docs/configuration.md index b76d7ed7..9875b9cd 100644 --- a/docs/configuration.md +++ b/docs/configuration.md @@ -55,12 +55,13 @@ OPENAI_API_KEY="your-api-key" ## Question Answering Providers -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. +Configure which LLM provider to use for question answering. Any provider and model supported by [Pydantic AI](https://ai.pydantic.dev/models/) can be used. ### Ollama (Default) ```bash -QA_PROVIDER="ollama:qwen3" +QA_PROVIDER="ollama" +QA_MODEL="qwen3" OLLAMA_BASE_URL="http://localhost:11434" ``` @@ -69,7 +70,8 @@ OLLAMA_BASE_URL="http://localhost:11434" OpenAI QA is included in the default installation. Simply configure: ```bash -QA_PROVIDER="openai:gpt-4o-mini" # or openai:gpt-4, openai:gpt-3.5-turbo, etc. +QA_PROVIDER="openai" +QA_MODEL="gpt-4o-mini" # or gpt-4, gpt-3.5-turbo, etc. OPENAI_API_KEY="your-api-key" ``` @@ -78,7 +80,8 @@ OPENAI_API_KEY="your-api-key" Anthropic QA is included in the default installation. Simply configure: ```bash -QA_PROVIDER="anthropic:claude-3-5-haiku-20241022" # or anthropic:claude-3-5-sonnet-20241022, etc. +QA_PROVIDER="anthropic" +QA_MODEL="claude-3-5-haiku-20241022" # or claude-3-5-sonnet-20241022, etc. ANTHROPIC_API_KEY="your-api-key" ``` @@ -88,13 +91,16 @@ Any provider supported by Pydantic AI can be used. Examples include: ```bash # Google Gemini -QA_PROVIDER="gemini:gemini-1.5-flash" +QA_PROVIDER="gemini" +QA_MODEL="gemini-1.5-flash" # Groq -QA_PROVIDER="groq:llama-3.3-70b-versatile" +QA_PROVIDER="groq" +QA_MODEL="llama-3.3-70b-versatile" # Mistral -QA_PROVIDER="mistral:mistral-small-latest" +QA_PROVIDER="mistral" +QA_MODEL="mistral-small-latest" ``` See the [Pydantic AI documentation](https://ai.pydantic.dev/models/) for the complete list of supported providers and models. diff --git a/src/haiku/rag/config.py b/src/haiku/rag/config.py index 478cf6e2..78328fe6 100644 --- a/src/haiku/rag/config.py +++ b/src/haiku/rag/config.py @@ -22,7 +22,8 @@ class AppConfig(BaseModel): RERANK_PROVIDER: str = "ollama" RERANK_MODEL: str = "qwen3" - QA_PROVIDER: str = "ollama:qwen3" + QA_PROVIDER: str = "ollama" + QA_MODEL: str = "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 d401597b..70ea8fe4 100644 --- a/src/haiku/rag/qa/__init__.py +++ b/src/haiku/rag/qa/__init__.py @@ -4,10 +4,12 @@ from haiku.rag.qa.agent import QuestionAnswerAgent def get_qa_agent(client: HaikuRAG, use_citations: bool = False) -> QuestionAnswerAgent: - provider_model = Config.QA_PROVIDER + provider = Config.QA_PROVIDER + model_name = Config.QA_MODEL return QuestionAnswerAgent( client=client, - provider_model=provider_model, + provider=provider, + model=model_name, use_citations=use_citations, ) diff --git a/src/haiku/rag/qa/agent.py b/src/haiku/rag/qa/agent.py index bb0028fa..b76112bc 100644 --- a/src/haiku/rag/qa/agent.py +++ b/src/haiku/rag/qa/agent.py @@ -23,14 +23,15 @@ class QuestionAnswerAgent: def __init__( self, client: HaikuRAG, - provider_model: str, + provider: str, + 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, @@ -57,21 +58,16 @@ class QuestionAnswerAgent: for chunk, score in expanded_results ] - 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) - + def _get_model(self, provider: str, model: str): + """Get the appropriate model object for the provider.""" if provider == "ollama": return OpenAIModel( model_name=model, provider=OllamaProvider(base_url=f"{Config.OLLAMA_BASE_URL}/v1"), ) else: - # For other providers, use the provider:model string directly - return provider_model + # For all other providers, use the provider:model format + return f"{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 f9233e6b..68ffe8c8 100644 --- a/tests/llm_judge.py +++ b/tests/llm_judge.py @@ -37,24 +37,12 @@ class LLMJudgeResponseSchema(BaseModel): class LLMJudge: """LLM-as-judge for evaluating answer equivalence using Pydantic AI.""" - def __init__(self, provider_model: str = Config.QA_PROVIDER): - self.provider_model = provider_model - - # 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 + def __init__(self, model: str = Config.QA_MODEL): + # 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( diff --git a/tests/test_qa.py b/tests/test_qa.py index 49348fdf..af8addd7 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, provider_model="ollama:qwen3") + qa = QuestionAnswerAgent(client, "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, provider_model="openai:gpt-4o-mini") + qa = QuestionAnswerAgent(client, "openai", "gpt-4o-mini") llm_judge = LLMJudge() doc = qa_corpus[1] @@ -63,9 +63,7 @@ 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, provider_model="anthropic:claude-3-5-haiku-20241022" - ) + qa = QuestionAnswerAgent(client, "anthropic", "claude-3-5-haiku-20241022") llm_judge = LLMJudge() doc = qa_corpus[1]