Add ask, research as tools to mcp

This commit is contained in:
Yiorgis Gozadinos 2025-10-08 09:13:21 +03:00
parent 8323cbedc3
commit e7aca6292a
No known key found for this signature in database

View file

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