Add ask, research as tools to mcp
This commit is contained in:
parent
8323cbedc3
commit
e7aca6292a
1 changed files with 99 additions and 0 deletions
|
|
@ -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
|
||||
|
|
|
|||
Loading…
Reference in a new issue