chat_with_agent.py

#!/usr/bin/env python3
"""
Interactive command-line chat with the deployed AgentCore Runtime agent.

Usage:
    python3 chat_with_agent.py
    
Commands:
    - Type your message and press Enter to chat
    - Type 'exit' or 'quit' to end the session
    - Type 'clear' to clear the screen
    - Type 'help' to see available commands
"""

import json
import base64
import os
import requests
import yaml
from bedrock_agentcore_starter_toolkit import Runtime

# ANSI color codes for better readability
class Colors:
    HEADER = '\033[95m'
    BLUE = '\033[94m'
    CYAN = '\033[96m'
    GREEN = '\033[92m'
    YELLOW = '\033[93m'
    RED = '\033[91m'
    ENDC = '\033[0m'
    BOLD = '\033[1m'

def print_header():
    """Print the chat header"""
    print("\n" + "="*80)
    print(f"{Colors.BOLD}{Colors.CYAN}AgentCore Returns Agent - Interactive Chat{Colors.ENDC}")
    print("="*80)
    print(f"{Colors.YELLOW}Type your message and press Enter. Type 'exit' to quit.{Colors.ENDC}")
    print("="*80 + "\n")

def print_help():
    """Print help information"""
    print(f"\n{Colors.BOLD}Available Commands:{Colors.ENDC}")
    print(f"  {Colors.GREEN}exit, quit{Colors.ENDC}  - End the chat session")
    print(f"  {Colors.GREEN}clear{Colors.ENDC}       - Clear the screen")
    print(f"  {Colors.GREEN}help{Colors.ENDC}        - Show this help message")
    print(f"  {Colors.GREEN}status{Colors.ENDC}      - Show agent configuration")
    print()

def clear_screen():
    """Clear the terminal screen"""
    os.system('clear' if os.name != 'nt' else 'cls')

def print_status(agent_name, region, actor_id):
    """Print current status"""
    print(f"\n{Colors.BOLD}Current Status:{Colors.ENDC}")
    print(f"  Agent: {Colors.CYAN}{agent_name}{Colors.ENDC}")
    print(f"  Region: {Colors.CYAN}{region}{Colors.ENDC}")
    print(f"  Actor ID: {Colors.CYAN}{actor_id}{Colors.ENDC}")
    print()

