Improviments

README.md

@@ -53,7 +53,7 @@ To set up and run the PlantUML Agent locally, follow these steps:
 2.  **Create a conda environment:**
 
     ```bash
-    conda create -n plantuml-agent python=3.10 --y
+    conda create -n plantuml-agent python=3.12 --y
     conda activate plantuml-agent
     ```
 

app.py

@@ -1,23 +1,21 @@
-import gradio as gr
-from smolagents import InferenceClientModel, CodeAgent
 from src.documents import get_processed_documents
-from src.retriever_tool import RetrieverTool
+from src.retriever import RetrieverTool
 from src.plantuml import render_plantuml
-from src.examples import get_plantuml_examples
-import re
-from dotenv import load_dotenv
-import os
+from src.utils import get_prompt, get_processed_response
+from src.interface import create_interface
 
-load_dotenv()
+from smolagents import InferenceClientModel, CodeAgent
+from cachetools import LRUCache, cached
 
-user_token = os.getenv("HF_TOKEN")
+import os
+from dotenv import load_dotenv
 
-PDF_PATH = "PlantUML_Language_Reference_Guide_en.pdf"
+load_dotenv()
 
-docs_processed = get_processed_documents(PDF_PATH)
+docs_processed = get_processed_documents("PlantUML_Language_Reference_Guide_en.pdf")
 retriever_tool = RetrieverTool(docs_processed)
 
