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()