06_memory_enabled_agent.py

"""
Standalone Strands Agent: returns_agent_with_memory
Generated by AgentCore MCP Server

Integrations:
- Memory: Enabled
- Gateway: Disabled
- Knowledge Base: Enabled
"""

import os
import json
from strands import Agent, tool
from strands.models import BedrockModel
from strands_tools import retrieve
from strands_tools import current_time
from bedrock_agentcore.memory.integrations.strands.config import AgentCoreMemoryConfig, RetrievalConfig
from bedrock_agentcore.memory.integrations.strands.session_manager import AgentCoreMemorySessionManager
from datetime import datetime

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

# Model configuration
bedrock_model = BedrockModel(model_id=MODEL_ID, temperature=0.3)

# ============================================================================
# LOAD CONFIGURATION FROM FILES
# ============================================================================

# Load memory configuration
try:
    with open('memory_config.json') as f:
        memory_config = json.load(f)
        MEMORY_ID = memory_config['memory_id']
        print(f"✓ Memory ID loaded: {MEMORY_ID}")
except FileNotFoundError:
    print("⚠️  memory_config.json not found - set MEMORY_ID environment variable")
    MEMORY_ID = os.environ.get("MEMORY_ID")

# Load knowledge base configuration
try:
    with open('kb_config.json') as f:
        kb_config = json.load(f)
        KB_ID = kb_config['knowledge_base_id']
        print(f"✓ Knowledge Base ID loaded: {KB_ID}")
except FileNotFoundError:
    print("⚠️  kb_config.json not found - set KNOWLEDGE_BASE_ID environment variable")
    KB_ID = os.environ.get("KNOWLEDGE_BASE_ID", "YOUR_KB_ID_HERE")

# ============================================================================
# MEMORY CONFIGURATION
# ============================================================================

def create_memory_session_manager(memory_id: str, session_id: str = SESSION_ID, actor_id: str = ACTOR_ID):
    """Create AgentCore Memory session manager with all three namespaces"""
    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),
        }
    )
    
    return AgentCoreMemorySessionManager(
        agentcore_memory_config=agentcore_memory_config,
        region_name=REGION
    )

# System prompt
system_prompt = f"""You are a personalized returns assistant who remembers customer preferences and history. 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 customer conversation history and preferences through memory. Use this to provide personalized service."""

# ============================================================================
# CUSTOM TOOLS
# ============================================================================

@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


def run_agent(user_input: str, session_id: str = SESSION_ID, actor_id: str = ACTOR_ID):
    """Run the agent with user input"""
    
    # Build tools list - includes all custom tools plus retrieve and current_time
    custom_tools = [
        retrieve, 
        current_time, 
        check_return_eligibility, 
        calculate_refund_amount, 
        format_policy_response
    ]
    
    # Configure memory
    if not MEMORY_ID:
        print("⚠️  MEMORY_ID not configured - agent will run without memory")
        session_manager = None
    else:
        session_manager = create_memory_session_manager(MEMORY_ID, session_id, actor_id)
        print(f"✓ Memory configured with all three namespaces:")
        print(f"  - Semantic: app/{actor_id}/semantic")
        print(f"  - Preferences: app/{actor_id}/preferences")
        print(f"  - Summary: app/{actor_id}/{session_id}/summary")
    
    # Create agent
    agent = Agent(
        model=bedrock_model,
        tools=custom_tools,
        system_prompt=system_prompt,
        session_manager=session_manager
    )
    
    response = agent(user_input)
    return response.message["content"][0]["text"]

if __name__ == "__main__":
    print("\n" + "="*80)
    print("MEMORY-ENABLED RETURNS AGENT")
    print("="*80)
    print(f"Agent: returns_agent_with_memory")
    print(f"Model: {MODEL_ID}")
    print(f"Region: {REGION}")
    print(f"Temperature: 0.3")
    print("="*80)
    print("\nIntegrations:")
    print(f"  ✓ Memory: Enabled (ID: {MEMORY_ID})")
    print(f"  ✓ Knowledge Base: Enabled (ID: {KB_ID})")
    print(f"  ✓ Custom Tools: 3 (check_return_eligibility, calculate_refund_amount, format_policy_response)")
    print(f"  ✓ Built-in Tools: 2 (retrieve, current_time)")
    print("="*80)
    
    # Example usage
    user_query = "Hello, I need help with a return"
    print(f"\nUser Query: {user_query}\n")
    result = run_agent(user_query, actor_id="user_001")
    print("\n" + "="*80)
    print("AGENT RESPONSE:")
    print("="*80)
    print(result)