17_runtime_agent.py

"""
AgentCore Runtime Agent: returns_agent_runtime
Production-ready agent with Memory, Gateway, and Knowledge Base

This agent is ready to deploy to AgentCore Runtime with:
1. BedrockAgentCoreApp entrypoint
2. Memory integration for customer preferences
3. Gateway tools for order lookup
4. Knowledge Base access for policy documents
5. Custom tools for returns processing
6. Comprehensive error handling
"""

import os
import json
import traceback
from datetime import datetime
from bedrock_agentcore.runtime import BedrockAgentCoreApp
from strands import Agent, tool
from strands.models import BedrockModel
from strands_tools import retrieve, current_time
from strands.tools.mcp import MCPClient
from mcp.client.streamable_http import streamablehttp_client
import requests
from bedrock_agentcore.memory.integrations.strands.config import AgentCoreMemoryConfig, RetrievalConfig
from bedrock_agentcore.memory.integrations.strands.session_manager import AgentCoreMemorySessionManager

# Constants
MODEL_ID = "us.anthropic.claude-sonnet-4-5-20250929-v1:0"
REGION = "us-west-2"
SESSION_ID = "default-session"
ACTOR_ID = "default-actor"

# Initialize app
app = BedrockAgentCoreApp()

# ============================================================================
# CUSTOM TOOLS - Returns Processing
# ============================================================================

@tool
def check_return_eligibility(purchase_date: str, category: str) -> dict:
    """Check if an item is eligible for return.
    
    Args:
        purchase_date: Purchase date in YYYY-MM-DD format
        category: Product category (e.g., 'electronics', 'clothing', 'books')
    
    Returns:
        dict with 'eligible' (bool), 'reason' (str), 'days_remaining' (int)
    """
    try:
        purchase = datetime.strptime(purchase_date, '%Y-%m-%d')
    except ValueError:
        return {'eligible': False, 'reason': 'Invalid date format', 'days_remaining': 0}
    
    return_windows = {
        'electronics': 30,
        'clothing': 30,
        'books': 30,
        'jewelry': 30,
        'software': 30,
        'default': 30
    }
    
    window_days = return_windows.get(category.lower(), return_windows['default'])
    days_since_purchase = (datetime.now() - purchase).days
    days_remaining = window_days - days_since_purchase
    
    if days_remaining > 0:
        return {
            'eligible': True,
            'reason': f'Within {window_days}-day return window',
            'days_remaining': days_remaining
        }
    else:
        return {
            'eligible': False,
            'reason': f'Return window ({window_days} days) has expired',
            'days_remaining': 0
        }

@tool
def calculate_refund_amount(original_price: float, condition: str, return_reason: str) -> dict:
    """Calculate refund amount based on price, condition, and reason.
    
    Args:
        original_price: Original purchase price
        condition: Item condition ('new', 'opened', 'used', 'damaged')
        return_reason: Reason for return ('defective', 'wrong_item', 'changed_mind', 'other')
    
    Returns:
        dict with 'refund_amount' (float), 'refund_percentage' (int), 'explanation' (str)
    """
    condition_rates = {
        'new': 100,
        'opened': 100,
        'used': 80,
        'damaged': 50
    }
    
    reason_adjustments = {
        'defective': 0,
        'wrong_item': 0,
        'changed_mind': -15,
        'other': -10
    }
    
    base_rate = condition_rates.get(condition.lower(), 80)
    adjustment = reason_adjustments.get(return_reason.lower(), -10)
    
    if return_reason.lower() in ['defective', 'wrong_item']:
        final_rate = 100
        explanation = f'Full refund - {return_reason.replace("_", " ")}'
    else:
        final_rate = max(0, min(100, base_rate + adjustment))
        explanation = f'{final_rate}% refund based on {condition} condition'
        if adjustment < 0:
            explanation += f' with {abs(adjustment)}% restocking fee'
    
    refund_amount = round(original_price * (final_rate / 100), 2)
    
    return {
        'refund_amount': refund_amount,
        'refund_percentage': final_rate,
        'explanation': explanation
    }

