atlas-chat 0.1.0__py3-none-any.whl

atlas/tests/test_tool_planner.py

@@ -0,0 +1,300 @@
+"""Tests for the tool_planner MCP server."""
+
+import base64
+import importlib.util
+from pathlib import Path
+from unittest.mock import AsyncMock, MagicMock
+
+import pytest
+
+# Load tool_planner/main.py as a uniquely-named module to avoid colliding
+# with backend/main.py which is already on sys.path.
+_tool_planner_path = Path(__file__).parent.parent / "mcp" / "tool_planner" / "main.py"
+_spec = importlib.util.spec_from_file_location("tool_planner_main", _tool_planner_path)
+_mod = importlib.util.module_from_spec(_spec)
+_spec.loader.exec_module(_mod)
+
+format_tools_for_llm = _mod.format_tools_for_llm
+build_planning_prompt = _mod.build_planning_prompt
+_build_artifact_response = _mod._build_artifact_response
+# plan_with_tools is wrapped by @mcp.tool into a FunctionTool; get the raw fn
+_plan_with_tools_fn = _mod.plan_with_tools.fn
+
+
+def _decode_artifact(result: dict) -> str:
+    """Decode the base64 script content from an artifact response."""
+    return base64.b64decode(result["artifacts"][0]["b64"]).decode("utf-8")
+
+
+async def _call_plan_with_tools(**kwargs):
+    """Call the unwrapped plan_with_tools async function."""
+    return await _plan_with_tools_fn(**kwargs)
+
+
+# ---------------------------------------------------------------------------
+# format_tools_for_llm tests
+# ---------------------------------------------------------------------------
+
+class TestFormatToolsForLlm:
+    """Tests for the format_tools_for_llm helper."""
+
+    def test_empty_data_returns_no_tools_message(self):
+        result = format_tools_for_llm({})
+        assert result == "(No tools available)"
+
+    def test_empty_servers_returns_no_tools_message(self):
+        result = format_tools_for_llm({"available_servers": []})
+        assert result == "(No tools available)"
+
+    def test_single_server_single_tool(self):
+        data = {
+            "available_servers": [
+                {
+                    "server_name": "calculator",
+                    "description": "Math calculator",
+                    "tools": [
+                        {
+                            "name": "calculator_evaluate",
+                            "description": "Evaluate math expressions",
+                            "parameters": {
+                                "type": "object",
+                                "properties": {
+                                    "expression": {
+                                        "type": "string",
+                                        "description": "Math expression",
+                                    }
+                                },
+                                "required": ["expression"],
+                            },
+                        }
+                    ],
+                }
+            ]
+        }
+        result = format_tools_for_llm(data)
+        assert "Server: calculator (Math calculator)" in result
+        assert "Tool: calculator_evaluate" in result
+        assert "Description: Evaluate math expressions" in result
+        assert "expression (string, required): Math expression" in result
+
+    def test_multiple_servers_and_tools(self):
+        data = {
+            "available_servers": [
+                {
+                    "server_name": "calc",
+                    "description": "",
+                    "tools": [
+                        {"name": "calc_add", "description": "Add numbers", "parameters": {}},
+                    ],
+                },
+                {
+                    "server_name": "pptx",
+                    "description": "Slides",
+                    "tools": [
+                        {"name": "pptx_create", "description": "Create slides", "parameters": {}},
+                        {"name": "pptx_export", "description": "Export slides", "parameters": {}},
+                    ],
+                },
+            ]
+        }
+        result = format_tools_for_llm(data)
+        assert "Server: calc" in result
+        assert "Server: pptx (Slides)" in result
+        assert "Tool: calc_add" in result
+        assert "Tool: pptx_create" in result
+        assert "Tool: pptx_export" in result
+
+    def test_optional_and_required_params(self):
+        data = {
+            "available_servers": [
+                {
+                    "server_name": "srv",
+                    "description": "",
+                    "tools": [
+                        {
+                            "name": "srv_tool",
+                            "description": "",
+                            "parameters": {
+                                "type": "object",
+                                "properties": {
+                                    "required_param": {"type": "string"},
+                                    "optional_param": {"type": "number"},
+                                },
+                                "required": ["required_param"],
+                            },
+                        }
+                    ],
+                }
+            ]
+        }
+        result = format_tools_for_llm(data)
+        assert "required_param (string, required)" in result
+        assert "optional_param (number, optional)" in result
+
+    def test_underscore_prefixed_params_excluded(self):
+        data = {
+            "available_servers": [
+                {
+                    "server_name": "srv",
+                    "description": "",
+                    "tools": [
+                        {
+                            "name": "srv_tool",
+                            "description": "",
+                            "parameters": {
+                                "type": "object",
+                                "properties": {
+                                    "task": {"type": "string"},
+                                    "_mcp_data": {"type": "object"},
+                                },
+                            },
+                        }
+                    ],
+                }
+            ]
+        }
+        result = format_tools_for_llm(data)
+        assert "_mcp_data" not in result
+        assert "task (string" in result
+
+
+# ---------------------------------------------------------------------------
+# build_planning_prompt tests
+# ---------------------------------------------------------------------------
+
+class TestBuildPlanningPrompt:
+    """Tests for the build_planning_prompt helper."""
+
+    def test_includes_task(self):
+        result = build_planning_prompt("create a presentation", "some tools")
+        assert "Task: create a presentation" in result
+
+    def test_includes_tools_reference(self):
+        result = build_planning_prompt("task", "Server: calc\n  Tool: calc_add")
+        assert "Server: calc" in result
+        assert "Tool: calc_add" in result
+
+    def test_includes_cli_usage_instructions(self):
+        result = build_planning_prompt("task", "tools")
+        assert "python atlas_chat_cli.py" in result
+        assert "--tools tool_name" in result
+        assert "-o result.txt" in result
+
+    def test_includes_loop_example(self):
+        result = build_planning_prompt("task", "tools")
+        assert "for item" in result
+
+
+# ---------------------------------------------------------------------------
+# plan_with_tools tests
+# ---------------------------------------------------------------------------
+
+class TestBuildArtifactResponse:
+    """Tests for the _build_artifact_response helper."""
+
+    def test_returns_required_keys(self):
+        result = _build_artifact_response("#!/bin/bash\necho hi", "greet user")
+        assert "results" in result
+        assert "artifacts" in result
+        assert "display" in result
+
+    def test_artifact_contains_base64_script(self):
+        script = "#!/bin/bash\nset -e\necho hello"
+        result = _build_artifact_response(script, "say hello")
+        artifact = result["artifacts"][0]
+        decoded = base64.b64decode(artifact["b64"]).decode("utf-8")
+        assert decoded == script
+
+    def test_artifact_mime_and_viewer(self):
+        result = _build_artifact_response("echo x", "task")
+        artifact = result["artifacts"][0]
+        assert artifact["mime"] == "application/x-sh"
+        assert artifact["viewer"] == "code"
+
+    def test_filename_derived_from_task(self):
+        result = _build_artifact_response("echo x", "Create a PowerPoint about dogs")
+        filename = result["artifacts"][0]["name"]
+        assert filename.endswith(".sh")
+        assert "create" in filename
+        assert " " not in filename
+
+    def test_display_opens_canvas_with_code_hint(self):
+        result = _build_artifact_response("echo x", "task")
+        display = result["display"]
+        assert display["open_canvas"] is True
+        assert display["viewer_hint"] == "code"
+        assert display["primary_file"].endswith(".sh")
+
+
+class TestPlanWithTools:
+    """Tests for the plan_with_tools tool function."""
+
+    @pytest.mark.asyncio
+    async def test_without_ctx_returns_artifact(self):
+        result = await _call_plan_with_tools(
+            task="test task", _mcp_data={"available_servers": []}
+        )
+        assert result["results"]["operation"] == "plan_with_tools"
+        script = _decode_artifact(result)
+        assert "Sampling unavailable" in script
+        assert "test task" in script
+
+    @pytest.mark.asyncio
+    async def test_without_mcp_data_still_works(self):
+        result = await _call_plan_with_tools(task="do something")
+        script = _decode_artifact(result)
+        assert "Sampling unavailable" in script
+        assert "No tools available" in script
+
+    @pytest.mark.asyncio
+    async def test_with_mocked_ctx_sample(self):
+        mock_ctx = MagicMock()
+        mock_result = MagicMock()
+        mock_result.text = "#!/bin/bash\nset -e\npython atlas_chat_cli.py 'hello' --tools calc_add"
+        mock_ctx.sample = AsyncMock(return_value=mock_result)
+        mock_ctx.report_progress = AsyncMock()
+
+        mcp_data = {
+            "available_servers": [
+                {
+                    "server_name": "calc",
+                    "description": "Calculator",
+                    "tools": [
+                        {
+                            "name": "calc_add",
+                            "description": "Add numbers",
+                            "parameters": {
+                                "type": "object",
+                                "properties": {"a": {"type": "number"}},
+                            },
+                        }
+                    ],
+                }
+            ]
+        }
+
+        result = await _call_plan_with_tools(
+            task="add two numbers", _mcp_data=mcp_data, ctx=mock_ctx
+        )
+
+        script = _decode_artifact(result)
+        assert "atlas_chat_cli.py" in script
+        assert result["artifacts"][0]["name"].endswith(".sh")
+        mock_ctx.sample.assert_awaited_once()
+
+        call_kwargs = mock_ctx.sample.call_args
+        assert call_kwargs.kwargs["temperature"] == 0.3
+        assert call_kwargs.kwargs["max_tokens"] == 10000
+        assert "task planner" in call_kwargs.kwargs["system_prompt"].lower()
+
+    @pytest.mark.asyncio
+    async def test_sample_returns_none_text(self):
+        mock_ctx = MagicMock()
+        mock_result = MagicMock()
+        mock_result.text = None
+        mock_ctx.sample = AsyncMock(return_value=mock_result)
+        mock_ctx.report_progress = AsyncMock()
+
+        result = await _call_plan_with_tools(task="test", _mcp_data={}, ctx=mock_ctx)
+        script = _decode_artifact(result)
+        assert "Unable to generate plan" in script