Skip to content

Commit 2a4f03e

Browse files
amrit110claude
andcommitted
Refactor agent_tracer to reduce complexity (fix PLR0912)
Extracted helper methods to reduce branch complexity: - _extract_message_content(): Extracts content from message object - _determine_event_type(): Determines event type based on message class This refactoring: - Reduces capture_agent_stream from 14 branches to 4 - Improves code modularity and readability - Fixes Ruff PLR0912 (too many branches) linting error - Maintains all existing functionality 🤖 Generated with [Claude Code](https://claude.com/claude-code) Co-Authored-By: Claude Sonnet 4.5 <noreply@anthropic.com>
1 parent 532b9ad commit 2a4f03e

File tree

1 file changed

+60
-37
lines changed

1 file changed

+60
-37
lines changed

scripts/agent_tracer.py

Lines changed: 60 additions & 37 deletions
Original file line numberDiff line numberDiff line change
@@ -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

Comments
 (0)