Spaces:
Running
Running
| 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): | |
| try: | |
| # Set the OpenAI API key dynamically based on user input | |
| openai.api_key = api_key | |
| response = openai.Completion.create( | |
| engine="gpt-4", | |
| prompt=f"Generate a {diagram_type} diagram in Mermaid.js syntax based on the following prompt: {prompt}", | |
| max_tokens=500 | |
| ) | |
| return response.choices[0].text.strip() | |
| except Exception as e: | |
| 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") | |
| # User input for OpenAI API key | |
| api_key = st.text_input("Enter your OpenAI API Key:", type="password") | |
| # Check if API key is provided | |
| if not api_key: | |
| st.warning("Please enter your OpenAI API Key to continue.") | |
| return | |
| # User prompt for diagram type | |
| prompt = st.text_area("Enter your prompt for the diagram:", "") | |
| diagram_types = ["UML Diagram", "ER Diagram", "State Diagram", "Class Diagram", "Sequence Diagram"] | |
| diagram_choice = st.selectbox("Select the type of diagram to generate:", diagram_types) | |
| if st.button("Generate Diagram"): | |
| if prompt: | |
| diagram_code = get_diagram_code(prompt, diagram_choice, api_key) | |
| if diagram_code: | |
| # 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}") | |
| else: | |
| st.error("Failed to generate diagram code.") | |
| else: | |
| st.error("Please enter a prompt.") | |
| if __name__ == "__main__": | |
| main() | |