Amzon Bedrock inline agent with return control. Run entire agent from one file. Session state stored in cloud.

amazon_bedrock_inline_return_control.py

import boto3
import json
import uuid
import datetime
from botocore.exceptions import ClientError

debug = False

def get_current_time():
    """Function to get the current time"""
    now = datetime.datetime.now()
    current_time = now.strftime("%H:%M:%S")
    timezone = datetime.datetime.now().astimezone().tzname()
    return {
        "time": current_time,
        "timezone": timezone
    }

def get_current_date():
    """Function to get the current date"""
    now = datetime.datetime.now()
    current_date = now.strftime("%Y-%m-%d")
    timezone = datetime.datetime.now().astimezone().tzname()
    return {
        "date": current_date,
        "timezone": timezone
    }

def add_two_numbers(a, b):
    """Function to add two numbers"""
    return a + b

def process_agent_response(response):
    """Process the agent's response and handle any return control events"""
    invocation_id = None
    response_text = ""
    function_info = None
    
    for event in response["completion"]:
        if "returnControl" in event:
            function_info = event["returnControl"]
            invocation_id = function_info["invocationId"]
            print("\nAgent is requesting information...")
            return invocation_id, response_text, function_info
        elif "chunk" in event and "bytes" in event["chunk"]:
            text = event["chunk"]["bytes"].decode('utf-8')
            response_text += text
            
    return invocation_id, response_text, function_info

def chat_with_agent():
    # Create Bedrock clients
    bedrock_runtime = boto3.client('bedrock-runtime')
    bedrock_agent_runtime = boto3.client('bedrock-agent-runtime')
    
    # Generate a unique session ID for the chat
    session_id = str(uuid.uuid4())
    print(f"\nStarting new chat session (ID: {session_id})")
    print("Type 'exit' or 'quit' to end the chat")
    print("-" * 50)
    
    # Define the action group with a time function
    action_groups = [
        {
            "actionGroupName": "TimeActions",
            "description": "Actions related to getting the current time",
            "actionGroupExecutor": {
                "customControl": "RETURN_CONTROL"
            },
            "functionSchema": {
                "functions": [
                    {
                        "name": "get_time",
                        "description": "Get the current time and date",
                        "parameters": {},
                        "requireConfirmation": "DISABLED"
                    },
                    {
                        "name": "get_date",
                        "description": "Get the current date",
                        "parameters": {},
                        "requireConfirmation": "DISABLED"
                    }
                ]
            }
        },
        {
            "actionGroupName": "MathActions",
            "description": "Actions related to mathematical operations",
            "actionGroupExecutor": {
                "customControl": "RETURN_CONTROL"
            },
            "functionSchema": {
                "functions": [
                    {
                        "name": "add_two_numbers",
                        "description": "Add two numbers together",  
                        "parameters": {
                            "a": {
                                "description": "The first number to add",
                                "required": True,
                                "type": "number"    
                            },
                            "b": {
                                "description": "The second number to add",
                                "required": True,
                                "type": "number"
                            }
                        }   
                    }
                ]
            }
        }
    ]
    
    # Define the agent instruction
    instruction = "You are a helpful assistant that can tell the time when asked. You can also engage in general conversation."
    
    # Specify the foundation model
    foundation_model = "us.amazon.nova-pro-v1:0"
    
    try:
        while True:
            # Get user input
            user_input = input("\nYou: ").strip()
            
            # Check if user wants to exit
            if user_input.lower() in ['exit', 'quit']:
                print("\nEnding chat session. Goodbye!")
                break
            
            # Skip empty inputs
            if not user_input:
                continue
            
            # First invocation - Send user input to agent
            response = bedrock_agent_runtime.invoke_inline_agent(
                sessionId=session_id,
                actionGroups=action_groups,
                instruction=instruction,
                foundationModel=foundation_model,
                inputText=user_input,
                enableTrace=True
            )
            
            # Process the response and check for return control
            invocation_id, response_text, function_info = process_agent_response(response)

            if invocation_id and function_info:
                # Extract function details from the invocationInputs
                function_input = function_info.get("invocationInputs", [])[0].get("functionInvocationInput", {})
                function_name = function_input.get("function")
                action_group = function_input.get("actionGroup")
                parameters = function_input.get("parameters", [])
                
                # Convert parameters list to dictionary with proper type conversion
                param_dict = {}
                for param in parameters:
                    name = param.get("name")
                    value = param.get("value")
                    param_type = param.get("type")
                    
                    # Convert value based on type
                    if param_type == "number":
                        try:
                            # Handle both integers and floats
                            value = float(value)
                            # Convert to int if it's a whole number
                            if value.is_integer():
                                value = int(value)
                        except ValueError:
                            print(f"\nWarning: Could not convert {value} to number")
                            continue
                    
                    param_dict[name] = value
                
                if debug:
                    print(f"\nRequested function: {function_name}")
                    print(f"Parameters: {param_dict}")
                
                # Call the appropriate function
                result = None
                if function_name == "get_time":
                    result = get_current_time()
                elif function_name == "get_date":
                    result = get_current_date()
                elif function_name == "add_two_numbers":
                    result = add_two_numbers(param_dict["a"], param_dict["b"])
                
                if result:
                    # Format the result for return control
                    return_control_result = {
                        "functionResult": {
                            "actionGroup": action_group,
                            "function": function_name,
                            "responseBody": {
                                "application/json": {
                                    "body": json.dumps(result)
                                }
                            }
                        }
                    }
                    
                    # Second invocation - Return the result to the agent
                    response = bedrock_agent_runtime.invoke_inline_agent(
                        sessionId=session_id,
                        actionGroups=action_groups,
                        instruction=instruction,
                        foundationModel=foundation_model,
                        inlineSessionState={
                            "invocationId": invocation_id,
                            "returnControlInvocationResults": [return_control_result]
                        }
                    )
                    
                    # Process the final response
                    _, response_text, _ = process_agent_response(response)
                else:
                    print(f"\nWarning: Function {function_name} did not return a result")
            
            # Print the agent's response
            print("\nAssistant:", response_text.strip())
            
    except ClientError as e:
        print(f"\nError: {e}")
    except KeyboardInterrupt:
        print("\n\nChat session interrupted. Goodbye!")
    except Exception as e:
        print(f"\nUnexpected error: {e}")

if __name__ == "__main__":
    chat_with_agent()