haiku.rag 0.5.5__py3-none-any.whl → 0.6.0__py3-none-any.whl

haiku/rag/embeddings/__init__.py

@@ -17,20 +17,14 @@ def get_embedder() -> EmbedderBase:
         except ImportError:
             raise ImportError(
                 "VoyageAI embedder requires the 'voyageai' package. "
-                "Please install haiku.rag with the 'voyageai' extra:"
+                "Please install haiku.rag with the 'voyageai' extra: "
                 "uv pip install haiku.rag[voyageai]"
             )
         return VoyageAIEmbedder(Config.EMBEDDINGS_MODEL, Config.EMBEDDINGS_VECTOR_DIM)
 
     if Config.EMBEDDINGS_PROVIDER == "openai":
-        try:
-            from haiku.rag.embeddings.openai import Embedder as OpenAIEmbedder
-        except ImportError:
-            raise ImportError(
-                "OpenAI embedder requires the 'openai' package. "
-                "Please install haiku.rag with the 'openai' extra:"
-                "uv pip install haiku.rag[openai]"
-            )
+        from haiku.rag.embeddings.openai import Embedder as OpenAIEmbedder
+
         return OpenAIEmbedder(Config.EMBEDDINGS_MODEL, Config.EMBEDDINGS_VECTOR_DIM)
 
     raise ValueError(f"Unsupported embedding provider: {Config.EMBEDDINGS_PROVIDER}")

haiku/rag/embeddings/openai.py

@@ -1,16 +1,13 @@
-try:
-    from openai import AsyncOpenAI
+from openai import AsyncOpenAI
 
-    from haiku.rag.embeddings.base import EmbedderBase
+from haiku.rag.embeddings.base import EmbedderBase
 
-    class Embedder(EmbedderBase):
-        async def embed(self, text: str) -> list[float]:
-            client = AsyncOpenAI()
-            response = await client.embeddings.create(
-                model=self._model,
-                input=text,
-            )
-            return response.data[0].embedding
 
-except ImportError:
-    pass
+class Embedder(EmbedderBase):
+    async def embed(self, text: str) -> list[float]:
+        client = AsyncOpenAI()
+        response = await client.embeddings.create(
+            model=self._model,
+            input=text,
+        )
+        return response.data[0].embedding

haiku/rag/qa/__init__.py

@@ -1,44 +1,15 @@
 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
+from haiku.rag.qa.agent import QuestionAnswerAgent
 
 
-def get_qa_agent(
-    client: HaikuRAG, model: str = "", use_citations: bool = False
-) -> 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, use_citations
-        )
+def get_qa_agent(client: HaikuRAG, use_citations: bool = False) -> QuestionAnswerAgent:
+    provider = Config.QA_PROVIDER
+    model_name = 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[openai]"
-            )
-        return QuestionAnswerOpenAIAgent(
-            client, model or Config.QA_MODEL, use_citations
-        )
-
-    if Config.QA_PROVIDER == "anthropic":
-        try:
-            from haiku.rag.qa.anthropic import QuestionAnswerAnthropicAgent
-        except ImportError:
-            raise ImportError(
-                "Anthropic QA agent requires the 'anthropic' package. "
-                "Please install haiku.rag with the 'anthropic' extra:"
-                "uv pip install haiku.rag[anthropic]"
-            )
-        return QuestionAnswerAnthropicAgent(
-            client, model or Config.QA_MODEL, use_citations
-        )
-
-    raise ValueError(f"Unsupported QA provider: {Config.QA_PROVIDER}")
+    return QuestionAnswerAgent(
+        client=client,
+        provider=provider,
+        model=model_name,
+        use_citations=use_citations,
+    )

haiku/rag/qa/agent.py

