Spaces:
Sleeping
Sleeping
File size: 5,521 Bytes
dc12f61 |
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 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 |
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() |