AWS Diagram MCP Server

Instructions

Implementation Reference

• aws_diagram_mcp_server/server.py:180-224 (handler)

  MCP tool handler function for 'get_diagram_examples'. Registered with @mcp.tool decorator. Receives diagram_type parameter and delegates to the core get_diagram_examples function, returning its model_dump().

  @mcp.tool(name='get_diagram_examples')
  async def mcp_get_diagram_examples(
      diagram_type: DiagramType = Field(
          default=DiagramType.ALL,
          description='Type of diagram example to return. Options: aws, sequence, flow, class, k8s, onprem, custom, all',
      ),
  ):
      """Get example code for different types of diagrams.
  
      This tool provides ready-to-use example code for various diagram types.
      Use these examples to understand the syntax and capabilities of the diagrams package
      before creating your own custom diagrams.
  
      USAGE INSTRUCTIONS:
      1. Select the diagram type you're interested in (or 'all' to see all examples)
      2. Study the returned examples to understand the structure and syntax
      3. Use these examples as templates for your own diagrams
      4. When ready, modify an example or write your own code and use generate_diagram
  
      EXAMPLE CATEGORIES:
      - aws: AWS cloud architecture diagrams (basic services, grouped workers, clustered web services, Bedrock)
      - sequence: Process and interaction flow diagrams
      - flow: Decision trees and workflow diagrams
      - class: Object relationship and inheritance diagrams
      - k8s: Kubernetes architecture diagrams
      - onprem: On-premises infrastructure diagrams
      - custom: Custom diagrams with custom icons
      - all: All available examples across categories
  
      Each example demonstrates different features of the diagrams package:
      - Basic connections between components
      - Grouping with Clusters
      - Advanced styling with Edge attributes
      - Different layout directions
      - Multiple component instances
      - Custom icons and nodes
  
      Parameters:
          diagram_type (str): Type of diagram example to return. Options: aws, sequence, flow, class, k8s, onprem, custom, all
  
      Returns:
          Dictionary with example code for the requested diagram type(s), organized by example name
      """
      result = get_diagram_examples(diagram_type)
      return result.model_dump()

