Update app.py via AI Editor

app.py

@@ -3,11 +3,17 @@ import io
 import base64
 import threading
 import logging
-from dash import Dash, dcc, html, Input, Output, State
+import tempfile
+import shutil
+import uuid
+import json
+from pathlib import Path
+from flask import request
+from dash import Dash, dcc, html, Input, Output, State, ctx, no_update
 import dash_bootstrap_components as dbc
 import openai
+import anthropic
 from google import generativeai as genai
-from anthropic import Anthropic
 import requests
 from diagrams import Diagram, Cluster
 from diagrams.aws.compute import EC2
@@ -20,10 +26,8 @@ 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.")
@@ -40,46 +44,74 @@ else:
     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
+if not anthropic_api_key:
     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])
+# Session management globals
+SESSION_DIR = Path(tempfile.gettempdir()) / "arch_diagram_sessions"
+SESSION_DIR.mkdir(parents=True, exist_ok=True)
+session_locks = {}
+session_data = {}
+
+def get_or_create_session_id():
+    sid = request.cookies.get("sid")
+    if not sid or not isinstance(sid, str) or len(sid) < 8:
+        sid = str(uuid.uuid4())
+    return sid
 
-# Global variable to store the generated diagram
-generated_diagram = None
+def get_session_dir(sid):
+    p = SESSION_DIR / sid
+    p.mkdir(parents=True, exist_ok=True)
+    return p
 
-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:
+def get_session_lock(sid):
+    if sid not in session_locks:
+        session_locks[sid] = threading.Lock()
+    return session_locks[sid]
+
+def get_session_data(sid):
+    if sid not in session_data:
+        session_data[sid] = {"diagram_base64": None, "description": "", "model": "gpt-3.5-turbo"}
+    return session_data[sid]
+
+def save_session_state(sid, data):
+    session_data[sid] = data
+    sdir = get_session_dir(sid)
+    with open(sdir / "session.json", "w") as f:
+        json.dump(data, f)
+
+def load_session_state(sid):
+    sdir = get_session_dir(sid)
+    sf = sdir / "session.json"
+    if sf.exists():
+        try:
+            with open(sf) as f:
+                session_data[sid] = json.load(f)
+        except Exception as e:
+            logging.error(f"Session load error: {e}")
+
+def generate_diagram(description, sid):
+    sdir = get_session_dir(sid)
+    img_path = sdir / "diagram.png"
+    with Diagram("Architecture", show=False, outformat="png", filename=str(img_path.with_suffix(""))):
         with Cluster("Cloud"):
             lb = ELB("Load Balancer")
             with Cluster("Web Tier"):
-                web = [EC2("Web Server 1"),
-                       EC2("Web Server 2")]
+                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
+    with open(img_path, "rb") as f:
+        img_bytes = f.read()
+    diagram_base64 = base64.b64encode(img_bytes).decode('utf-8')
+    return diagram_base64
 
 def get_ai_response(model, prompt):
     if model == 'gpt-3.5-turbo':
@@ -89,17 +121,25 @@ def get_ai_response(model, 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)
+        if genai is None:
+            raise Exception("Gemini API not configured.")
+        model_obj = genai.GenerativeModel('gemini-1.5-pro-latest')
+        response = model_obj.generate_content(prompt)
         return response.text
     elif model == 'claude-3-5-haiku-20241022':
