feat: add mcp command group with create-mcp-proxy functionality

.bumpversion.cfg

@@ -1,5 +1,5 @@
 [bumpversion]
-current_version = 0.0.18
+current_version = 0.0.19
 commit = True
 tag = True
 

.github/workflows/github-actions.yml

@@ -19,42 +19,6 @@ jobs:
             toxpython: 'python3.11'
             tox_env: 'docs'
             os: 'ubuntu-latest'
-          - name: 'py39-pydantic28-cover (ubuntu)'
-            python: '3.9'
-            toxpython: 'python3.9'
-            python_arch: 'x64'
-            tox_env: 'py39-pydantic28-cover'
-            os: 'ubuntu-latest'
-          - name: 'py39-pydantic28-cover (windows)'
-            python: '3.9'
-            toxpython: 'python3.9'
-            python_arch: 'x64'
-            tox_env: 'py39-pydantic28-cover'
-            os: 'windows-latest'
-          - name: 'py39-pydantic28-cover (macos)'
-            python: '3.9'
-            toxpython: 'python3.9'
-            python_arch: 'x64'
-            tox_env: 'py39-pydantic28-cover'
-            os: 'macos-13'
-          - name: 'py39-pydantic210-cover'
-            python: '3.9'
-            toxpython: 'python3.9'
-            python_arch: 'x64'
-            tox_env: 'py39-pydantic210-cover'
-            os: 'ubuntu-latest'
-          - name: 'py39-pydantic28-nocov'
-            python: '3.9'
-            toxpython: 'python3.9'
-            python_arch: 'x64'
-            tox_env: 'py39-pydantic28-nocov'
-            os: 'ubuntu-latest'
-          - name: 'py39-pydantic210-nocov'
-            python: '3.9'
-            toxpython: 'python3.9'
-            python_arch: 'x64'
-            tox_env: 'py39-pydantic210-nocov'
-            os: 'ubuntu-latest'
           - name: 'py310-pydantic28-cover'
             python: '3.10'
             toxpython: 'python3.10'
@@ -127,30 +91,6 @@ jobs:
             python_arch: 'x64'
             tox_env: 'py312-pydantic210-nocov'
             os: 'ubuntu-latest'
-          - name: 'pypy39-pydantic28-cover'
-            python: 'pypy-3.9'
-            toxpython: 'pypy3.9'
-            python_arch: 'x64'
-            tox_env: 'pypy39-pydantic28-cover'
-            os: 'ubuntu-latest'
-          - name: 'pypy39-pydantic210-cover'
-            python: 'pypy-3.9'
-            toxpython: 'pypy3.9'
-            python_arch: 'x64'
-            tox_env: 'pypy39-pydantic210-cover'
-            os: 'ubuntu-latest'
-          - name: 'pypy39-pydantic28-nocov'
-            python: 'pypy-3.9'
-            toxpython: 'pypy3.9'
-            python_arch: 'x64'
-            tox_env: 'pypy39-pydantic28-nocov'
-            os: 'ubuntu-latest'
-          - name: 'pypy39-pydantic210-nocov'
-            python: 'pypy-3.9'
-            toxpython: 'pypy3.9'
-            python_arch: 'x64'
-            tox_env: 'pypy39-pydantic210-nocov'
-            os: 'ubuntu-latest'
           - name: 'pypy310-pydantic28-cover'
             python: 'pypy-3.10'
             toxpython: 'pypy3.10'

.gitignore

@@ -72,3 +72,4 @@ docs/_build
 
 # Mypy Cache
 .mypy_cache/
+.aider*

docs/conf.py

@@ -15,7 +15,7 @@
 year = "2024"
 author = "Altimate Inc."
 copyright = f"{year}, {author}"
-version = release = "0.0.18"
+version = release = "0.0.19"
 
 pygments_style = "trac"
 templates_path = ["."]

setup.py

