Skip to content

Text to diagram

Jump to bottom
Somkiat Puisungnoen edited this page Oct 24, 2024 · 8 revisions

Text to diagram

0. Setup environment

$python -m venv ./demo/venv
$source ./demo/venv/bin/activate

1. Install dependencies

  • File requirements.txt 
streamlit
Pillow
openai

Run

$pip install -r requirements.txt

2. Create file mermaid.py 

import base64
from PIL import Image
import requests
import io

def generate_diagram(graph):
    graphbytes = graph.encode("ascii")
    base64_bytes = base64.b64encode(graphbytes)
    base64_string = base64_bytes.decode("ascii")
    img = Image.open(io.BytesIO(requests.get(
        'https://mermaid.ink/img/' + base64_string).content))
    return img

3. Create file prompt.py

  • Create prompt template
  • Send question to OpenAI
from openai import OpenAI
import streamlit as st

def generate_prompt(prompt, chart_type, direction):
    # Preset instruction messages for the model
    messages = [
        {
            "role": "user",
            "content": "You are a bot that only communicates in Mermaid.js formatted markdown.",
        },
        {
            "role": "user",
            "content": "Do not provide any additional information or notes, ONLY markdown.",
        },
    ]

    # Generate prompt using OpenAI model
    prompt_formatted = f"""
    Generate markdown for mermaid.js for a chart of type {chart_type}
    with the following details and only return the markdown that can be pasted into a mermaid.js viewer:
    {prompt}
    """

    # Add prompt to messages
    messages.append({"role": "user", "content": prompt_formatted})

    return messages


# Generate response using OpenAI model
def SendChatRequest(prompt, chart_type, direction):

    # Assemble the prompt
    full_prompt = generate_prompt(prompt, chart_type, direction)

    # Send prompt to OpenAI
    openai_api_key = st.session_state.get("OPENAI_API_KEY")
    client = OpenAI(api_key=openai_api_key)
    response = client.chat.completions.create(
        model="gpt-4o", messages=full_prompt, max_tokens=1000
    )
    graph = response.choices[0].message.content
    print(graph)

    # Remove ```mermaid and ``` from the response
    graph = graph.replace("```mermaid", "")
    graph = graph.replace("```", "")

    print(graph)
    return graph

4. Create file main.py

  • Main program of Streamlit
import streamlit as st
from components.sidebar import sidebar
from mermaid import generate_diagram
from prompt import SendChatRequest

def clear_submit():
    st.session_state["submit"] = False

# Set general streamlit config with title and header
st.set_page_config(page_title="MermaidGPT", page_icon="🧜‍♀️", layout="wide")
st.header("🧜‍♀️MermaidGPT")

# Initialize sidebar
sidebar()

# Chart options
col1, col2 = st.columns(2)
with col1:
   chart_type = st.selectbox('Please select a chart type', ['Flowchart', 'Sequence Diagram', 'Class Diagram', 'State Diagram', 'Entity Relationship Diagram', 'User Journey', 'Gantt', 'Pie Chart', 'Quadrant Chart', 'Requirement Diagram'])
with col2:
   orientation = st.selectbox('Please select an orientation', ['Horizontal', 'Vertical'])

# Text input for prompt
chart_prompt = st.text_area('Enter your chart/diagram details (be as specific as possible): ', height=200)

button = st.button("Generate")
if button:
    if not st.session_state.get("OPENAI_API_KEY"):
        st.error("Please configure your OpenAI API key!")
    elif not chart_prompt:
        st.error("Please enter a chart prompt!")
    else:
        st.session_state["submit"] = True

        # Run OpenAI API call and receive md response
        result = SendChatRequest(chart_prompt, chart_type, orientation)

        # Display the image using streamlit
        try:
            diagram_img = generate_diagram(result)
            st.image(diagram_img, use_column_width=True)
        except:
            st.error("Something went wrong. Please try again or slightly rephrase your prompt.")
        st.write(f'Resulting markdown: {result}')

Create ./components/sidebar.py 

import streamlit as st

def set_openai_api_key(api_key: str):
    st.session_state["OPENAI_API_KEY"] = api_key

def sidebar():
    with st.sidebar:
        st.markdown(
            "## How to use\n"
            "1. Enter your [OpenAI API key](https://platform.openai.com/account/api-keys) below🔑\n"
            "2. Choose the chart type you wish to generate 📈 \n"
            "3. Give a detailed description of all components and transitions in your chart✍️ \n"
        )
        api_key_input = st.text_input(
            "OpenAI API Key",
            type="password",
            placeholder="Paste your OpenAI API key here (sk-...)",
            help="You can get your API key from https://platform.openai.com/account/api-keys.",
            value=st.session_state.get("OPENAI_API_KEY", ""),
        )

        if api_key_input:
            set_openai_api_key(api_key_input)

Run program with streamlit

streamlit run main.py

Open application in web browser

Sample prompt to generate sequence diagram

Step to login
1. login with username and password
2. Submit from to api gateway and redirect to auth service
3. Auth service check data in LDAP
4. Send result back to user

Reference

https://github.com/jgordley/MermaidGPT

Clone this wiki locally