|
import streamlit as st |
|
import openai |
|
import urllib.parse |
|
import io |
|
from PIL import Image |
|
|
|
# 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-4o", # 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 |
|
|
|
# Function to convert SVG to PNG |
|
def convert_svg_to_png(svg_file): |
|
try: |
|
# Read the uploaded SVG file and convert it to PNG using cairosvg |
|
png_output = io.BytesIO() |
|
cairosvg.svg2png(file_obj=svg_file, write_to=png_output) |
|
png_output.seek(0) |
|
return png_output |
|
except Exception as e: |
|
st.error(f"Error during SVG to PNG conversion: {e}") |
|
return None |
|
|
|
# Streamlit App UI |
|
def main(): |
|
st.title("Generate Diagrams using AI and MermaidFlow") |
|
|
|
# Display a logo or icon |
|
image = Image.open("11.png") # Path to your image file |
|
st.image(image, width=300) # You can adjust the width as needed |
|
|
|
# Display sample prompt |
|
with st.expander("Sample Prompt For Users"): |
|
st.markdown(""" |
|
Create a UML diagram for a **Library Management System** that includes classes such as **Book** (with attributes like `bookID`, `title`, `author`, `publisher`, `genre`, `availabilityStatus` and methods like `checkAvailability()`, `updateAvailability()`), **Member** (with attributes such as `memberID`, `name`, `email`, `phoneNumber`, `membershipDate` and methods like `borrowBook()`, `returnBook()`, `reserveBook()`), **Staff** (with attributes such as `staffID`, `name`, `role`, `email`, `phoneNumber` and methods like `addBook()`, `removeBook()`, `updateBookInfo()`), **Transaction** (with attributes like `transactionID`, `transactionDate`, `dueDate`, `fineAmount` and methods like `calculateFine()`, `generateReceipt()`), and **Reservation** (with attributes like `reservationID`, `reservationDate`, `expirationDate` and methods like `cancelReservation()`, `checkReservationStatus()`). Define the relationships where a **Member** can borrow many **Books**, a **Staff** can manage many **Books**, a **Transaction** is linked to one **Book** and one **Member**, and a **Reservation** is made by a **Member** for a **Book**. Ensure appropriate visibility markers for methods and attributes, and include necessary relationships such as associations, multiplicities, and inheritance. Optionally, consider adding sequence or activity diagrams for processes like book borrowing. |
|
""") |
|
|
|
# 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') |
|
|
|
# Encode the Mermaid code for use in the MermaidFlow editor URL |
|
encoded_code = urllib.parse.quote(diagram_code) # URL encode the Mermaid code |
|
|
|
# MermaidFlow Editor URL with the encoded Mermaid code |
|
mermaidflow_url = f"https://www.mermaidflow.app/editor?code={encoded_code}" |
|
|
|
# Embed MermaidFlow Editor in iframe |
|
st.markdown(f'<iframe src="{mermaidflow_url}" width="100%" height="600px"></iframe>', unsafe_allow_html=True) |
|
|
|
# Provide a direct link to MermaidFlow editor for convenience |
|
st.markdown(f"Copy and paste the Generated Mermaid Code in the Live Editor and download the generated diagram by clicking on Export as SVG format to get a high-resolution generated image.") |
|
|
|
else: |
|
st.error("Failed to generate diagram code.") |
|
else: |
|
st.error("Please enter a prompt.") |
|
|
|
if __name__ == "__main__": |
|
main() |