@@ -13,7 +13,7 @@ def read(*names, **kwargs):
 
 setup(
     name="altimate-datapilot-cli",
-    version="0.0.18",
+    version="0.0.19",
     license="MIT",
     description="Assistant for Data Teams",
     long_description="{}\n{}".format(
@@ -68,6 +68,8 @@ def read(*names, **kwargs):
         "tabulate~=0.9.0",
         "requests>=2.31",
         "sqlglot~=25.30.0",
+        "mcp~=1.9.0",
+        "pyperclip~=1.8.2",
     ],
     extras_require={
         # eg:

src/datapilot/__init__.py

@@ -1 +1 @@
-__version__ = "0.0.18"
+__version__ = "0.0.19"

src/datapilot/cli/main.py

@@ -1,5 +1,6 @@
 import click
 
+from datapilot.core.mcp_utils.mcp import mcp
 from datapilot.core.platforms.dbt.cli.cli import dbt
 
 
@@ -9,3 +10,4 @@ def datapilot():
 
 
 datapilot.add_command(dbt)
+datapilot.add_command(mcp)

src/datapilot/core/mcp_utils/__init__.py

@@ -0,0 +1,2 @@
+DBT = []
+SQL = []

src/datapilot/core/mcp_utils/mcp.py

@@ -0,0 +1,176 @@
+import asyncio
+import json
+import logging
+import shutil
+from dataclasses import dataclass
+
+import click
+import pyperclip
+from mcp import ClientSession
+from mcp import StdioServerParameters
+from mcp.client.stdio import stdio_client
+
+logging.basicConfig(level=logging.INFO)
+
+
+@dataclass
+class InputParameter:
+    name: str
+    type: str
+    required: bool
+    key: str
+    description: str
+
+
+def find_input_tokens(data):
+    tokens = set()
+    if isinstance(data, list):
+        for item in data:
+            tokens.update(find_input_tokens(item))
+    elif isinstance(data, dict):
+        for value in data.values():
+            tokens.update(find_input_tokens(value))
+    elif isinstance(data, str) and data.startswith("${input:"):
+        tokens.add(data[8:-1].strip())
+    return tokens
+
+
+# New mcp group
+@click.group()
+def mcp():
+    """mcp specific commands."""
+
+
+@mcp.command("inspect-mcp-server")
+def create_mcp_proxy():
+    content = click.edit()
+    if content is None:
+        click.echo("No input provided.")
+        return
+
+    try:
+        config = json.loads(content)
+    except json.JSONDecodeError:
+        click.echo("Invalid JSON content.")
+        return
+
+    inputs = {}
+    mcp_config = config.get("mcp", {})
+
+    # Select server
+    # Support both "servers" and "mcpServers" naming conventions
+    servers = mcp_config.get("mcpServers", mcp_config.get("servers", {}))
+    server_names = list(servers.keys())
+
+    if not server_names:
+        ctx = click.get_current_context()
+        click.secho("Error: No servers configured in mcp config (tried keys: 'mcpServers' and 'servers')", fg="red")
+        ctx.exit(1)
+
+    if len(server_names) > 1:
+        server_name = click.prompt("Choose a server", type=click.Choice(server_names), show_choices=True)
+    else:
+        server_name = server_names[0]
+
+    if server_name in servers:
+        server_config = servers[server_name]
+
+        # Collect input tokens ONLY from this server's config
+        input_ids = find_input_tokens(server_config.get("args", []))
+        input_ids.update(find_input_tokens(server_config.get("env", {})))
+
+        # Create prompt definitions using BOTH discovered tokens AND configured inputs
+        existing_input_ids = {i["id"] for i in mcp_config.get("inputs", [])}
+        inputs_to_prompt = input_ids.intersection(existing_input_ids)
+        inputs_to_prompt.update(input_ids)  # Add any undiscovered-by-config inputs
+
+        input_configs = []
+        for input_id in inputs_to_prompt:
+            input_def = next((d for d in mcp_config.get("inputs", []) if d["id"] == input_id), {})
+            inputs[input_id] = click.prompt(
+                input_def.get("description", input_id),
+                hide_input=True,
+            )
+            # Create InputParameters config entry
+            input_configs.append(
+                InputParameter(
+                    name=input_def.get("name", input_id),
+                    type="password",
+                    required=True,
+                    key=input_id,
+                    description=input_def.get("description", ""),
+                ).__dict__
+            )
+
+        # 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(server_config=server_config, command=server_config["command"], args=processed_args, env=processed_env)
+        )
+        # Add processed parameters to output
+        output_with_name = {
+            "name": server_name,
+            "config": input_configs,
+            "command": server_config["command"],
+            "args": [arg.replace("${input:", "${") if isinstance(arg, str) else arg for arg in server_config.get("args", [])],
+            "env": [
+                {"key": k, "value": v.replace("${input:", "${") if isinstance(v, str) else v}
+                for k, v in server_config.get("env", {}).items()
+            ],
+            **output,
+        }
+        output_json = json.dumps(output_with_name, indent=2)
+        click.echo(output_json)
+        try:
+            pyperclip.copy(output_json)
+            click.secho("\nOutput copied to clipboard!", fg="green")
+        except pyperclip.PyperclipException as e:
+            click.secho(f"\nFailed to copy to clipboard: {e!s}", fg="yellow")
+
+
+async def list_tools(server_config: dict, 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}")
+
+    try:
+        # Only support stdio server type
+        server_type = server_config.get("type", "stdio")
+        if server_type != "stdio":
+            raise click.UsageError(f"Only stdio MCP servers are supported. Found type: {server_type}")
+
+        server_params = StdioServerParameters(
+            command=command_path,
+            args=args,
+            env=env,
+        )
+
+        async with stdio_client(server_params) as (read, write):
+            async with ClientSession(read, write) as session:
+                await session.initialize()
+                tools = await session.list_tools()
+                mcp_tools = [
+                    {
+                        "name": tool.name,
+                        "description": tool.description,
+                        "inputSchema": tool.inputSchema,
+                    }
+                    for tool in tools.tools
+                ]
+
+                return {
+                    "tools": mcp_tools,
+                }
+    except Exception as e:
+        raise click.UsageError("Could not connect to MCP server: " + str(e)) from e