diff --git a/.isort.cfg b/.isort.cfg index 8112570b..10d71835 100644 --- a/.isort.cfg +++ b/.isort.cfg @@ -2,4 +2,4 @@ line_length = 120 multi_line_output = 3 include_trailing_comma = True -known_third_party = graphviz,jinja2 +known_third_party = graphviz,inspection,jinja2,mcp diff --git a/mcp-server/GEMINI.md b/mcp-server/GEMINI.md index 10944df3..5d85c146 100644 --- a/mcp-server/GEMINI.md +++ b/mcp-server/GEMINI.md @@ -78,7 +78,7 @@ The server will define three main tools. It should use the `mcp.server.fastmcp` - `code` (string, required): The Python code using `diagrams` DSL. - `filename` (string, optional): Desired output filename. -* **Returns**: +* **Returns**: - `status`: "success" or "error" - `message`: Path to file or error message. diff --git a/mcp-server/README.md b/mcp-server/README.md index 0174fe96..3263918d 100644 --- a/mcp-server/README.md +++ b/mcp-server/README.md @@ -14,12 +14,12 @@ graph TD subgraph Container ["Docker Container"] MCPServer["MCP Server (server.py)"] NodeRegistry["(Node Registry)"] - + subgraph Logic ["Core Logic"] Inspector["src/inspection.py"] Executor["generate_diagram"] end - + subgraph Libs ["Dependencies"] DiagramsLib["diagrams package"] Graphviz["Graphviz Binary"] @@ -34,12 +34,12 @@ graph TD %% Tool Flows Client -- "list_icons()" --> MCPServer MCPServer -- "Query" --> NodeRegistry - + Client -- "generate_diagram(code)" --> MCPServer MCPServer -- "Pass Code" --> Executor Executor -- "exec()" --> DiagramsLib DiagramsLib -- "Render" --> Graphviz - + %% Output Graphviz -- "Generates PNG" --> Executor Executor -- "Writes File (Volume Mount)" --> HostFS @@ -109,7 +109,7 @@ Generates a diagram from Python code. ```python from diagrams import Diagram from diagrams.aws.compute import EC2 - + with Diagram("Simple", show=False): EC2("web") ``` diff --git a/mcp-server/src/inspection.py b/mcp-server/src/inspection.py index 222029f4..3300f75e 100644 --- a/mcp-server/src/inspection.py +++ b/mcp-server/src/inspection.py @@ -1,15 +1,17 @@ -import pkgutil import importlib import inspect +import pkgutil import sys from collections import defaultdict + import diagrams from diagrams import Node + def get_all_nodes(): """ Dynamically inspects the diagrams package and returns a dictionary of all available Nodes. - + Returns: dict: A nested dictionary structure: { @@ -20,7 +22,7 @@ def get_all_nodes(): """ # Initialize the structure icons = defaultdict(lambda: defaultdict(list)) - + # We also keep a flat map for the execution context: Name -> Class # This handles potential name collisions by favoring the last seen or explicit logic if needed. node_registry = {} @@ -33,14 +35,14 @@ def get_all_nodes(): for _, provider_name, ispkg in pkgutil.iter_modules(path, prefix): if not ispkg: continue - + # e.g., provider_name = "diagrams.aws" short_provider = provider_name.split(".")[-1] - + # Skip internal modules if any (base, etc are actually useful, but we focus on providers) - if short_provider in ['base', 'custom']: - # 'custom' and 'base' might be treated differently, but for now we scan them - pass + if short_provider in ['base', 'custom']: + # 'custom' and 'base' might be treated differently, but for now we scan them + pass try: provider_module = importlib.import_module(provider_name) @@ -63,13 +65,14 @@ def get_all_nodes(): # or at least is defined in the diagrams package if obj.__module__.startswith("diagrams"): # Ensure the object actually belongs to this service (or a submodule of it) - # This prevents listing imported classes from other services (e.g. Trace in operations vs devtools) + # This prevents listing imported classes from other services (e.g. Trace + # in operations vs devtools) if not obj.__module__.startswith(service_name): continue - + icons[short_provider][short_service].append(name) node_registry[name] = obj - + except ImportError: continue except Exception: diff --git a/mcp-server/src/server.py b/mcp-server/src/server.py index e1cb0c78..7f75e966 100644 --- a/mcp-server/src/server.py +++ b/mcp-server/src/server.py @@ -1,14 +1,15 @@ +import base64 +import contextlib import os import sys import tempfile -import contextlib -import base64 from pathlib import Path -from mcp.server.fastmcp import FastMCP -from diagrams import Diagram, Cluster, Edge, Node # Import our helper from inspection import get_all_nodes +from mcp.server.fastmcp import FastMCP + +from diagrams import Cluster, Diagram, Edge, Node # Initialize FastMCP mcp = FastMCP("diagrams-mcp") @@ -18,11 +19,12 @@ print("Loading diagram nodes...", file=sys.stderr) ALL_ICONS, NODE_REGISTRY = get_all_nodes() print(f"Loaded {len(NODE_REGISTRY)} nodes.", file=sys.stderr) + @mcp.tool() def list_icons(provider_filter: str = None, service_filter: str = None): """ List available icons from the diagrams package, with optional filtering. - + Args: provider_filter: Filter icons by provider name (e.g., "aws", "gcp", "k8s") service_filter: Filter icons by service name (e.g., "compute", "database") @@ -30,26 +32,30 @@ def list_icons(provider_filter: str = None, service_filter: str = None): if not provider_filter: # Return list of providers return {"providers": list(ALL_ICONS.keys())} - + if provider_filter not in ALL_ICONS: return {"error": f"Provider '{provider_filter}' not found. Available: {list(ALL_ICONS.keys())}"} - + provider_data = ALL_ICONS[provider_filter] - + if not service_filter: # Return all services for this provider return provider_data - + if service_filter not in provider_data: - return {"error": f"Service '{service_filter}' not found in '{provider_filter}'. Available: {list(provider_data.keys())}"} - + return { + "error": f"Service '{service_filter}' not found in '{provider_filter}'. Available: { + list( + provider_data.keys())}"} + return {service_filter: provider_data[service_filter]} + @mcp.tool() def get_diagram_examples(diagram_type: str = "all"): """ Get example code for different types of diagrams. - + Args: diagram_type: Type of diagram example to return (aws, k8s, flow, etc. or 'all') """ @@ -70,11 +76,11 @@ from diagrams.k8s.network import Ingress, Service with Diagram("K8s Cluster", show=False): ingress = Ingress("domain.com") - + with Cluster("App"): svc = Service("svc") pods = [Pod("pod1"), Pod("pod2")] - + ingress >> svc >> pods """, "custom": """ @@ -86,28 +92,29 @@ with Diagram("Custom", show=False): Custom("Label", "./my-icon.png") """ } - + if diagram_type == "all": return examples - + return {diagram_type: examples.get(diagram_type, "No example found for this type.")} + @mcp.tool() def generate_diagram(code: str, filename: str = None, timeout: int = 90): """ Generate a diagram from Python code using the diagrams package. - + Args: code: Python code using the diagrams package DSL. filename: Optional filename to save the diagram to. timeout: Execution timeout in seconds. """ - + # Create a temporary directory for execution with tempfile.TemporaryDirectory() as temp_dir: original_cwd = os.getcwd() os.chdir(temp_dir) - + try: # Prepare the execution context # We inject Diagram, Cluster, Edge, and ALL discovered nodes (EC2, Pod, etc.) @@ -120,57 +127,60 @@ def generate_diagram(code: str, filename: str = None, timeout: int = 90): "Node": Node, **NODE_REGISTRY } - + # Execute the code # We wrap it in a try/except block within the exec to catch runtime errors try: exec(code, exec_globals) except Exception as e: return {"status": "error", "message": f"Runtime error: {str(e)}"} - + # Find the generated file # Diagrams generates files based on the name passed to Diagram() class # We look for any .png file created in the temp dir generated_files = list(Path(".").glob("*.png")) - + if not generated_files: - return {"status": "error", "message": "No diagram image was generated. Did you call with Diagram(..., show=False)?"} - + return { + "status": "error", + "message": "No diagram image was generated. Did you call with Diagram(..., show=False)?"} + # Use the most recently modified file or the first one generated_files.sort(key=lambda f: f.stat().st_mtime, reverse=True) output_file = generated_files[0] - + # If a filename was requested, we might want to rename it? - # For now, we return the path. + # For now, we return the path. # In a real MCP setup, we might copy this to a mounted volume. - + # Copy the generated file back to the original working directory # This ensures that if the user mounted their project to the working directory, # the file appears in their project. import shutil - + target_dir = Path(original_cwd) target_filename = filename if filename else output_file.name target_path = target_dir / target_filename - + # Ensure extension if not target_path.suffix: target_path = target_path.with_suffix(".png") - + shutil.copy2(output_file, target_path) final_path = str(target_path) - + return { - "status": "success", + "status": "success", "path": final_path, "filename": target_path.name } except Exception as e: return {"status": "error", "message": f"System error: {str(e)}"} - + finally: os.chdir(original_cwd) + if __name__ == "__main__": mcp.run()