• aws_diagram_mcp_server/diagrams_tools.py:349-584 (helper)

  Core helper function implementing the logic to generate and return example diagram code snippets based on the requested diagram_type. Constructs a dictionary of examples and wraps in DiagramExampleResponse.

  def get_diagram_examples(diagram_type: DiagramType = DiagramType.ALL) -> DiagramExampleResponse:
      """Get example code for different types of diagrams.
  
      Args:
          diagram_type: Type of diagram example to return.
  
      Returns:
          DiagramExampleResponse: Dictionary with example code for the requested diagram type(s)
      """
      examples = {}
  
      # Basic examples
      if diagram_type in [DiagramType.AWS, DiagramType.ALL]:
          examples['aws_basic'] = """with Diagram("Web Service Architecture", show=False):
      ELB("lb") >> EC2("web") >> RDS("userdb")
  """
  
      if diagram_type in [DiagramType.SEQUENCE, DiagramType.ALL]:
          examples['sequence'] = """with Diagram("User Authentication Flow", show=False):
      user = User("User")
      login = InputOutput("Login Form")
      auth = Decision("Authenticated?")
      success = Action("Access Granted")
      failure = Action("Access Denied")
  
      user >> login >> auth
      auth >> success
      auth >> failure
  """
  
      if diagram_type in [DiagramType.FLOW, DiagramType.ALL]:
          examples['flow'] = """with Diagram("Order Processing Flow", show=False):
      start = Predefined("Start")
      order = InputOutput("Order Received")
      check = Decision("In Stock?")
      process = Action("Process Order")
      wait = Delay("Backorder")
      ship = Action("Ship Order")
      end = Predefined("End")
  
      start >> order >> check
      check >> process >> ship >> end
      check >> wait >> process
  """
  
      if diagram_type in [DiagramType.CLASS, DiagramType.ALL]:
          examples['class'] = """with Diagram("Simple Class Diagram", show=False):
      base = Python("BaseClass")
      child1 = Python("ChildClass1")
      child2 = Python("ChildClass2")
  
      base >> child1
      base >> child2
  """
  
      # Advanced examples from the documentation
      if diagram_type in [DiagramType.AWS, DiagramType.ALL]:
          examples[
              'aws_grouped_workers'
          ] = """with Diagram("Grouped Workers", show=False, direction="TB"):
      ELB("lb") >> [EC2("worker1"),
                    EC2("worker2"),
                    EC2("worker3"),
                    EC2("worker4"),
                    EC2("worker5")] >> RDS("events")
  """
  
          examples[
              'aws_clustered_web_services'
          ] = """with Diagram("Clustered Web Services", show=False):
      dns = Route53("dns")
      lb = ELB("lb")
  
      with Cluster("Services"):
          svc_group = [ECS("web1"),
                       ECS("web2"),
                       ECS("web3")]
  
      with Cluster("DB Cluster"):
          db_primary = RDS("userdb")
          db_primary - [RDS("userdb ro")]
  
      memcached = ElastiCache("memcached")
  
      dns >> lb >> svc_group
      svc_group >> db_primary
      svc_group >> memcached
  """
  
          examples['aws_event_processing'] = """with Diagram("Event Processing", show=False):
      source = EKS("k8s source")
  
      with Cluster("Event Flows"):
          with Cluster("Event Workers"):
              workers = [ECS("worker1"),
                         ECS("worker2"),
                         ECS("worker3")]
  
          queue = SQS("event queue")
  
          with Cluster("Processing"):
              handlers = [Lambda("proc1"),
                          Lambda("proc2"),
                          Lambda("proc3")]
  
      store = S3("events store")
      dw = Redshift("analytics")
  
      source >> workers >> queue >> handlers
      handlers >> store
      handlers >> dw
  """
  
          examples[
              'aws_bedrock'
          ] = """with Diagram("S3 Image Processing with Bedrock", show=False, direction="LR"):
      user = User("User")
  
      with Cluster("Amazon S3 Bucket"):
          input_folder = S3("Input Folder")
          output_folder = S3("Output Folder")
  
      lambda_function = Lambda("Image Processor Function")
      bedrock = Bedrock("Claude Sonnet 3.7")
  
      user >> Edge(label="Upload Image") >> input_folder
      input_folder >> Edge(label="Trigger") >> lambda_function
      lambda_function >> Edge(label="Process Image") >> bedrock
      bedrock >> Edge(label="Return Bounding Box") >> lambda_function
      lambda_function >> Edge(label="Upload Processed Image") >> output_folder
      output_folder >> Edge(label="Download Result") >> user
  """
  
      if diagram_type in [DiagramType.K8S, DiagramType.ALL]:
          examples['k8s_exposed_pod'] = """with Diagram("Exposed Pod with 3 Replicas", show=False):
      net = Ingress("domain.com") >> Service("svc")
      net >> [Pod("pod1"),
              Pod("pod2"),
              Pod("pod3")] << ReplicaSet("rs") << Deployment("dp") << HPA("hpa")
  """
  
          examples['k8s_stateful'] = """with Diagram("Stateful Architecture", show=False):
      with Cluster("Apps"):
          svc = Service("svc")
          sts = StatefulSet("sts")
  
          apps = []
          for _ in range(3):
              pod = Pod("pod")
              pvc = PVC("pvc")
              pod - sts - pvc
              apps.append(svc >> pod >> pvc)
  
      apps << PV("pv") << StorageClass("sc")
  """
  
      if diagram_type in [DiagramType.ONPREM, DiagramType.ALL]:
          examples[
              'onprem_web_service'
          ] = """with Diagram("Advanced Web Service with On-Premises", show=False):
      ingress = Nginx("ingress")
  
      metrics = Prometheus("metric")
      metrics << Grafana("monitoring")
  
      with Cluster("Service Cluster"):
          grpcsvc = [
              Server("grpc1"),
              Server("grpc2"),
              Server("grpc3")]
  
      with Cluster("Sessions HA"):
          primary = Redis("session")
          primary - Redis("replica") << metrics
          grpcsvc >> primary
  
      with Cluster("Database HA"):
          primary = PostgreSQL("users")
          primary - PostgreSQL("replica") << metrics
          grpcsvc >> primary
  
      aggregator = Fluentd("logging")
      aggregator >> Kafka("stream") >> Spark("analytics")
  
      ingress >> grpcsvc >> aggregator
  """
  
          examples[
              'onprem_web_service_colored'
          ] = """with Diagram(name="Advanced Web Service with On-Premise (colored)", show=False):
      ingress = Nginx("ingress")
  
      metrics = Prometheus("metric")
      metrics << Edge(color="firebrick", style="dashed") << Grafana("monitoring")
  
      with Cluster("Service Cluster"):
          grpcsvc = [
              Server("grpc1"),
              Server("grpc2"),
              Server("grpc3")]
  
      with Cluster("Sessions HA"):
          primary = Redis("session")
          primary - Edge(color="brown", style="dashed") - Redis("replica") << Edge(label="collect") << metrics
          grpcsvc >> Edge(color="brown") >> primary
  
      with Cluster("Database HA"):
          primary = PostgreSQL("users")
          primary - Edge(color="brown", style="dotted") - PostgreSQL("replica") << Edge(label="collect") << metrics
          grpcsvc >> Edge(color="black") >> primary
  
      aggregator = Fluentd("logging")
      aggregator >> Edge(label="parse") >> Kafka("stream") >> Edge(color="black", style="bold") >> Spark("analytics")
  
      ingress >> Edge(color="darkgreen") << grpcsvc >> Edge(color="darkorange") >> aggregator
  """
  
      if diagram_type in [DiagramType.CUSTOM, DiagramType.ALL]:
          examples['custom_rabbitmq'] = """# Download an image to be used into a Custom Node class
  rabbitmq_url = "https://jpadilla.github.io/rabbitmqapp/assets/img/icon.png"
  rabbitmq_icon = "rabbitmq.png"
  urlretrieve(rabbitmq_url, rabbitmq_icon)
  
  with Diagram("Broker Consumers", show=False):
      with Cluster("Consumers"):
          consumers = [
              Pod("worker"),
              Pod("worker"),
              Pod("worker")]
  
      queue = Custom("Message queue", rabbitmq_icon)
  
      queue >> consumers >> Aurora("Database")
  """
  
      return DiagramExampleResponse(examples=examples)

• aws_diagram_mcp_server/models.py:22-33 (schema)

  Pydantic Enum defining the possible diagram_type input values for the get_diagram_examples tool.

  class DiagramType(str, Enum):
      """Enum for supported diagram types."""
  
      AWS = 'aws'
      SEQUENCE = 'sequence'
      FLOW = 'flow'
      CLASS = 'class'
      K8S = 'k8s'
      ONPREM = 'onprem'
      CUSTOM = 'custom'
      ALL = 'all'

• aws_diagram_mcp_server/models.py:76-80 (schema)

  Pydantic model defining the response structure for get_diagram_examples, containing a dictionary of example names to code snippets.

  class DiagramExampleResponse(BaseModel):
      """Response model for diagram examples."""
  
      examples: Dict[str, str]