1818from typing import Any
1919from typing import Optional
2020
21- from ...agents .base_agent import BaseAgent
2221from ...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
3849def 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 ]
0 commit comments