danvics/haiku.rag
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions

OpenAI Question/Answer agent

Browse source
This commit is contained in:
Yiorgis Gozadinos 2025-06-27 09:14:16 +03:00
parent 357460da57
commit 85b106c461
No known key found for this signature in database
8 changed files with 202 additions and 53 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

8
src/haiku/rag/config.py
View file

@ -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

26
src/haiku/rag/qa/__init__.py
View file

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

27
src/haiku/rag/qa/base.py
View file

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

30
src/haiku/rag/qa/ollama.py
View file

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

101
src/haiku/rag/qa/openai.py Normal file
View file

@ -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

7
tests/conftest.py
View file

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

4
tests/generate_benchmark_db.py
View file

@ -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

52
tests/test_qa.py
View file

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