8000 chore: Fix formatting · google/adk-python@eb2b9ef · GitHub
[go: up one dir, main page]
Skip to content

Commit eb2b9ef

Browse files
selcukguncopybara-github
selcukgun
authored and
copybara-github
committed
chore: Fix formatting
PiperOrigin-RevId: 766407362
1 parent 174afb3 commit eb2b9ef

File tree

1 file changed

+57
-29
lines changed

1 file changed

+57
-29
lines changed

‎src/google/adk/cli/agent_graph.py

Lines changed: 57 additions & 29 deletions
< 10000 td data-grid-cell-id="diff-f74a0b28d97172ba6fbc7779764ecb5a9f3a2ffa68690670a522544b348667b6-22-23-1" data-selected="false" role="gridcell" style="background-color:var(--diffBlob-additionNum-bgColor, var(--diffBlob-addition-bgColor-num));text-align:center" tabindex="-1" valign="top" class="focusable-grid-cell diff-line-number position-relative left-side">23
Original file line numberDiff line numberDiff line change
@@ -19,7 +19,10 @@
1919

2020
import graphviz
2121

22-
from ..agents import BaseAgent, SequentialAgent, LoopAgent, ParallelAgent
22+
from ..agents import BaseAgent
+
from ..agents import LoopAgent
24+
from ..agents import ParallelAgent
25+
from ..agents import SequentialAgent
2326
from ..agents.llm_agent import LlmAgent
2427
from ..tools.agent_tool import AgentTool
2528
from ..tools.base_tool import BaseTool
@@ -35,15 +38,20 @@
3538
retrieval_tool_module_loaded = True
3639

3740

38-
async def build_graph(graph: graphviz.Digraph, agent: BaseAgent, highlight_pairs, parent_agent=None):
41+
async def build_graph(
42+
graph: graphviz.Digraph,
43+
agent: BaseAgent,
44+
highlight_pairs,
45+
parent_agent=None,
46+
):
3947
"""
4048
Build a graph of the agent and its sub-agents.
4149
Args:
4250
graph: The graph to build on.
4351
agent: The agent to build the graph for.
4452
highlight_pairs: A list of pairs of nodes to highlight.
4553
parent_agent: The parent agent of the current agent. This is specifically used when building Workflow Agents to directly connect a node to nodes inside a Workflow Agent.
46-
54+
4755
Returns:
4856
None
4957
"""
@@ -56,11 +64,11 @@ def get_node_name(tool_or_agent: Union[BaseAgent, BaseTool]):
5664
if isinstance(tool_or_agent, BaseAgent):
5765
# Added Workflow Agent checks for different agent types
5866
if isinstance(tool_or_agent, SequentialAgent):
59-
return tool_or_agent.name + f" (Sequential Agent)"
67+
return tool_or_agent.name + f' (Sequential Agent)'
6068
elif isinstance(tool_or_agent, LoopAgent):
61-
return tool_or_agent.name + f" (Loop Agent)"
69+
return tool_or_agent.name + f' (Loop Agent)'
6270
elif isinstance(tool_or_agent, ParallelAgent):
63-
return tool_or_agent.name + f" (Parallel Agent)"
71+
return tool_or_agent.name + f' (Parallel Agent)'
6472
else:
6573
return tool_or_agent.name
6674
elif isinstance(tool_or_agent, BaseTool):
@@ -93,7 +101,7 @@ def get_node_caption(tool_or_agent: Union[BaseAgent, BaseTool]):
93101
def get_node_shape(tool_or_agent: Union[BaseAgent, BaseTool]):
94102
if isinstance(tool_or_agent, BaseAgent):
95103
return 'ellipse'
96-
104+
97105
elif retrieval_tool_module_loaded and isinstance(
98106
tool_or_agent, BaseRetrievalTool
99107
):
@@ -109,7 +117,7 @@ def get_node_shape(tool_or_agent: Union[BaseAgent, BaseTool]):
109117
tool_or_agent,
110118
)
111119
return 'cylinder'
112-
120+
113121
def should_build_agent_cluster(tool_or_agent: Union[BaseAgent, BaseTool]):
114122
if isinstance(tool_or_agent, BaseAgent):
115123
if isinstance(tool_or_agent, SequentialAgent):
@@ -135,19 +143,24 @@ def should_build_agent_cluster(tool_or_agent: Union[BaseAgent, BaseTool]):
135143
tool_or_agent,
136144
)
137145
return False
138-
146+
139147
def build_cluster(child: graphviz.Digraph, agent: BaseAgent, name: str):
140-
if isinstance(agent, LoopAgent):
148+
if isinstance(agent, LoopAgent):
141149
# Draw the edge from the parent agent to the first sub-agent
142150
draw_edge(parent_agent.name, agent.sub_agents[0].name)
143151
length = len(agent.sub_agents)
144152
currLength = 0
145153
# Draw the edges between the sub-agents
146-
for sub_agent_int_sequential in agent.sub_agents:
154+
for sub_agent_int_sequential in agent.sub_agents:
147155
build_graph(child, sub_agent_int_sequential, highlight_pairs)
148156
# Draw the edge between the current sub-agent and the next one
149157
# If it's the last sub-agent, draw an edge to the first one to indicating a loop
150-
draw_edge(agent.sub_agents[currLength].name, agent.sub_agents[0 if currLength == length - 1 else currLength+1 ].name)
158+
draw_edge(
159+
agent.sub_agents[currLength].name,
160+
agent.sub_agents[
161+
0 if currLength == length - 1 else currLength + 1
162+
].name,
163+
)
151164
currLength += 1
152165
elif isinstance(agent, SequentialAgent):
153166
# Draw the edge from the parent agent to the first sub-agent
@@ -160,10 +173,13 @@ def build_cluster(child: graphviz.Digraph, agent: BaseAgent, name: str):
160173
build_graph(child, sub_agent_int_sequential, highlight_pairs)
161174
# Draw the edge between the current sub-agent and the next one
162175
# If it's the last sub-agent, don't draw an edge to avoid a loop
163-
draw_edge(agent.sub_agents[currLength].name, agent.sub_agents[currLength+1].name) if currLength != length - 1 else None
176+
draw_edge(
177+
agent.sub_agents[currLength].name,
178+
agent.sub_agents[currLength + 1].name,
179+
) if currLength != length - 1 else None
164180
currLength += 1
165181

