Added the option to specify the maximum number of agent loops. (#1309)

Author: knziiy

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

@@ -9,24 +9,40 @@
 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_max_iterations, get_system_prompt
 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)
 
+    def iteration_limit_handler(self, **ev):
+        if ev.get("init_event_loop"):
+            self.iteration_count = 0
+        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 +80,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,12 @@ 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

packages/cdk/test/__snapshots__/generative-ai-use-cases.test.ts.snap

@@ -4992,7 +4992,7 @@ exports[`GenerativeAiUseCases matches the snapshot (closed network mode) 4`] = `
         "AgentCoreRuntimeName": "GenericAgentCoreRuntime",
         "CustomConfig": {
           "containerImageUri": {
-            "Fn::Sub": "123456890123.dkr.ecr.us-east-1.\${AWS::URLSuffix}/cdk-hnb659fds-container-assets-123456890123-us-east-1:736d91ac33c9bb64c93f49d1258dbf8aff8d6ec20441291f20f42076efb65f00",
+            "Fn::Sub": "123456890123.dkr.ecr.us-east-1.\${AWS::URLSuffix}/cdk-hnb659fds-container-assets-123456890123-us-east-1:89f9ea87f3ee318598f0b63b57f12a3762e5abe60ad4c925b317d2cf2c737d0d",
           },
           "environmentVariables": {
             "FILE_BUCKET": {
@@ -26205,7 +26205,7 @@ exports[`GenerativeAiUseCases matches the snapshot 4`] = `
         "AgentCoreRuntimeName": "GenericAgentCoreRuntime",
         "CustomConfig": {
           "containerImageUri": {
-            "Fn::Sub": "123456890123.dkr.ecr.us-east-1.\${AWS::URLSuffix}/cdk-hnb659fds-container-assets-123456890123-us-east-1:736d91ac33c9bb64c93f49d1258dbf8aff8d6ec20441291f20f42076efb65f00",
+            "Fn::Sub": "123456890123.dkr.ecr.us-east-1.\${AWS::URLSuffix}/cdk-hnb659fds-container-assets-123456890123-us-east-1:89f9ea87f3ee318598f0b63b57f12a3762e5abe60ad4c925b317d2cf2c737d0d",
           },
           "environmentVariables": {
             "FILE_BUCKET": {