44"""
55
66import asyncio
7+ import logging
78import warnings
89from abc import ABC , abstractmethod
910from concurrent .futures import ThreadPoolExecutor
1516from ..types .content import ContentBlock
1617from ..types .event_loop import Metrics , Usage
1718
19+ logger = logging .getLogger (__name__ )
20+
1821
1922class Status (Enum ):
2023 """Execution status for both graphs and nodes."""
@@ -59,6 +62,54 @@ def get_agent_results(self) -> list[AgentResult]:
5962 flattened .extend (nested_node_result .get_agent_results ())
6063 return flattened
6164
65+ def to_dict (self ) -> dict [str , Any ]:
66+ """Convert NodeResult to JSON-serializable dict, ignoring state field."""
67+ if isinstance (self .result , Exception ):
68+ result_data : dict [str , Any ] = {"type" : "exception" , "message" : str (self .result )}
69+ elif isinstance (self .result , AgentResult ):
70+ result_data = self .result .to_dict ()
71+ else :
72+ # MultiAgentResult case
73+ result_data = self .result .to_dict ()
74+
75+ return {
76+ "result" : result_data ,
77+ "execution_time" : self .execution_time ,
78+ "status" : self .status .value ,
79+ "accumulated_usage" : self .accumulated_usage ,
80+ "accumulated_metrics" : self .accumulated_metrics ,
81+ "execution_count" : self .execution_count ,
82+ }
83+
84+ @classmethod
85+ def from_dict (cls , data : dict [str , Any ]) -> "NodeResult" :
86+ """Rehydrate a NodeResult from persisted JSON."""
87+ if "result" not in data :
88+ raise TypeError ("NodeResult.from_dict: missing 'result'" )
89+ raw = data ["result" ]
90+
91+ result : Union [AgentResult , "MultiAgentResult" , Exception ]
92+ if isinstance (raw , dict ) and raw .get ("type" ) == "agent_result" :
93+ result = AgentResult .from_dict (raw )
94+ elif isinstance (raw , dict ) and raw .get ("type" ) == "exception" :
95+ result = Exception (str (raw .get ("message" , "node failed" )))
96+ elif isinstance (raw , dict ) and raw .get ("type" ) == "multiagent_result" :
97+ result = MultiAgentResult .from_dict (raw )
98+ else :
99+ raise TypeError (f"NodeResult.from_dict: unsupported result payload: { raw !r} " )
100+
101+ usage = _parse_usage (data .get ("accumulated_usage" , {}))
102+ metrics = _parse_metrics (data .get ("accumulated_metrics" , {}))
103+
104+ return cls (
105+ result = result ,
106+ execution_time = int (data .get ("execution_time" , 0 )),
107+ status = Status (data .get ("status" , "pending" )),
108+ accumulated_usage = usage ,
109+ accumulated_metrics = metrics ,
110+ execution_count = int (data .get ("execution_count" , 0 )),
111+ )
112+
62113
63114@dataclass
64115class MultiAgentResult :
@@ -76,6 +127,38 @@ class MultiAgentResult:
76127 execution_count : int = 0
77128 execution_time : int = 0
78129
130+ @classmethod
131+ def from_dict (cls , data : dict [str , Any ]) -> "MultiAgentResult" :
132+ """Rehydrate a MultiAgentResult from persisted JSON."""
133+ if data .get ("type" ) != "multiagent_result" :
134+ raise TypeError (f"MultiAgentResult.from_dict: unexpected type { data .get ('type' )!r} " )
135+
136+ results = {k : NodeResult .from_dict (v ) for k , v in data .get ("results" , {}).items ()}
137+ usage = _parse_usage (data .get ("accumulated_usage" , {}))
138+ metrics = _parse_metrics (data .get ("accumulated_metrics" , {}))
139+
140+ multiagent_result = cls (
141+ status = Status (data .get ("status" , Status .PENDING .value )),
142+ results = results ,
143+ accumulated_usage = usage ,
144+ accumulated_metrics = metrics ,
145+ execution_count = int (data .get ("execution_count" , 0 )),
146+ execution_time = int (data .get ("execution_time" , 0 )),
147+ )
148+ return multiagent_result
149+
150+ def to_dict (self ) -> dict [str , Any ]:
151+ """Convert MultiAgentResult to JSON-serializable dict."""
152+ return {
153+ "type" : "multiagent_result" ,
154+ "status" : self .status .value ,
155+ "results" : {k : v .to_dict () for k , v in self .results .items ()},
156+ "accumulated_usage" : self .accumulated_usage ,
157+ "accumulated_metrics" : self .accumulated_metrics ,
158+ "execution_count" : self .execution_count ,
159+ "execution_time" : self .execution_time ,
160+ }
161+
79162
80163class MultiAgentBase (ABC ):
81164 """Base class for multi-agent helpers.
@@ -122,3 +205,34 @@ def execute() -> MultiAgentResult:
122205 with ThreadPoolExecutor () as executor :
123206 future = executor .submit (execute )
124207 return future .result ()
208+
209+ def serialize_state (self ) -> dict [str , Any ]:
210+ """Return a JSON-serializable snapshot of the orchestrator state."""
211+ raise NotImplementedError
212+
213+ def deserialize_state (self , payload : dict [str , Any ]) -> None :
214+ """Restore orchestrator state from a session dict."""
215+ raise NotImplementedError
216+
217+
218+ # Private helper function to avoid duplicate code
219+
220+
221+ def _parse_usage (usage_data : dict [str , Any ]) -> Usage :
222+ """Parse Usage from dict data."""
223+ usage = Usage (
224+ inputTokens = usage_data .get ("inputTokens" , 0 ),
225+ outputTokens = usage_data .get ("outputTokens" , 0 ),
226+ totalTokens = usage_data .get ("totalTokens" , 0 ),
227+ )
228+ # Add optional fields if they exist
229+ if "cacheReadInputTokens" in usage_data :
230+ usage ["cacheReadInputTokens" ] = usage_data ["cacheReadInputTokens" ]
231+ if "cacheWriteInputTokens" in usage_data :
232+ usage ["cacheWriteInputTokens" ] = usage_data ["cacheWriteInputTokens" ]
233+ return usage
234+
235+
236+ def _parse_metrics (metrics_data : dict [str , Any ]) -> Metrics :
237+ """Parse Metrics from dict data."""
238+ return Metrics (latencyMs = metrics_data .get ("latencyMs" , 0 ))
0 commit comments