Update app.py

app.py

@@ -1,8 +1,5 @@
 import streamlit as st
 import openai
-import tempfile
-import os
-from pymermaid import mermaid
 
 # Function to call OpenAI GPT model for prompt processing
 def get_diagram_code(prompt, diagram_type, api_key):
@@ -11,7 +8,7 @@ def get_diagram_code(prompt, diagram_type, api_key):
         openai.api_key = api_key
 
         response = openai.Completion.create(
-            engine="gpt-4o",
+            engine="gpt-4",
             prompt=f"Generate a {diagram_type} diagram in Mermaid.js syntax based on the following prompt: {prompt}",
             max_tokens=500
         )
@@ -20,19 +17,6 @@ def get_diagram_code(prompt, diagram_type, api_key):
         st.error(f"Error: {e}")
         return None
 
-# Function to render Mermaid.js diagram using pymermaid
-def render_mermaid(diagram_code):
-    # Create a temporary file to save the mermaid code
-    with tempfile.NamedTemporaryFile(delete=False, suffix=".mmd") as mmd_file:
-        mmd_file.write(diagram_code.encode('utf-8'))
-        mmd_file.close()
-
-        # Generate diagram image using pymermaid
-        output_image_path = mmd_file.name.replace(".mmd", ".png")
-        mermaid.render(mmd_file.name, output_image_path)
-        
-        return output_image_path
-
 # Streamlit App UI
 def main():
     st.title("Generate Diagrams using GPT-4 and Mermaid.js")
@@ -59,24 +43,10 @@ def main():
                 # Render Mermaid code to Streamlit
                 st.code(diagram_code, language='mermaid')
 
-                # Save diagram as image
-                try:
-                    image_path = render_mermaid(diagram_code)
-                    st.image(image_path, caption="Generated Diagram", use_column_width=True)
-                    
-                    # Provide download option
-                    with open(image_path, "rb") as img_file:
-                        st.download_button(
-                            label="Download Diagram Image",
-                            data=img_file,
-                            file_name="diagram.png",
-                            mime="image/png"
-                        )
-                    
-                    # Clean up the temporary files
-                    os.remove(image_path)
-                except Exception as e:
-                    st.error(f"Error generating diagram image: {e}")
+                # Display Mermaid diagram using iframe (mermaid.live)
+                diagram_url = f"https://mermaid-js.github.io/mermaid-live-editor/#/edit/{diagram_code}"
+                st.markdown(f'<iframe src="{diagram_url}" width="800" height="600"></iframe>', unsafe_allow_html=True)
+
             else:
                 st.error("Failed to generate diagram code.")
         else: