Spaces:
Sleeping
Sleeping
| import zlib | |
| import requests | |
| import gradio as gr | |
| from groq import Groq | |
| from PIL import Image | |
| import io | |
| # Validate the API key | |
| def validate_api_key(api_key): | |
| try: | |
| client = Groq(api_key=api_key) | |
| response = client.chat.completions.create( | |
| model="llama3-70b-8192", | |
| messages=[ | |
| {"role": "system", "content": "Validation test for API key."}, | |
| {"role": "user", "content": "Hello, world!"} | |
| ], | |
| temperature=0.1, | |
| max_tokens=5, | |
| ) | |
| if response and response.choices: | |
| return True, "API key is valid!" | |
| else: | |
| return False, "API key validation failed: Empty response." | |
| except Exception as e: | |
| return False, f"Invalid API key. Error: {str(e)}" | |
| def encode_plantuml(uml_code): | |
| compressed = zlib.compress(uml_code.encode('utf-8'))[2:-4] | |
| base64_chars = "0123456789ABCDEFGHIJKLMNOPQRSTUVWXYZabcdefghijklmnopqrstuvwxyz-_" | |
| result = [] | |
| buffer = 0 | |
| buffer_size = 0 | |
| for byte in compressed: | |
| buffer = (buffer << 8) | byte | |
| buffer_size += 8 | |
| while buffer_size >= 6: | |
| buffer_size -= 6 | |
| result.append(base64_chars[(buffer >> buffer_size) & 0x3F]) | |
| if buffer_size > 0: | |
| result.append(base64_chars[(buffer << (6 - buffer_size)) & 0x3F]) | |
| return ''.join(result) | |
| def generate_plantuml_url(uml_code): | |
| encoded_uml = encode_plantuml(uml_code) | |
| base_url = "http://www.plantuml.com/plantuml/png/" | |
| return f"{base_url}{encoded_uml}" | |
| def download_image(url): | |
| try: | |
| response = requests.get(url) | |
| response.raise_for_status() | |
| return Image.open(io.BytesIO(response.content)) | |
| except Exception as e: | |
| print(f"Error downloading image: {e}") | |
| return None | |
| def generate_uml_diagram(prompt, api_key): | |
| try: | |
| client = Groq(api_key=api_key) | |
| if not prompt.startswith("@startuml"): | |
| response = client.chat.completions.create( | |
| model="llama3-70b-8192", | |
| messages=[ | |
| {"role": "system", "content": "You are an expert at creating PlantUML diagrams. Always return ONLY the raw PlantUML code without markdown code block fences or explanatory text."}, | |
| {"role": "user", "content": f"Create a PlantUML diagram for the following description: {prompt}"} | |
| ], | |
| max_tokens=300 | |
| ) | |
| uml_code = response.choices[0].message.content.strip() | |
| if "```" in uml_code: | |
| uml_code = uml_code.split("```")[1].strip() | |
| else: | |
| uml_code = prompt.strip() | |
| plantuml_url = generate_plantuml_url(uml_code) | |
| image = download_image(plantuml_url) | |
| if image: | |
| return image | |
| else: | |
| return "Failed to generate image. Please check the input or try again." | |
| except Exception as e: | |
| print(f"Error: {e}") | |
| return f"Error generating diagram: {str(e)}" | |
| def on_login(api_key): | |
| is_valid, message = validate_api_key(api_key) | |
| if is_valid: | |
| return api_key, gr.update(visible=False), gr.update(visible=True), gr.Markdown(f"**Success!** {message}") | |
| else: | |
| return "", gr.update(visible=True), gr.update(visible=False), gr.Markdown(f"**Error!** {message}", visible=True) | |
| # Gradio UI | |
| with gr.Blocks(title="DiagramCraft AI") as demo: | |
| # Login Section | |
| with gr.Column(visible=True) as login_section: | |
| gr.Markdown("## Login to Access the DiagramCraft AI") | |
| api_key_input = gr.Textbox(label="Enter your API Key:", type="password") | |
| login_button = gr.Button("Login") | |
| login_feedback = gr.Markdown(visible=False) | |
| gr.Markdown(""" | |
| ## Welcome to the PlantUML Diagram Generator! | |
| ### Login Instructions: | |
| To access this tool, you will need a **Groq API key**. Follow these steps to get one: | |
| 1. **Sign Up / Log In**: | |
| - Visit [Groq API Portal](https://console.groq.com/keys). | |
| - Create an account or log in | |
| 2. **Generate Your API Key**: | |
| - Navigate to the **API Keys** section after logging in. | |
| - Click on **Create New Key** and give it a name (e.g., "PlantUML Tool"). | |
| - Copy the generated key. | |
| 3. **Paste Your API Key**: | |
| - Enter your API key in the text box above and click **Login**. | |
| - Once logged in, you can start generating UML diagrams instantly! | |
| ### Note: | |
| - If you encounter any issues during login, double-check your API key. | |
| - For help, visit the [Groq Support Page](https://support.groq.com/). | |
| """) | |
| # UML Diagram Generator Section | |
| with gr.Column(visible=False) as diagram_section: | |
| gr.Markdown(""" | |
| ## DiagramCraft AI | |
| """) | |
| with gr.Row(): | |
| with gr.Column(scale=3): | |
| gr.Markdown(""" | |
| ### About This Tool: | |
| This project helps you generate UML diagrams quickly and easily! It supports various diagram types such as: | |
| - **Class Diagrams** | |
| - **Sequence Diagrams** | |
| - **ER Diagrams** | |
| - **Use Case Diagrams** | |
| - **Data Flow Diagrams (DFDs)** | |
| Simply provide a **description** or **PlantUML code** to visualize your diagram instantly. | |
| ### How to Use: | |
| 1. **Enter a Description**: | |
| - Describe your system or diagram in plain language (e.g., "A library with books, authors, and members"). | |
| 2. **Click Generate**: | |
| - The tool will process your input and create a UML diagram. | |
| 3. **View the Diagram**: | |
| - The generated image will appear below. | |
| ### Troubleshooting: | |
| - **No image appears?** Click **Generate** again or try after a few moments. | |
| - **Still not working?** Double-check your input for clarity or correctness. | |
| ### Example Use Cases: | |
| - Visualize system architectures for projects. | |
| - Create technical documentation quickly. | |
| - Design and share models with your team! | |
| Happy diagramming! | |
| """) | |
| with gr.Column(scale=6): | |
| gr.Markdown(""" | |
| Enter a description or PlantUML code to generate your diagram. | |
| If no image appears, click "Generate" again or try after some time. | |
| """) | |
| prompt_input = gr.Textbox(label="Describe the diagram or paste PlantUML code") | |
| generate_button = gr.Button("Generate Diagram") | |
| with gr.Row(): | |
| uml_image = gr.Image(label="Generated UML Diagram", height=500) | |
| state = gr.State() # Manage API key state | |
| login_button.click( | |
| on_login, | |
| inputs=api_key_input, | |
| outputs=[state, login_section, diagram_section, login_feedback] | |
| ) | |
| generate_button.click( | |
| fn=generate_uml_diagram, | |
| inputs=[prompt_input, state], | |
| outputs=[uml_image] | |
| ) | |
| demo.launch() | |