feat: add JSON config parsing with input token replacement

Author: mdesmet

src/datapilot/core/mcp_utils/mcp.py

@@ -21,32 +21,65 @@ def create_mcp_proxy():
     content = click.edit()
     if content is None:
         click.echo("No input provided.")
+        return
 
-    output = asyncio.run(list_tools())
-    click.echo(json.dumps(output, indent=2))
+    try:
+        config = json.loads(content)
+    except json.JSONDecodeError:
+        click.echo("Invalid JSON content.")
+        return
 
-async def list_tools(command: str, args: list[str], env: dict[str, str]) -> str:
-    command = shutil.which(command)
+    inputs = {}
+    mcp_config = config.get("mcp", {})
+    
+    # Process inputs first
+    for input_def in mcp_config.get("inputs", []):
+        input_id = input_def["id"]
+        inputs[input_id] = click.prompt(
+            input_def.get("description", input_id),
+            hide_input=input_def.get("password", False)
+        )
+
+    # Process servers
+    servers = mcp_config.get("servers", {})
+    for server_name, server_config in servers.items():
+        # Replace input tokens in args
+        processed_args = [
+            inputs.get(arg[8:-1], arg) if isinstance(arg, str) and arg.startswith("${input:") else arg 
+            for arg in server_config.get("args", [])
+        ]
+
+        # Replace input tokens in environment variables
+        processed_env = {
+            k: inputs.get(v[8:-1], v) if isinstance(v, str) and v.startswith("${input:") else v
+            for k, v in server_config.get("env", {}).items()
+        }
+
+        # Execute with processed parameters
+        output = asyncio.run(list_tools(
+            command=server_config["command"],
+            args=processed_args,
+            env=processed_env
+        ))
+        click.echo(f"\nServer: {server_name}")
+        click.echo(json.dumps(output, indent=2))
+
+async def list_tools(command: str, args: list[str], env: dict[str, str]):
+    command_path = shutil.which(command)
+    if not command_path:
+        raise click.UsageError(f"Command not found: {command}")
 
-    # Create server parameters for stdio connection
     server_params = StdioServerParameters(
-        command=command,  # Executable
-        args=args,  # Optional command line arguments
-        env=None,  # Optional environment variables
+        command=command_path,
+        args=args,
+        env=env,  # Now using processed env
     )
-
+    
     async with stdio_client(server_params) as (read, write):
-        async with ClientSession(
-            read, write
-        ) as session:
-            # Initialize the connection
+        async with ClientSession(read, write) as session:
             await session.initialize()
-
-            # List available tools
             tools = await session.list_tools()
-
-            # print as json
-            tools_list = [
+            return [
                 {
                     "name": tool.name,
                     "description": tool.description,
@@ -55,5 +88,3 @@ async def list_tools(command: str, args: list[str], env: dict[str, str]) -> str:
                 for tool in tools.tools
             ]
 
-            return tools_list
-