@@ -0,0 +1,76 @@
+from pydantic import BaseModel, Field
+from pydantic_ai import Agent, RunContext
+from pydantic_ai.models.openai import OpenAIModel
+from pydantic_ai.providers.ollama import OllamaProvider
+
+from haiku.rag.client import HaikuRAG
+from haiku.rag.config import Config
+from haiku.rag.qa.prompts import SYSTEM_PROMPT, SYSTEM_PROMPT_WITH_CITATIONS
+
+
+class SearchResult(BaseModel):
+    content: str = Field(description="The document text content")
+    score: float = Field(description="Relevance score (higher is more relevant)")
+    document_uri: str = Field(description="Source URI/path of the document")
+
+
+class Dependencies(BaseModel):
+    model_config = {"arbitrary_types_allowed": True}
+    client: HaikuRAG
+
+
+class QuestionAnswerAgent:
+    def __init__(
+        self,
+        client: HaikuRAG,
+        provider: str,
+        model: str,
+        use_citations: bool = False,
+        q: float = 0.0,
+    ):
+        self._client = client
+
+        system_prompt = SYSTEM_PROMPT_WITH_CITATIONS if use_citations else SYSTEM_PROMPT
+        model_obj = self._get_model(provider, model)
+
+        self._agent = Agent(
+            model=model_obj,
+            deps_type=Dependencies,
+            system_prompt=system_prompt,
+        )
+
+        @self._agent.tool
+        async def search_documents(
+            ctx: RunContext[Dependencies],
+            query: str,
+            limit: int = 3,
+        ) -> list[SearchResult]:
+            """Search the knowledge base for relevant documents."""
+            search_results = await ctx.deps.client.search(query, limit=limit)
+            expanded_results = await ctx.deps.client.expand_context(search_results)
+
+            return [
+                SearchResult(
+                    content=chunk.content,
+                    score=score,
+                    document_uri=chunk.document_uri or "",
+                )
+                for chunk, score in expanded_results
+            ]
+
+    def _get_model(self, provider: str, model: str):
+        """Get the appropriate model object for the provider."""
+        if provider == "ollama":
+            return OpenAIModel(
+                model_name=model,
+                provider=OllamaProvider(base_url=f"{Config.OLLAMA_BASE_URL}/v1"),
+            )
+        else:
+            # For all other providers, use the provider:model format
+            return f"{provider}:{model}"
+
+    async def answer(self, question: str) -> str:
+        """Answer a question using the RAG system."""
+        deps = Dependencies(client=self._client)
+        result = await self._agent.run(question, deps=deps)
+        return result.output

haiku/rag/qa/prompts.py

@@ -18,6 +18,7 @@ Guidelines:
 - Stick to the answer, do not ellaborate or provide context unless explicitly asked for it.
 
 Be concise, and always maintain accuracy over completeness. Prefer short, direct answers that are well-supported by the documents.
+/no_think
 """
 
 SYSTEM_PROMPT_WITH_CITATIONS = """
@@ -55,4 +56,5 @@ Citations:
 - /path/to/document2.pdf: "The manual provides guidance on military procedures and..."
 
 Be concise, and always maintain accuracy over completeness. Prefer short, direct answers that are well-supported by the documents.
