Spaces:
Sleeping
Sleeping
File size: 3,269 Bytes
f6662f9 98fd32e f6662f9 98fd32e f6662f9 98fd32e f6662f9 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 |
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() |