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 pydantic import BaseModel
|
||||||
|
|
||||||
from haiku.rag.client import HaikuRAG
|
from haiku.rag.client import HaikuRAG
|
||||||
|
from haiku.rag.config import Config
|
||||||
|
from haiku.rag.research.models import ResearchReport
|
||||||
|
|
||||||
|
|
||||||
class SearchResult(BaseModel):
|
class SearchResult(BaseModel):
|
||||||
|
|
@ -153,4 +155,101 @@ def create_mcp_server(db_path: Path) -> FastMCP:
|
||||||
except Exception:
|
except Exception:
|
||||||
return False
|
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
|
return mcp
|
||||||
|
|
|
||||||
Loading…
Reference in a new issue