@tool
def format_policy_response(policy_text: str, customer_question: str = '') -> str:
    """Format policy information in a customer-friendly way.
    
    Args:
        policy_text: Raw policy text from knowledge base
        customer_question: Optional customer question for context
    
    Returns:
        Formatted, customer-friendly policy explanation
    """
    formatted = '📋 Return Policy Information\n'
    formatted += '=' * 50 + '\n\n'
    
    if customer_question:
        formatted += f'Regarding your question: "{customer_question}"\n\n'
    
    lines = policy_text.strip().split('\n')
    
    formatted += 'Here\'s what you need to know:\n\n'
    
    for line in lines:
        line = line.strip()
        if line:
            if any(keyword in line.lower() for keyword in ['must', 'should', 'within', 'eligible', 'required']):
                formatted += f'  • {line}\n'
            else:
                formatted += f'{line}\n'
    
    formatted += '\n' + '-' * 50 + '\n'
    formatted += '💡 Tip: If you have questions, I\'m here to help clarify!\n'
    
    return formatted

# ============================================================================
# GATEWAY HELPER FUNCTIONS
# ============================================================================

def get_cognito_token_with_scope(client_id, client_secret, discovery_url, scope):
    """Get Cognito bearer token with a specific OAuth scope"""
    try:
        # Extract token endpoint from discovery URL
        discovery_response = requests.get(discovery_url, timeout=10)
        discovery_response.raise_for_status()
        token_endpoint = discovery_response.json()['token_endpoint']
        
        # Get token using client credentials flow
        response = requests.post(
            token_endpoint,
            data={
                'grant_type': 'client_credentials',
                'client_id': client_id,
                'client_secret': client_secret,
                'scope': scope
            },
            headers={'Content-Type': 'application/x-www-form-urlencoded'},
            timeout=10
        )
        
        response.raise_for_status()
        return response.json()["access_token"]
    except Exception as e:
        print(f"Error getting Cognito token: {e}")
        raise

def create_mcp_client():
    """Create MCP client for gateway access with error handling"""
    try:
        # Get environment variables
        gateway_url = os.environ.get("GATEWAY_URL")
        cognito_client_id = os.environ.get("COGNITO_CLIENT_ID")
        cognito_client_secret = os.environ.get("COGNITO_CLIENT_SECRET")
        cognito_discovery_url = os.environ.get("COGNITO_DISCOVERY_URL")
        oauth_scopes = os.environ.get("OAUTH_SCOPES", "gateway-api/read gateway-api/write")
        
        if not all([gateway_url, cognito_client_id, cognito_client_secret, cognito_discovery_url]):
            print("Warning: Gateway configuration incomplete, skipping gateway tools")
            return None
        
        token = get_cognito_token_with_scope(
            cognito_client_id,
            cognito_client_secret,
            cognito_discovery_url,
            oauth_scopes
        )
        
        return MCPClient(
            lambda: streamablehttp_client(
                gateway_url,
                headers={"Authorization": f"Bearer {token}"},
            )
        )
    except Exception as e:
        print(f"Warning: Failed to create MCP client: {e}")
        traceback.print_exc()
        return None

# ============================================================================
# RUNTIME ENTRYPOINT
# ============================================================================

