app.py

import gradio as gr
from smolagents import InferenceClientModel, CodeAgent
from src.documents import get_processed_documents
from src.retriever_tool import RetrieverTool
from src.plantuml import render_plantuml
from src.examples import get_plantuml_examples
import re
from dotenv import load_dotenv
import os

load_dotenv()

user_token = os.getenv("HF_TOKEN")

PDF_PATH = "PlantUML_Language_Reference_Guide_en.pdf"

docs_processed = get_processed_documents(PDF_PATH)
retriever_tool = RetrieverTool(docs_processed)

model = InferenceClientModel('microsoft/phi-4', api_key=user_token)

agent = CodeAgent(
    tools=[retriever_tool],
    model=model,
    max_steps=4,
    verbosity_level=2,
)

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.")
    
    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"

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 PlantUML diagrams using LLMs and view them in real-time.</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 request 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="")

    submit.click(respond, [msg], [msg, diagram_text_output, diagram_output])
    msg.submit(respond, [msg], [msg, diagram_text_output, diagram_output])

demo.launch()