Create app.py

app.py

@@ -0,0 +1,189 @@
+import os
+import io
+import base64
+import threading
+import logging
+from dash import Dash, dcc, html, Input, Output, State
+import dash_bootstrap_components as dbc
+import openai
+from google import generativeai as genai
+from anthropic import Anthropic
+import requests
+from diagrams import Diagram, Cluster
+from diagrams.aws.compute import EC2
+from diagrams.aws.database import RDS
+from diagrams.aws.network import ELB
+from diagrams.gcp.compute import ComputeEngine
+from diagrams.gcp.database import SQL
+from diagrams.gcp.network import LoadBalancing
+from diagrams.onprem.compute import Server
+from diagrams.onprem.database import PostgreSQL
+from diagrams.onprem.network import Nginx
+
+# Configure logging
+logging.basicConfig(level=logging.INFO)
+
+# Initialize API clients
+openai.api_key = os.getenv("OPENAI_API_KEY")
+if not openai.api_key:
+    logging.warning("OPENAI_API_KEY not set. GPT-3.5 model will not be available.")
+
+google_api_key = os.getenv("GOOGLE_API_KEY")
+if google_api_key:
+    try:
+        genai.configure(api_key=google_api_key)
+    except Exception as e:
+        logging.error(f"Failed to configure Google Gemini: {e}")
+        genai = None
+else:
+    genai = None
+    logging.warning("GOOGLE_API_KEY not set. Gemini model will not be available.")
+
+anthropic_api_key = os.getenv("ANTHROPIC_API_KEY")
+if anthropic_api_key:
+    try:
+        anthropic = Anthropic(api_key=anthropic_api_key)
+    except Exception as e:
+        logging.error(f"Failed to initialize Anthropic client: {e}")
+        anthropic = None
+else:
+    anthropic = None
+    logging.warning("ANTHROPIC_API_KEY not set. Claude model will not be available.")
+
+grok_api_key = os.getenv("GROK_API_KEY")
+if not grok_api_key:
+    logging.warning("GROK_API_KEY not set. Groq model will not be available.")
+
+app = Dash(__name__, external_stylesheets=[dbc.themes.BOOTSTRAP])
+
+# Global variable to store the generated diagram
+generated_diagram = None
+
+def generate_diagram(description):
+    global generated_diagram
+    # This is a placeholder function. In a real implementation, you would parse the AI model's
+    # output and create a diagram based on that. For now, we'll create a simple diagram.
+    with Diagram("Architecture", show=False) as diag:
+        with Cluster("Cloud"):
+            lb = ELB("Load Balancer")
+            with Cluster("Web Tier"):
+                web = [EC2("Web Server 1"),
+                       EC2("Web Server 2")]
+            with Cluster("Database Tier"):
+                db_master = RDS("Master")
+                db_slave = RDS("Slave")
+        
+        lb >> web
+        web >> db_master
+        db_master - db_slave
+
+    img_bytes = diag.dot.pipe(format='png')
+    generated_diagram = base64.b64encode(img_bytes).decode('utf-8')
+    return generated_diagram
+
+def get_ai_response(model, prompt):
+    if model == 'gpt-3.5-turbo':
+        response = openai.ChatCompletion.create(
+            model="gpt-3.5-turbo",
+            messages=[{"role": "user", "content": prompt}]
+        )
+        return response.choices[0].message.content
+    elif model == 'gemini-1.5-flash-latest':
+        model = genai.GenerativeModel('gemini-1.5-pro')
+        response = model.generate_content(prompt)
+        return response.text
+    elif model == 'claude-3-5-haiku-20241022':
+        response = anthropic.messages.create(
+            model="claude-3-5-haiku-20241022",
+            max_tokens=2000,
+            messages=[{"role": "user", "content": prompt}]
+        )
+        return response.content[0].text
+    elif model == 'grok-3-mini-fast-beta':
+        groq_url = "https://api.groq.com/openai/v1/chat/completions"
+        headers = {
+            "Authorization": f"Bearer {grok_api_key}",
+            "Content-Type": "application/json"
+        }
+        data = {
+            "model": "grok-3-mini-fast-beta",
+            "messages": [{"role": "user", "content": prompt}]
+        }
+        response = requests.post(groq_url, json=data, headers=headers)
+        return response.json()['choices'][0]['message']['content']
+    else:
+        return "Error: Invalid model selected."
+
+app.layout = dbc.Container([
+    html.H1("Architecture Diagram Generator", className="my-4"),
+    dbc.Row([
+        dbc.Col([
+            dcc.Dropdown(
+                id='model-dropdown',
+                options=[
+                    {'label': 'GPT-3.5 Turbo', 'value': 'gpt-3.5-turbo'},
+                    {'label': 'Gemini 1.5 Flash', 'value': 'gemini-1.5-flash-latest'},
+                    {'label': 'Claude 3.5 Haiku', 'value': 'claude-3-5-haiku-20241022'},
+                    {'label': 'Grok 3 Mini Fast', 'value': 'grok-3-mini-fast-beta'}
+                ],
+                value='gpt-3.5-turbo',
+                className="mb-3"
+            ),
+            dbc.Textarea(id='description-input', placeholder="Enter architecture description here...", style={'height': '200px'}, className="mb-3"),
+            dbc.Button("Generate Diagram", id='generate-button', color="primary", className="mb-3"),
+            dbc.Button("Download Diagram", id='download-button', color="secondary", className="mb-3", disabled=True),
+            dcc.Download(id="download-diagram"),
+            html.Div(id='status-message', className="mb-3")
+        ], width=4),
+        dbc.Col([
+            html.Img(id='diagram-output', style={'width': '100%'})
+        ], width=8)
+    ])
+], fluid=True)
+
+@app.callback(
+    [Output('diagram-output', 'src'),
+     Output('status-message', 'children'),
+     Output('download-button', 'disabled')],
+    [Input('generate-button', 'n_clicks')],
+    [State('description-input', 'value'),
+     State('model-dropdown', 'value')],
+    prevent_initial_call=True
+)
+def update_diagram(n_clicks, description, model):
+    if n_clicks is None:
+        return dash.no_update, dash.no_update, True
+    
+    if not description:
+        return dash.no_update, "Please enter a description.", True
+
+    try:
+        # Get AI response
+        ai_response = get_ai_response(model, description)
+        
+        # Generate diagram based on AI response
+        diagram_base64 = generate_diagram(ai_response)
+        
+        return f'data:image/png;base64,{diagram_base64}', "Diagram generated successfully!", False
+    except Exception as e:
+        logging.error(f"Error generating diagram: {str(e)}")
+        return dash.no_update, f"Error: {str(e)}", True
+
+@app.callback(
+    Output("download-diagram", "data"),
+    Input("download-button", "n_clicks"),
+    prevent_initial_call=True
+)
+def download_diagram(n_clicks):
+    if n_clicks is None:
+        return dash.no_update
+    
+    if generated_diagram is None:
+        return dash.no_update
+
+    return dcc.send_bytes(base64.b64decode(generated_diagram), "architecture_diagram.png")
+
+if __name__ == '__main__':
+    print("Starting the Dash application...")
+    app.run(debug=True, host='0.0.0.0', port=7860)
+    print("Dash application has finished running.")