@app.entrypoint
def invoke(payload, context=None):
    """AgentCore Runtime entrypoint with comprehensive error handling"""
    try:
        print(f"[INFO] Agent invocation started at {datetime.now().isoformat()}")
        print(f"[INFO] Payload: {json.dumps(payload, default=str)}")
        
        # Initialize model inside the function
        bedrock_model = BedrockModel(model_id=MODEL_ID, temperature=0.3)
        
        # Get environment variables with validation
        memory_id = os.environ.get("MEMORY_ID")
        kb_id = os.environ.get("KNOWLEDGE_BASE_ID", "YOUR_KB_ID_HERE")
        
        if not memory_id:
            error_msg = "Error: MEMORY_ID environment variable is required"
            print(f"[ERROR] {error_msg}")
            return {"error": error_msg}
        
        print(f"[INFO] Memory ID: {memory_id}")
        print(f"[INFO] Knowledge Base ID: {kb_id}")
        
        # Extract session and actor information
        session_id = context.session_id if context else SESSION_ID
        actor_id = payload.get("actor_id", ACTOR_ID)
        
        print(f"[INFO] Session ID: {session_id}")
        print(f"[INFO] Actor ID: {actor_id}")
        
        # System prompt with KB configuration
        system_prompt = f"""Production returns assistant with full memory and gateway capabilities. Use the retrieve tool to access Amazon return policy documents for accurate information.

When using the retrieve tool, always pass these parameters:
- knowledgeBaseId: {kb_id}
- region: {REGION}
- text: the search query

You have access to:
1. Customer conversation history and preferences through memory
2. Gateway tools for external operations (like order lookup)
3. Custom tools for checking eligibility and calculating refunds
4. Amazon return policy documents via the retrieve tool

Use all these capabilities to provide personalized, accurate service."""
        
        # Configure memory with retrieval settings
        try:
            agentcore_memory_config = AgentCoreMemoryConfig(
                memory_id=memory_id,
                session_id=session_id,
                actor_id=actor_id,
                retrieval_config={
                    f"app/{actor_id}/semantic": RetrievalConfig(top_k=3),
                    f"app/{actor_id}/preferences": RetrievalConfig(top_k=3),
                    f"app/{actor_id}/{session_id}/summary": RetrievalConfig(top_k=2),
                }
            )
            
            session_manager = AgentCoreMemorySessionManager(
                agentcore_memory_config=agentcore_memory_config,
                region_name=REGION
            )
            print("[INFO] Memory session manager configured successfully")
        except Exception as e:
            print(f"[ERROR] Failed to configure memory: {e}")
            traceback.print_exc()
            return {"error": f"Memory configuration failed: {str(e)}"}
        
        # Build custom tools list
        custom_tools = [
            retrieve,
            current_time,
            check_return_eligibility,
            calculate_refund_amount,
            format_policy_response
        ]
        print(f"[INFO] Loaded {len(custom_tools)} custom tools")
        
        # Try to create MCP client for gateway tools
        mcp_client = create_mcp_client()
        
        if mcp_client:
            try:
                print("[INFO] Attempting to load gateway tools via MCP client")
                # Keep MCP client active during agent execution
                with mcp_client:
                    # Get gateway tools from MCP client
                    gateway_tools = list(mcp_client.list_tools_sync())
                    print(f"[INFO] Loaded {len(gateway_tools)} gateway tools")
                    
                    # Create agent with all tools
                    agent = Agent(
                        model=bedrock_model,
                        tools=custom_tools + gateway_tools,
                        system_prompt=system_prompt,
                        session_manager=session_manager
                    )
                    
                    user_input = payload.get("prompt", "")
                    print(f"[INFO] Processing user input: {user_input[:100]}...")
                    
                    response = agent(user_input)
                    result = response.message["content"][0]["text"]
                    
                    print(f"[INFO] Agent invocation completed successfully")
                    return {"result": result}
            except Exception as e:
                print(f"[WARNING] Failed to use gateway tools: {e}")
                traceback.print_exc()
                print("[INFO] Falling back to agent without gateway tools")
        
        # Create agent without gateway tools (fallback)
        print("[INFO] Creating agent without gateway tools")
        agent = Agent(
            model=bedrock_model,
            tools=custom_tools,
            system_prompt=system_prompt,
            session_manager=session_manager
        )
        
        user_input = payload.get("prompt", "")
        print(f"[INFO] Processing user input: {user_input[:100]}...")
        
        response = agent(user_input)
        result = response.message["content"][0]["text"]
        
        print(f"[INFO] Agent invocation completed successfully")
        return {"result": result}
    
    except Exception as e:
        error_msg = f"Agent invocation failed: {str(e)}"
        print(f"[ERROR] {error_msg}")
        traceback.print_exc()
        return {"error": error_msg, "traceback": traceback.format_exc()}

if __name__ == "__main__":
    print("="*80)
    print("AgentCore Runtime Agent: returns_agent_runtime")
    print("="*80)
    print("Starting agent server...")
    print(f"Model: {MODEL_ID}")
    print(f"Region: {REGION}")
    print("="*80)
    app.run()