Update app.py

app.py

@@ -1,7 +1,12 @@
 import os
 import streamlit as st
 import requests
+import re
 from dotenv import load_dotenv
+from PIL import Image
+import matplotlib.pyplot as plt
+from io import BytesIO
+import tempfile
 
 # Load environment variables (for FLOWISE_API token)
 load_dotenv()
@@ -15,8 +20,16 @@ def query(payload):
     response = requests.post(API_URL, headers=headers, json=payload)
     return response.json()
 
+# Extract only Python code from text
+def extract_python_code(text):
+    code_pattern = r"```python(.*?)```"
+    match = re.search(code_pattern, text, re.DOTALL)
+    if match:
+        return match.group(1).strip()
+    return "No Python code found in response."
+
 # Streamlit UI
-st.title("AWS Architecture Diagram Generator")
+st.title("Architecture Diagram Generator")
 
 # Get input from the user
 user_input = st.text_area("Describe your architecture:", 
@@ -27,11 +40,31 @@ if st.button("Generate Diagram"):
         # Send the user's input to the API
         payload = {"question": user_input}
         output = query(payload)
-        print(output)
         
-        # Check and display the API response
-        if "diagram_url" in output:
-            st.image(output["diagram_url"], caption="Generated Architecture Diagram")
+        # Check and display the Python code from the API response
+        if "text" in output:
+            python_code = extract_python_code(output["text"])
+            st.code(python_code, language="python")
+
+            # Save the generated code to a temporary file and execute it
+            with tempfile.TemporaryDirectory() as tmpdirname:
+                temp_file = os.path.join(tmpdirname, "architecture_diagram.py")
+                with open(temp_file, "w") as f:
+                    f.write(python_code)
+
+                # Execute the code
+                try:
+                    exec(open(temp_file).read())
+                    # Assuming the diagram was saved as "WordPress_Architecture.png" in the script
+                    diagram_path = os.path.join(tmpdirname, "WordPress_Architecture.png")
+
+                    # Display the generated image
+                    if os.path.exists(diagram_path):
+                        st.image(diagram_path, caption="Generated Architecture Diagram")
+                    else:
+                        st.error("No diagram image was generated.")
+                except Exception as e:
+                    st.error(f"Error executing the Python code: {e}")
         else:
             st.write("Response:", output)
     else: