07_test_memory_agent.py

#!/usr/bin/env python3
"""
Test script for memory-enabled returns agent.

This script tests the agent's ability to recall customer preferences
and history from AgentCore Memory.
"""

import json
import sys
import os

# Add current directory to path to import the agent
sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))

print("=" * 80)
print("MEMORY-ENABLED AGENT TEST")
print("=" * 80)

# Load 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("❌ Error: memory_config.json not found")
    print("   Run 03_create_memory.py first")
    sys.exit(1)

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("❌ Error: kb_config.json not found")
    sys.exit(1)

# Set environment variables (in case they're not already set)
os.environ['MEMORY_ID'] = memory_id
os.environ['KNOWLEDGE_BASE_ID'] = kb_id

print("\n" + "=" * 80)
print("IMPORTING MEMORY-ENABLED AGENT")
print("=" * 80)

try:
    # Import the memory-enabled agent
    from agent_06_memory_enabled_agent import run_agent
    print("✓ Agent imported successfully")
except ImportError:
    try:
        # Try alternative import name
        import importlib.util
        spec = importlib.util.spec_from_file_location("memory_agent", "06_memory_enabled_agent.py")
        memory_agent = importlib.util.module_from_spec(spec)
        spec.loader.exec_module(memory_agent)
        run_agent = memory_agent.run_agent
        print("✓ Agent imported successfully")
    except Exception as e:
        print(f"❌ Error importing agent: {e}")
        print("   Make sure 06_memory_enabled_agent.py exists")
        sys.exit(1)

# Test configuration
print("\n" + "=" * 80)
print("TEST CONFIGURATION")
print("=" * 80)
print(f"Customer ID: user_001")
print(f"Session ID: test_session_001")
print(f"Memory ID: {memory_id}")
print(f"Knowledge Base ID: {kb_id}")

# Test query
test_query = "Hi! I'm thinking about returning something. What do you remember about my preferences?"

print("\n" + "=" * 80)
print("RUNNING TEST")
print("=" * 80)
print(f"\nUser (user_001): {test_query}\n")
print("Agent is thinking...")
print("-" * 80)

try:
    # Run the agent with user_001 as the actor
    response = run_agent(
        user_input=test_query,
        session_id="test_session_001",
        actor_id="user_001"
    )
    
    print("\n" + "=" * 80)
    print("AGENT RESPONSE")
    print("=" * 80)
    print(response)
    
    # Analyze the response
    print("\n" + "=" * 80)
    print("MEMORY RECALL ANALYSIS")
    print("=" * 80)
    
    response_lower = response.lower()
    
    # Check for preference recall
    email_mentioned = any(keyword in response_lower for keyword in ['email', 'notification', 'communicate'])
    laptop_mentioned = any(keyword in response_lower for keyword in ['laptop', 'defective', 'return'])
    
    print("\n✓ Memory Recall Check:")
    if email_mentioned:
        print("  ✓ Email notification preference: RECALLED")
    else:
        print("  ⚠️  Email notification preference: NOT MENTIONED")
    
    if laptop_mentioned:
        print("  ✓ Defective laptop return: RECALLED")
    else:
        print("  ⚠️  Defective laptop return: NOT MENTIONED")
    
    print("\n📊 Expected Memories:")
    print("  - Preference: Email notifications over SMS")
    print("  - History: Recently returned a defective laptop")
    print("  - History: Asked about return windows for electronics")
    
    if email_mentioned and laptop_mentioned:
        print("\n✅ SUCCESS: Agent successfully recalled customer preferences and history!")
    elif email_mentioned or laptop_mentioned:
        print("\n⚠️  PARTIAL: Agent recalled some but not all customer information")
    else:
        print("\n⚠️  NOTE: Agent may not have explicitly mentioned stored memories")
        print("   This could be normal if the agent is waiting for more specific questions")
    
except Exception as e:
    print(f"\n❌ Error running agent: {e}")
    import traceback
    traceback.print_exc()
    sys.exit(1)

print("\n" + "=" * 80)
print("TEST COMPLETE")
print("=" * 80)
print("\n💡 Tips:")
print("  - The agent has access to all memories from user_001")
print("  - Try asking more specific questions about returns or preferences")
print("  - The agent will use memory to personalize responses")
print("  - Memory includes: preferences, semantic facts, and conversation summaries")