-        response = anthropic.messages.create(
+        if not anthropic_api_key:
+            raise Exception("Anthropic API not configured.")
+        client = anthropic.Anthropic(api_key=anthropic_api_key)
+        response = client.messages.create(
             model="claude-3-5-haiku-20241022",
-            max_tokens=2000,
+            system=None,
+            max_tokens=2048,
             messages=[{"role": "user", "content": prompt}]
         )
         return response.content[0].text
     elif model == 'grok-3-mini-fast-beta':
+        if not grok_api_key:
+            raise Exception("Groq API not configured.")
         groq_url = "https://api.groq.com/openai/v1/chat/completions"
         headers = {
             "Authorization": f"Bearer {grok_api_key}",
@@ -112,62 +152,121 @@ def get_ai_response(model, 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)
+        raise Exception("Error: Invalid model selected.")
+
+app = Dash(__name__, external_stylesheets=[dbc.themes.BOOTSTRAP])
+
+def serve_layout():
+    sid = get_or_create_session_id()
+    load_session_state(sid)
+    sdata = get_session_data(sid)
+    diagram = sdata.get("diagram_base64", None)
+    resp = dbc.Container(fluid=True, children=[
+        dbc.Row([
+            dbc.Col([
+                html.H1("Architecture Diagram Generator", className="mb-4"),
+            ], width=12)
+        ], className="mt-2"),
+        dbc.Row([
+            dbc.Col([
+                dbc.Card([
+                    dbc.CardHeader("Controls"),
+                    dbc.CardBody([
+                        html.Div([
+                            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=sdata.get("model", "gpt-3.5-turbo"),
+                                className="mb-3"
+                            ),
+                            dbc.Textarea(
+                                id='description-input', 
+                                placeholder="Enter architecture description here...",
+                                value=sdata.get("description", ""),
+                                style={'height': '200px', 'whiteSpace': 'pre-wrap', 'wordWrap': 'break-word'},
+                                className="mb-3"
+                            ),
+                            dbc.Button("Generate Diagram", 
+                                       id='generate-button', 
+                                       color="primary", 
+                                       className="me-2 mb-2"),
+                            dbc.Button("Download Diagram", 
+                                       id='download-button', 
+                                       color="secondary", 
+                                       className="mb-2", 
+                                       disabled=(diagram is None)),
+                            dcc.Download(id="download-diagram"),
+                            html.Div(id='status-message', className="mb-3"),
+                        ]),
+                    ])
+                ])
+            ], width=4, style={'backgroundColor': '#f8f9fa', 'minHeight': '100vh'}),
+            dbc.Col([
+                dbc.Card([
+                    dbc.CardHeader("Diagram"),
+                    dbc.CardBody([
+                        dcc.Loading(
+                            id="loading",
+                            type="default",
+                            fullscreen=False,
+                            children=[
+                                html.Img(
+                                    id='diagram-output',
+                                    src=f'data:image/png;base64,{diagram}' if diagram else "",
+                                    style={'width': '100%', 'minHeight': '300px'}
+                                ),
+                            ]
+                        )
+                    ])
+                ]),
+            ], width=8, style={'backgroundColor': '#fff', 'minHeight': '100vh'}),
+        ])
     ])
-], fluid=True)
+    return resp
+
+app.layout = serve_layout
+
+@app.server.after_request
+def set_cookie(response):
+    sid = get_or_create_session_id()
+    response.set_cookie("sid", sid, max_age=60*60*24*30, httponly=False, samesite="Lax")
+    return response
 
 @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')],
+    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):
+    sid = get_or_create_session_id()
+    lock = get_session_lock(sid)
     if n_clicks is None:
-        return dash.no_update, dash.no_update, True
-    
+        return no_update, 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
+        return no_update, "Please enter a description.", True
+    with lock:
+        try:
+            ai_response = get_ai_response(model, description)
+            diagram_base64 = generate_diagram(ai_response, sid)
+            sdata = get_session_data(sid)
+            sdata["diagram_base64"] = diagram_base64
+            sdata["description"] = description
+            sdata["model"] = model
+            save_session_state(sid, sdata)
+            logging.info(f"Session {sid}: Diagram generated successfully.")
+            return f'data:image/png;base64,{diagram_base64}', "Diagram generated successfully!", False
+        except Exception as e:
+            logging.error(f"Session {sid}: Error generating diagram: {str(e)}")
+            return no_update, f"Error: {str(e)}", True
 
 @app.callback(
     Output("download-diagram", "data"),
@@ -175,15 +274,20 @@ def update_diagram(n_clicks, description, model):
     prevent_initial_call=True
 )
 def download_diagram(n_clicks):
+    sid = get_or_create_session_id()
+    lock = get_session_lock(sid)
     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")
+        return no_update
+    with lock:
+        sdata = get_session_data(sid)
+        diagram_base64 = sdata.get("diagram_base64", None)
+        if not diagram_base64:
+            return no_update
+        diagram_bytes = base64.b64decode(diagram_base64)
+        logging.info(f"Session {sid}: Diagram downloaded.")
+        return dcc.send_bytes(diagram_bytes, "architecture_diagram.png")
 
 if __name__ == '__main__':
     print("Starting the Dash application...")
-    app.run(debug=True, host='0.0.0.0', port=7860)
+    app.run(debug=True, host='0.0.0.0', port=7860, threaded=True)
     print("Dash application has finished running.")