Skip to content

Commit 1d61969

Browse files
committed
fix: collect eval state from workflow nodes
1 parent 59f7bdf commit 1d61969

2 files changed

Lines changed: 65 additions & 6 deletions

File tree

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

Lines changed: 17 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -18,13 +18,24 @@
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
2322

2423

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

2940
if (
3041
isinstance(agent, LlmAgent)
@@ -36,11 +47,11 @@ def _create_empty_state(agent: BaseAgent, all_state: dict[str, Any]):
3647

3748

3849
def create_empty_state(
39-
agent: BaseAgent, initialized_states: Optional[dict[str, Any]] = None
50+
agent: Any, initialized_states: Optional[dict[str, Any]] = None
4051
) -> dict[str, Any]:
4152
"""Creates empty str for non-initialized states."""
4253
non_initialized_states = {}
43-
_create_empty_state(agent, non_initialized_states)
54+
_create_empty_state(agent, non_initialized_states, set())
4455
for key in initialized_states or {}:
4556
if key in non_initialized_states:
4657
del non_initialized_states[key]
Lines changed: 48 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,48 @@
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 google.adk.agents.llm_agent import LlmAgent
16+
from google.adk.cli.utils.state import create_empty_state
17+
from google.adk.workflow import START
18+
from google.adk.workflow._workflow import Workflow
19+
20+
21+
def test_create_empty_state_reads_agent_tree():
22+
child = LlmAgent(name='child', instruction='Use {child_key}')
23+
root = LlmAgent(
24+
name='root',
25+
instruction='Use {root_key}',
26+
sub_agents=[child],
27+
)
28+
29+
assert create_empty_state(root) == {
30+
'child_key': '',
31+
'root_key': '',
32+
}
33+
34+
35+
def test_create_empty_state_reads_workflow_graph_nodes():
36+
node = LlmAgent(name='node', instruction='Use {workflow_key}')
37+
workflow = Workflow(name='workflow', edges=[(START, node)])
38+
39+
assert create_empty_state(workflow) == {'workflow_key': ''}
40+
41+
42+
def test_create_empty_state_skips_initialized_workflow_state():
43+
node = LlmAgent(name='node', instruction='Use {workflow_key} and {fresh_key}')
44+
workflow = Workflow(name='workflow', edges=[(START, node)])
45+
46+
assert create_empty_state(workflow, {'workflow_key': 'set'}) == {
47+
'fresh_key': ''
48+
}

0 commit comments

Comments
 (0)