main.py

import streamlit as st
import subprocess
import os
import sys
import re
from PIL import Image
from groq import Groq

# Set page configuration for a better experience
st.set_page_config(page_title="Diagram Generator", page_icon=":art:", layout="wide")

# Custom CSS for a polished look
st.markdown("""

    <style>

    .main { background-color: #f5f5f5; }

    .stButton>button {

        background-color: #007bff;

        color: white;

        font-weight: bold;

    }

    </style>

    """, unsafe_allow_html=True)

# Sidebar: let the user choose which type of diagram they want to generate
diagram_type = st.sidebar.selectbox("Select Diagram Type", ["AWS Architecture Diagram", "Flowchart Diagram"])

if diagram_type == "AWS Architecture Diagram":
    st.sidebar.header("Instructions for AWS Diagrams")
    st.sidebar.write("""

    1. Enter a detailed AWS architecture description in the text area.

    2. Click **Generate Diagram** to create the diagram.

    3. Once generated, the diagram image will be shown along with a download option.

    4. Expand the **Show Generated Code** section to view the code.

    """)
else:
    st.sidebar.header("Instructions for Flowchart Diagrams")
    st.sidebar.write("""

    1. Enter a detailed flowchart description in the text area.

    2. Click **Generate Diagram** to create the flowchart.

    3. Once generated, the flowchart image will be shown along with a download option.

    4. Expand the **Show Generated Code** section to view the code and error details.

    """)

# ----- Common Imports Block -----
common_imports = (
    "class NodeList(list):\n"
    "    def __rshift__(self, other):\n"
    "        for item in self:\n"
    "            item >> other\n"
    "        return other\n"
    "\n"
    "    def __lshift__(self, other):\n"
    "        for item in self:\n"
    "            other >> item\n"
    "        return other\n"
    "\n"
    "def list_rshift(self, other):\n"
    "    return NodeList(self).__rshift__(other)\n"
    "\n"
    "def list_lshift(self, other):\n"
    "    return NodeList(self).__lshift__(other)\n"
    "\n"
    "try:\n"
    "    list.__rshift__ = list_rshift\n"
    "    list.__lshift__ = list_lshift\n"
    "except Exception as e:\n"
    "    pass"
)

# ----- AWS Imports Block -----
aws_imports = (
    common_imports +
    "\n" +
    "from diagrams.aws.compute import *\n"
    "from diagrams.aws.database import *\n"
    "from diagrams.aws.analytics import *\n"
    "from diagrams.aws.integration import *\n"
    "from diagrams.aws.management import *\n"
    "from diagrams.aws.ml import *\n"
    "from diagrams.aws.network import *\n"
    "from diagrams.aws.storage import *\n"
    "from diagrams import Diagram, Cluster, Edge\n"
    "from diagrams.aws.security import *"
)

# ----- Flowchart Imports Block -----
flowchart_imports = (
    common_imports +
    "\n" +
    "from diagrams.programming.flowchart import *\n"
    "from diagrams import Diagram, Cluster, Edge"
)

# ----- Groq API Call Function -----
def call_groq_api(prompt):
    client = Groq(api_key=os.getenv("GROQ_API_KEY"))
    chat_completion = client.chat.completions.create(
        messages=[{"role": "user", "content": prompt}],
        model="deepseek-r1-distill-llama-70b",
        stream=False,
    )
    return chat_completion.choices[0].message.content

# ----- Diagram Generation Function -----
import re
import subprocess
import sys
import os

def generate_diagram(diagram_code, import_block):
    # Extract content between the three backticks (```diagram and closing backticks)
    match = re.search(r'```diagram\n(.*?)\n```', diagram_code, re.DOTALL)
    if not match:
        raise ValueError("Invalid diagram code format. No content found between the diagram code backticks.")

    diagram_code = match.group(1)  # Extract content after '```diagram' and before closing '```'

    # Force the diagram name and filename to "diagram"
    diagram_code = re.sub(r'with Diagram\(".*?"', 'with Diagram("diagram"', diagram_code, count=1)
    diagram_code = re.sub(r'filename="[^"]*"', 'filename="diagram"', diagram_code)


    diagram_code = re.sub(r'DynamoDB', 'Dynamodb', diagram_code)
    # Combine imports and diagram code
    full_code = import_block + "\n" + diagram_code

    # Write the full code to a script file
    script_path = "diagram_app.py"
    with open(script_path, "w") as f:
        f.write(full_code)
    
    # Execute the script to generate the diagram
    try:
        result = subprocess.run(
            [sys.executable, script_path],
            check=True,
            stdout=subprocess.PIPE,
            stderr=subprocess.PIPE
        )
    except subprocess.CalledProcessError as e:
        error_details = e.stderr.decode('utf-8') if e.stderr else str(e)
        print("Error generating the diagram. Details below:")
        print(error_details)
        return None

    # Look for the generated diagram image
    image_path = "diagram.png"
    if os.path.exists(image_path):
        return image_path
    else:
        print("Diagram image not found.")
        return None
