diff --git a/haiku_rag_slim/haiku/rag/agents/rlm/docker_sandbox.py b/haiku_rag_slim/haiku/rag/agents/rlm/docker_sandbox.py index 9f77fe51..7d91f7ca 100644 --- a/haiku_rag_slim/haiku/rag/agents/rlm/docker_sandbox.py +++ b/haiku_rag_slim/haiku/rag/agents/rlm/docker_sandbox.py @@ -24,7 +24,7 @@ class SandboxResult: success: bool -class DockerSandbox: +class DockerSandbox: # pragma: no cover """Execute code in a persistent Docker container. Use as an async context manager to manage container lifecycle: diff --git a/tests/test_app.py b/tests/test_app.py index 92c5db2c..d34d83b1 100644 --- a/tests/test_app.py +++ b/tests/test_app.py @@ -669,3 +669,57 @@ def test_migrate_closes_store_on_exception(tmp_path): app.migrate() mock_store.close.assert_called_once() + + +@pytest.mark.asyncio +async def test_rlm(app: HaikuRAGApp, monkeypatch): + """Test rlm method calls client.rlm and prints results.""" + from haiku.rag.agents.rlm.models import RLMResult + + mock_result = RLMResult( + answer="The total is 42.", + program="result = sum(values)\nprint(result)", + ) + + mock_client = AsyncMock() + mock_client.rlm = AsyncMock(return_value=mock_result) + mock_client.__aenter__.return_value = mock_client + + mock_print = MagicMock() + monkeypatch.setattr(app.console, "print", mock_print) + + with patch("haiku.rag.app.HaikuRAG", return_value=mock_client): + await app.rlm("What is the total?") + + mock_client.rlm.assert_called_once_with( + "What is the total?", documents=None, filter=None + ) + calls = [str(c) for c in mock_print.call_args_list] + assert any("Question" in c for c in calls) + assert any("Program" in c for c in calls) + assert any("Answer" in c for c in calls) + + +@pytest.mark.asyncio +async def test_rlm_with_document_and_filter(app: HaikuRAGApp, monkeypatch): + """Test rlm method passes document and filter to client.""" + from haiku.rag.agents.rlm.models import RLMResult + + mock_result = RLMResult( + answer="Answer with filter", + program="print('filtered')", + ) + + mock_client = AsyncMock() + mock_client.rlm = AsyncMock(return_value=mock_result) + mock_client.__aenter__.return_value = mock_client + + mock_print = MagicMock() + monkeypatch.setattr(app.console, "print", mock_print) + + with patch("haiku.rag.app.HaikuRAG", return_value=mock_client): + await app.rlm("What is it?", document="doc-123", filter="uri LIKE '%test%'") + + mock_client.rlm.assert_called_once_with( + "What is it?", documents=["doc-123"], filter="uri LIKE '%test%'" + ) diff --git a/tests/test_mcp.py b/tests/test_mcp.py index dee34e18..0f8a56bb 100644 --- a/tests/test_mcp.py +++ b/tests/test_mcp.py @@ -307,3 +307,71 @@ async def test_mcp_research_question(): assert result.title == "Research Title" assert result.executive_summary == "Summary" mock_graph.run.assert_called_once() + + +@pytest.mark.asyncio +async def test_mcp_rlm_question(): + """Test rlm_question 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.agents.rlm.models import RLMResult + + mock_result = RLMResult( + answer="The total is 42.", + program="result = sum(values)\nprint(result)", + ) + + with patch("haiku.rag.mcp.HaikuRAG") as mock_rag_class: + mock_rag = AsyncMock() + mock_rag.rlm = AsyncMock(return_value=mock_result) + 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() + rlm_tool = next(t for t in tools.values() if t.name == "rlm_question") + + result = await rlm_tool.fn(question="What is the total?") + + assert result == "The total is 42." + mock_rag.rlm.assert_called_once_with( + "What is the total?", documents=None, filter=None + ) + + +@pytest.mark.asyncio +async def test_mcp_rlm_question_with_document_and_filter(): + """Test rlm_question tool with document and filter parameters.""" + with tempfile.TemporaryDirectory() as temp_dir: + db_path = Path(temp_dir) / "test.lancedb" + mcp = create_mcp_server(db_path) + + from haiku.rag.agents.rlm.models import RLMResult + + mock_result = RLMResult( + answer="Filtered answer", + program="print('filtered')", + ) + + with patch("haiku.rag.mcp.HaikuRAG") as mock_rag_class: + mock_rag = AsyncMock() + mock_rag.rlm = AsyncMock(return_value=mock_result) + 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() + rlm_tool = next(t for t in tools.values() if t.name == "rlm_question") + + result = await rlm_tool.fn( + question="Analyze this", + document="doc-123", + filter="uri LIKE '%test%'", + ) + + assert result == "Filtered answer" + mock_rag.rlm.assert_called_once_with( + "Analyze this", + documents=["doc-123"], + filter="uri LIKE '%test%'", + )