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 AI") | |
| # 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') | |
| # Prepare Mermaid code for the Mermaid AI platform | |
| encoded_code = diagram_code.replace("\n", "%0A").replace(" ", "%20") # Encode the Mermaid code | |
| mermaid_ai_url = f"https://www.mermaidchart.com/app/projects/1179e788-52af-40a6-a788-6f15474a9c04/diagrams/f787a6a7-4922-4063-ad40-29bdec71a7ab/version/v0.1/ai?mermaidCode={encoded_code}" | |
| # Provide a link to the Mermaid AI platform | |
| st.markdown(f"Click the link to generate and save the diagram: [Mermaid AI - Generate Diagram]({mermaid_ai_url})") | |
| else: | |
| st.error("Failed to generate diagram code.") | |
| else: | |
| st.error("Please enter a prompt.") | |
| if __name__ == "__main__": | |
| main() | |