From e7aca6292a4b06f2cb33df9e9fdbb24ecc1c30bd Mon Sep 17 00:00:00 2001 From: Yiorgis Gozadinos Date: Wed, 8 Oct 2025 09:13:21 +0300 Subject: [PATCH 1/2] 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 From 43dac018c00b825933a3704970431de6c1d998a6 Mon Sep 17 00:00:00 2001 From: Yiorgis Gozadinos Date: Wed, 8 Oct 2025 09:13:43 +0300 Subject: [PATCH 2/2] Test mcp tools --- tests/test_mcp.py | 322 ++++++++++++++++++++++++++++++++++++++++++++++ 1 file changed, 322 insertions(+) create mode 100644 tests/test_mcp.py diff --git a/tests/test_mcp.py b/tests/test_mcp.py new file mode 100644 index 00000000..02b1b275 --- /dev/null +++ b/tests/test_mcp.py @@ -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()