File size: 2,600 Bytes
5105400
 
 
 
 
 
 
 
 
6fe0a96
 
 
 
 
 
 
5105400
 
6fe0a96
 
 
5105400
 
 
 
 
 
56f5f92
5105400
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
56f5f92
b5df131
56f5f92
5fd0c51
56f5f92
 
b5df131
5105400
 
 
 
 
 
 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
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()