NodeLayout

  1# NodeLayout.py
  2
  3from PySide6.QtWidgets import QGraphicsItem, QLineEdit, QGraphicsProxyWidget, QVBoxLayout, QWidget, QLabel, QComboBox, QTextEdit, QHBoxLayout
  4from PySide6.QtCore import QRectF, Qt
  5from PySide6.QtGui import QPainter, QBrush, QPen
  6
  7class Slot:
  8    def __init__(self, data_attr, parent, layout, widget_class=QTextEdit):
  9        self.data_attr = data_attr
 10        self.parent = parent
 11        self.layout = layout
 12        self.widget_class = widget_class
 13
 14        self.create_widgets()
 15        self.add_to_layout()
 16        self.connect_signals()
 17
 18    def create_widgets(self):
 19        self.label = QLabel(self.data_attr.capitalize())
 20        self.edit = self.widget_class(getattr(self.parent.data, self.data_attr))
 21
 22    def add_to_layout(self):
 23        self.layout.addWidget(self.label)
 24        self.layout.addWidget(self.edit)
 25
 26    def connect_signals(self):
 27        if isinstance(self.edit, QTextEdit):
 28            self.edit.textChanged.connect(self.update_data)
 29        elif isinstance(self.edit, QLineEdit):
 30            self.edit.textChanged.connect(self.update_data)
 31
 32    def update_data(self):
 33        if isinstance(self.edit, QTextEdit):
 34            setattr(self.parent.data, self.data_attr, self.edit.toPlainText())
 35        elif isinstance(self.edit, QLineEdit):
 36            setattr(self.parent.data, self.data_attr, self.edit.text())
 37
 38    def show(self):
 39        self.label.show()
 40        self.edit.show()
 41
 42    def hide(self):
 43        self.label.hide()
 44        self.edit.hide()
 45
 46class NodeLayout(QGraphicsItem):
 47    def __init__(self, parent):
 48        super().__init__(parent)
 49        self.parent = parent  # Save the reference to the parent node
 50
 51        # Create a widget to hold the QLineEdit widgets
 52        self.container_widget = QWidget()
 53        self.layout = QVBoxLayout()
 54        self.container_widget.setLayout(self.layout)
 55
 56        # Type Label and Combo Box
 57        self.type_layout = QHBoxLayout()
 58        self.type_label = QLabel("Type:")
 59        self.type_combo = QComboBox()
 60        self.type_combo.addItems(["None", "Start", "Agent", "Task", "Step", "Team"])
 61        self.type_combo.setCurrentText(self.parent.data.type)
 62        self.type_layout.addWidget(self.type_label)
 63        self.type_layout.addWidget(self.type_combo)
 64        self.layout.addLayout(self.type_layout)
 65        self.type_combo.currentTextChanged.connect(self.update_data_type)
 66
 67        # Initialize slots
 68        self.slots = {}
 69        self.add_slot("name", QLineEdit)
 70
 71        # Agent
 72        self.add_slot("role", QTextEdit)
 73        self.add_slot("goal", QTextEdit)
 74        self.add_slot("backstory", QTextEdit)
 75        
 76        # Task
 77        self.add_slot("agent", QTextEdit)
 78        self.add_slot("description", QTextEdit)
 79        self.add_slot("expected_output", QTextEdit)
 80        
 81        # Step
 82        self.add_slot("tool", QTextEdit)
 83        self.add_slot("arg", QTextEdit)
 84        self.add_slot("output_var", QTextEdit)
 85        
 86        self.proxy_widget = QGraphicsProxyWidget(self)
 87        self.proxy_widget.setWidget(self.container_widget)
 88
 89        # Initialize visibility based on the current type
 90        self.update_field_visibility()
 91
 92        # Call update_proxy_widget_geometry once to ensure layout is correct
 93        self.update_proxy_widget_geometry()
 94
 95    def add_slot(self, data_attr, widget_class):
 96        slot = Slot(data_attr, self.parent, self.layout, widget_class)
 97        self.slots[data_attr] = slot
 98
 99    def boundingRect(self):
100        return self.parent.rect
101
102    def paint(self, painter, option, widget):
103        rect = self.boundingRect()
104        painter.setBrush(QBrush(Qt.white))
105        painter.setPen(QPen(Qt.black))
106        painter.drawRect(rect)
107
108        # Draw right-bottom square
109        painter.setBrush(QBrush(Qt.black))
110        painter.drawRect(rect.right() - 10, rect.bottom() - 10, 10, 10)
111
112    def update_proxy_widget_geometry(self):
113        rect = self.boundingRect()
114        self.proxy_widget.setGeometry(rect.adjusted(5, 5, -5, -5))
115
116    def update_data_type(self):
117        self.parent.data.type = self.type_combo.currentText()
118        self.update_field_visibility()  # Update field visibility when type changes
119
120    def update_field_visibility(self):
121        node_type = self.parent.data.type
122        for slot in self.slots.values():
123            slot.hide()
124
125        if node_type in ["Start", "None"]:
126            pass
127        elif node_type == "Team":
128            self.slots["name"].show()
129        elif node_type == "Agent":
130            self.slots["name"].show()
131            self.slots["role"].show()
132            self.slots["goal"].show()
133            self.slots["backstory"].show()
134        elif node_type == "Task":
135            self.slots["agent"].show()
136            self.slots["description"].show()
137            self.slots["expected_output"].show()
138        elif node_type == "Step":
139            self.slots["tool"].show()
140            self.slots["arg"].show()
141            self.slots["output_var"].show()
142
143
144
145        # Update geometry after setting visibility
146        self.update_proxy_widget_geometry()