fix layout

app.py

@@ -2,22 +2,25 @@ from src.agent import PlantUMLAgent
 from src.documents import get_processed_documents, get_diagrams_names
 from src.interface import create_interface
 from src.preview import render_plantuml, sanitize_plantuml_code
+import gradio as gr
 
 from cachetools import LRUCache, cached
 
-from dotenv import load_dotenv
-
-load_dotenv()
-
 diagrams_names = get_diagrams_names()
 docs_processed = get_processed_documents(diagrams_names)
 
-agent = PlantUMLAgent('mistral-small-latest', docs_processed)
-
 cache = LRUCache(maxsize=100)
 
 @cached(cache)
-def respond(message, diagram_name):
+def respond(message, diagram_name, api_key):
+    if not api_key:
+        gr.Warning("Please provide a valid Mistral API key to generate diagrams.")
+        return message, "", ""
+    if not message:
+        gr.Warning("Please provide a description for the UML diagram.")
+        return "", "", ""
+    
+    agent = PlantUMLAgent('mistral-small-latest', docs_processed, api_key)
     response = agent.predict(message, diagram_name)
 
     plantuml_code = sanitize_plantuml_code(response)    

requirements.txt

@@ -3,7 +3,6 @@ langchain
 langchain_community
 rank_bm25
 numpy==1.26.4
-python-dotenv
 cachetools
 beautifulsoup4
 docling

src/agent.py

@@ -1,12 +1,12 @@
 from langchain_community.retrievers import BM25Retriever
 from mistralai import Mistral
-import os
 
 class PlantUMLAgent:
-    def __init__(self, model, docs_processed):
+    def __init__(self, model, docs_processed, api_key):
         self.model = model
         self.docs_processed = docs_processed
-        self.client = Mistral(api_key=os.getenv("MISTRAL_API_KEY"))
+        self.client = Mistral(api_key=api_key)
+
 
     def predict(self, message, diagram_name):
         diagram_name = diagram_name.replace(" ", "-").lower()

src/interface.py

@@ -26,6 +26,16 @@ def create_interface(diagram_types, respond):
 
         with gr.Row():
             with gr.Column(scale=1):
+                
+                gr.Markdown("## 🔑 Mistral API Key Input")
+                api_key_input = gr.Textbox(
+                    label="🔐 Mistral API Key",
+                    placeholder="Input your Mistral API key here",
+                    type="password",
+                    lines=1,
+                    interactive=True
+                )
+
                 gr.Markdown("## 📝 Input Description")
                 msg = gr.Textbox(
                     label="Describe your UML diagram",
@@ -73,7 +83,7 @@ def create_interface(diagram_types, respond):
             """
         )
 
-        submit.click(respond, [msg, diagram_dropdown], [msg, diagram_text_output, diagram_output])
-        msg.submit(respond, [msg, diagram_dropdown], [msg, diagram_text_output, diagram_output])
+        submit.click(respond, [msg, diagram_dropdown, api_key_input], [msg, diagram_text_output, diagram_output])
+        msg.submit(respond, [msg, diagram_dropdown, api_key_input], [msg, diagram_text_output, diagram_output])
 
     return demo