haiku.rag 0.3.1__tar.gz → 0.3.3__tar.gz

{haiku_rag-0.3.1 → haiku_rag-0.3.3}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: haiku.rag
-Version: 0.3.1
+Version: 0.3.3
 Summary: Retrieval Augmented Generation (RAG) with SQLite
 Author-email: Yiorgis Gozadinos <ggozadinos@gmail.com>
 License: MIT
@@ -116,3 +116,4 @@ Full documentation at: https://ggozad.github.io/haiku.rag/
 - [Configuration](https://ggozad.github.io/haiku.rag/configuration/) - Environment variables
 - [CLI](https://ggozad.github.io/haiku.rag/cli/) - Command reference
 - [Python API](https://ggozad.github.io/haiku.rag/python/) - Complete API docs
+- [Benchmarks](https://ggozad.github.io/haiku.rag/benchmarks/) - Performance Benchmarks

{haiku_rag-0.3.1 → haiku_rag-0.3.3}/README.md

@@ -77,3 +77,4 @@ Full documentation at: https://ggozad.github.io/haiku.rag/
 - [Configuration](https://ggozad.github.io/haiku.rag/configuration/) - Environment variables
 - [CLI](https://ggozad.github.io/haiku.rag/cli/) - Command reference
 - [Python API](https://ggozad.github.io/haiku.rag/python/) - Complete API docs
+- [Benchmarks](https://ggozad.github.io/haiku.rag/benchmarks/) - Performance Benchmarks

haiku_rag-0.3.3/docs/benchmarks.md

@@ -0,0 +1,27 @@
+# Benchmarks
+
+We use the [repliqa](https://huggingface.co/datasets/ServiceNow/repliqa) dataset for the evaluation of `haiku.rag`.
+
+You can perform your own evaluations using as example the script found at
+`tests/generate_benchmark_db.py`.
+
+## Recall
+
+In order to calculate recall, we load the `News Stories` from `repliqa_3` which is 1035 documents and index them in a sqlite db. Subsequently, we run a search over the `question` field for each row of the dataset and check whether we match the document that answers the question.
+
+
+The recall obtained is ~0.73 for matching in the top result, raising to ~0.75 for the top 3 results.
+
+| Model                                 | Document in top 1 | Document in top 3 |
+|---------------------------------------|-------------------|-------------------|
+| Ollama / `mxbai-embed-large`          | 0.73              | 0.75              |
+| OpenAI / `text-embeddings-3-small`    | 0.75              | 0.88              |
+
+## Question/Answer evaluation
+
+Again using the same dataset, we use a QA agent to answer the question. In addition we use an LLM judge (using the Ollama `qwen3`) to evaluate whether the answer is correct or not. The obtained accuracy is as follows:
+
+| Embedding Model              | QA Model                          | Accuracy  |
+|------------------------------|-----------------------------------|-----------|
+| Ollama / `mxbai-embed-large` | Ollama / `qwen3`                  | 0.64      |
+| Ollama / `mxbai-embed-large` | Anthropic / `Claude Sonnet 3.7`   | 0.79      |

{haiku_rag-0.3.1 → haiku_rag-0.3.3}/mkdocs.yml

@@ -63,6 +63,7 @@ nav:
       - Server: server.md
       - MCP: mcp.md
       - Python: python.md
+      - Benchmarks: benchmarks.md
 markdown_extensions:
   - admonition
   - attr_list

{haiku_rag-0.3.1 → haiku_rag-0.3.3}/pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "haiku.rag"
-version = "0.3.1"
+version = "0.3.3"
 description = "Retrieval Augmented Generation (RAG) with SQLite"
 authors = [{ name = "Yiorgis Gozadinos", email = "ggozadinos@gmail.com" }]
 license = { text = "MIT" }

{haiku_rag-0.3.1 → haiku_rag-0.3.3}/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}")

haiku_rag-0.3.3/src/haiku/rag/qa/anthropic.py

@@ -0,0 +1,106 @@
+from collections.abc import Sequence
+
+try:
+    from anthropic import AsyncAnthropic
+    from anthropic.types import MessageParam, TextBlock, ToolParam, ToolUseBlock
+
+    from haiku.rag.client import HaikuRAG
+    from haiku.rag.qa.base import QuestionAnswerAgentBase
+
+    class QuestionAnswerAnthropicAgent(QuestionAnswerAgentBase):
+        def __init__(self, client: HaikuRAG, model: str = "claude-3-5-haiku-20241022"):
+            super().__init__(client, model or self._model)
+            self.tools: Sequence[ToolParam] = [
+                ToolParam(
+                    name="search_documents",
+                    description="Search the knowledge base for relevant documents",
+                    input_schema={
+                        "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"],
+                    },
+                )
+            ]
+
+        async def answer(self, question: str) -> str:
+            anthropic_client = AsyncAnthropic()
+
+            messages: list[MessageParam] = [{"role": "user", "content": question}]
+
+            max_rounds = 5  # Prevent infinite loops
+
+            for _ in range(max_rounds):
+                response = await anthropic_client.messages.create(
+                    model=self._model,
+                    max_tokens=4096,
+                    system=self._system_prompt,
+                    messages=messages,
+                    tools=self.tools,
+                    temperature=0.0,
+                )
+
+                if response.stop_reason == "tool_use":
+                    messages.append({"role": "assistant", "content": response.content})
+
+                    # Process tool calls
+                    tool_results = []
+                    for content_block in response.content:
+                        if isinstance(content_block, ToolUseBlock):
+                            if content_block.name == "search_documents":
+                                args = content_block.input
+                                query = (
+                                    args.get("query", question)
+                                    if isinstance(args, dict)
+                                    else question
+                                )
+                                limit = (
+                                    int(args.get("limit", 3))
+                                    if isinstance(args, dict)
+                                    else 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)
+
+                                tool_results.append(
+                                    {
+                                        "type": "tool_result",
+                                        "tool_use_id": content_block.id,
+                                        "content": context,
+                                    }
+                                )
+
+                    if tool_results:
+                        messages.append({"role": "user", "content": tool_results})
+                else:
+                    # No tool use, return the response
+                    if response.content:
+                        first_content = response.content[0]
+                        if isinstance(first_content, TextBlock):
+                            return first_content.text
+                    return ""
+
+            # If we've exhausted max rounds, return empty string
+            return ""
+
+except ImportError:
+    pass

haiku_rag-0.3.3/src/haiku/rag/qa/ollama.py

@@ -0,0 +1,64 @@
+from ollama import AsyncClient
+
+from haiku.rag.client import HaikuRAG
+from haiku.rag.config import Config
+from haiku.rag.qa.base import QuestionAnswerAgentBase
+
+OLLAMA_OPTIONS = {"temperature": 0.0, "seed": 42, "num_ctx": 64000}
+
+
+class QuestionAnswerOllamaAgent(QuestionAnswerAgentBase):
+    def __init__(self, client: HaikuRAG, model: str = Config.QA_MODEL):
+        super().__init__(client, model or self._model)
+
+    async def answer(self, question: str) -> str:
+        ollama_client = AsyncClient(host=Config.OLLAMA_BASE_URL)
+
+        messages = [
+            {"role": "system", "content": self._system_prompt},
+            {"role": "user", "content": question},
+        ]
+
+        max_rounds = 5  # Prevent infinite loops
+
+        for _ in range(max_rounds):
+            response = await ollama_client.chat(
+                model=self._model,
+                messages=messages,
+                tools=self.tools,
+                options=OLLAMA_OPTIONS,
+                think=False,
+            )
+
+            if response.get("message", {}).get("tool_calls"):
+                messages.append(response["message"])
+
+                for tool_call in response["message"]["tool_calls"]:
+                    if tool_call["function"]["name"] == "search_documents":
+                        args = 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(
+                            {
+                                "role": "tool",
+                                "content": context,
+                                "tool_call_id": tool_call.get("id", "search_tool"),
+                            }
+                        )
+            else:
+                # No tool calls, return the response
+                return response["message"]["content"]
+
+        # If we've exhausted max rounds, return empty string
+        return ""

haiku_rag-0.3.3/src/haiku/rag/qa/openai.py

@@ -0,0 +1,100 @@
+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()
+
+            messages: list[ChatCompletionMessageParam] = [
+                ChatCompletionSystemMessageParam(
+                    role="system", content=self._system_prompt
+                ),
+                ChatCompletionUserMessageParam(role="user", content=question),
+            ]
+
+            max_rounds = 5  # Prevent infinite loops
+
+            for _ in range(max_rounds):
+                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,
+                                )
+                            )
+                else:
+                    # No tool calls, return the response
+                    return response_message.content or ""
+
+            # If we've exhausted max rounds, return empty string
+            return ""
+
+except ImportError:
+    pass

haiku_rag-0.3.3/src/haiku/rag/qa/prompts.py

@@ -0,0 +1,20 @@
+SYSTEM_PROMPT = """
+You are a knowledgeable assistant that helps users find information from a document knowledge base.
+
+Your process:
+1. When a user asks a question, use the search_documents tool to find relevant information
+2. Search with specific keywords and phrases from the user's question
+3. Review the search results and their relevance scores
+4. If you need additional context, perform follow-up searches with different keywords
+5. Provide a comprehensive answer based only on the retrieved documents
+
+Guidelines:
+- Base your answers strictly on the provided document content
+- Quote or reference specific information when possible
+- If multiple documents contain relevant information, synthesize them coherently
+- Indicate when information is incomplete or when you need to search for additional context
+- If the retrieved documents don't contain sufficient information, clearly state: "I cannot find enough information in the knowledge base to answer this question."
+- For complex questions, consider breaking them down and performing multiple searches
+
+Be concise, and always maintain accuracy over completeness. Prefer short, direct answers that are well-supported by the documents.
+"""

haiku_rag-0.3.3/tests/generate_benchmark_db.py

@@ -0,0 +1,151 @@
+import asyncio
+from pathlib import Path
+
+from datasets import Dataset, load_dataset
+from llm_judge import LLMJudge
+from rich.console import Console
+from rich.progress import Progress
+
+from haiku.rag.client import HaikuRAG
+from haiku.rag.qa import get_qa_agent
+
+console = Console()
+
+db_path = Path(__file__).parent / "data" / "benchmark.sqlite"
+
+
+async def populate_db():
+    ds: Dataset = load_dataset("ServiceNow/repliqa")["repliqa_3"]  # type: ignore
+    corpus = ds.filter(lambda doc: doc["document_topic"] == "News Stories")
+
+    with Progress() as progress:
+        task = progress.add_task("[green]Populating database...", total=len(corpus))
+
+        async with HaikuRAG(db_path) as rag:
+            for doc in corpus:
+                uri = doc["document_id"]  # type: ignore
+                existing_doc = await rag.get_document_by_uri(uri)
+                if existing_doc is not None:
+                    progress.advance(task)
+                    continue
+
+                await rag.create_document(
+                    content=doc["document_extracted"],  # type: ignore
+                    uri=uri,
+                )
+                progress.advance(task)
+
+
+async def run_match_benchmark():
+    ds: Dataset = load_dataset("ServiceNow/repliqa")["repliqa_3"]  # type: ignore
+    corpus = ds.filter(lambda doc: doc["document_topic"] == "News Stories")
+
+    correct_at_1 = 0
+    correct_at_2 = 0
+    correct_at_3 = 0
+    total_queries = 0
+
+    with Progress() as progress:
+        task = progress.add_task(
+            "[blue]Running retrieval benchmark...", total=len(corpus)
+        )
+
+        async with HaikuRAG(db_path) as rag:
+            for doc in corpus:
+                doc_id = doc["document_id"]  # type: ignore
+                matches = await rag.search(
+                    query=doc["question"],  # type: ignore
+                    limit=3,
+                )
+
+                total_queries += 1
+
+                # Check position of correct document in results
+                for position, (chunk, _) in enumerate(matches):
+                    retrieved = await rag.get_document_by_id(chunk.document_id)
+                    if retrieved and retrieved.uri == doc_id:
+                        if position == 0:  # First position
+                            correct_at_1 += 1
+                            correct_at_2 += 1
+                            correct_at_3 += 1
+                        elif position == 1:  # Second position
+                            correct_at_2 += 1
+                            correct_at_3 += 1
+                        elif position == 2:  # Third position
+                            correct_at_3 += 1
+                        break
+
+                progress.advance(task)
+
+    # Calculate recall metrics
+    recall_at_1 = correct_at_1 / total_queries
+    recall_at_2 = correct_at_2 / total_queries
+    recall_at_3 = correct_at_3 / total_queries
+
+    console.print("\n=== Retrieval Benchmark Results ===", style="bold cyan")
+    console.print(f"Total queries: {total_queries}")
+    console.print(f"Recall@1: {recall_at_1:.4f}")
+    console.print(f"Recall@2: {recall_at_2:.4f}")
+    console.print(f"Recall@3: {recall_at_3:.4f}")
+
+    return {"recall@1": recall_at_1, "recall@2": recall_at_2, "recall@3": recall_at_3}
+
+
+async def run_qa_benchmark(k: int | None = None):
+    """Run QA benchmarking on the corpus."""
+    ds: Dataset = load_dataset("ServiceNow/repliqa")["repliqa_3"]  # type: ignore
+    corpus = ds.filter(lambda doc: doc["document_topic"] == "News Stories")
+
+    if k is not None:
+        corpus = corpus.select(range(min(k, len(corpus))))
+
+    judge = LLMJudge()
+    correct_answers = 0
+    total_questions = 0
+
+    with Progress() as progress:
+        task = progress.add_task("[yellow]Running QA benchmark...", total=len(corpus))
+
+        async with HaikuRAG(db_path) as rag:
+            qa = get_qa_agent(rag)
+
+            for doc in corpus:
+                question = doc["question"]  # type: ignore
+                expected_answer = doc["answer"]  # type: ignore
+
+                generated_answer = await qa.answer(question)
+                is_equivalent = await judge.judge_answers(
+                    question, generated_answer, expected_answer
+                )
+                console.print(f"Question: {question}")
+                console.print(f"Expected: {expected_answer}")
+                console.print(f"Generated: {generated_answer}")
+                console.print(f"Equivalent: {is_equivalent}\n")
+
+                if is_equivalent:
+                    correct_answers += 1
+                total_questions += 1
+                console.print("Current score:", correct_answers, "/", total_questions)
+
+                progress.advance(task)
+
+    accuracy = correct_answers / total_questions if total_questions > 0 else 0
+
+    console.print("\n=== QA Benchmark Results ===", style="bold cyan")
+    console.print(f"Total questions: {total_questions}")
+    console.print(f"Correct answers: {correct_answers}")
+    console.print(f"QA Accuracy: {accuracy:.4f} ({accuracy * 100:.2f}%)")
+
+
+async def main():
+    await populate_db()
+
+    console.print("Running retrieval benchmarks...", style="bold blue")
+    await run_match_benchmark()
+
+    console.print("\nRunning QA benchmarks...", style="bold yellow")
+    await run_qa_benchmark()
+
+
+if __name__ == "__main__":
+    asyncio.run(main())

{haiku_rag-0.3.1 → haiku_rag-0.3.3}/tests/llm_judge.py

@@ -49,6 +49,7 @@ class LLMJudge:
         1. Do both answers provide the same answer?
         2. Do both answers directly address the question asked?
         3. Minor differences in wording or style are acceptable if the meaning of the answer is the same.
+        4. If one answer is more detailed but the other is correct, they can still be considered equivalent.
 
         Be strict but fair in your evaluation. Focus on factual correctness and whether both answers would satisfy someone asking the question."""
 

{haiku_rag-0.3.1 → haiku_rag-0.3.3}/uv.lock

@@ -816,7 +816,7 @@ wheels = [
 
 [[package]]
 name = "haiku-rag"
-version = "0.3.1"
+version = "0.3.3"
 source = { editable = "." }
 dependencies = [
     { name = "fastmcp" },

haiku_rag-0.3.1/BENCHMARKS.md

@@ -1,13 +0,0 @@
-# `haiku.rag` benchmarks
-
-We use [repliqa](https://huggingface.co/datasets/ServiceNow/repliqa) for the evaluation of `haiku.rag`
-
-* Recall
-
-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.

haiku_rag-0.3.1/src/haiku/rag/qa/anthropic.py

@@ -1,112 +0,0 @@
-from collections.abc import Sequence
-
-try:
-    from anthropic import AsyncAnthropic
-    from anthropic.types import MessageParam, TextBlock, ToolParam, ToolUseBlock
-
-    from haiku.rag.client import HaikuRAG
-    from haiku.rag.qa.base import QuestionAnswerAgentBase
-
-    class QuestionAnswerAnthropicAgent(QuestionAnswerAgentBase):
-        def __init__(self, client: HaikuRAG, model: str = "claude-3-5-haiku-20241022"):
-            super().__init__(client, model or self._model)
-            self.tools: Sequence[ToolParam] = [
-                ToolParam(
-                    name="search_documents",
-                    description="Search the knowledge base for relevant documents",
-                    input_schema={
-                        "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"],
-                    },
-                )
-            ]
-
-        async def answer(self, question: str) -> str:
-            anthropic_client = AsyncAnthropic()
-
-            messages: list[MessageParam] = [{"role": "user", "content": question}]
-
-            response = await anthropic_client.messages.create(
-                model=self._model,
-                max_tokens=4096,
-                system=self._system_prompt,
-                messages=messages,
-                tools=self.tools,
-                temperature=0.0,
-            )
-
-            if response.stop_reason == "tool_use":
-                messages.append({"role": "assistant", "content": response.content})
-
-                # Process tool calls
-                tool_results = []
-                for content_block in response.content:
-                    if isinstance(content_block, ToolUseBlock):
-                        if content_block.name == "search_documents":
-                            args = content_block.input
-                            query = (
-                                args.get("query", question)
-                                if isinstance(args, dict)
-                                else question
-                            )
-                            limit = (
-                                int(args.get("limit", 3))
-                                if isinstance(args, dict)
-                                else 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)
-
-                            tool_results.append(
-                                {
-                                    "type": "tool_result",
-                                    "tool_use_id": content_block.id,
-                                    "content": context,
-                                }
-                            )
-
-                if tool_results:
-                    messages.append({"role": "user", "content": tool_results})
-
-                    final_response = await anthropic_client.messages.create(
-                        model=self._model,
-                        max_tokens=4096,
-                        system=self._system_prompt,
-                        messages=messages,
-                        temperature=0.0,
-                    )
-                    if final_response.content:
-                        first_content = final_response.content[0]
-                        if isinstance(first_content, TextBlock):
-                            return first_content.text
-                    return ""
-
-            if response.content:
-                first_content = response.content[0]
-                if isinstance(first_content, TextBlock):
-                    return first_content.text
-            return ""
-
-except ImportError:
-    pass

haiku_rag-0.3.1/src/haiku/rag/qa/ollama.py

@@ -1,67 +0,0 @@
-from ollama import AsyncClient
-
-from haiku.rag.client import HaikuRAG
-from haiku.rag.config import Config
-from haiku.rag.qa.base import QuestionAnswerAgentBase
-
-OLLAMA_OPTIONS = {"temperature": 0.0, "seed": 42, "num_ctx": 64000}
-
-
-class QuestionAnswerOllamaAgent(QuestionAnswerAgentBase):
-    def __init__(self, client: HaikuRAG, model: str = Config.QA_MODEL):
-        super().__init__(client, model or self._model)
-
-    async def answer(self, question: str) -> str:
-        ollama_client = AsyncClient(host=Config.OLLAMA_BASE_URL)
-
-        # Define the search tool
-
-        messages = [
-            {"role": "system", "content": self._system_prompt},
-            {"role": "user", "content": question},
-        ]
-
-        # Initial response with tool calling
-        response = await ollama_client.chat(
-            model=self._model,
-            messages=messages,
-            tools=self.tools,
-            options=OLLAMA_OPTIONS,
-            think=False,
-        )
-
-        if response.get("message", {}).get("tool_calls"):
-            for tool_call in response["message"]["tool_calls"]:
-                if tool_call["function"]["name"] == "search_documents":
-                    args = 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(response["message"])
-                    messages.append(
-                        {
-                            "role": "tool",
-                            "content": context,
-                            "tool_call_id": tool_call.get("id", "search_tool"),
-                        }
-                    )
-
-            final_response = await ollama_client.chat(
-                model=self._model,
-                messages=messages,
-                think=False,
-                options=OLLAMA_OPTIONS,
-            )
-            return final_response["message"]["content"]
-        else:
-            return response["message"]["content"]

haiku_rag-0.3.1/src/haiku/rag/qa/openai.py

@@ -1,101 +0,0 @@
-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

haiku_rag-0.3.1/src/haiku/rag/qa/prompts.py

@@ -1,7 +0,0 @@
-SYSTEM_PROMPT = """
-You are a helpful assistant that uses a RAG library to answer the user's prompt.
-Your task is to provide a concise and accurate answer based on the provided context.
-You should ask the provided tools to find relevant documents and then use the content of those documents to answer the question.
-Never make up information, always use the context to answer the question.
-If the context does not contain enough information to answer the question, respond with "I cannot answer that based on the provided context."
-"""

haiku_rag-0.3.1/tests/generate_benchmark_db.py

@@ -1,129 +0,0 @@
-import asyncio
-from pathlib import Path
-
-from datasets import Dataset, load_dataset
-from llm_judge import LLMJudge
-from tqdm import tqdm
-
-from haiku.rag.client import HaikuRAG
-from haiku.rag.qa.ollama import QuestionAnswerOllamaAgent
-
-db_path = Path(__file__).parent / "data" / "benchmark.sqlite"
-
-
-async def populate_db():
-    if (db_path).exists():
-        print("Benchmark database already exists. Skipping creation.")
-        return
-
-    ds: Dataset = load_dataset("ServiceNow/repliqa")["repliqa_3"]  # type: ignore
-    corpus = ds.filter(lambda doc: doc["document_topic"] == "News Stories")
-
-    async with HaikuRAG(db_path) as rag:
-        for i, doc in enumerate(tqdm(corpus)):
-            await rag.create_document(
-                content=doc["document_extracted"],  # type: ignore
-                uri=doc["document_id"],  # type: ignore
-            )
-
-
-async def run_match_benchmark():
-    ds: Dataset = load_dataset("ServiceNow/repliqa")["repliqa_3"]  # type: ignore
-    corpus = ds.filter(lambda doc: doc["document_topic"] == "News Stories")
-
-    correct_at_1 = 0
-    correct_at_2 = 0
-    correct_at_3 = 0
-    total_queries = 0
-
-    async with HaikuRAG(db_path) as rag:
-        for i, doc in enumerate(tqdm(corpus)):
-            doc_id = doc["document_id"]  # type: ignore
-            matches = await rag.search(
-                query=doc["question"],  # type: ignore
-                limit=3,
-            )
-
-            total_queries += 1
-
-            # Check position of correct document in results
-            for position, (chunk, _) in enumerate(matches):
-                retrieved = await rag.get_document_by_id(chunk.document_id)
-                if retrieved and retrieved.uri == doc_id:
-                    if position == 0:  # First position
-                        correct_at_1 += 1
-                        correct_at_2 += 1
-                        correct_at_3 += 1
-                    elif position == 1:  # Second position
-                        correct_at_2 += 1
-                        correct_at_3 += 1
-                    elif position == 2:  # Third position
-                        correct_at_3 += 1
-                    break
-
-    # Calculate recall metrics
-    recall_at_1 = correct_at_1 / total_queries
-    recall_at_2 = correct_at_2 / total_queries
-    recall_at_3 = correct_at_3 / total_queries
-
-    print("\n=== Retrieval Benchmark Results ===")
-    print(f"Total queries: {total_queries}")
-    print(f"Recall@1: {recall_at_1:.4f}")
-    print(f"Recall@2: {recall_at_2:.4f}")
-    print(f"Recall@3: {recall_at_3:.4f}")
-
-    return {"recall@1": recall_at_1, "recall@2": recall_at_2, "recall@3": recall_at_3}
-
-
-async def run_qa_benchmark(k: int | None = None):
-    """Run QA benchmarking on the corpus."""
-    ds: Dataset = load_dataset("ServiceNow/repliqa")["repliqa_3"]  # type: ignore
-    corpus = ds.filter(lambda doc: doc["document_topic"] == "News Stories")
-
-    if k is not None:
-        corpus = corpus.select(range(min(k, len(corpus))))
-
-    judge = LLMJudge()
-    correct_answers = 0
-    total_questions = 0
-
-    async with HaikuRAG(db_path) as rag:
-        qa = QuestionAnswerOllamaAgent(rag)
-
-        for i, doc in enumerate(tqdm(corpus, desc="QA Benchmarking")):
-            question = doc["question"]  # type: ignore
-            expected_answer = doc["answer"]  # type: ignore
-
-            generated_answer = await qa.answer(question)
-            is_equivalent = await judge.judge_answers(
-                question, generated_answer, expected_answer
-            )
-            print(f"Question: {question}")
-            print(f"Expected: {expected_answer}")
-            print(f"Generated: {generated_answer}")
-            print(f"Equivalent: {is_equivalent}\n")
-
-            if is_equivalent:
-                correct_answers += 1
-            total_questions += 1
-
-    accuracy = correct_answers / total_questions if total_questions > 0 else 0
-
-    print("\n=== QA Benchmark Results ===")
-    print(f"Total questions: {total_questions}")
-    print(f"Correct answers: {correct_answers}")
-    print(f"QA Accuracy: {accuracy:.4f} ({accuracy * 100:.2f}%)")
-
-
-async def main():
-    await populate_db()
-
-    print("Running retrieval benchmarks...")
-    await run_match_benchmark()
-
-    print("\nRunning QA benchmarks...")
-    await run_qa_benchmark()
-
-
-if __name__ == "__main__":
-    asyncio.run(main())