Update app.py

app.py

@@ -1,50 +1,42 @@
 import gradio as gr
 import openai
-import plantuml
+from plantuml import PlantUML
+import tempfile
 import os
 
-def generate_plantuml_code(api_key, prompt):
+def generate_plantuml_code(api_key, input_text):
     openai.api_key = api_key
-    messages = [
-        {'role': 'system', 'content': 'You are a helpful assistant.'},
-        {'role': 'user', 'content': prompt}
-    ]
-    
-    completions = openai.ChatCompletion.create(
-        model="gpt-3.5-turbo",
-        messages=messages,
-        max_tokens=512,
-        n=1,
-        stop=None,
-        temperature=0.7,
+    response = openai.ChatCompletion.create(
+        model='gpt-3.5-turbo',
+        messages=[{"role": "user", "content": f"Generate PlantUML code for the following use case: {input_text}"}],
+        max_tokens=150,
+        temperature=0,
     )
-    message = completions.choices[0].message['content']
-    return message
-
-def plantuml_code_to_svg(plantuml_code):
-    with open("diagram.txt", "w") as f:
-        f.write(plantuml_code)
-
-    os.system(f"java -jar plantuml.jar -tsvg diagram.txt")
-
-    with open("diagram.svg", "r") as f:
-        svg = f.read()
-
-    return svg
-
-def create_diagram(api_key, description):
-    plantuml_code = generate_plantuml_code(api_key, description)
-    svg_image = plantuml_code_to_svg(plantuml_code)
-    return plantuml_code, svg_image
-
-api_key_input = gr.Textbox(label="Enter your OpenAI API key")
-description_input = gr.Textbox(lines=5, label="Describe the diagram in natural language")
-plantuml_code_output = gr.Textbox(lines=10, label="Generated PlantUML code")
-diagram_output = gr.Image(plot=True, label="Diagram")
-
-iface = gr.Interface(
-    fn=create_diagram,
-    inputs=[api_key_input, description_input],
-    outputs=[plantuml_code_output, diagram_output],
-)
-iface.launch()
+    plantuml_code = response["choices"][0]["message"]["content"]
+    return plantuml_code
+
+def render_plantuml_diagram(plantuml_code):
+    with tempfile.TemporaryDirectory() as tempdir:
+        puml = PlantUML(f"file://{tempdir}")
+        puml.processes(plantuml_code.encode())
+        with open(os.path.join(tempdir, "output.svg"), "r") as f:
+            svg_diagram = f.read()
+    return svg_diagram
+
+def app(api_key, input_text="", plantuml_code=""):
+    if input_text:
+        plantuml_code = generate_plantuml_code(api_key, input_text)
+    svg_diagram = render_plantuml_diagram(plantuml_code)
+    return plantuml_code, svg_diagram
+
+inputs = [
+    gr.inputs.Textbox(label="Enter OpenAI API Key"),
+    gr.inputs.Textbox(lines=5, label="Enter use case in natural language (optional)"),
+    gr.inputs.Textbox(lines=5, label="Paste PlantUML code (optional)")
+]
+outputs = [
+    gr.outputs.Textbox(label="Generated PlantUML Code"),
+    gr.outputs.HTML(label="Rendered Diagram")
+]
+
+gr.Interface(fn=app, inputs=inputs, outputs=outputs).launch()