Update app.py via AI Editor

app.py

@@ -96,9 +96,6 @@ def clear_session(sid):
         logging.error(f"Error clearing session {sid}: {e}")
 
 def parse_ai_response_to_graphviz(ai_response):
-    # Simple heuristic to parse AI response to Graphviz digraph elements.
-    # If the AI response is in markdown or text, try to extract nodes and edges.
-    # This is a minimal approach; for more complex parsing, extend as needed.
     lines = [line.strip() for line in ai_response.split("\n") if line.strip()]
     edges = []
     nodes = set()
@@ -124,20 +121,19 @@ def parse_ai_response_to_graphviz(ai_response):
     return nodes, edges
 
 def generate_diagram(description, sid):
-    sdir = get_session_dir(sid)
-    img_path = sdir / "diagram.png"
-    # Parse the description into nodes and edges
     nodes, edges = parse_ai_response_to_graphviz(description)
     dot = graphviz.Digraph(format="png")
     for node in nodes:
         dot.node(node)
     for src, dst in edges:
         dot.edge(src, dst)
-    dot.render(str(img_path.with_suffix('')), format='png', cleanup=True)
-    with open(img_path, "rb") as f:
-        img_bytes = f.read()
-    diagram_base64 = base64.b64encode(img_bytes).decode('utf-8')
-    return diagram_base64
+    try:
+        img_bytes = dot.pipe(format="png")
+        diagram_base64 = base64.b64encode(img_bytes).decode('utf-8')
+        return diagram_base64
+    except Exception as e:
+        logging.error(f"Graphviz error: {e}")
+        raise Exception("Graphviz rendering failed. Please check your description or try again.")
 
 def get_ai_response(model, prompt):
     if model == 'gpt-3.5-turbo':