fix: sevefal important fixes to MCP compiler

Author: evilsocket

nerve/tools/mcp/client.py

@@ -10,6 +10,8 @@
 
 
 class Client:
+    name: str = ""
+
     def __init__(self, name: str, server: Configuration.MCPServer):
         self.name = name
         self.server = server

nerve/tools/mcp/compiler.py

@@ -3,6 +3,7 @@
 import typing as t
 
 import jinja2
+from mcp import Tool
 from pydantic import create_model
 
 from nerve.models import Configuration
@@ -81,36 +82,44 @@ def _get_python_type(
         return None, type_mapping.get(schema_type, t.Any)  # type: ignore
 
 
+async def create_function_body(client: Client, mcp_tool: Tool) -> tuple[str, dict[str, t.Any]]:
+    typed_args = []
+    client_name = f"nerve_mcp_{client.name}_client"
+    type_defs = {client_name: client}
+
+    for name, arg_props in mcp_tool.inputSchema.get("properties", {}).items():
+        args_def, arg_type = _get_python_type(name, arg_props)
+        typed_args.append(
+            {"name": name, "type": _stringify_type(arg_type), "description": arg_props.get("description", "")}
+        )
+        if args_def:
+            type_defs.update(args_def)
+
+    # load the template from the same directory as this script
+    template_path = os.path.join(os.path.dirname(__file__), "body.j2")
+    with open(template_path) as f:
+        template_content = f.read()
+
+    return (
+        jinja2.Environment()
+        .from_string(template_content)
+        .render(client_name=client_name, tool=mcp_tool, arguments=typed_args),
+        type_defs,
+    )
+
+
 async def get_tools_from_mcp(name: str, server: Configuration.MCPServer) -> list[t.Callable[..., t.Any]]:
     # connect and list tools
     client = Client(name, server)
     mpc_tools = await client.tools()
-    client_name = f"nerve_mcp_{name}_client"
     compiled_tools = []
 
     for mcp_tool in mpc_tools:
-        typed_args = []
-        type_defs = {client_name: client}
-
-        for name, arg_props in mcp_tool.inputSchema.get("properties", {}).items():
-            args_def, arg_type = _get_python_type(name, arg_props)
-            typed_args.append({"name": name, "type": _stringify_type(arg_type)})
-            if args_def:
-                type_defs.update(args_def)
-
-        # load the template from the same directory as this script
-        template_path = os.path.join(os.path.dirname(__file__), "body.j2")
-        with open(template_path) as f:
-            template_content = f.read()
-
-        func_body = (
-            jinja2.Environment()
-            .from_string(template_content)
-            .render(client_name=client_name, tool=mcp_tool, arguments=typed_args)
-        )
+        func_body, type_defs = await create_function_body(client, mcp_tool)
+
         # print(func_body)
         exec(func_body, type_defs)
 
-        tool_fn = wrap_tool_function(type_defs[mcp_tool.name])  # type: ignore
+        tool_fn = wrap_tool_function(type_defs[mcp_tool.name])
         compiled_tools.append(tool_fn)
     return compiled_tools

nerve/tools/mcp/compilter_test.py

@@ -0,0 +1,120 @@
+import unittest
+from unittest.mock import MagicMock
+
+from mcp import Tool
+
+from nerve.tools.mcp.client import Client
+from nerve.tools.mcp.compiler import create_function_body
+
+
+class TestCreateFunctionBody(unittest.IsolatedAsyncioTestCase):
+    async def test_create_function_body_with_string_argument(self) -> None:
+        # Mock the client
+        client = MagicMock(spec=Client)
+
+        # Create a mock tool with a simple string argument
+        mock_tool = Tool(
+            name="hello_world",
+            description="A simple hello world function",
+            inputSchema={
+                "type": "object",
+                "properties": {"name": {"type": "string", "description": "The name to greet"}},
+                "required": ["name"],
+            },
+        )
+
+        func_body, type_defs = await create_function_body(client, mock_tool)
+
+        self.assertIn(
+            '''
+async def hello_world(name: Annotated[str, "The name to greet"]) -> Any:
+    """A simple hello world function"""
+'''.strip(),
+            func_body,
+        )
+
+    async def test_create_function_body_with_string_and_int_arguments(self) -> None:
+        # Mock the client
+        client = MagicMock(spec=Client)
+
+        # Create a mock tool with string and int arguments
+        mock_tool = Tool(
+            name="calculate",
+            description="A function that performs a calculation",
+            inputSchema={
+                "type": "object",
+                "properties": {
+                    "operation": {"type": "string", "description": "The operation to perform"},
+                    "value": {"type": "integer", "description": "The value to calculate with"},
+                },
+                "required": ["operation", "value"],
+            },
+        )
+
+        func_body, type_defs = await create_function_body(client, mock_tool)
+
+        self.assertIn(
+            '''
+async def calculate(operation: Annotated[str, "The operation to perform"], value: Annotated[int, "The value to calculate with"]) -> Any:
+    """A function that performs a calculation"""
+'''.strip(),
+            func_body,
+        )
+
+    async def test_create_function_body_with_complex_nested_arguments(self) -> None:
+        # Mock the client
+        client = MagicMock(spec=Client)
+
+        # Create a mock tool with complex nested arguments
+        mock_tool = Tool(
+            name="process_data",
+            description="A function that processes complex data structures",
+            inputSchema={
+                "type": "object",
+                "properties": {
+                    "user": {
+                        "type": "object",
+                        "properties": {
+                            "name": {"type": "string", "description": "User's name"},
+                            "age": {"type": "integer", "description": "User's age"},
+                            "preferences": {
+                                "type": "array",
+                                "items": {"type": "string"},
+                                "description": "User's preferences",
+                            },
+                        },
+                        "description": "User information",
+                    },
+                    "settings": {
+                        "type": "object",
+                        "properties": {
+                            "enabled": {"type": "boolean", "description": "Whether processing is enabled"},
+                            "options": {
+                                "type": "object",
+                                "properties": {
+                                    "mode": {"type": "string", "description": "Processing mode"},
+                                    "priority": {"type": "integer", "description": "Processing priority"},
+                                },
+                                "description": "Processing options",
+                            },
+                        },
+                        "description": "Processing settings",
+                    },
+                },
+                "required": ["user", "settings"],
+            },
+        )
+
+        func_body, type_defs = await create_function_body(client, mock_tool)
+
+        print()
+        print(func_body)
+        print()
+
+        self.assertIn(
+            '''
+async def process_data(user: Annotated[user_0, "User information"], settings: Annotated[settings_1, "Processing settings"]) -> Any:
+    """A function that processes complex data structures"""
+'''.strip(),
+            func_body,
+        )