+/no_think
 """

haiku/rag/reranking/ollama.py

@@ -1,14 +1,12 @@
-import json
-
-from ollama import AsyncClient
 from pydantic import BaseModel
+from pydantic_ai import Agent
+from pydantic_ai.models.openai import OpenAIModel
+from pydantic_ai.providers.ollama import OllamaProvider
 
 from haiku.rag.config import Config
 from haiku.rag.reranking.base import RerankerBase
 from haiku.rag.store.models.chunk import Chunk
 
-OLLAMA_OPTIONS = {"temperature": 0.0, "seed": 42, "num_ctx": 16384}
-
 
 class RerankResult(BaseModel):
     """Individual rerank result with index and relevance score."""
@@ -26,7 +24,28 @@ class RerankResponse(BaseModel):
 class OllamaReranker(RerankerBase):
     def __init__(self, model: str = Config.RERANK_MODEL):
         self._model = model
-        self._client = AsyncClient(host=Config.OLLAMA_BASE_URL)
+
+        # Create the reranking prompt
+        system_prompt = """You are a document reranking assistant. Given a query and a list of document chunks, you must rank them by relevance to the query.
+
+Return your response as a JSON object with a "results" array. Each result should have:
+- "index": the original index of the document (integer)
+- "relevance_score": a score between 0.0 and 1.0 indicating relevance (float, where 1.0 is most relevant)
+
+Only return the top documents up to the requested limit, ordered by decreasing relevance score.
+/no_think
+"""
+
+        model_obj = OpenAIModel(
+            model_name=model,
+            provider=OllamaProvider(base_url=f"{Config.OLLAMA_BASE_URL}/v1"),
+        )
+
+        self._agent = Agent(
+            model=model_obj,
+            output_type=RerankResponse,
+            system_prompt=system_prompt,
+        )
 
     async def rerank(
         self, query: str, chunks: list[Chunk], top_n: int = 10
@@ -38,15 +57,6 @@ class OllamaReranker(RerankerBase):
         for i, chunk in enumerate(chunks):
             documents.append({"index": i, "content": chunk.content})
 
-        # Create the prompt for reranking
-        system_prompt = """You are a document reranking assistant. Given a query and a list of document chunks, you must rank them by relevance to the query.
-
-Return your response as a JSON object with a "results" array. Each result should have:
-- "index": the original index of the document (integer)
-- "relevance_score": a score between 0.0 and 1.0 indicating relevance (float, where 1.0 is most relevant)
-
-Only return the top documents up to the requested limit, ordered by decreasing relevance score."""
-
         documents_text = ""
         for doc in documents:
             documents_text += f"Index {doc['index']}: {doc['content']}\n\n"
@@ -56,27 +66,14 @@ Only return the top documents up to the requested limit, ordered by decreasing r
 Documents to rerank:
 {documents_text.strip()}
 
-Please rank these documents by relevance to the query and return the top {top_n} results as JSON."""
-
-        messages = [
-            {"role": "system", "content": system_prompt},
-            {"role": "user", "content": user_prompt},
-        ]
+Rank these documents by relevance to the query and return the top {top_n} results as JSON."""
 
         try:
-            response = await self._client.chat(
-                model=self._model,
-                messages=messages,
-                format=RerankResponse.model_json_schema(),
-                options=OLLAMA_OPTIONS,
-            )
-
-            content = response["message"]["content"]
+            result = await self._agent.run(user_prompt)
 
-            parsed_response = RerankResponse.model_validate(json.loads(content))
             return [
-                (chunks[result.index], result.relevance_score)
-                for result in parsed_response.results[:top_n]
+                (chunks[result_item.index], result_item.relevance_score)
+                for result_item in result.output.results[:top_n]
             ]
 
         except Exception:

{haiku_rag-0.5.5.dist-info → haiku_rag-0.6.0.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: haiku.rag
-Version: 0.5.5
+Version: 0.6.0
 Summary: Retrieval Augmented Generation (RAG) with SQLite
 Author-email: Yiorgis Gozadinos <ggozadinos@gmail.com>
 License: MIT
@@ -22,6 +22,7 @@ Requires-Dist: docling>=2.15.0
 Requires-Dist: fastmcp>=2.8.1
 Requires-Dist: httpx>=0.28.1
 Requires-Dist: ollama>=0.5.3
+Requires-Dist: pydantic-ai>=0.7.2
 Requires-Dist: pydantic>=2.11.7
 Requires-Dist: python-dotenv>=1.1.0
 Requires-Dist: rich>=14.0.0
@@ -29,14 +30,8 @@ Requires-Dist: sqlite-vec>=0.1.6
 Requires-Dist: tiktoken>=0.9.0
 Requires-Dist: typer>=0.16.0
 Requires-Dist: watchfiles>=1.1.0
-Provides-Extra: anthropic
-Requires-Dist: anthropic>=0.56.0; extra == 'anthropic'
-Provides-Extra: cohere
-Requires-Dist: cohere>=5.16.1; extra == 'cohere'
 Provides-Extra: mxbai
 Requires-Dist: mxbai-rerank>=0.1.6; extra == 'mxbai'
-Provides-Extra: openai
-Requires-Dist: openai>=1.0.0; extra == 'openai'
 Provides-Extra: voyageai
 Requires-Dist: voyageai>=0.3.2; extra == 'voyageai'
 Description-Content-Type: text/markdown
@@ -51,7 +46,7 @@ Retrieval-Augmented Generation (RAG) library on SQLite.
 
 - **Local SQLite**: No external servers required
 - **Multiple embedding providers**: Ollama, VoyageAI, OpenAI
-- **Multiple QA providers**: Ollama, OpenAI, Anthropic
+- **Multiple QA providers**: Any provider/model supported by Pydantic AI
 - **Hybrid search**: Vector + full-text search with Reciprocal Rank Fusion
 - **Reranking**: Default search result reranking with MixedBread AI or Cohere
 - **Question answering**: Built-in QA agents on your documents

{haiku_rag-0.5.5.dist-info → haiku_rag-0.6.0.dist-info}/RECORD

@@ -9,22 +9,19 @@ haiku/rag/mcp.py,sha256=tMN6fNX7ZtAER1R6DL1GkC9HZozTC4HzuQs199p7icI,4551
 haiku/rag/monitor.py,sha256=r386nkhdlsU8UECwIuVwnrSlgMk3vNIuUZGNIzkZuec,2770
 haiku/rag/reader.py,sha256=qkPTMJuQ_o4sK-8zpDl9WFYe_MJ7aL_gUw6rczIpW-g,3274
 haiku/rag/utils.py,sha256=g-uNTG60iBLgkeHHuah6eVZEkX3NFLs-LZU1YnzJzLQ,2967
-haiku/rag/embeddings/__init__.py,sha256=yFBlxS0jBiVHl_rWz5kb43t6Ha132U1ZGdlIPfhzPdg,1491
+haiku/rag/embeddings/__init__.py,sha256=n7aHW3BxHlpGxU4ze4YYDOsljzFpEep8dwVE2n45JoE,1218
 haiku/rag/embeddings/base.py,sha256=NTQvuzbZPu0LBo5wAu3qGyJ4xXUaRAt1fjBO0ygWn_Y,465
 haiku/rag/embeddings/ollama.py,sha256=y6-lp0XpbnyIjoOEdtSzMdEVkU5glOwnWQ1FkpUZnpI,370
-haiku/rag/embeddings/openai.py,sha256=i4Ui5hAJkcKqJkH9L3jJo7fuGYHn07td532w-ksg_T8,431
+haiku/rag/embeddings/openai.py,sha256=iA-DewCOSip8PLU_RhEJHFHBle4DtmCCIGNfGs58Wvk,357
 haiku/rag/embeddings/voyageai.py,sha256=0hiRTIqu-bpl-4OaCtMHvWfPdgbrzhnfZJowSV8pLRA,415
-haiku/rag/qa/__init__.py,sha256=vC9S6cvZtPz-UfA_v4DMwI7eam6567BXNrUwHsMo_i8,1633
-haiku/rag/qa/anthropic.py,sha256=o0RVn7lcdYvoCUGXh551jeuoB3ANJSZ7uz2R_h_pZ2w,4321
-haiku/rag/qa/base.py,sha256=dCX14ifJW4QMCNFP_pmss9SYWM9Qm1cSWZrMl6A-2C8,3541
-haiku/rag/qa/ollama.py,sha256=3T9ciKWpCIY7jejvdrsMC_wIvGRWQEWA0AwKjOlX35M,2131
-haiku/rag/qa/openai.py,sha256=4BFc8pzFI-CTDxxKMskMxMKkacvUoRTVWI8kKntl3Jw,3718
-haiku/rag/qa/prompts.py,sha256=WTA66brySfzIkuDZ_hRQQKGx12ngIu9nUDKMNGg2-Bg,3321
+haiku/rag/qa/__init__.py,sha256=Sl7Kzrg9CuBOcMF01wc1NtQhUNWjJI0MhIHfCWrb8V4,434
+haiku/rag/qa/agent.py,sha256=r6tYKvOW4W1HxBRHH1kmzlzb1bIJcQSuHd6cG9ANqXY,2594
+haiku/rag/qa/prompts.py,sha256=xdT4cyrOrAK9UDgVqyev1wHF49jD57Bh40gx2sH4NPI,3341
 haiku/rag/reranking/__init__.py,sha256=fwC3pauteJwh9Ulm2270QvwAdwr4NMr4RUEuolC-wKU,1063
 haiku/rag/reranking/base.py,sha256=LM9yUSSJ414UgBZhFTgxGprlRqzfTe4I1vgjricz2JY,405
 haiku/rag/reranking/cohere.py,sha256=1iTdiaa8vvb6oHVB2qpWzUOVkyfUcimVSZp6Qr4aq4c,1049
 haiku/rag/reranking/mxbai.py,sha256=46sVTsTIkzIX9THgM3u8HaEmgY7evvEyB-N54JTHvK8,867
-haiku/rag/reranking/ollama.py,sha256=tCrLlNNDBCZu7J3to1gvBq-sOvN1flYEA7E3H3Jq0mU,2790
+haiku/rag/reranking/ollama.py,sha256=Q3dJepxFyB9CRCtrZvcwX-Drrpa2-8TMO7YGhxD1Qcs,2610
 haiku/rag/store/__init__.py,sha256=hq0W0DAC7ysqhWSP2M2uHX8cbG6kbr-sWHxhq6qQcY0,103
 haiku/rag/store/engine.py,sha256=cOMBToLilI1Di1qQrFzGLqtRMsuvtiX0Q5RNIEzQy9w,6232
 haiku/rag/store/models/__init__.py,sha256=s0E72zneGlowvZrFWaNxHYjOAUjgWdLxzdYsnvNRVlY,88
@@ -37,8 +34,8 @@ haiku/rag/store/repositories/document.py,sha256=ki8LiDukwU1469Yw51i0rQFvBzUQeYkF
 haiku/rag/store/repositories/settings.py,sha256=qZLXvLsErnCWL0nBQQNfRnatHzCKhtUDLvUK9k-W_fU,2463
 haiku/rag/store/upgrades/__init__.py,sha256=kKS1YWT_P-CYKhKtokOLTIFNKf9jlfjFFr8lyIMeogM,100
 haiku/rag/store/upgrades/v0_3_4.py,sha256=GLogKZdZ40NX1vBHKdOJju7fFzNUCHoEnjSZg17Hm2U,663
-haiku_rag-0.5.5.dist-info/METADATA,sha256=rponlCmspT548_0Z_YbYSp8Q2c1QQlCEXzRMx5sxPfs,4455
-haiku_rag-0.5.5.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
-haiku_rag-0.5.5.dist-info/entry_points.txt,sha256=G1U3nAkNd5YDYd4v0tuYFbriz0i-JheCsFuT9kIoGCI,48
-haiku_rag-0.5.5.dist-info/licenses/LICENSE,sha256=eXZrWjSk9PwYFNK9yUczl3oPl95Z4V9UXH7bPN46iPo,1065
-haiku_rag-0.5.5.dist-info/RECORD,,
+haiku_rag-0.6.0.dist-info/METADATA,sha256=oLxNtf0pFMyLwc9sVsiztYbrpiyVNkg0wsX0TZdUYFw,4283
+haiku_rag-0.6.0.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
+haiku_rag-0.6.0.dist-info/entry_points.txt,sha256=G1U3nAkNd5YDYd4v0tuYFbriz0i-JheCsFuT9kIoGCI,48
+haiku_rag-0.6.0.dist-info/licenses/LICENSE,sha256=eXZrWjSk9PwYFNK9yUczl3oPl95Z4V9UXH7bPN46iPo,1065
+haiku_rag-0.6.0.dist-info/RECORD,,

haiku/rag/qa/anthropic.py

@@ -1,108 +0,0 @@
-from collections.abc import Sequence
-
-try:
-    from anthropic import AsyncAnthropic  # type: ignore
-    from anthropic.types import (  # type: ignore
-        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",
-            use_citations: bool = False,
-        ):
-            super().__init__(client, model or self._model, use_citations)
-            self.tools: Sequence[ToolParam] = [
-                ToolParam(
-                    name="search_documents",
-                    description="Search the knowledge base for relevant documents. Returns a JSON array with content, score, and document_uri for each result.",
-                    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
-                                )
-
-                                context = await self._search_and_expand(
-                                    query, limit=limit
-                                )
-
-                                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/qa/base.py

@@ -1,89 +0,0 @@
-import json
-
-from haiku.rag.client import HaikuRAG
-from haiku.rag.qa.prompts import SYSTEM_PROMPT, SYSTEM_PROMPT_WITH_CITATIONS
-
-
-class QuestionAnswerAgentBase:
-    _model: str = ""
-    _system_prompt: str = SYSTEM_PROMPT
-
-    def __init__(self, client: HaikuRAG, model: str = "", use_citations: bool = False):
-        self._model = model
-        self._client = client
-        self._system_prompt = (
-            SYSTEM_PROMPT_WITH_CITATIONS if use_citations else SYSTEM_PROMPT
-        )
-
-    async def answer(self, question: str) -> str:
-        raise NotImplementedError(
-            "QABase is an abstract class. Please implement the answer method in a subclass."
-        )
-
-    async def _search_and_expand(self, query: str, limit: int = 3) -> str:
-        """Search for documents and expand context, then format as JSON"""
-        search_results = await self._client.search(query, limit=limit)
-        expanded_results = await self._client.expand_context(search_results)
-        return self._format_search_results(expanded_results)
-
-    def _format_search_results(self, search_results) -> str:
-        """Format search results as JSON list of {content, score, document_uri}"""
-        formatted_results = []
-        for chunk, score in search_results:
-            formatted_results.append(
-                {
-                    "content": chunk.content,
-                    "score": score,
-                    "document_uri": chunk.document_uri,
-                }
-            )
-        return json.dumps(formatted_results, indent=2)
-
-    tools = [
-        {
-            "type": "function",
-            "function": {
-                "name": "search_documents",
-                "description": "Search the knowledge base for relevant documents. Returns a JSON array of search results.",
-                "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"],
-                },
-                "returns": {
-                    "type": "string",
-                    "description": "JSON array of search results",
-                    "schema": {
-                        "type": "array",
-                        "items": {
-                            "type": "object",
-                            "properties": {
-                                "content": {
-                                    "type": "string",
-                                    "description": "The document text content",
-                                },
-                                "score": {
-                                    "type": "number",
-                                    "description": "Relevance score (higher is more relevant)",
-                                },
-                                "document_uri": {
-                                    "type": "string",
-                                    "description": "Source URI/path of the document",
-                                },
-                            },
-                        },
-                    },
-                },
-            },
-        }
-    ]

haiku/rag/qa/ollama.py

@@ -1,60 +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": 16384}
-
-
-class QuestionAnswerOllamaAgent(QuestionAnswerAgentBase):
-    def __init__(
-        self,
-        client: HaikuRAG,
-        model: str = Config.QA_MODEL,
-        use_citations: bool = False,
-    ):
-        super().__init__(client, model or self._model, use_citations)
-
-    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))
-
-                        context = await self._search_and_expand(query, limit=limit)
-                        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/qa/openai.py

@@ -1,97 +0,0 @@
-from collections.abc import Sequence
-
-try:
-    from openai import AsyncOpenAI  # type: ignore
-    from openai.types.chat import (  # type: ignore
-        ChatCompletionAssistantMessageParam,
-        ChatCompletionMessageParam,
-        ChatCompletionSystemMessageParam,
-        ChatCompletionToolMessageParam,
-        ChatCompletionUserMessageParam,
-    )
-    from openai.types.chat.chat_completion_tool_param import (  # type: ignore
-        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",
-            use_citations: bool = False,
-        ):
-            super().__init__(client, model or self._model, use_citations)
-            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))
-
-                            context = await self._search_and_expand(query, limit=limit)
-
-                            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