Spaces:
Sleeping
Sleeping
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 | |
# ------------------------------ | |
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() |