From 1705e7a21cff15b93179f1287c5fa116359f6170 Mon Sep 17 00:00:00 2001 From: Yiorgis Gozadinos Date: Fri, 4 Jul 2025 12:14:13 +0300 Subject: [PATCH] Honour QA_MODEL for anthropic and OpenAI --- src/haiku/rag/qa/__init__.py | 7 ++----- tests/generate_benchmark_db.py | 4 ++-- 2 files changed, 4 insertions(+), 7 deletions(-) diff --git a/src/haiku/rag/qa/__init__.py b/src/haiku/rag/qa/__init__.py index bcf53380..f9047a91 100644 --- a/src/haiku/rag/qa/__init__.py +++ b/src/haiku/rag/qa/__init__.py @@ -8,7 +8,6 @@ def get_qa_agent(client: HaikuRAG, model: str = "") -> QuestionAnswerAgentBase: """ Factory function to get the appropriate QA agent based on the configuration. """ - if Config.QA_PROVIDER == "ollama": return QuestionAnswerOllamaAgent(client, model or Config.QA_MODEL) @@ -21,7 +20,7 @@ def get_qa_agent(client: HaikuRAG, model: str = "") -> QuestionAnswerAgentBase: "Please install haiku.rag with the 'openai' extra:" "uv pip install haiku.rag --extra openai" ) - return QuestionAnswerOpenAIAgent(client, model or "gpt-4o-mini") + return QuestionAnswerOpenAIAgent(client, model or Config.QA_MODEL) if Config.QA_PROVIDER == "anthropic": try: @@ -32,8 +31,6 @@ def get_qa_agent(client: HaikuRAG, model: str = "") -> QuestionAnswerAgentBase: "Please install haiku.rag with the 'anthropic' extra:" "uv pip install haiku.rag --extra anthropic" ) - return QuestionAnswerAnthropicAgent( - client, model or "claude-3-5-haiku-20241022" - ) + return QuestionAnswerAnthropicAgent(client, model or Config.QA_MODEL) raise ValueError(f"Unsupported QA provider: {Config.QA_PROVIDER}") diff --git a/tests/generate_benchmark_db.py b/tests/generate_benchmark_db.py index a736317a..b1ade523 100644 --- a/tests/generate_benchmark_db.py +++ b/tests/generate_benchmark_db.py @@ -6,7 +6,7 @@ from llm_judge import LLMJudge from tqdm import tqdm from haiku.rag.client import HaikuRAG -from haiku.rag.qa.ollama import QuestionAnswerOllamaAgent +from haiku.rag.qa import get_qa_agent db_path = Path(__file__).parent / "data" / "benchmark.sqlite" @@ -88,7 +88,7 @@ async def run_qa_benchmark(k: int | None = None): total_questions = 0 async with HaikuRAG(db_path) as rag: - qa = QuestionAnswerOllamaAgent(rag) + qa = get_qa_agent(rag) for i, doc in enumerate(tqdm(corpus, desc="QA Benchmarking")): question = doc["question"] # type: ignore