def main():
    # Ensure we're in the project root directory
    script_dir = os.path.dirname(os.path.abspath(__file__))
    os.chdir(script_dir)
    
    # Check if runtime config exists
    if not os.path.exists('runtime_config.json'):
        print(f"{Colors.RED}❌ Error: Agent not deployed yet{Colors.ENDC}")
        print("Please run 19_deploy_agent.py first")
        exit(1)

    # Load configuration files
    try:
        with open('runtime_config.json') as f:
            runtime_config = json.load(f)
        with open('cognito_config.json') as f:
            cognito_config = json.load(f)
        with open('runtime_execution_role_config.json') as f:
            role_config = json.load(f)
        with open('.bedrock_agentcore.yaml') as f:
            bedrock_config = yaml.safe_load(f)
    except FileNotFoundError as e:
        print(f"{Colors.RED}❌ Error: Required config file not found: {e}{Colors.ENDC}")
        exit(1)
        exit(1)

    # Get configuration
    agent_name = runtime_config['agent_name']
    region = runtime_config['region']
    
    default_agent = bedrock_config.get('default_agent')
    agent_config = bedrock_config.get('agents', {}).get(default_agent, {})
    entrypoint = agent_config.get('entrypoint')

    # Initialize Runtime
    print(f"{Colors.CYAN}Initializing agent...{Colors.ENDC}")
    runtime = Runtime()

    # Build authorizer configuration
    auth_config = {
        "customJWTAuthorizer": {
            "allowedClients": [cognito_config["client_id"]],
            "discoveryUrl": cognito_config["discovery_url"]
        }
    }

    # Configure runtime
    try:
        runtime.configure(
            entrypoint=entrypoint,
            agent_name=agent_name,
            execution_role=role_config["role_arn"],
            auto_create_ecr=True,
            memory_mode="NO_MEMORY",
            requirements_file="requirements.txt",
            region=region,
            authorizer_configuration=auth_config
        )
    except Exception as e:
        print(f"{Colors.YELLOW}⚠ Configuration warning: {e}{Colors.ENDC}")

    # Get OAuth token function
    def get_oauth_token():
        credentials = f"{cognito_config['client_id']}:{cognito_config['client_secret']}"
        encoded_credentials = base64.b64encode(credentials.encode()).decode()
        
        response = requests.post(
            cognito_config["token_endpoint"],
            headers={
                "Authorization": f"Basic {encoded_credentials}",
                "Content-Type": "application/x-www-form-urlencoded"
            },
            data={
                "grant_type": "client_credentials",
                "scope": "gateway-api/read gateway-api/write"
            },
            timeout=10
        )
        
        if response.status_code == 200:
            return response.json()["access_token"]
        return None

    # Generate unique actor ID for this session
    import uuid
    actor_id = f"cli_user_{str(uuid.uuid4())[:8]}"

    # Print header
    clear_screen()
    print_header()
    print(f"{Colors.GREEN}✓ Agent ready!{Colors.ENDC}")
    print_status(agent_name, region, actor_id)

    # Chat loop
    message_count = 0
    while True:
        try:
            # Get user input
            user_input = input(f"{Colors.BOLD}{Colors.BLUE}You: {Colors.ENDC}").strip()
            
            if not user_input:
                continue
            
            # Handle commands
            if user_input.lower() in ['exit', 'quit']:
                print(f"\n{Colors.CYAN}Goodbye! Chat session ended.{Colors.ENDC}\n")
                break
            
            if user_input.lower() == 'clear':
                clear_screen()
                print_header()
                continue
            
            if user_input.lower() == 'help':
                print_help()
                continue
            
            if user_input.lower() == 'status':
                print_status(agent_name, region, actor_id)
                continue
            
            # Get OAuth token
            token = get_oauth_token()
            if not token:
                print(f"{Colors.RED}❌ Failed to authenticate{Colors.ENDC}")
                continue
            
            # Prepare payload
            payload = {
                "prompt": user_input,
                "actor_id": actor_id
            }
            
            # Show thinking indicator
            print(f"{Colors.YELLOW}Agent is thinking...{Colors.ENDC}", end='\r')
            
            # Invoke agent
            try:
                response = runtime.invoke(payload, bearer_token=token)
                
                # Clear thinking indicator
                print(" " * 50, end='\r')
                
                # Parse and display response
                result_text = None
                if isinstance(response, dict):
                    if 'response' in response:
                        response_data = json.loads(response['response']) if isinstance(response['response'], str) else response['response']
                        result_text = response_data.get('result', str(response_data))
                    elif 'result' in response:
                        result_text = response['result']
                    else:
                        result_text = str(response)
                else:
                    result_text = str(response)
                
                # Print agent response
                print(f"{Colors.BOLD}{Colors.GREEN}Agent: {Colors.ENDC}{result_text}\n")
                message_count += 1
                
            except Exception as e:
                print(f"{Colors.RED}❌ Error: {e}{Colors.ENDC}\n")
        
        except KeyboardInterrupt:
            print(f"\n\n{Colors.CYAN}Chat interrupted. Type 'exit' to quit or continue chatting.{Colors.ENDC}\n")
            continue
        except EOFError:
            print(f"\n\n{Colors.CYAN}Goodbye! Chat session ended.{Colors.ENDC}\n")
            break

    # Print summary
    if message_count > 0:
        print(f"{Colors.CYAN}Session summary: {message_count} message(s) exchanged{Colors.ENDC}\n")

if __name__ == "__main__":
    main()