From 357460da577cf13ca9a1a9ffb02fa5f95e162ae3 Mon Sep 17 00:00:00 2001 From: Yiorgis Gozadinos Date: Thu, 26 Jun 2025 11:56:49 +0300 Subject: [PATCH 1/4] Update benchmarks --- BENCHMARKS.md | 4 ++++ 1 file changed, 4 insertions(+) diff --git a/BENCHMARKS.md b/BENCHMARKS.md index bc66cab3..977e1927 100644 --- a/BENCHMARKS.md +++ b/BENCHMARKS.md @@ -7,3 +7,7 @@ We use [repliqa](https://huggingface.co/datasets/ServiceNow/repliqa) for the eva We load the `News Stories` from `repliqa_3` which is 1035 documents, using `tests/generate_benchmark_db.py`, using the `mxbai-embed-large` Ollama embeddings. Subsequently, we run a search over the `question` for each row of the dataset and check whether we match the document that answers the question. The recall obtained is ~0.75 for matching in the top result, raising to ~0.75 for the top 3 results. + +* Question/Answer evaluation + +We use the `News Stories` from `repliqa_3` using the `mxbai-embed-large` Ollama embeddings, with a QA agent also using Ollama with the `qwen3` model (8b). For each story we ask the `question` and use an LLM judge (also `qwen3`) to evaluate whether the answer is correct or not. Thus we obtain accuracy of ~0.54. From 85b106c4614b41b4508b2298dca7a8b317819fe1 Mon Sep 17 00:00:00 2001 From: Yiorgis Gozadinos Date: Fri, 27 Jun 2025 09:14:16 +0300 Subject: [PATCH 2/4] 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}" ) From 4bbc23dbd40925bcf29b8c094546e43ea4a75c50 Mon Sep 17 00:00:00 2001 From: Yiorgis Gozadinos Date: Sat, 28 Jun 2025 09:21:45 +0300 Subject: [PATCH 3/4] Add QA to client, cli --- src/haiku/rag/app.py | 11 +++++++++++ src/haiku/rag/cli.py | 15 +++++++++++++++ src/haiku/rag/client.py | 16 +++++++++++++++- 3 files changed, 41 insertions(+), 1 deletion(-) diff --git a/src/haiku/rag/app.py b/src/haiku/rag/app.py index 6db14c83..7e33ba5b 100644 --- a/src/haiku/rag/app.py +++ b/src/haiku/rag/app.py @@ -61,6 +61,17 @@ class HaikuRAGApp: for chunk, score in results: self._rich_print_search_result(chunk, score) + async def ask(self, question: str): + async with HaikuRAG(db_path=self.db_path) as self.client: + try: + answer = await self.client.ask(question) + self.console.print(f"[bold blue]Question:[/bold blue] {question}") + self.console.print() + self.console.print("[bold green]Answer:[/bold green]") + self.console.print(Markdown(answer)) + except Exception as e: + self.console.print(f"[red]Error: {e}[/red]") + def _rich_print_document(self, doc: Document, truncate: bool = False): """Format a document for display.""" if truncate: diff --git a/src/haiku/rag/cli.py b/src/haiku/rag/cli.py index 71e2c8b9..2e012cf1 100644 --- a/src/haiku/rag/cli.py +++ b/src/haiku/rag/cli.py @@ -113,6 +113,21 @@ def search( event_loop.run_until_complete(app.search(query=query, limit=limit, k=k)) +@cli.command("ask", help="Ask a question using the QA agent") +def ask( + question: str = typer.Argument( + help="The question to ask", + ), + db: Path = typer.Option( + get_default_data_dir() / "haiku.rag.sqlite", + "--db", + help="Path to the SQLite database file", + ), +): + app = HaikuRAGApp(db_path=db) + event_loop.run_until_complete(app.ask(question=question)) + + @cli.command( "serve", help="Start the haiku.rag MCP server (by default in streamable HTTP mode)" ) diff --git a/src/haiku/rag/client.py b/src/haiku/rag/client.py index 920f262a..0f24b3b9 100644 --- a/src/haiku/rag/client.py +++ b/src/haiku/rag/client.py @@ -36,7 +36,7 @@ class HaikuRAG: """Async context manager entry.""" return self - async def __aexit__(self, exc_type, exc_val, exc_tb): + async def __aexit__(self, exc_type, exc_val, exc_tb): # noqa: ARG002 """Async context manager exit.""" self.close() return False @@ -256,6 +256,20 @@ class HaikuRAG: """ return await self.chunk_repository.search_chunks_hybrid(query, limit, k) + async def ask(self, question: str) -> str: + """Ask a question using the configured QA agent. + + Args: + question: The question to ask + + Returns: + The generated answer as a string + """ + from haiku.rag.qa import get_qa_agent + + qa_agent = get_qa_agent(self) + return await qa_agent.answer(question) + def close(self): """Close the underlying store connection.""" self.store.close() From 273a3bda4f7899c73ebddb86ab2018b54b633cba Mon Sep 17 00:00:00 2001 From: Yiorgis Gozadinos Date: Sat, 28 Jun 2025 09:37:59 +0300 Subject: [PATCH 4/4] Document QA --- README.md | 8 ++++++++ docs/cli.md | 9 +++++++++ docs/configuration.md | 47 +++++++++++++++++++++++++++++++++++++++++++ docs/index.md | 6 ++++++ docs/python.md | 15 ++++++++++++++ 5 files changed, 85 insertions(+) diff --git a/README.md b/README.md index 8e20b3e1..cf86ae99 100644 --- a/README.md +++ b/README.md @@ -9,6 +9,7 @@ Retrieval-Augmented Generation (RAG) library on SQLite. - **Local SQLite**: No external servers required - **Multiple embedding providers**: Ollama, VoyageAI, OpenAI - **Hybrid search**: Vector + full-text search with Reciprocal Rank Fusion +- **Question answering**: Built-in QA agents on your documents - **File monitoring**: Auto-index files when run as server - **40+ file formats**: PDF, DOCX, HTML, Markdown, audio, URLs - **MCP server**: Expose as tools for AI assistants @@ -27,6 +28,9 @@ haiku-rag add-src document.pdf # Search haiku-rag search "query" +# Ask questions +haiku-rag ask "Who is the author of haiku.rag?" + # Start server with file monitoring export MONITOR_DIRECTORIES="/path/to/docs" haiku-rag serve @@ -45,6 +49,10 @@ async with HaikuRAG("database.db") as client: results = await client.search("query") for chunk, score in results: print(f"{score:.3f}: {chunk.content}") + + # Ask questions + answer = await client.ask("Who is the author of haiku.rag?") + print(answer) ``` ## MCP Server diff --git a/docs/cli.md b/docs/cli.md index e5100062..fae3db8a 100644 --- a/docs/cli.md +++ b/docs/cli.md @@ -47,6 +47,15 @@ With options: haiku-rag search "python programming" --limit 10 --k 100 ``` +## Question Answering + +Ask questions about your documents: +```bash +haiku-rag ask "Who is the author of haiku.rag?" +``` + +The QA agent will search your documents for relevant information and provide a comprehensive answer. + ## Server Start the MCP server: diff --git a/docs/configuration.md b/docs/configuration.md index 8ba6cc7e..cae8a1dd 100644 --- a/docs/configuration.md +++ b/docs/configuration.md @@ -55,3 +55,50 @@ EMBEDDINGS_MODEL="text-embedding-3-small" # or text-embedding-3-large EMBEDDINGS_VECTOR_DIM=1536 OPENAI_API_KEY="your-api-key" ``` + +## Question Answering Providers + +Configure which LLM provider to use for question answering. + +### Ollama (Default) + +```bash +QA_PROVIDER="ollama" +QA_MODEL="qwen3" +OLLAMA_BASE_URL="http://localhost:11434" +``` + +### OpenAI + +For OpenAI QA, you need to install haiku.rag with OpenAI extras: + +```bash +uv pip install haiku.rag --extra openai +``` + +Then configure: + +```bash +QA_PROVIDER="openai" +QA_MODEL="gpt-4o-mini" # or gpt-4, gpt-3.5-turbo, etc. +OPENAI_API_KEY="your-api-key" +``` + +## Other Settings + +### Database and Storage + +```bash +# Default data directory (where SQLite database is stored) +DEFAULT_DATA_DIR="/path/to/data" +``` + +### Document Processing + +```bash +# Chunk size for document processing +CHUNK_SIZE=256 + +# Chunk overlap for better context +CHUNK_OVERLAP=32 +``` diff --git a/docs/index.md b/docs/index.md index 404ae8b4..66117c63 100644 --- a/docs/index.md +++ b/docs/index.md @@ -8,6 +8,7 @@ - **Local SQLite**: No need to run additional servers - **Support for various embedding providers**: Ollama, VoyageAI, OpenAI or add your own - **Hybrid Search**: Vector search using `sqlite-vec` combined with full-text search `FTS5`, using Reciprocal Rank Fusion +- **Question Answering**: Built-in QA agents using Ollama or OpenAI. - **File monitoring**: Automatically index files when run as a server - **Extended file format support**: Parse 40+ file formats including PDF, DOCX, HTML, Markdown, audio and more. Or add a URL! - **MCP server**: Exposes functionality as MCP tools @@ -31,12 +32,16 @@ async with HaikuRAG("database.db") as client: # Search documents results = await client.search("query") + + # Ask questions + answer = await client.ask("Who is the author of haiku.rag?") ``` Or use the CLI: ```bash haiku-rag add "Your document content" haiku-rag search "query" +haiku-rag ask "Who is the author of haiku.rag?" ``` ## Documentation @@ -44,6 +49,7 @@ haiku-rag search "query" - [Installation](installation.md) - Install haiku.rag with different providers - [Configuration](configuration.md) - Environment variables and settings - [CLI](cli.md) - Command line interface usage +- [Question Answering](qa.md) - QA agents and natural language queries - [Server](server.md) - File monitoring and server mode - [MCP](mcp.md) - Model Context Protocol integration - [Python](python.md) - Python API reference diff --git a/docs/python.md b/docs/python.md index 9908210e..ebc87f4c 100644 --- a/docs/python.md +++ b/docs/python.md @@ -91,4 +91,19 @@ for chunk, relevance_score in results: print(f"Relevance: {relevance_score:.3f}") print(f"Content: {chunk.content}") print(f"From document: {chunk.document_id}") + print(f"Document URI: {chunk.document_uri}") + print(f"Document metadata: {chunk.document_meta}") ``` + +## Question Answering + +Ask questions about your documents: + +```python +answer = await client.ask("Who is the author of haiku.rag?") +print(answer) +``` + +The QA agent will search your documents for relevant information and use the configured LLM to generate a comprehensive answer. + +The QA provider and model can be configured via environment variables (see [Configuration](configuration.md)).