@@ -60,32 +60,22 @@ def __init__(self):
6060 self .previous_step_type = None
6161
6262 def yield_printable_events (self , chunk ):
63- for printable_event in self ._yield_printable_events (
64- chunk , self .previous_event_type , self .previous_step_type
65- ):
63+ for printable_event in self ._yield_printable_events (chunk , self .previous_event_type , self .previous_step_type ):
6664 yield printable_event
6765
68- self .previous_event_type , self .previous_step_type = (
69- self ._get_event_type_step_type (chunk )
70- )
66+ self .previous_event_type , self .previous_step_type = self ._get_event_type_step_type (chunk )
7167
72- def _yield_printable_events (
73- self , chunk , previous_event_type = None , previous_step_type = None
74- ):
68+ def _yield_printable_events (self , chunk , previous_event_type = None , previous_step_type = None ):
7569 if hasattr (chunk , "error" ):
76- yield TurnStreamPrintableEvent (
77- role = None , content = chunk .error ["message" ], color = "red"
78- )
70+ yield TurnStreamPrintableEvent (role = None , content = chunk .error ["message" ], color = "red" )
7971 return
8072
8173 if not hasattr (chunk , "event" ):
8274 # Need to check for custom tool first
8375 # since it does not produce event but instead
8476 # a Message
8577 if isinstance (chunk , ToolResponseMessage ):
86- yield TurnStreamPrintableEvent (
87- role = "CustomTool" , content = chunk .content , color = "green"
88- )
78+ yield TurnStreamPrintableEvent (role = "CustomTool" , content = chunk .content , color = "green" )
8979 return
9080
9181 event = chunk .event
@@ -101,9 +91,7 @@ def _yield_printable_events(
10191 if step_type == "shield_call" and event_type == "step_complete" :
10292 violation = event .payload .step_details .violation
10393 if not violation :
104- yield TurnStreamPrintableEvent (
105- role = step_type , content = "No Violation" , color = "magenta"
106- )
94+ yield TurnStreamPrintableEvent (role = step_type , content = "No Violation" , color = "magenta" )
10795 else :
10896 yield TurnStreamPrintableEvent (
10997 role = step_type ,
@@ -114,9 +102,7 @@ def _yield_printable_events(
114102 # handle inference
115103 if step_type == "inference" :
116104 if event_type == "step_start" :
117- yield TurnStreamPrintableEvent (
118- role = step_type , content = "" , end = "" , color = "yellow"
119- )
105+ yield TurnStreamPrintableEvent (role = step_type , content = "" , end = "" , color = "yellow" )
120106 elif event_type == "step_progress" :
121107 if event .payload .delta .type == "tool_call" :
122108 if isinstance (event .payload .delta .tool_call , str ):
@@ -167,13 +153,9 @@ def _yield_printable_events(
167153
168154 def _get_event_type_step_type (self , chunk ):
169155 if hasattr (chunk , "event" ):
170- previous_event_type = (
171- chunk .event .payload .event_type if hasattr (chunk , "event" ) else None
172- )
156+ previous_event_type = chunk .event .payload .event_type if hasattr (chunk , "event" ) else None
173157 previous_step_type = (
174- chunk .event .payload .step_type
175- if previous_event_type not in {"turn_start" , "turn_complete" }
176- else None
158+ chunk .event .payload .step_type if previous_event_type not in {"turn_start" , "turn_complete" } else None
177159 )
178160 return previous_event_type , previous_step_type
179161 return None , None
0 commit comments