166-
elif isinstance(agent, ParallelAgent):
182+
elif isinstance(agent, ParallelAgent):
167183
# Draw the edge from the parent agent to every sub-agent
168184
for sub_agent in agent.sub_agents:
169185
build_graph(child, sub_agent, highlight_pairs)
@@ -191,7 +207,9 @@ def draw_node(tool_or_agent: Union[BaseAgent, BaseTool]):
191207
if name in highlight_tuple:
192208
# if in highlight, draw highlight node
193209
if asCluster:
194-
cluster = graphviz.Digraph(name='cluster_' + name) # adding "cluster_" to the name makes the graph render as a cluster subgraph
210+
cluster = graphviz.Digraph(
211+
name='cluster_' + name
212+
) # adding "cluster_" to the name makes the graph render as a cluster subgraph
195213
build_cluster(cluster, agent, name)
196214
graph.subgraph(cluster)
197215
else:
@@ -207,19 +225,21 @@ def draw_node(tool_or_agent: Union[BaseAgent, BaseTool]):
207225
return
208226
# if not in highlight, draw non-highlight node
209227
if asCluster:
210-
211-
cluster = graphviz.Digraph(name='cluster_' + name) # adding "cluster_" to the name makes the graph render as a cluster subgraph
228+
229+
cluster = graphviz.Digraph(
230+
name='cluster_' + name
231+
) # adding "cluster_" to the name makes the graph render as a cluster subgraph
212232
build_cluster(cluster, agent, name)
213233
graph.subgraph(cluster)
214-
234+
215235
else:
216236
graph.node(
217-
name,
218-
caption,
219-
shape=shape,
220-
style='rounded',
221-
color=light_gray,
222-
fontcolor=light_gray,
237+
name,
238+
caption,
239+
shape=shape,
240+
style='rounded',
241+
color=light_gray,
242+
fontcolor=light_gray,
223243
)
224244

225245
return
@@ -234,17 +254,25 @@ def draw_edge(from_name, to_name):
234254
graph.edge(from_name, to_name, color=light_green, dir='back')
235255
return
236256
# if no need to highlight, color gray
237-
if (should_build_agent_cluster(agent)):
238-
239-
graph.edge(from_name, to_name, color=light_gray, )
257+
if should_build_agent_cluster(agent):
258+
259+
graph.edge(
260+
from_name,
261+
to_name,
262+
color=light_gray,
263+
)
240264
else:
241265
graph.edge(from_name, to_name, arrowhead='none', color=light_gray)
242266

243267
draw_node(agent)
244268
for sub_agent in agent.sub_agents:
245-
269+
246270
build_graph(graph, sub_agent, highlight_pairs, agent)
247-
if (not should_build_agent_cluster(sub_agent) and not should_build_agent_cluster(agent)): # This is to avoid making a node for a Workflow Agent
271+
if not should_build_agent_cluster(
272+
sub_agent
273+
) and not should_build_agent_cluster(
274+
agent
275+
): # This is to avoid making a node for a Workflow Agent
248276
draw_edge(agent.name, sub_agent.name)
249277
if isinstance(agent, LlmAgent):
250278
for tool in await agent.canonical_tools():

0 commit comments

Comments
 (0)
0