OpenAI Question/Answer agent

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

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}")

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"],
+                },
+            },
+        }
+    ]

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

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

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()

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

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:
+try:
-    import sys
+    from haiku.rag.qa.openai import QuestionAnswerOpenAIAgent
-    from pathlib import Path
 
-    sys.path.append(str(Path(__file__).parent))
+    OPENAI_AVAILABLE = True
-    from llm_judge import LLMJudge
+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}"
     )