From e7aca6292a4b06f2cb33df9e9fdbb24ecc1c30bd Mon Sep 17 00:00:00 2001
From: Yiorgis Gozadinos <ggozadinos@gmail.com>
Date: Wed, 8 Oct 2025 09:13:21 +0300
Subject: [PATCH] Add ask, research as tools to mcp

---
 src/haiku/rag/mcp.py | 99 ++++++++++++++++++++++++++++++++++++++++++++
 1 file changed, 99 insertions(+)

diff --git a/src/haiku/rag/mcp.py b/src/haiku/rag/mcp.py
index 654cf8e8..379086c8 100644
--- a/src/haiku/rag/mcp.py
+++ b/src/haiku/rag/mcp.py
@@ -5,6 +5,8 @@ from fastmcp import FastMCP
 from pydantic import BaseModel
 
 from haiku.rag.client import HaikuRAG
+from haiku.rag.config import Config
+from haiku.rag.research.models import ResearchReport
 
 
 class SearchResult(BaseModel):
@@ -153,4 +155,101 @@ def create_mcp_server(db_path: Path) -> FastMCP:
         except Exception:
             return False
 
+    @mcp.tool()
+    async def ask_question(
+        question: str,
+        cite: bool = False,
+        deep: bool = False,
+    ) -> str:
+        """Ask a question using the QA agent.
+
+        Args:
+            question: The question to ask.
+            cite: Whether to include citations in the response.
+            deep: Use deep multi-agent QA for complex questions that require decomposition.
+
+        Returns:
+            The answer as a string.
+        """
+        try:
+            async with HaikuRAG(db_path) as rag:
+                if deep:
+                    from haiku.rag.config import Config
+                    from haiku.rag.qa.deep.dependencies import DeepQAContext
+                    from haiku.rag.qa.deep.graph import build_deep_qa_graph
+                    from haiku.rag.qa.deep.nodes import DeepQAPlanNode
+                    from haiku.rag.qa.deep.state import DeepQADeps, DeepQAState
+
+                    graph = build_deep_qa_graph()
+                    context = DeepQAContext(
+                        original_question=question, use_citations=cite
+                    )
+                    state = DeepQAState(context=context)
+                    deps = DeepQADeps(client=rag)
+
+                    start_node = DeepQAPlanNode(
+                        provider=Config.QA_PROVIDER,
+                        model=Config.QA_MODEL,
+                    )
+
+                    result = await graph.run(
+                        start_node=start_node, state=state, deps=deps
+                    )
+                    answer = result.output.answer
+                else:
+                    answer = await rag.ask(question, cite=cite)
+                return answer
+        except Exception as e:
+            return f"Error answering question: {e!s}"
+
+    @mcp.tool()
+    async def research_question(
+        question: str,
+        max_iterations: int = 3,
+        confidence_threshold: float = 0.8,
+        max_concurrency: int = 1,
+    ) -> ResearchReport | None:
+        """Run multi-agent research to investigate a complex question.
+
+        The research process uses multiple agents to plan, search, evaluate, and synthesize
+        information iteratively until confidence threshold is met or max iterations reached.
+
+        Args:
+            question: The research question to investigate.
+            max_iterations: Maximum search/analyze iterations (default: 3).
+            confidence_threshold: Minimum confidence score (0-1) to stop early (default: 0.8).
+            max_concurrency: Maximum concurrent searches per iteration (default: 1).
+
+        Returns:
+            A research report with findings, or None if an error occurred.
+        """
+        try:
+            from haiku.rag.graph.nodes.plan import PlanNode
+            from haiku.rag.research.dependencies import ResearchContext
+            from haiku.rag.research.graph import build_research_graph
+            from haiku.rag.research.state import ResearchDeps, ResearchState
+
+            async with HaikuRAG(db_path) as rag:
+                graph = build_research_graph()
+                state = ResearchState(
+                    context=ResearchContext(original_question=question),
+                    max_iterations=max_iterations,
+                    confidence_threshold=confidence_threshold,
+                    max_concurrency=max_concurrency,
+                )
+                deps = ResearchDeps(client=rag)
+
+                result = await graph.run(
+                    PlanNode(
+                        provider=Config.RESEARCH_PROVIDER or Config.QA_PROVIDER,
+                        model=Config.RESEARCH_MODEL or Config.QA_MODEL,
+                    ),
+                    state=state,
+                    deps=deps,
+                )
+
+                return result.output
+        except Exception:
+            return None
+
     return mcp