@@ -195,6 +195,64 @@ def extract_tool_info(self, content: str, event_type: str) -> dict[str, Any] | N
195195
196196 return None
197197
198+ def _extract_message_content (self , message : Any ) -> str :
199+ """Extract content string from message object.
200+
201+ Args:
202+ message: Agent SDK message object
203+
204+ Returns:
205+ Extracted content string
206+
207+ """
208+ if not hasattr (message , "content" ):
209+ return ""
210+
211+ if isinstance (message .content , str ):
212+ return message .content
213+
214+ if isinstance (message .content , list ):
215+ # Handle content blocks
216+ return " " .join (
217+ block .get ("text" , "" ) if isinstance (block , dict ) else str (block )
218+ for block in message .content
219+ )
220+
221+ return str (message .content )
222+
223+ def _determine_event_type (self , message : Any , msg_class : str , content : str ) -> str :
224+ """Determine event type based on message class and content.
225+
226+ Args:
227+ message: Agent SDK message object
228+ msg_class: Class name of message
229+ content: Extracted content string
230+
231+ Returns:
232+ Event type string (ERROR, TOOL_RESULT, TOOL_CALL, etc.)
233+
234+ """
235+ if msg_class == "ToolResultBlock" :
236+ # Check if it's an actual error
237+ is_error_attr = getattr (message , "is_error" , None )
238+ # If attribute not accessible, parse from string representation
239+ if is_error_attr is None :
240+ msg_str = str (message )
241+ if "is_error=True" in msg_str :
242+ is_error_attr = True
243+ elif "is_error=False" in msg_str or "is_error=None" in msg_str :
244+ is_error_attr = False
245+ return "ERROR" if is_error_attr is True else "TOOL_RESULT"
246+
247+ if msg_class == "ToolUseBlock" :
248+ return "TOOL_CALL"
249+
250+ if msg_class == "TextBlock" :
251+ return self .classify_message (content if content else str (message ))
252+
253+ # Fallback to content-based classification
254+ return self .classify_message (content if content else str (message ))
255+
198256 async def capture_agent_stream (
199257 self , agent_stream : AsyncIterator [Any ]
200258 ) -> AsyncIterator [Any ]:
@@ -208,44 +266,9 @@ async def capture_agent_stream(
208266
209267 """
210268 async for message in agent_stream :
211- # Determine message type from class name
212269 msg_class = message .__class__ .__name__
213-
214- # Extract content
215- content = ""
216- if hasattr (message , "content" ):
217- if isinstance (message .content , str ):
218- content = message .content
219- elif isinstance (message .content , list ):
220- # Handle content blocks
221- content = " " .join (
222- block .get ("text" , "" ) if isinstance (block , dict ) else str (block )
223- for block in message .content
224- )
225- else :
226- content = str (message .content )
227-
228- # Determine event type based on message class
229- if msg_class == "ToolResultBlock" :
230- # Check if it's an actual error
231- # Try to access is_error attribute, checking both direct attribute and string repr
232- is_error_attr = getattr (message , "is_error" , None )
233- # If attribute not accessible, parse from string representation
234- if is_error_attr is None :
235- msg_str = str (message )
236- if "is_error=True" in msg_str :
237- is_error_attr = True
238- elif "is_error=False" in msg_str or "is_error=None" in msg_str :
239- is_error_attr = False
240- event_type = "ERROR" if is_error_attr is True else "TOOL_RESULT"
241- elif msg_class == "ToolUseBlock" :
242- event_type = "TOOL_CALL"
243- elif msg_class == "TextBlock" :
244- # Classify text content
245- event_type = self .classify_message (content if content else str (message ))
246- else :
247- # Fallback to content-based classification
248- event_type = self .classify_message (content if content else str (message ))
270+ content = self ._extract_message_content (message )
271+ event_type = self ._determine_event_type (message , msg_class , content )
249272
250273 if content or str (message ):
251274 self .event_sequence += 1
0 commit comments