CmdForge/scripts/test_nodegraph.py

128 lines
3.5 KiB
Python

#!/usr/bin/env python3
"""Test script for NodeGraphQt integration."""
import sys
from PySide6.QtWidgets import QApplication, QMainWindow, QVBoxLayout, QWidget
from NodeGraphQt import NodeGraph, BaseNode
class InputNode(BaseNode):
"""Node representing tool input."""
__identifier__ = 'cmdforge.nodes'
NODE_NAME = 'Input'
def __init__(self):
super().__init__()
self.set_color(100, 100, 180) # Blue-ish
self.add_output('input', color=(180, 180, 250))
self.add_output('file', color=(180, 180, 250))
class PromptNode(BaseNode):
"""Node representing a prompt step."""
__identifier__ = 'cmdforge.nodes'
NODE_NAME = 'Prompt'
def __init__(self):
super().__init__()
self.set_color(80, 120, 180) # Indigo-ish
self.add_input('in', color=(180, 180, 250))
self.add_output('response', color=(180, 250, 180))
# Add properties
self.add_text_input('provider', 'Provider', text='claude')
self.add_text_input('output_var', 'Output Var', text='response')
class CodeNode(BaseNode):
"""Node representing a code step."""
__identifier__ = 'cmdforge.nodes'
NODE_NAME = 'Code'
def __init__(self):
super().__init__()
self.set_color(80, 160, 100) # Green-ish
self.add_input('in', color=(180, 250, 180))
self.add_output('result', color=(250, 220, 180))
# Add properties
self.add_text_input('output_var', 'Output Var', text='result')
class OutputNode(BaseNode):
"""Node representing tool output."""
__identifier__ = 'cmdforge.nodes'
NODE_NAME = 'Output'
def __init__(self):
super().__init__()
self.set_color(180, 120, 80) # Orange-ish
self.add_input('in', color=(250, 220, 180))
def main():
app = QApplication(sys.argv)
# Create main window
window = QMainWindow()
window.setWindowTitle('CmdForge Flow Visualization - Prototype')
window.setGeometry(100, 100, 1000, 700)
# Create node graph
graph = NodeGraph()
# Register our custom nodes
graph.register_node(InputNode)
graph.register_node(PromptNode)
graph.register_node(CodeNode)
graph.register_node(OutputNode)
# Get the graph widget
graph_widget = graph.widget
# Set as central widget
window.setCentralWidget(graph_widget)
# Create a sample tool visualization
input_node = graph.create_node('cmdforge.nodes.InputNode', name='Input')
input_node.set_pos(-300, 0)
prompt_node = graph.create_node('cmdforge.nodes.PromptNode', name='Summarize')
prompt_node.set_pos(-50, 0)
prompt_node.set_property('provider', 'claude')
prompt_node.set_property('output_var', 'response')
code_node = graph.create_node('cmdforge.nodes.CodeNode', name='Format')
code_node.set_pos(200, 0)
code_node.set_property('output_var', 'formatted')
output_node = graph.create_node('cmdforge.nodes.OutputNode', name='Output')
output_node.set_pos(450, 0)
# Connect nodes
input_node.output(0).connect_to(prompt_node.input(0))
prompt_node.output(0).connect_to(code_node.input(0))
code_node.output(0).connect_to(output_node.input(0))
# Auto-layout
graph.auto_layout_nodes()
# Fit view
graph.fit_to_selection()
window.show()
print("Flow visualization prototype running...")
print("Nodes created: Input -> Prompt -> Code -> Output")
print("Close the window to exit.")
sys.exit(app.exec())
if __name__ == '__main__':
main()