Applying Black to increase readability and Modularizing Event Loop

Author: fede-dash

src/strands/event_loop/error_handler.py

@@ -13,7 +13,10 @@
 from ..types.exceptions import ContextWindowOverflowException, ModelThrottledException
 from ..types.models import Model
 from ..types.streaming import StopReason
-from .message_processor import find_last_message_with_tool_results, truncate_tool_results
+from .message_processor import (
+    find_last_message_with_tool_results,
+    truncate_tool_results,
+)
 
 logger = logging.getLogger(__name__)
 
@@ -100,7 +103,10 @@ def handle_input_too_long_error(
 
     # If we found a message with toolResult
     if last_message_with_tool_results is not None:
-        logger.debug("message_index=<%s> | found message with tool results at index", last_message_with_tool_results)
+        logger.debug(
+            "message_index=<%s> | found message with tool results at index",
+            last_message_with_tool_results,
+        )
 
         # Truncate the tool results in this message
         truncate_tool_results(messages, last_message_with_tool_results)

src/strands/event_loop/event_loop.py

@@ -18,7 +18,11 @@
 from ..tools.executor import run_tools, validate_and_prepare_tools
 from ..types.content import Message, Messages
 from ..types.event_loop import ParallelToolExecutorInterface
-from ..types.exceptions import ContextWindowOverflowException, EventLoopException, ModelThrottledException
+from ..types.exceptions import (
+    ContextWindowOverflowException,
+    EventLoopException,
+    ModelThrottledException,
+)
 from ..types.models import Model
 from ..types.streaming import Metrics, StopReason
 from ..types.tools import ToolConfig, ToolHandler, ToolResult, ToolUse
@@ -28,8 +32,12 @@
 
 logger = logging.getLogger(__name__)
 
+MAX_ATTEMPTS = 6
+INITIAL_DELAY = 4
+MAX_DELAY = 240  # 4 minutes
 
-def initialize_state(**kwargs: Any) -> Any:
+
+def initialize_state(kwargs: Dict[str, Any]) -> Dict[str, Any]:
     """Initialize the request state if not present.
 
     Creates an empty request_state dictionary if one doesn't already exist in the
@@ -51,7 +59,7 @@ def event_loop_cycle(
     system_prompt: Optional[str],
     messages: Messages,
     tool_config: Optional[ToolConfig],
-    callback_handler: Any,
+    callback_handler: Callable[..., Any],
     tool_handler: Optional[ToolHandler],
     tool_execution_handler: Optional[ParallelToolExecutorInterface] = None,
     **kwargs: Any,
@@ -100,10 +108,12 @@ def event_loop_cycle(
     # Initialize cycle state
     kwargs["event_loop_cycle_id"] = uuid.uuid4()
 
-    event_loop_metrics: EventLoopMetrics = kwargs.get("event_loop_metrics", EventLoopMetrics())
+    event_loop_metrics: EventLoopMetrics = kwargs.get(
+        "event_loop_metrics", EventLoopMetrics()
+    )
 
     # Initialize state and get cycle trace
-    kwargs = initialize_state(**kwargs)
+    kwargs = initialize_state(kwargs)
     cycle_start_time, cycle_trace = event_loop_metrics.start_cycle()
     kwargs["event_loop_cycle_trace"] = cycle_trace
 
@@ -130,9 +140,9 @@ def event_loop_cycle(
     stop_reason: StopReason
     usage: Any
     metrics: Metrics
-    max_attempts = 6
-    initial_delay = 4
-    max_delay = 240  # 4 minutes
+    max_attempts = MAX_ATTEMPTS
+    initial_delay = INITIAL_DELAY
+    max_delay = MAX_DELAY
     current_delay = initial_delay
 
     # Retry loop for handling throttling exceptions
@@ -145,13 +155,15 @@ def event_loop_cycle(
         )
 
         try:
-            stop_reason, message, usage, metrics, kwargs["request_state"] = stream_messages(
-                model,
-                system_prompt,
-                messages,
-                tool_config,
-                callback_handler,
-                **kwargs,
+            stop_reason, message, usage, metrics, kwargs["request_state"] = (
+                stream_messages(
+                    model,
+                    system_prompt,
+                    messages,
+                    tool_config,
+                    callback_handler,
+                    **kwargs,
+                )
             )
             if model_invoke_span:
                 tracer.end_model_invoke_span(model_invoke_span, message, usage)
@@ -177,7 +189,13 @@ def event_loop_cycle(
 
             # Handle throttling errors with exponential backoff
             should_retry, current_delay = handle_throttling_error(
-                e, attempt, max_attempts, current_delay, max_delay, callback_handler, kwargs
+                e,
+                attempt,
+                max_attempts,
+                current_delay,
+                max_delay,
+                callback_handler,
+                kwargs,
             )
             if should_retry:
                 continue
@@ -204,80 +222,29 @@ def event_loop_cycle(
 
         # If the model is requesting to use tools
         if stop_reason == "tool_use":
-            tool_uses: List[ToolUse] = []
-            tool_results: List[ToolResult] = []
-            invalid_tool_use_ids: List[str] = []
-
-            # Extract and validate tools
-            validate_and_prepare_tools(message, tool_uses, tool_results, invalid_tool_use_ids)
-
-            # Check if tools are available for execution
-            if tool_uses:
-                if tool_handler is None:
-                    raise ValueError("toolUse present but tool handler not set")
-                if tool_config is None:
-                    raise ValueError("toolUse present but tool config not set")
-
-                # Create the tool handler process callable
-                tool_handler_process: Callable[[ToolUse], ToolResult] = partial(
-                    tool_handler.process,
-                    messages=messages,
-                    model=model,
-                    system_prompt=system_prompt,
-                    tool_config=tool_config,
-                    callback_handler=callback_handler,
-                    **kwargs,
+            if not tool_handler:
+                raise EventLoopException(
+                    "Model requested tool use but no tool handler provided",
+                    kwargs["request_state"],
                 )
 
-                # Execute tools (parallel or sequential)
-                run_tools(
-                    handler=tool_handler_process,
-                    tool_uses=tool_uses,
-                    event_loop_metrics=event_loop_metrics,
-                    request_state=cast(Any, kwargs["request_state"]),
-                    invalid_tool_use_ids=invalid_tool_use_ids,
-                    tool_results=tool_results,
-                    cycle_trace=cycle_trace,
-                    parent_span=cycle_span,
-                    parallel_tool_executor=tool_execution_handler,
-                )
-
-                # Update state for the next cycle
-                kwargs = prepare_next_cycle(kwargs, event_loop_metrics)
-
-                # Create the tool result message
-                tool_result_message: Message = {
-                    "role": "user",
-                    "content": [{"toolResult": result} for result in tool_results],
-                }
-                messages.append(tool_result_message)
-                callback_handler(message=tool_result_message)
-
-                if cycle_span:
-                    tracer.end_event_loop_cycle_span(
-                        span=cycle_span, message=message, tool_result_message=tool_result_message
-                    )
-
-                # Check if we should stop the event loop
-                if kwargs["request_state"].get("stop_event_loop"):
-                    event_loop_metrics.end_cycle(cycle_start_time, cycle_trace)
-                    return (
-                        stop_reason,
-                        message,
-                        event_loop_metrics,
-                        kwargs["request_state"],
-                    )
-
-                # Recursive call to continue the conversation
-                return recurse_event_loop(
-                    model=model,
-                    system_prompt=system_prompt,
-                    messages=messages,
-                    tool_config=tool_config,
-                    callback_handler=callback_handler,
-                    tool_handler=tool_handler,
-                    **kwargs,
-                )
+            # Handle tool execution
+            return _handle_tool_execution(
+                stop_reason,
+                message,
+                model,
+                system_prompt,
+                messages,
+                tool_config,
+                tool_handler,
+                callback_handler,
+                tool_execution_handler,
+                event_loop_metrics,
+                cycle_trace,
+                cycle_span,
+                cycle_start_time,
+                kwargs,
+            )
 
         # End the cycle and return results
         event_loop_metrics.end_cycle(cycle_start_time, cycle_trace)
@@ -359,7 +326,9 @@ def recurse_event_loop(
     )
 
 
-def prepare_next_cycle(kwargs: Dict[str, Any], event_loop_metrics: EventLoopMetrics) -> Dict[str, Any]:
+def prepare_next_cycle(
+    kwargs: Dict[str, Any], event_loop_metrics: EventLoopMetrics
+) -> Dict[str, Any]:
     """Prepare state for the next event loop cycle.
 
     Updates the keyword arguments with the current event loop metrics and stores the current cycle ID as the parent
@@ -377,3 +346,81 @@ def prepare_next_cycle(kwargs: Dict[str, Any], event_loop_metrics: EventLoopMetr
     kwargs["event_loop_parent_cycle_id"] = kwargs["event_loop_cycle_id"]
 
     return kwargs
+
+
+def _handle_tool_execution(
+    stop_reason: StopReason,
+    message: Message,
+    model: Model,
+    system_prompt: Optional[str],
+    messages: Messages,
+    tool_config: ToolConfig,
+    tool_handler: ToolHandler,
+    callback_handler: Callable[..., Any],
+    tool_execution_handler: Optional[ParallelToolExecutorInterface],
+    event_loop_metrics: EventLoopMetrics,
+    cycle_trace: Trace,
+    cycle_span: Any,
+    cycle_start_time: float,
+    kwargs: Dict[str, Any],
+) -> Tuple[StopReason, Message, EventLoopMetrics, Dict[str, Any]]:
+    tool_uses: List[ToolUse] = []
+    tool_results: List[ToolResult] = []
+    invalid_tool_use_ids: List[str] = []
+
+    validate_and_prepare_tools(message, tool_uses, tool_results, invalid_tool_use_ids)
+
+    if not tool_uses:
+        return stop_reason, message, event_loop_metrics, kwargs["request_state"]
+
+    tool_handler_process = partial(
+        tool_handler.process,
+        messages=messages,
+        model=model,
+        system_prompt=system_prompt,
+        tool_config=tool_config,
+        callback_handler=callback_handler,
+        **kwargs,
+    )
+
+    run_tools(
+        handler=tool_handler_process,
+        tool_uses=tool_uses,
+        event_loop_metrics=event_loop_metrics,
+        request_state=cast(Any, kwargs["request_state"]),
+        invalid_tool_use_ids=invalid_tool_use_ids,
+        tool_results=tool_results,
+        cycle_trace=cycle_trace,
+        parent_span=cycle_span,
+        parallel_tool_executor=tool_execution_handler,
+    )
+
+    kwargs = prepare_next_cycle(kwargs, event_loop_metrics)
+
+    tool_result_message: Message = {
+        "role": "user",
+        "content": [{"toolResult": result} for result in tool_results],
+    }
+
+    messages.append(tool_result_message)
+    callback_handler(message=tool_result_message)
+
+    if cycle_span:
+        tracer = get_tracer()
+        tracer.end_event_loop_cycle_span(
+            span=cycle_span, message=message, tool_result_message=tool_result_message
+        )
+
+    if kwargs["request_state"].get("stop_event_loop", False):
+        event_loop_metrics.end_cycle(cycle_start_time, cycle_trace)
+        return stop_reason, message, event_loop_metrics, kwargs["request_state"]
+
+    return recurse_event_loop(
+        model=model,
+        system_prompt=system_prompt,
+        messages=messages,
+        tool_config=tool_config,
+        callback_handler=callback_handler,
+        tool_handler=tool_handler,
+        **kwargs,
+    )

src/strands/event_loop/message_processor.py

@@ -68,14 +68,20 @@ def clean_orphaned_empty_tool_uses(messages: Messages) -> bool:
                     tool_results.add(tool_id)
 
     # Filter for orphaned empty toolUses (no corresponding toolResult)
-    orphaned_tool_uses = {tool_id: info for tool_id, info in empty_tool_uses.items() if tool_id not in tool_results}
+    orphaned_tool_uses = {
+        tool_id: info
+        for tool_id, info in empty_tool_uses.items()
+        if tool_id not in tool_results
+    }
 
     # Apply fixes in reverse order of occurrence (to avoid index shifting)
     if not orphaned_tool_uses:
         return False
 
     # Sort by message index and content index in reverse order
-    sorted_orphaned = sorted(orphaned_tool_uses.items(), key=lambda x: (x[1][0], x[1][1]), reverse=True)
+    sorted_orphaned = sorted(
+        orphaned_tool_uses.items(), key=lambda x: (x[1][0], x[1][1]), reverse=True
+    )
 
     # Apply fixes
     for tool_id, (msg_idx, content_idx, tool_name) in sorted_orphaned:
@@ -91,13 +97,22 @@ def clean_orphaned_empty_tool_uses(messages: Messages) -> bool:
             # Check if this is the sole content in the message
             if len(messages[msg_idx]["content"]) == 1:
                 # Replace with a message indicating the attempted tool
-                messages[msg_idx]["content"] = [{"text": f"[Attempted to use {tool_name}, but operation was canceled]"}]
-                logger.debug("message_index=<%s> | replaced content with context message", msg_idx)
+                messages[msg_idx]["content"] = [
+                    {
+                        "text": f"[Attempted to use {tool_name}, but operation was canceled]"
+                    }
+                ]
+                logger.debug(
+                    "message_index=<%s> | replaced content with context message",
+                    msg_idx,
+                )
             else:
                 # Simply remove the orphaned toolUse entry
                 messages[msg_idx]["content"].pop(content_idx)
                 logger.debug(
-                    "message_index=<%s>, content_index=<%s> | removed content item from message", msg_idx, content_idx
+                    "message_index=<%s>, content_index=<%s> | removed content item from message",
+                    msg_idx,
+                    content_idx,
                 )
         except Exception as e:
             logger.warning("failed to fix orphaned tool use | %s", e)
@@ -156,7 +171,9 @@ def truncate_tool_results(messages: Messages, msg_idx: int) -> bool:
         if isinstance(content, dict) and "toolResult" in content:
             # Update status to error with informative message
             message["content"][i]["toolResult"]["status"] = "error"
-            message["content"][i]["toolResult"]["content"] = [{"text": "The tool result was too large!"}]
+            message["content"][i]["toolResult"]["content"] = [
+                {"text": "The tool result was too large!"}
+            ]
             changes_made = True
 
     return changes_made