Spaces:
Running
Running
| import streamlit as st | |
| import openai | |
| # 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 | |
| # For GPT-4 and other chat models, use the v1/chat/completions endpoint | |
| response = openai.ChatCompletion.create( | |
| model="gpt-4", # Make sure you're using the right model | |
| messages=[ | |
| {"role": "system", "content": "You are a helpful assistant."}, | |
| {"role": "user", "content": f"Generate a {diagram_type} diagram in Mermaid.js syntax based on the following prompt: {prompt}"} | |
| ], | |
| max_tokens=500 | |
| ) | |
| return response['choices'][0]['message']['content'].strip() | |
| except Exception as e: | |
| st.error(f"Error: {e}") | |
| return None | |
| # 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') | |
| # 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: | |
| st.error("Please enter a prompt.") | |
| if __name__ == "__main__": | |
| main() | |