# ----- Main UI -----
st.title(diagram_type)
if diagram_type == "AWS Architecture Diagram":
    st.write("Generate a clear and organized AWS architecture diagram simply by describing your architecture below.")
    placeholder_text = "Enter your detailed AWS architecture description here..."
else:
    st.write("Generate an easy-to-understand flowchart diagram by describing your process or system below.")
    placeholder_text = "Enter your detailed flowchart description here..."

user_input = st.text_area("Diagram Description", height=200, placeholder=placeholder_text)

if st.button("Generate Diagram"):
    if not user_input.strip():
        st.error("Please enter a diagram description.")
    else:
        # Build the prompt and choose the proper import block based on selection
        if diagram_type == "AWS Architecture Diagram":
            full_prompt = (
                "Your task is to generate code using a diagram DSL similar to the Diagrams library. Always make an easy diagram to understand the architecture.\n"
                "Client class doesnt exist use Route53 from the 20 list only and related pickup Follow these guidelines:\n"
                "- Start with a with Diagram(\"diagram\", show=False, direction=\"TB\") block.\n"
                "- Use clusters with with Cluster(\"Cluster Name\") to group related nodes.\n"
                "- Instantiate nodes using functions named after AWS services.\n"
                "- Connect nodes using arrow operators (>> or <<) to indicate data flow.\n"
                "- Nested clusters can be used to represent hierarchical architectures.\n"
                "- Strictly always Use only the following 20 AWS services: and these 20 classes only given below as these are only supported and use same naming for them \n"
                "    1. EC2\n"
                "    2. Lambda\n"
                "    3. Fargate\n"
                "    4. ECS\n"
                "    5. RDS\n"
                "    6. Dynamodb\n"
                "    7. Redshift\n"
                "    8. Athena\n"
                "    9. ELB\n"
                "    10. APIGateway\n"
                "    11. CloudFront\n"
                "    12. Route53\n"
                "    13. VPC\n"
                "    14. S3\n"
                "    15. EBS\n"
                "    16. SQS\n"
                "    17. SNS\n"
                "    18. Cloudwatch\n"
                "    19. Sagemaker\n"
                "    20. ElasticBeanstalk\n"
                "Do not include any extra services, explanation, or markdown formatting. Do not include any formatting tags such as <think>, </think>, triple backticks (```diagram or ```), or the word 'python' anywhere in your output.\n"
                "Ensure that only one diagram is generated. \n"
                "Examples of proper diagram code structure (do not include these examples in your output):\n"
                "1. **Connect nodes using arrow operators like `>>` or `<<` to indicate data flow:**\n"
                "```python\n"
                "with Diagram(\"Grouped Workers\", show=False, direction=\"TB\"):\n"
                "    ELB(\"lb\") >> [EC2(\"worker1\"),\n"
                "                  EC2(\"worker2\"),\n"
                "                  EC2(\"worker3\"),\n"
                "                  EC2(\"worker4\"),\n"
                "                  EC2(\"worker5\")] >> RDS(\"events\")\n"
                "```\n"
                "2. **Nested Clusters:**\n"
                "```python\n"
                "with Diagram(\"Event Processing\", show=False):\n"
                "    source = EKS(\"k8s source\")\n\n"
                "    with Cluster(\"Event Flows\"):\n"
                "        with Cluster(\"Event Workers\"):\n"
                "            workers = [ECS(\"worker1\"), ECS(\"worker2\"), ECS(\"worker3\")]\n\n"
                "    queue = SQS(\"event queue\")\n\n"
                "    with Cluster(\"Processing\"):\n"
                "        handlers = [Lambda(\"proc1\"), Lambda(\"proc2\"), Lambda(\"proc3\")]\n\n"
                "    store = S3(\"events store\")\n"
                "    dw = Redshift(\"analytics\")\n\n"
                "    source >> workers >> queue >> handlers\n"
                "    handlers >> store\n"
                "    handlers >> dw\n"
                "```\n"
                "3. **Custom Edges and Styling:**\n"
                "```python\n"
                "with Diagram(name=\"Advanced Web Service with On-Premises (colored)\", show=False):\n"
                "    ingress = Nginx(\"ingress\")\n\n"
                "    metrics = Prometheus(\"metric\")\n"
                "    metrics << Edge(color=\"firebrick\", style=\"dashed\") << Grafana(\"monitoring\")\n\n"
                "    with Cluster(\"Service Cluster\"):\n"
                "        grpcsvc = [Server(\"grpc1\"), Server(\"grpc2\"), Server(\"grpc3\")]\n\n"
                "    with Cluster(\"Sessions HA\"):\n"
                "        primary = Redis(\"session\")\n"
                "        primary - Edge(color=\"brown\", style=\"dashed\") - Redis(\"replica\") << Edge(label=\"collect\") << metrics\n"
                "        grpcsvc >> Edge(color=\"brown\") >> primary\n\n"
                "    with Cluster(\"Database HA\"):\n"
                "        primary = PostgreSQL(\"users\")\n"
                "        primary - Edge(color=\"brown\", style=\"dotted\") - PostgreSQL(\"replica\") << Edge(label=\"collect\") << metrics\n"
                "        grpcsvc >> Edge(color=\"black\") >> primary\n\n"
                "    aggregator = Fluentd(\"logging\")\n"
                "    aggregator >> Edge(label=\"parse\") >> Kafka(\"stream\") >> Edge(color=\"black\", style=\"bold\") >> Spark(\"analytics\")\n\n"
                "    ingress >> Edge(color=\"darkgreen\") << grpcsvc >> Edge(color=\"darkorange\") >> aggregator\n"
                "```\n"
                "Architecture Description:\n"
                " it is Dynamodb always the function name and not DynamoDB Keep the naming in detail and explanation with good names and strictly keep the names same as given as those are library names. ALWAYS GIVE CODE WITH BACKTICKS AND 'diagram' as heading and only one diagram\n" + user_input
            )
            import_block = aws_imports
        else:
            full_prompt = (
                "Your task is to generate code using a diagram DSL similar to the Diagrams library. "
                "Always create an easy-to-understand flowchart diagram that clearly represents the process or system.\n"
                "Follow these guidelines:\n"
                "- Start with a with Diagram(\"diagram\", show=False, direction=\"TB\") block.\n"
                "- Use clusters (with Cluster(\"Cluster Name\")) to group related nodes if needed.\n"
                "- Use only the following flowchart node functions for your diagram: and same naming\n"
                "    • Action\n"
                "    • Collate\n"
                "    • Database\n"
                "    • Decision\n"
                "    • Delay\n"
                "    • Display\n"
                "    • Document\n"
                "    • InputOutput\n"
                "    • Inspection\n"
                "    • InternalStorage\n"
                "    • LoopLimit\n"
                "    • ManualInput\n"
                "    • ManualLoop\n"
                "    • Merge\n"
                "    • MultipleDocuments\n"
                "    • OffPageConnectorLeft\n"
                "    • OffPageConnectorRight\n"
                "    • Or\n"
                "    • PredefinedProcess\n"
                "    • Preparation\n"
                "    • Sort\n"
                "    • StartEnd\n"
                "    • StoredData\n"
                "    • SummingJunction\n"
                "Connect nodes using arrow operators (>> or <<) to indicate the flow.\n"
                "Do not include any extra nodes, explanation, or markdown formatting. Do not include any formatting tags such as <think>, </think>, triple backticks (```diagram or ```), or the word 'python' anywhere in your output.\n"
                "Ensure that only one diagram is generated.\n"
                "Examples of proper diagram code structure (do not include these examples in your output):\n"
                "1. **Connect nodes using arrow operators like `>>` or `<<` to indicate data flow:**\n"
                "```python\n"
                "with Diagram(\"Grouped Workers\", show=False, direction=\"TB\"):\n"
                "    Action(\"Start\") >> [Action(\"Task1\"), Action(\"Task2\"), Action(\"Task3\")] >> Action(\"End\")\n"
                "```\n"
                "2. **Nested Clusters:**\n"
                "```python\n"
                "with Diagram(\"Event Processing\", show=False):\n"
                "    source = Action(\"Data Ingest\")\n\n"
                "    with Cluster(\"Data Flow\"):\n"
                "        source >> Action(\"Processing\") >> Action(\"Storage\")\n\n"
                "    storage = Action(\"Store Data\")\n"
                "    analytics = Action(\"Analytics\")\n\n"
                "    source >> storage\n"
                "    storage >> analytics\n"
                "```\n"
                "3. **Custom Edges and Styling:**\n"
                "```python\n"
                "with Diagram(\"Custom Flowchart\", show=False):\n"
                "    start = Action(\"Start\")\n"
                "    task1 = Action(\"Task 1\")\n"
                "    task2 = Action(\"Task 2\")\n"
                "    task3 = Action(\"Task 3\")\n\n"
                "    start >> Edge(color=\"blue\") >> task1\n"
                "    task1 >> Edge(color=\"red\", style=\"dashed\") >> task2\n"
                "    task2 >> task3\n"
                "```\n"
                "Flowchart Description: Strictly only use the classes given above and nothing else. ALWAYS GIVE CODE WITH BACKTICKS AND 'diagram' as heading and only one diagram\n" + user_input
            )
            import_block = flowchart_imports
        
        st.info("Generating diagram code from your description...")
        try:
            diagram_code = call_groq_api(full_prompt)
        except Exception as ex:
            st.error("An error occurred while calling the Groq API:")
            
            diagram_code = None

        # Always show the generated diagram code in an expander for debugging,
        # even if errors occur later.

        if diagram_code:
            with st.spinner("Generating diagram image..."):
                image_path = generate_diagram(diagram_code, import_block)
            if image_path:
                st.success("Diagram generated successfully!")
                st.image(image_path, caption="Generated Diagram", use_container_width=True)
                with open(image_path, "rb") as file:
                    st.download_button("Download Diagram", file, file_name="diagram.png", mime="image/png")
            else:
                st.error("Failed to generate the diagram image. Please try again.")
        else:
            st.error("Failed to generate diagram code. Please try again.")