diff --git a/packages/cdk/lambda-python/generic-agent-core-runtime/src/agent.py b/packages/cdk/lambda-python/generic-agent-core-runtime/src/agent.py index 6775ffc69..950757045 100644 --- a/packages/cdk/lambda-python/generic-agent-core-runtime/src/agent.py +++ b/packages/cdk/lambda-python/generic-agent-core-runtime/src/agent.py @@ -9,7 +9,7 @@ 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 @@ -17,16 +17,32 @@ 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): diff --git a/packages/cdk/lambda-python/generic-agent-core-runtime/src/config.py b/packages/cdk/lambda-python/generic-agent-core-runtime/src/config.py index 26aec7fc6..4d546c3b4 100644 --- a/packages/cdk/lambda-python/generic-agent-core-runtime/src/config.py +++ b/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 diff --git a/packages/cdk/test/__snapshots__/generative-ai-use-cases.test.ts.snap b/packages/cdk/test/__snapshots__/generative-ai-use-cases.test.ts.snap index 50a28d781..2d92004ba 100644 --- a/packages/cdk/test/__snapshots__/generative-ai-use-cases.test.ts.snap +++ b/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": {