danvics/haiku.rag
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions

Merge pull request #93 from ggozad/feat/mcp-enhancements

Browse source 
Add ask, deep ask, research as MCP tools.
This commit is contained in:
Yiorgis Gozadinos 2025-10-08 09:22:04 +03:00 committed by GitHub
commit 09528221b5
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
2 changed files with 421 additions and 0 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

99
src/haiku/rag/mcp.py
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

322
tests/test_mcp.py Normal file
View file

@ -0,0 +1,322 @@
import tempfile
from pathlib import Path
from unittest.mock import AsyncMock, patch
import pytest
from haiku.rag.mcp import create_mcp_server
from haiku.rag.research.models import ResearchReport
from haiku.rag.store.models.document import Document
@pytest.mark.asyncio
async def test_mcp_add_document_from_file():
"""Test add_document_from_file tool is properly wired."""
with tempfile.TemporaryDirectory() as temp_dir:
db_path = Path(temp_dir) / "test.lancedb"
mcp = create_mcp_server(db_path)
mock_doc = Document(content="test", uri="file:///test.txt")
mock_doc.id = "doc123"
with patch("haiku.rag.mcp.HaikuRAG") as mock_rag_class:
mock_rag = AsyncMock()
mock_rag.create_document_from_source = AsyncMock(return_value=mock_doc)
mock_rag_class.return_value.__aenter__ = AsyncMock(return_value=mock_rag)
mock_rag_class.return_value.__aexit__ = AsyncMock(return_value=None)
tools = await mcp.get_tools()
add_file_tool = next(
t for t in tools.values() if t.name == "add_document_from_file"
)
result = await add_file_tool.fn(file_path="/test.txt") # type: ignore[attr-defined]
assert result == "doc123"
mock_rag.create_document_from_source.assert_called_once()
@pytest.mark.asyncio
async def test_mcp_ask_question():
"""Test ask_question tool is properly wired."""
with tempfile.TemporaryDirectory() as temp_dir:
db_path = Path(temp_dir) / "test.lancedb"
mcp = create_mcp_server(db_path)
with patch("haiku.rag.mcp.HaikuRAG") as mock_rag_class:
mock_rag = AsyncMock()
mock_rag.ask = AsyncMock(return_value="This is the answer")
mock_rag_class.return_value.__aenter__ = AsyncMock(return_value=mock_rag)
mock_rag_class.return_value.__aexit__ = AsyncMock(return_value=None)
tools = await mcp.get_tools()
ask_tool = next(t for t in tools.values() if t.name == "ask_question")
result = await ask_tool.fn( # type: ignore[attr-defined]
question="What is this?", cite=False, deep=False
)
assert result == "This is the answer"
mock_rag.ask.assert_called_once_with("What is this?", cite=False)
@pytest.mark.asyncio
async def test_mcp_add_document_from_url():
"""Test add_document_from_url tool is properly wired."""
with tempfile.TemporaryDirectory() as temp_dir:
db_path = Path(temp_dir) / "test.lancedb"
mcp = create_mcp_server(db_path)
mock_doc = Document(content="test", uri="https://example.com")
mock_doc.id = "doc456"
with patch("haiku.rag.mcp.HaikuRAG") as mock_rag_class:
mock_rag = AsyncMock()
mock_rag.create_document_from_source = AsyncMock(return_value=mock_doc)
mock_rag_class.return_value.__aenter__ = AsyncMock(return_value=mock_rag)
mock_rag_class.return_value.__aexit__ = AsyncMock(return_value=None)
tools = await mcp.get_tools()
add_url_tool = next(
t for t in tools.values() if t.name == "add_document_from_url"
)
result = await add_url_tool.fn(url="https://example.com") # type: ignore[attr-defined]
assert result == "doc456"
mock_rag.create_document_from_source.assert_called_once()
@pytest.mark.asyncio
async def test_mcp_add_document_from_text():
"""Test add_document_from_text tool is properly wired."""
with tempfile.TemporaryDirectory() as temp_dir:
db_path = Path(temp_dir) / "test.lancedb"
mcp = create_mcp_server(db_path)
mock_doc = Document(content="test content", uri="text://test")
mock_doc.id = "doc789"
with patch("haiku.rag.mcp.HaikuRAG") as mock_rag_class:
mock_rag = AsyncMock()
mock_rag.create_document = AsyncMock(return_value=mock_doc)
mock_rag_class.return_value.__aenter__ = AsyncMock(return_value=mock_rag)
mock_rag_class.return_value.__aexit__ = AsyncMock(return_value=None)
tools = await mcp.get_tools()
add_text_tool = next(
t for t in tools.values() if t.name == "add_document_from_text"
)
result = await add_text_tool.fn( # type: ignore[attr-defined]
content="test content", uri="text://test"
)
assert result == "doc789"
mock_rag.create_document.assert_called_once()
@pytest.mark.asyncio
async def test_mcp_search_documents():
"""Test search_documents tool is properly wired."""
with tempfile.TemporaryDirectory() as temp_dir:
db_path = Path(temp_dir) / "test.lancedb"
mcp = create_mcp_server(db_path)
from haiku.rag.store.models.chunk import Chunk
mock_chunk1 = Chunk(content="Result 1", document_id="doc1")
mock_chunk2 = Chunk(content="Result 2", document_id="doc2")
mock_results = [(mock_chunk1, 0.9), (mock_chunk2, 0.8)]
with patch("haiku.rag.mcp.HaikuRAG") as mock_rag_class:
mock_rag = AsyncMock()
mock_rag.search = AsyncMock(return_value=mock_results)
mock_rag_class.return_value.__aenter__ = AsyncMock(return_value=mock_rag)
mock_rag_class.return_value.__aexit__ = AsyncMock(return_value=None)
tools = await mcp.get_tools()
search_tool = next(
t for t in tools.values() if t.name == "search_documents"
)
result = await search_tool.fn(query="test query", limit=5) # type: ignore[attr-defined]
assert len(result) == 2
assert result[0].document_id == "doc1"
assert result[0].content == "Result 1"
assert result[0].score == 0.9
assert result[1].document_id == "doc2"
mock_rag.search.assert_called_once_with("test query", 5)
@pytest.mark.asyncio
async def test_mcp_get_document():
"""Test get_document tool is properly wired."""
with tempfile.TemporaryDirectory() as temp_dir:
db_path = Path(temp_dir) / "test.lancedb"
mcp = create_mcp_server(db_path)
from datetime import UTC, datetime
mock_doc = Document(content="test", uri="file:///test.txt", title="Test Doc")
mock_doc.id = "doc123"
mock_doc.created_at = datetime(2024, 1, 1, tzinfo=UTC)
mock_doc.updated_at = datetime(2024, 1, 2, tzinfo=UTC)
with patch("haiku.rag.mcp.HaikuRAG") as mock_rag_class:
mock_rag = AsyncMock()
mock_rag.get_document_by_id = AsyncMock(return_value=mock_doc)
mock_rag_class.return_value.__aenter__ = AsyncMock(return_value=mock_rag)
mock_rag_class.return_value.__aexit__ = AsyncMock(return_value=None)
tools = await mcp.get_tools()
get_tool = next(t for t in tools.values() if t.name == "get_document")
result = await get_tool.fn(document_id="doc123") # type: ignore[attr-defined]
assert result is not None
assert result.id == "doc123"
assert result.content == "test"
assert result.title == "Test Doc"
mock_rag.get_document_by_id.assert_called_once_with("doc123")
@pytest.mark.asyncio
async def test_mcp_list_documents():
"""Test list_documents tool is properly wired."""
with tempfile.TemporaryDirectory() as temp_dir:
db_path = Path(temp_dir) / "test.lancedb"
mcp = create_mcp_server(db_path)
from datetime import UTC, datetime
mock_doc1 = Document(content="test1", uri="file:///test1.txt")
mock_doc1.id = "doc1"
mock_doc1.created_at = datetime(2024, 1, 1, tzinfo=UTC)
mock_doc1.updated_at = datetime(2024, 1, 1, tzinfo=UTC)
mock_doc2 = Document(content="test2", uri="file:///test2.txt")
mock_doc2.id = "doc2"
mock_doc2.created_at = datetime(2024, 1, 2, tzinfo=UTC)
mock_doc2.updated_at = datetime(2024, 1, 2, tzinfo=UTC)
with patch("haiku.rag.mcp.HaikuRAG") as mock_rag_class:
mock_rag = AsyncMock()
mock_rag.list_documents = AsyncMock(return_value=[mock_doc1, mock_doc2])
mock_rag_class.return_value.__aenter__ = AsyncMock(return_value=mock_rag)
mock_rag_class.return_value.__aexit__ = AsyncMock(return_value=None)
tools = await mcp.get_tools()
list_tool = next(t for t in tools.values() if t.name == "list_documents")
result = await list_tool.fn(limit=10, offset=0) # type: ignore[attr-defined]
assert len(result) == 2
assert result[0].id == "doc1"
assert result[1].id == "doc2"
mock_rag.list_documents.assert_called_once_with(10, 0)
@pytest.mark.asyncio
async def test_mcp_delete_document():
"""Test delete_document tool is properly wired."""
with tempfile.TemporaryDirectory() as temp_dir:
db_path = Path(temp_dir) / "test.lancedb"
mcp = create_mcp_server(db_path)
with patch("haiku.rag.mcp.HaikuRAG") as mock_rag_class:
mock_rag = AsyncMock()
mock_rag.delete_document = AsyncMock(return_value=True)
mock_rag_class.return_value.__aenter__ = AsyncMock(return_value=mock_rag)
mock_rag_class.return_value.__aexit__ = AsyncMock(return_value=None)
tools = await mcp.get_tools()
delete_tool = next(t for t in tools.values() if t.name == "delete_document")
result = await delete_tool.fn(document_id="doc123") # type: ignore[attr-defined]
assert result is True
mock_rag.delete_document.assert_called_once_with("doc123")
@pytest.mark.asyncio
async def test_mcp_ask_question_deep():
"""Test ask_question tool with deep=True is properly wired."""
with tempfile.TemporaryDirectory() as temp_dir:
db_path = Path(temp_dir) / "test.lancedb"
mcp = create_mcp_server(db_path)
with (
patch("haiku.rag.mcp.HaikuRAG") as mock_rag_class,
patch("haiku.rag.qa.deep.graph.build_deep_qa_graph") as mock_graph_builder,
):
mock_rag = AsyncMock()
mock_rag_class.return_value.__aenter__ = AsyncMock(return_value=mock_rag)
mock_rag_class.return_value.__aexit__ = AsyncMock(return_value=None)
mock_graph = AsyncMock()
mock_result = AsyncMock()
mock_result.output.answer = "Deep answer"
mock_graph.run = AsyncMock(return_value=mock_result)
mock_graph_builder.return_value = mock_graph
tools = await mcp.get_tools()
ask_tool = next(t for t in tools.values() if t.name == "ask_question")
result = await ask_tool.fn( # type: ignore[attr-defined]
question="Deep question?", cite=True, deep=True
)
assert result == "Deep answer"
mock_graph.run.assert_called_once()
@pytest.mark.asyncio
async def test_mcp_research_question():
"""Test research_question tool is properly wired."""
with tempfile.TemporaryDirectory() as temp_dir:
db_path = Path(temp_dir) / "test.lancedb"
mcp = create_mcp_server(db_path)
mock_report = ResearchReport(
title="Research Title",
executive_summary="Summary",
main_findings=["Finding 1"],
conclusions=["Conclusion 1"],
recommendations=["Recommendation 1"],
sources_summary="Sources used",
)
with (
patch("haiku.rag.mcp.HaikuRAG") as mock_rag_class,
patch(
"haiku.rag.research.graph.build_research_graph"
) as mock_graph_builder,
):
mock_rag = AsyncMock()
mock_rag_class.return_value.__aenter__ = AsyncMock(return_value=mock_rag)
mock_rag_class.return_value.__aexit__ = AsyncMock(return_value=None)
mock_graph = AsyncMock()
mock_result = AsyncMock()
mock_result.output = mock_report
mock_graph.run = AsyncMock(return_value=mock_result)
mock_graph_builder.return_value = mock_graph
tools = await mcp.get_tools()
research_tool = next(
t for t in tools.values() if t.name == "research_question"
)
result = await research_tool.fn( # type: ignore[attr-defined]
question="Research question?",
max_iterations=1,
confidence_threshold=0.5,
max_concurrency=1,
)
assert result is not None
assert result.title == "Research Title"
assert result.executive_summary == "Summary"
mock_graph.run.assert_called_once()