21_invoke_agent.py

#!/usr/bin/env python3
"""
Script to invoke deployed AgentCore Runtime agent.

This script:
1. Loads Cognito credentials
2. Gets OAuth token for authentication
3. Invokes the agent with a test query
4. Displays the full response
"""

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

print("="*80)
print("AGENTCORE RUNTIME AGENT INVOCATION")
print("="*80)

# Check if runtime config exists
if not os.path.exists('runtime_config.json'):
    print("\n❌ Error: Agent not deployed yet")
    print("Please run 19_deploy_agent.py first")
    exit(1)

# Load runtime config
with open('runtime_config.json') as f:
    runtime_deployed_config = json.load(f)
    agent_arn = runtime_deployed_config.get('agent_arn')
    print(f"\nAgent ARN: {agent_arn}")

# Load configuration files
try:
    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)
except FileNotFoundError as e:
    print(f"\n❌ Error: Required config file not found: {e}")
    exit(1)

# Load .bedrock_agentcore.yaml to get agent name and entrypoint
if not os.path.exists('.bedrock_agentcore.yaml'):
    print("\n❌ Error: .bedrock_agentcore.yaml not found")
    print("Please run 19_deploy_agent.py first")
    exit(1)

import yaml
with open('.bedrock_agentcore.yaml') as f:
    runtime_config = yaml.safe_load(f)

default_agent = runtime_config.get('default_agent')
agent_config = runtime_config.get('agents', {}).get(default_agent, {})
agent_name = agent_config.get('name')
entrypoint = agent_config.get('entrypoint')

print(f"Agent Name: {agent_name}")

# ============================================================================
# STEP 1: Generate OAuth Bearer Token
# ============================================================================

print("\n" + "="*80)
print("STEP 1: GENERATING OAUTH TOKEN")
print("="*80)

try:
    # Use token endpoint from config
    token_endpoint = cognito_config["token_endpoint"]
    
    # Prepare credentials for Basic Auth
    credentials = f"{cognito_config['client_id']}:{cognito_config['client_secret']}"
    encoded_credentials = base64.b64encode(credentials.encode()).decode()
    
    # Get OAuth scopes - use gateway-api scopes
    oauth_scopes = "gateway-api/read gateway-api/write"
    
    print(f"Token Endpoint: {token_endpoint}")
    print(f"Client ID: {cognito_config['client_id']}")
    print(f"Scopes: {oauth_scopes}")
    
    # Request token
    response = requests.post(
        token_endpoint,
        headers={
            "Authorization": f"Basic {encoded_credentials}",
            "Content-Type": "application/x-www-form-urlencoded"
        },
        data={
            "grant_type": "client_credentials",
            "scope": oauth_scopes
        },
        timeout=10
    )
    
    if response.status_code != 200:
        print(f"\n❌ Failed to get OAuth token")
        print(f"Status Code: {response.status_code}")
        print(f"Response: {response.text}")
        exit(1)
    
    bearer_token = response.json()["access_token"]
    print("\n✓ OAuth token obtained successfully")
    
except Exception as e:
    print(f"\n❌ Error getting OAuth token: {e}")
    import traceback
    traceback.print_exc()
    exit(1)

# ============================================================================
# STEP 2: Configure Runtime
# ============================================================================

print("\n" + "="*80)
print("STEP 2: CONFIGURING RUNTIME")
print("="*80)

# Initialize Runtime
runtime = Runtime()

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

# Configure runtime (to load existing configuration)
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="us-west-2",
        authorizer_configuration=auth_config
    )
    print("✓ Runtime configured successfully")
except Exception as e:
    print(f"❌ Error configuring runtime: {e}")
    exit(1)

# ============================================================================
# STEP 3: Invoke Agent
# ============================================================================

print("\n" + "="*80)
print("STEP 3: INVOKING AGENT")
print("="*80)

# Test query
test_query = "Can you look up my order ORD-001 and help me with a return?"
actor_id = "user_001"

print(f"\nActor ID: {actor_id}")
print(f"Query: {test_query}")

payload = {
    "prompt": test_query,
    "actor_id": actor_id
}

print("\nSending request to agent...")

try:
    response = runtime.invoke(
        payload,
        bearer_token=bearer_token
    )
    
    print("\n" + "="*80)
    print("✅ AGENT RESPONSE")
    print("="*80)
    
    # Parse response if it's a dict with 'result' key
    if isinstance(response, dict):
        if 'result' in response:
            print(response['result'])
        elif 'error' in response:
            print(f"❌ Agent returned error: {response['error']}")
            if 'traceback' in response:
                print("\nTraceback:")
                print(response['traceback'])
        else:
            print(json.dumps(response, indent=2))
    else:
        print(response)
    
    print("="*80)
    
    # Display verification checklist
    print("\n" + "="*80)
    print("VERIFICATION CHECKLIST")
    print("="*80)
    
    response_text = str(response).lower()
    
    # Check if agent used memory
    memory_used = any(word in response_text for word in ['remember', 'preference', 'email', 'previously'])
    print(f"  {'✓' if memory_used else '✗'} Used memory (customer preferences): {memory_used}")
    
    # Check if agent looked up order
    order_lookup = 'ord-001' in response_text or 'laptop' in response_text
    print(f"  {'✓' if order_lookup else '✗'} Looked up order ORD-001: {order_lookup}")
    
    # Check if agent provided return info
    return_info = any(word in response_text for word in ['return', 'eligible', 'refund', 'days'])
    print(f"  {'✓' if return_info else '✗'} Provided return information: {return_info}")
    
    print("="*80)
    
    if memory_used and order_lookup and return_info:
        print("\n✅ SUCCESS: Agent demonstrated full integration!")
        print("   - Memory: Recalled customer preferences")
        print("   - Gateway: Looked up order via Lambda")
        print("   - Knowledge Base: Accessed return policies")
        print("   - Custom Tools: Processed return eligibility")
    else:
        print("\n⚠️  PARTIAL: Some capabilities may not have been demonstrated")
        print("   This could be normal depending on the agent's response strategy")
    
    print("="*80)
    
except Exception as e:
    print("\n" + "="*80)
    print("❌ ERROR INVOKING AGENT")
    print("="*80)
    print(f"\nError: {e}")
    import traceback
    traceback.print_exc()
    
    print("\n" + "="*80)
    print("TROUBLESHOOTING")
    print("="*80)
    print("\n1. Check agent status:")
    print("   python3 20_check_status.py")
    print("\n2. Verify agent is in READY state")
    print("\n3. Check CloudWatch logs:")
    print(f"   aws logs tail /aws/bedrock-agentcore/runtimes/{agent_name.replace('_', '-')}-* --since 10m")
    print("\n4. Verify OAuth token is valid")
    print("\n5. Check IAM role permissions")
    print("="*80)
    exit(1)

print("\n✅ Invocation completed successfully!")