-model = InferenceClientModel('microsoft/phi-4', api_key=user_token)
+model = InferenceClientModel('microsoft/phi-4', api_key=os.getenv("HF_TOKEN"))
 
 agent = CodeAgent(
     tools=[retriever_tool],
@@ -26,76 +24,20 @@ agent = CodeAgent(
     verbosity_level=2,
 )
 
+@cached(LRUCache(maxsize=100))
 def get_plantuml_diagram(message):
-    prompt = ("Generate a PlantUML diagram based on the following message:\n"
-              f"{message}\n"
-              "Use the PlantUML syntax and include all necessary components.\n"
-              "Return the PlantUML code wrapped in @startuml and @enduml tags.")
-    
+    prompt = get_prompt(message)
     response = agent.run(prompt)
-    matches = re.findall(r"@startuml(.*?)@enduml", response, re.DOTALL)
-    if matches:
-        return f"@startuml\n{matches[0].strip()}\n@enduml"
-    return "@startuml\n@enduml"
+    
+    return get_processed_response(response)
 
 def respond(message):
     plantuml_code = get_plantuml_diagram(message)
     svg_output = render_plantuml(plantuml_code)
     return "", plantuml_code, svg_output
 
-def set_example_input(example_name):
-    return PLANTUML_EXAMPLES[example_name]["input"]
-
-PLANTUML_EXAMPLES = get_plantuml_examples()
-example_names = list(PLANTUML_EXAMPLES.keys())
-
-with gr.Blocks(title="PlantUML Agent") as demo:
-    with gr.Row():
-        gr.HTML(
-            """
-            <div style="display: flex; align-items: center; justify-content: center; gap: 10px;">
-                <img src="https://repository-images.githubusercontent.com/553868400/2f7375d8-c2ee-44fe-81d0-a8a29ec284da" alt="PlantUML Logo" style="height: 50px;">
-                <h1>PlantUML Agent</h1>
-            </div>
-            <p style="font-size: 1.2em;">Generate UML diagrams using <a href="https://plantuml.com" target="_blank">PlantUML</a> syntax with the help of an AI agent.</p>
-            <p style="font-size: 0.9em;">Diagrams supported: Sequence, Use Case, Class, Object, Activity, Component, Deployment, State, and Timing.</p>
-            """
-        )
+if __name__ == "__main__":
+    demo = create_interface(respond)
+    demo.launch()
     
-    with gr.Row():
-        with gr.Column(scale=1):
-            gr.Markdown("### Input")
-            msg = gr.Textbox(label="Type your PlantUML description here...", placeholder="e.g., create a basic use case class", lines=3)
-            gr.Markdown("### Examples:")
-            example_dropdown = gr.Dropdown(
-                choices=example_names,
-                label="Choose a diagram description example",
-                interactive=True
-            )
-
-            example_dropdown.change(
-                fn=set_example_input,
-                inputs=example_dropdown,
-                outputs=msg
-            )
-
-            submit = gr.Button("Generate Diagram")
-
-        
-    with gr.Row():
-        with gr.Column(scale=1):
-            gr.Markdown("### PlantUML Code")
-            diagram_text_output = gr.Code(label="", lines=15, interactive=False)
-        with gr.Column(scale=1):
-            gr.Markdown("### PlantUML Preview")
-            diagram_output = gr.HTML(label="")
-
-    with gr.Row():
-        gr.Markdown(
-            "Powered by SmolAgents, Phi-4, BM25Retriever, LangChain, Gradio, and PlantUML."
-        )
-
-    submit.click(respond, [msg], [msg, diagram_text_output, diagram_output])
-    msg.submit(respond, [msg], [msg, diagram_text_output, diagram_output])
-
 demo.launch()

requirements.txt

@@ -7,4 +7,5 @@ smolagents[gradio]
 plantweb
 rank_bm25
 numpy==1.26.4
-python-dotenv
+python-dotenv
+cachetools

src/examples.py

@@ -1,8 +1,8 @@
 def get_plantuml_examples():
     return {
-        'Unselected': {
+        'Default example': {
             "description": "Select a diagram type to see examples.",
-            "input": ""
+            "input": "Generate a simple use case diagram."
         },
         "Use Case Diagram": {
             "description": "A use case diagram showing how users (actors) interact with a system and what functionalities (use cases) they use.",

src/interface.py

@@ -0,0 +1,60 @@
+import gradio as gr
+from src.examples import get_plantuml_examples
+
+def create_interface(respond):
+    PLANTUML_EXAMPLES = get_plantuml_examples()
+    example_names = list(PLANTUML_EXAMPLES.keys())
+
+    def set_example_input(example_name):
+        return PLANTUML_EXAMPLES.get(example_name, "")
+    
+    with gr.Blocks(title="PlantUML Agent") as demo:
+        with gr.Row():
+            gr.HTML(
+                """
+                <div style="display: flex; align-items: center; justify-content: center; gap: 10px;">
+                    <img src="https://repository-images.githubusercontent.com/553868400/2f7375d8-c2ee-44fe-81d0-a8a29ec284da" alt="PlantUML Logo" style="height: 50px;">
+                    <h1>PlantUML Agent</h1>
+                </div>
+                <p style="font-size: 1.2em;">Generate UML diagrams using <a href="https://plantuml.com" target="_blank">PlantUML</a> syntax with the help of an AI agent.</p>
+                <p style="font-size: 0.9em;">Diagrams supported: Sequence, Use Case, Class, Object, Activity, Component, Deployment, State, and Timing.</p>
+                """
+            )
+        
+        with gr.Row():
+            with gr.Column(scale=1):
+                gr.Markdown("### Input")
+                msg = gr.Textbox(label="Type your PlantUML description here...", placeholder="e.g., create a basic use case class", lines=3)
+                gr.Markdown("### Examples:")
+                example_dropdown = gr.Dropdown(
+                    choices=example_names,
+                    label="Choose a diagram description example",
+                    interactive=True
+                )
+
+                example_dropdown.change(
+                    fn=set_example_input,
+                    inputs=example_dropdown,
+                    outputs=msg
+                )
+
+                submit = gr.Button("Generate Diagram")
+                
+
+        with gr.Row():
+            with gr.Column(scale=1):
+                gr.Markdown("### PlantUML Code")
+                diagram_text_output = gr.Code(label="", lines=15, interactive=False)
+            with gr.Column(scale=1):
+                gr.Markdown("### PlantUML Preview")
+                diagram_output = gr.HTML(label="")
+
+        with gr.Row():
+            gr.Markdown(
+                "Powered by SmolAgents, Phi-4, BM25Retriever, LangChain, LRUCache, Gradio, and PlantUML."
+            )
+
+        submit.click(respond, [msg], [msg, diagram_text_output, diagram_output])
+        msg.submit(respond, [msg], [msg, diagram_text_output, diagram_output])
+
+    return demo

src/plantuml.py

@@ -1,5 +1,5 @@
 from plantweb import render
 
 def render_plantuml(uml_code):
-    svg_output = render.render(uml_code, engine='plantuml', format='svg')
+    svg_output, _, _, _ = render.render(uml_code, engine='plantuml', format='svg')
     return svg_output

src/utils.py

@@ -0,0 +1,14 @@
+import re
+
+def get_prompt(message):
+    return ("Generate a PlantUML diagram based on the following message:\n"
+            f"{message}\n"
+            "Use the PlantUML syntax and include all necessary components.\n"
+            "Return the PlantUML code wrapped in @startuml and @enduml tags."
+            "Do not include any additional text outside these tags without commenting it out.")
+    
+def get_processed_response(response):
+    matches = re.findall(r"@startuml(.*?)@enduml", response, re.DOTALL)
+    if matches:
+        return f"@startuml\n{matches[0].strip()}\n@enduml"
+    return "@startuml\n@enduml"