diff --git a/.bumpversion.cfg b/.bumpversion.cfg index 6ad930e1..90cda220 100644 --- a/.bumpversion.cfg +++ b/.bumpversion.cfg @@ -1,5 +1,5 @@ [bumpversion] -current_version = 0.0.18 +current_version = 0.0.19 commit = True tag = True diff --git a/.github/workflows/github-actions.yml b/.github/workflows/github-actions.yml index 1478e850..57ec4278 100644 --- a/.github/workflows/github-actions.yml +++ b/.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' diff --git a/.gitignore b/.gitignore index 77973dd9..23bb0a16 100644 --- a/.gitignore +++ b/.gitignore @@ -72,3 +72,4 @@ docs/_build # Mypy Cache .mypy_cache/ +.aider* diff --git a/docs/conf.py b/docs/conf.py index afea67e9..68234685 100644 --- a/docs/conf.py +++ b/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 = ["."] diff --git a/setup.py b/setup.py index 28859c96..2da85465 100644 --- a/setup.py +++ b/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: diff --git a/src/datapilot/__init__.py b/src/datapilot/__init__.py index f18e5d09..a11f0b41 100644 --- a/src/datapilot/__init__.py +++ b/src/datapilot/__init__.py @@ -1 +1 @@ -__version__ = "0.0.18" +__version__ = "0.0.19" diff --git a/src/datapilot/cli/main.py b/src/datapilot/cli/main.py index 9e0cfb39..f0f17965 100644 --- a/src/datapilot/cli/main.py +++ b/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) diff --git a/src/datapilot/core/mcp_utils/__init__.py b/src/datapilot/core/mcp_utils/__init__.py new file mode 100644 index 00000000..5216d6d1 --- /dev/null +++ b/src/datapilot/core/mcp_utils/__init__.py @@ -0,0 +1,2 @@ +DBT = [] +SQL = [] diff --git a/src/datapilot/core/mcp_utils/mcp.py b/src/datapilot/core/mcp_utils/mcp.py new file mode 100644 index 00000000..7c2481a7 --- /dev/null +++ b/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