From 85b106c4614b41b4508b2298dca7a8b317819fe1 Mon Sep 17 00:00:00 2001 From: Yiorgis Gozadinos Date: Fri, 27 Jun 2025 09:14:16 +0300 Subject: [PATCH] OpenAI Question/Answer agent --- src/haiku/rag/config.py | 8 +++ src/haiku/rag/qa/__init__.py | 26 +++++++++ src/haiku/rag/qa/base.py | 27 ++++++++- src/haiku/rag/qa/ollama.py | 30 +--------- src/haiku/rag/qa/openai.py | 101 +++++++++++++++++++++++++++++++++ tests/conftest.py | 7 --- tests/generate_benchmark_db.py | 4 +- tests/test_qa.py | 52 +++++++++++------ 8 files changed, 202 insertions(+), 53 deletions(-) create mode 100644 src/haiku/rag/qa/openai.py diff --git a/src/haiku/rag/config.py b/src/haiku/rag/config.py index 270c0aea..c388064f 100644 --- a/src/haiku/rag/config.py +++ b/src/haiku/rag/config.py @@ -27,6 +27,10 @@ class AppConfig(BaseModel): OLLAMA_BASE_URL: str = "http://localhost:11434" + # Provider keys + VOYAGE_API_KEY: str = "" + OPENAI_API_KEY: str = "" + @field_validator("MONITOR_DIRECTORIES", mode="before") @classmethod def parse_monitor_directories(cls, v): @@ -41,3 +45,7 @@ class AppConfig(BaseModel): # Expose Config object for app to import Config = AppConfig.model_validate(os.environ) +if Config.OPENAI_API_KEY: + os.environ["OPENAI_API_KEY"] = Config.OPENAI_API_KEY +if Config.VOYAGE_API_KEY: + os.environ["VOYAGE_API_KEY"] = Config.VOYAGE_API_KEY diff --git a/src/haiku/rag/qa/__init__.py b/src/haiku/rag/qa/__init__.py index e69de29b..69ff4523 100644 --- a/src/haiku/rag/qa/__init__.py +++ b/src/haiku/rag/qa/__init__.py @@ -0,0 +1,26 @@ +from haiku.rag.client import HaikuRAG +from haiku.rag.config import Config +from haiku.rag.qa.base import QuestionAnswerAgentBase +from haiku.rag.qa.ollama import QuestionAnswerOllamaAgent + + +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) + + if Config.QA_PROVIDER == "openai": + try: + from haiku.rag.qa.openai import QuestionAnswerOpenAIAgent + except ImportError: + raise ImportError( + "OpenAI QA agent requires the 'openai' package. " + "Please install haiku.rag with the 'openai' extra:" + "uv pip install haiku.rag --extra openai" + ) + return QuestionAnswerOpenAIAgent(client, model or "gpt-4o-mini") + + raise ValueError(f"Unsupported QA provider: {Config.QA_PROVIDER}") diff --git a/src/haiku/rag/qa/base.py b/src/haiku/rag/qa/base.py index 6c8f8359..0ff2a55b 100644 --- a/src/haiku/rag/qa/base.py +++ b/src/haiku/rag/qa/base.py @@ -2,7 +2,7 @@ from haiku.rag.client import HaikuRAG from haiku.rag.qa.prompts import SYSTEM_PROMPT -class QABase: +class QuestionAnswerAgentBase: _model: str = "" _system_prompt: str = SYSTEM_PROMPT @@ -14,3 +14,28 @@ class QABase: raise NotImplementedError( "QABase is an abstract class. Please implement the answer method in a subclass." ) + + tools = [ + { + "type": "function", + "function": { + "name": "search_documents", + "description": "Search the knowledge base for relevant documents", + "parameters": { + "type": "object", + "properties": { + "query": { + "type": "string", + "description": "The search query to find relevant documents", + }, + "limit": { + "type": "integer", + "description": "Maximum number of results to return", + "default": 3, + }, + }, + "required": ["query"], + }, + }, + } + ] diff --git a/src/haiku/rag/qa/ollama.py b/src/haiku/rag/qa/ollama.py index 021190ca..c8cac4ce 100644 --- a/src/haiku/rag/qa/ollama.py +++ b/src/haiku/rag/qa/ollama.py @@ -2,12 +2,12 @@ from ollama import AsyncClient from haiku.rag.client import HaikuRAG from haiku.rag.config import Config -from haiku.rag.qa.base import QABase +from haiku.rag.qa.base import QuestionAnswerAgentBase OLLAMA_OPTIONS = {"temperature": 0.0, "seed": 42, "num_ctx": 64000} -class QA(QABase): +class QuestionAnswerOllamaAgent(QuestionAnswerAgentBase): def __init__(self, client: HaikuRAG, model: str = Config.QA_MODEL): super().__init__(client, model or self._model) @@ -15,30 +15,6 @@ class QA(QABase): ollama_client = AsyncClient(host=Config.OLLAMA_BASE_URL) # Define the search tool - tools = [ - { - "type": "function", - "function": { - "name": "search_documents", - "description": "Search the knowledge base for relevant documents", - "parameters": { - "type": "object", - "properties": { - "query": { - "type": "string", - "description": "The search query to find relevant documents", - }, - "limit": { - "type": "integer", - "description": "Maximum number of results to return", - "default": 3, - }, - }, - "required": ["query"], - }, - }, - } - ] messages = [ {"role": "system", "content": self._system_prompt}, @@ -49,7 +25,7 @@ class QA(QABase): response = await ollama_client.chat( model=self._model, messages=messages, - tools=tools, + tools=self.tools, options=OLLAMA_OPTIONS, think=False, ) diff --git a/src/haiku/rag/qa/openai.py b/src/haiku/rag/qa/openai.py new file mode 100644 index 00000000..f75a7396 --- /dev/null +++ b/src/haiku/rag/qa/openai.py @@ -0,0 +1,101 @@ +from collections.abc import Sequence + +try: + from openai import AsyncOpenAI + from openai.types.chat import ( + ChatCompletionAssistantMessageParam, + ChatCompletionMessageParam, + ChatCompletionSystemMessageParam, + ChatCompletionToolMessageParam, + ChatCompletionUserMessageParam, + ) + from openai.types.chat.chat_completion_tool_param import ChatCompletionToolParam + + from haiku.rag.client import HaikuRAG + from haiku.rag.qa.base import QuestionAnswerAgentBase + + class QuestionAnswerOpenAIAgent(QuestionAnswerAgentBase): + def __init__(self, client: HaikuRAG, model: str = "gpt-4o-mini"): + super().__init__(client, model or self._model) + self.tools: Sequence[ChatCompletionToolParam] = [ + ChatCompletionToolParam(tool) for tool in self.tools + ] + + async def answer(self, question: str) -> str: + openai_client = AsyncOpenAI() + + # Define the search tool + + messages: list[ChatCompletionMessageParam] = [ + ChatCompletionSystemMessageParam( + role="system", content=self._system_prompt + ), + ChatCompletionUserMessageParam(role="user", content=question), + ] + + # Initial response with tool calling + response = await openai_client.chat.completions.create( + model=self._model, + messages=messages, + tools=self.tools, + temperature=0.0, + ) + + response_message = response.choices[0].message + + if response_message.tool_calls: + messages.append( + ChatCompletionAssistantMessageParam( + role="assistant", + content=response_message.content, + tool_calls=[ + { + "id": tc.id, + "type": "function", + "function": { + "name": tc.function.name, + "arguments": tc.function.arguments, + }, + } + for tc in response_message.tool_calls + ], + ) + ) + + for tool_call in response_message.tool_calls: + if tool_call.function.name == "search_documents": + import json + + args = json.loads(tool_call.function.arguments) + query = args.get("query", question) + limit = int(args.get("limit", 3)) + + search_results = await self._client.search(query, limit=limit) + + context_chunks = [] + for chunk, score in search_results: + context_chunks.append( + f"Content: {chunk.content}\nScore: {score:.4f}" + ) + + context = "\n\n".join(context_chunks) + + messages.append( + ChatCompletionToolMessageParam( + role="tool", + content=context, + tool_call_id=tool_call.id, + ) + ) + + final_response = await openai_client.chat.completions.create( + model=self._model, + messages=messages, + temperature=0.0, + ) + return final_response.choices[0].message.content or "" + else: + return response_message.content or "" + +except ImportError: + pass diff --git a/tests/conftest.py b/tests/conftest.py index 2dcea549..31ed812d 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -3,8 +3,6 @@ from pathlib import Path import pytest from datasets import Dataset, load_dataset, load_from_disk -from .llm_judge import LLMJudge - @pytest.fixture(scope="session") def qa_corpus() -> Dataset: @@ -18,8 +16,3 @@ def qa_corpus() -> Dataset: corpus = ds.filter(lambda doc: doc["document_topic"] == "News Stories") corpus.save_to_disk(ds_path) return corpus - - -@pytest.fixture(scope="session") -def llm_judge() -> LLMJudge: - return LLMJudge() diff --git a/tests/generate_benchmark_db.py b/tests/generate_benchmark_db.py index bbebb5eb..a736317a 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 QA +from haiku.rag.qa.ollama import QuestionAnswerOllamaAgent 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 = QA(rag) + qa = QuestionAnswerOllamaAgent(rag) for i, doc in enumerate(tqdm(corpus, desc="QA Benchmarking")): question = doc["question"] # type: ignore diff --git a/tests/test_qa.py b/tests/test_qa.py index 18496df7..41312a75 100644 --- a/tests/test_qa.py +++ b/tests/test_qa.py @@ -1,29 +1,52 @@ -from typing import TYPE_CHECKING - import pytest from datasets import Dataset from haiku.rag.client import HaikuRAG -from haiku.rag.qa.ollama import QA +from haiku.rag.qa.ollama import QuestionAnswerOllamaAgent -if TYPE_CHECKING: - import sys - from pathlib import Path +try: + from haiku.rag.qa.openai import QuestionAnswerOpenAIAgent - sys.path.append(str(Path(__file__).parent)) - from llm_judge import LLMJudge + OPENAI_AVAILABLE = True +except ImportError: + QuestionAnswerOpenAIAgent = None + OPENAI_AVAILABLE = False + +from .llm_judge import LLMJudge @pytest.mark.asyncio -async def test_qa_with_dataset_question(qa_corpus: Dataset, llm_judge: "LLMJudge"): +async def test_qa_ollama_with_dataset_question(qa_corpus: Dataset): """Test QA with actual question from the dataset using LLM judge.""" client = HaikuRAG(":memory:") - qa = QA(client) + qa = QuestionAnswerOllamaAgent(client) + llm_judge = LLMJudge() + + doc = qa_corpus[1] + await client.create_document( + content=doc["document_extracted"], uri=doc["document_id"] + ) + + question = doc["question"] + expected_answer = doc["answer"] + + answer = await qa.answer(question) + is_equivalent = await llm_judge.judge_answers(question, answer, expected_answer) + + assert is_equivalent, ( + f"Generated answer not equivalent to expected answer.\nQuestion: {question}\nGenerated: {answer}\nExpected: {expected_answer}" + ) + + +@pytest.mark.asyncio +@pytest.mark.skipif(not OPENAI_AVAILABLE, reason="OpenAI not available") +async def test_qa_openai_basic(qa_corpus: Dataset): + """Test OpenAI QA basic functionality.""" + client = HaikuRAG(":memory:") + qa = QuestionAnswerOpenAIAgent(client) # type: ignore + llm_judge = LLMJudge() - # Use the first document from the corpus doc = qa_corpus[1] - - # Add the document to database await client.create_document( content=doc["document_extracted"], uri=doc["document_id"] ) @@ -32,11 +55,8 @@ async def test_qa_with_dataset_question(qa_corpus: Dataset, llm_judge: "LLMJudge expected_answer = doc["answer"] answer = await qa.answer(question) - # Use LLM judge to evaluate answer equivalence is_equivalent = await llm_judge.judge_answers(question, answer, expected_answer) - assert isinstance(answer, str) - assert len(answer) > 0 assert is_equivalent, ( f"Generated answer not equivalent to expected answer.\nQuestion: {question}\nGenerated: {answer}\nExpected: {expected_answer}" )