Update app.py via AI Editor

app.py

@@ -16,7 +16,7 @@ import openai
 import anthropic
 from google import generativeai as genai
 import requests
-import pygraphviz as pgv
+import graphviz
 
 logging.basicConfig(level=logging.INFO)
 
@@ -108,13 +108,12 @@ def strip_markdown_codeblock(text):
 def generate_diagram(dot_code, sid):
     dot_code = dot_code.strip()
     try:
-        G = pgv.AGraph(string=dot_code)
-        out_img = io.BytesIO()
-        G.draw(out_img, format="png", prog="dot")
-        diagram_base64 = base64.b64encode(out_img.getvalue()).decode('utf-8')
+        src = graphviz.Source(dot_code)
+        out_img = src.pipe(format="png")
+        diagram_base64 = base64.b64encode(out_img).decode('utf-8')
         return diagram_base64
     except Exception as e:
-        logging.error(f"PyGraphviz error: {e}")
+        logging.error(f"Graphviz error: {e}")
         raise Exception("Graphviz rendering failed. Please check your description or try again.")
 
 def get_ai_response(model, user_prompt):