test / app.py
FILALIHicham's picture
Add initial implementation of tools and requirements
dc12f61
import os
import logging
import json
import yaml
import base64
import tempfile
from smolagents import CodeAgent, HfApiModel, load_tool, tool
from tools.final_answer import FinalAnswerTool
from Gradio_UI import GradioUI
from plantuml import PlantUML
from smolagents.agent_types import AgentImage
from PIL import Image
# ------------------------------
# UML Generation Tool
# ------------------------------
@tool
def generate_uml_class_diagram(description: str, extra_details: str = "") -> object:
"""
Generates a UML CLASS diagram image (in PlantUML syntax) from a system description.
The tool instructs the LLM to assess the provided system description and,
if necessary, add reasonable assumptions to generate complete UML code.
It then enters a loop, re-calling the model until the output is clean—that is,
it starts with '@startuml' and ends with '@enduml' with no extra text.
Once a clean code block is obtained, it sends the code to a PlantUML server to produce
a PNG diagram, saves the image to a temporary file, opens it with PIL, and returns the image
wrapped as an AgentImage so that the Gradio UI displays it directly.
Args:
description: A natural language description (or topic) of the system.
extra_details: (Optional) Additional details to supplement the description.
Returns:
An AgentImage object wrapping the generated UML diagram image if successful,
otherwise a string with an error message.
"""
# Combine description and extra details.
full_description = description.strip()
if extra_details.strip():
full_description += "\n" + extra_details.strip()
# Build the initial prompt.
prompt = (
"You are an expert in UML class diagramming. Given the following system description, "
"generate complete PlantUML code for a UML class diagram that represents the system. "
"Do not include any extra commentary—only output the UML code, starting with '@startuml' "
"and ending with '@enduml'.\n\n"
f"System Description:\n{full_description}\n\nPlantUML Code:\n@startuml"
)
max_attempts = 5
attempt = 0
plantuml_code = ""
messages = [{"role": "user", "content": prompt}]
# Loop until the output is clean.
while attempt < max_attempts:
try:
response_obj = model(messages, stop_sequences=["@enduml"], max_tokens=2096)
raw_output = response_obj.content if hasattr(response_obj, "content") else response_obj
except Exception as e:
logging.error("Error calling model: " + str(e))
raw_output = ""
start_idx = raw_output.find("@startuml")
end_idx = raw_output.rfind("@enduml")
if start_idx != -1 and end_idx != -1:
candidate_code = raw_output[start_idx:end_idx + len("@enduml")]
else:
candidate_code = raw_output
# Check for any extra text.
before = raw_output[:start_idx].strip() if start_idx != -1 else ""
after = raw_output[end_idx + len("@enduml"):].strip() if end_idx != -1 else ""
if before == "" and after == "":
plantuml_code = candidate_code
break
else:
prompt = (
"Please provide only the UML code in PlantUML syntax with no extra text. "
"Your answer must start with '@startuml' and end with '@enduml'.\n\n"
f"System Description:\n{full_description}\n\nPlantUML Code:\n@startuml"
)
messages = [{"role": "user", "content": prompt}]
attempt += 1
if not plantuml_code:
plantuml_code = candidate_code
# Process the generated PlantUML code to generate a PNG diagram.
pl = PlantUML(url="http://www.plantuml.com/plantuml/img/")
try:
diagram_data = pl.processes(plantuml_code)
except Exception as e:
logging.error("Error generating diagram: " + str(e))
diagram_data = None
if diagram_data:
try:
# Save the diagram data to a temporary PNG file.
tmp_file = tempfile.NamedTemporaryFile(delete=False, suffix=".png", mode="wb")
tmp_file.write(diagram_data)
tmp_file.close()
# Open the image with PIL.
pil_image = Image.open(tmp_file.name)
# Wrap the image as an AgentImage and store the image object.
agent_image = AgentImage(pil_image)
agent_image.image = pil_image # Ensure the image attribute is set
return agent_image
except Exception as e:
logging.error("Error saving diagram to file: " + str(e))
return f"Error saving diagram to file: {str(e)}"
else:
return f"Error generating diagram. Generated PlantUML code was:\n{plantuml_code}"
# ------------------------------
# Setup Model and Agent
# ------------------------------
model = HfApiModel(
max_tokens=2096,
temperature=0.5,
model_id='Qwen/Qwen2.5-Coder-32B-Instruct',
custom_role_conversions=None,
)
final_answer = FinalAnswerTool()
with open("prompts.yaml", 'r') as stream:
prompt_templates = yaml.safe_load(stream)
agent = CodeAgent(
model=model,
tools=[final_answer, generate_uml_class_diagram],
max_steps=6,
verbosity_level=1,
grammar=None,
planning_interval=None,
name=None,
description=None,
prompt_templates=prompt_templates
)
GradioUI(agent).launch()