Skip to content

Commit cc27387

Browse files
committed
fix: collect eval state from workflow nodes
1 parent e90b119 commit cc27387

2 files changed

Lines changed: 68 additions & 6 deletions

File tree

src/google/adk/cli/utils/state.py

Lines changed: 18 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -18,13 +18,25 @@
1818
from typing import Any
1919
from typing import Optional
2020

21-
from ...agents.base_agent import BaseAgent
2221
from ...agents.llm_agent import LlmAgent
22+
from ...workflow import BaseNode
2323

2424

25-
def _create_empty_state(agent: BaseAgent, all_state: dict[str, Any]):
26-
for sub_agent in agent.sub_agents:
27-
_create_empty_state(sub_agent, all_state)
25+
def _create_empty_state(
26+
agent: BaseNode, all_state: dict[str, Any], visited: set[int]
27+
):
28+
agent_id = id(agent)
29+
if agent_id in visited:
30+
return
31+
visited.add(agent_id)
32+
33+
for sub_agent in getattr(agent, 'sub_agents', []) or []:
34+
_create_empty_state(sub_agent, all_state, visited)
35+
36+
graph = getattr(agent, 'graph', None)
37+
if graph is not None:
38+
for graph_node in getattr(graph, 'nodes', []) or []:
39+
_create_empty_state(graph_node, all_state, visited)
2840

2941
if (
3042
isinstance(agent, LlmAgent)
@@ -36,11 +48,11 @@ def _create_empty_state(agent: BaseAgent, all_state: dict[str, Any]):
3648

3749

3850
def create_empty_state(
39-
agent: BaseAgent, initialized_states: Optional[dict[str, Any]] = None
51+
agent: BaseNode, initialized_states: Optional[dict[str, Any]] = None
4052
) -> dict[str, Any]:
4153
"""Creates empty str for non-initialized states."""
4254
non_initialized_states = {}
43-
_create_empty_state(agent, non_initialized_states)
55+
_create_empty_state(agent, non_initialized_states, set())
4456
for key in initialized_states or {}:
4557
if key in non_initialized_states:
4658
del non_initialized_states[key]
Lines changed: 50 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,50 @@
1+
# Copyright 2026 Google LLC
2+
#
3+
# Licensed under the Apache License, Version 2.0 (the "License");
4+
# you may not use this file except in compliance with the License.
5+
# You may obtain a copy of the License at
6+
#
7+
# http://www.apache.org/licenses/LICENSE-2.0
8+
#
9+
# Unless required by applicable law or agreed to in writing, software
10+
# distributed under the License is distributed on an "AS IS" BASIS,
11+
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12+
# See the License for the specific language governing permissions and
13+
# limitations under the License.
14+
15+
from __future__ import annotations
16+
17+
from google.adk.agents.llm_agent import LlmAgent
18+
from google.adk.cli.utils.state import create_empty_state
19+
from google.adk.workflow import START
20+
from google.adk.workflow._workflow import Workflow
21+
22+
23+
def test_create_empty_state_reads_agent_tree():
24+
child = LlmAgent(name='child', instruction='Use {child_key}')
25+
root = LlmAgent(
26+
name='root',
27+
instruction='Use {root_key}',
28+
sub_agents=[child],
29+
)
30+
31+
assert create_empty_state(root) == {
32+
'child_key': '',
33+
'root_key': '',
34+
}
35+
36+
37+
def test_create_empty_state_reads_workflow_graph_nodes():
38+
node = LlmAgent(name='node', instruction='Use {workflow_key}')
39+
workflow = Workflow(name='workflow', edges=[(START, node)])
40+
41+
assert create_empty_state(workflow) == {'workflow_key': ''}
42+
43+
44+
def test_create_empty_state_skips_initialized_workflow_state():
45+
node = LlmAgent(name='node', instruction='Use {workflow_key} and {fresh_key}')
46+
workflow = Workflow(name='workflow', edges=[(START, node)])
47+
48+
assert create_empty_state(workflow, {'workflow_key': 'set'}) == {
49+
'fresh_key': ''
50+
}

0 commit comments

Comments
 (0)