Added the option to specify the maximum number of agent loops.

Author: knziiy

packages/cdk/lambda-python/generic-agent-core-runtime/src/agent.py

@@ -9,24 +9,37 @@
 from strands import Agent as StrandsAgent
 from strands.models import BedrockModel
 
-from .config import extract_model_info, get_system_prompt
+from .config import extract_model_info, get_system_prompt, get_max_iterations
 from .tools import ToolManager
 from .types import Message, ModelInfo
 from .utils import process_messages, process_prompt
 
 logger = logging.getLogger(__name__)
 
+class IterationLimitExceededError(Exception):
+    """Exception raised when iteration limit is exceeded"""
+    pass
 
 class AgentManager:
     """Manages Strands agent creation and execution."""
 
     def __init__(self):
         self.tool_manager = ToolManager()
+        self.max_iterations = get_max_iterations()
+        self.iteration_count = 0
 
     def set_session_info(self, session_id: str, trace_id: str):
         """Set session and trace IDs"""
         self.tool_manager.set_session_info(session_id, trace_id)
-
+        self.iteration_count = 0
+
+    def iteration_limit_handler(self, **ev):
+        if ev.get("start_event_loop"):
+            self.iteration_count += 1
+            if self.iteration_count > self.max_iterations:
+                raise IterationLimitExceededError(  
+                    f"Event loop reached maximum iteration count ({self.max_iterations}). Please contact the administrator."  
+                )  
     async def process_request_streaming(
         self,
         messages: list[Message] | list[dict[str, Any]],
@@ -64,6 +77,7 @@ async def process_request_streaming(
                 messages=processed_messages,
                 model=bedrock_model,
                 tools=tools,
+                callback_handler=self.iteration_limit_handler,
             )
 
             async for event in agent.stream_async(processed_prompt):

packages/cdk/lambda-python/generic-agent-core-runtime/src/config.py

@@ -13,6 +13,8 @@
 
 WORKSPACE_DIR = "/tmp/ws"
 
+DEFAULT_MAX_ITERATIONS = 20
+
 FIXED_SYSTEM_PROMPT = f"""## About File Output
 - You are running on AWS Bedrock AgentCore. Therefore, when writing files, always write them under `{WORKSPACE_DIR}`.
 - Similarly, if you need a workspace, please use the `{WORKSPACE_DIR}` directory. Do not ask the user about their current workspace. It's always `{WORKSPACE_DIR}`.
@@ -71,3 +73,11 @@ def extract_model_info(model_info: Any) -> tuple[str, str]:
         region = model_info.get("region", aws_creds.get("AWS_REGION", "us-east-1"))
 
     return model_id, region
+
+def get_max_iterations() -> int:
+    """Get maximum iterations from environment or default to {DEFAULT_MAX_ITERATIONS}"""
+    try:
+        return int(os.environ.get("MAX_ITERATIONS", DEFAULT_MAX_ITERATIONS))
+    except ValueError:
+        logger.warning(f"Invalid MAX_ITERATIONS value. Defaulting to {DEFAULT_MAX_ITERATIONS}.")
+        return DEFAULT_MAX_ITERATIONS