1+ """
2+ Flow Display for PraisonAI Agents
3+
4+ Visual display with agents in center and tools on sides.
5+ """
6+
7+ from typing import Dict , List , Set , Tuple
8+ from collections import defaultdict
9+ import threading
10+ from rich .console import Console
11+ from rich .panel import Panel
12+ from rich .text import Text
13+ from rich .align import Align
14+ from rich .columns import Columns
15+ from rich .table import Table
16+ from rich import box
17+
18+ class FlowDisplay :
19+ """Displays agent workflow with agents centered and tools on sides."""
20+
21+ def __init__ (self ):
22+ self .console = Console ()
23+ self .agents = [] # List of agents in order
24+ self .agent_tools = defaultdict (list ) # agent -> [tools]
25+ self .tracking = False
26+ self .lock = threading .Lock ()
27+
28+ def start (self ):
29+ """Start tracking workflow."""
30+ self .tracking = True
31+ self .agents .clear ()
32+ self .agent_tools .clear ()
33+
34+ # Register callbacks
35+ try :
36+ from praisonaiagents .main import register_display_callback
37+
38+ def on_interaction (** kwargs ):
39+ if self .tracking :
40+ agent = kwargs .get ('agent_name' , 'Unknown' )
41+ self ._add_agent (agent )
42+
43+ def on_tool_call (message , ** kwargs ):
44+ if self .tracking and "called function" in message :
45+ parts = message .split ("'" )
46+ if len (parts ) > 1 :
47+ tool_name = parts [1 ]
48+ agent = kwargs .get ('agent_name' , 'Unknown' )
49+ self ._add_tool (agent , tool_name )
50+
51+ register_display_callback ('interaction' , on_interaction )
52+ register_display_callback ('tool_call' , on_tool_call )
53+
54+ except ImportError :
55+ pass
56+
57+ def stop (self ):
58+ """Stop tracking and display the flow."""
59+ self .tracking = False
60+ self .display ()
61+
62+ def _add_agent (self , name : str ):
63+ """Add an agent if not already present."""
64+ with self .lock :
65+ if name not in self .agents :
66+ self .agents .append (name )
67+
68+ def _add_tool (self , agent_name : str , tool_name : str ):
69+ """Add a tool to an agent."""
70+ with self .lock :
71+ if agent_name not in self .agents :
72+ self .agents .append (agent_name )
73+ if tool_name not in self .agent_tools [agent_name ]:
74+ self .agent_tools [agent_name ].append (tool_name )
75+
76+ def display (self ):
77+ """Display the flow chart with agents in center and tools on sides."""
78+ if not self .agents :
79+ return
80+
81+ self .console .print ("\n [bold cyan]🔄 Agent Workflow Flow[/bold cyan]\n " )
82+
83+ # Display start
84+ self ._display_centered_node ("── start ──" , "grey35" )
85+ self ._display_arrow_down ()
86+
87+ # Display each agent with their tools
88+ for i , agent in enumerate (self .agents ):
89+ self ._display_agent_with_tools (agent )
90+
91+ # Add arrow to next agent or end
92+ if i < len (self .agents ) - 1 :
93+ self ._display_arrow_down ()
94+
95+ # Display end
96+ self ._display_arrow_down ()
97+ self ._display_centered_node ("── end ──" , "grey35" )
98+
99+ def _display_agent_with_tools (self , agent : str ):
100+ """Display agent with tools on the sides."""
101+ tools = self .agent_tools .get (agent , [])
102+
103+ if not tools :
104+ # No tools - just agent
105+ self ._display_centered_node (agent , "purple" )
106+ else :
107+ # Split tools between left and right
108+ left_tools = tools [::2 ] # Even indices
109+ right_tools = tools [1 ::2 ] # Odd indices
110+
111+ # Create the layout
112+ table = Table (show_header = False , show_edge = False , box = None , padding = 0 )
113+ table .add_column (justify = "center" , min_width = 20 ) # Left tools
114+ table .add_column (justify = "center" , min_width = 5 ) # Space
115+ table .add_column (justify = "center" , min_width = 20 ) # Agent
116+ table .add_column (justify = "center" , min_width = 5 ) # Space
117+ table .add_column (justify = "center" , min_width = 20 ) # Right tools
118+
119+ # Create panels
120+ left_panel = self ._create_tools_panel (left_tools ) if left_tools else ""
121+ agent_panel = Panel (
122+ Text (agent , style = "white on purple" , justify = "center" ),
123+ style = "white on purple" ,
124+ box = box .ROUNDED ,
125+ padding = (0 , 2 )
126+ )
127+ right_panel = self ._create_tools_panel (right_tools ) if right_tools else ""
128+
129+ # Add row
130+ table .add_row (left_panel , "" , agent_panel , "" , right_panel )
131+
132+ # Display centered
133+ self .console .print (Align .center (table ))
134+
135+ # Show arrows
136+ if left_tools or right_tools :
137+ arrow_parts = []
138+ if left_tools :
139+ arrow_parts .append ("←→" )
140+ else :
141+ arrow_parts .append (" " )
142+
143+ arrow_parts .append (" " ) # Center space
144+
145+ if right_tools :
146+ arrow_parts .append ("←→" )
147+ else :
148+ arrow_parts .append (" " )
149+
150+ self .console .print (Align .center (Text ("" .join (arrow_parts ))))
151+
152+ def _create_tools_panel (self , tools : List [str ]) -> Panel :
153+ """Create a panel for tools."""
154+ if not tools :
155+ return ""
156+
157+ if len (tools ) == 1 :
158+ return Panel (
159+ Text (tools [0 ], style = "black on yellow" , justify = "center" ),
160+ style = "black on yellow" ,
161+ box = box .ROUNDED ,
162+ padding = (0 , 1 )
163+ )
164+ else :
165+ # Multiple tools
166+ content = "\n " .join (tools )
167+ return Panel (
168+ Text (content , style = "black on yellow" , justify = "center" ),
169+ style = "black on yellow" ,
170+ box = box .ROUNDED ,
171+ padding = (0 , 1 )
172+ )
173+
174+ def _display_centered_node (self , label : str , color : str ):
175+ """Display a centered node."""
176+ panel = Panel (
177+ Text (label , style = f"white on { color } , justify = "center" ),
178+ style = f"white on { color } ,
179+ box = box .ROUNDED ,
180+ padding = (0 , 2 )
181+ )
182+ self .console .print (Align .center (panel ))
183+
184+ def _display_arrow_down (self ):
185+ """Display a downward arrow."""
186+ self .console .print ()
187+ self .console .print (Align .center ("↓" ))
188+ self .console .print ()
189+
190+
191+ # Simple function to create and use
192+ def track_workflow ():
193+ """Create a flow display tracker."""
194+ return FlowDisplay ()
0 commit comments