Amazon Bedrock AgentCore Memory demonstration - Basic memory management, conversation storage, and hierarchical event retrieval for AI agents

basic_memory_example.py

#!/usr/bin/env python3
"""
Basic Amazon Bedrock AgentCore Memory Example

This script demonstrates the fundamental concepts of using AgentCore Memory:
- Creating memory instances
- Storing conversation events
- Retrieving memories with semantic search
"""

from datetime import datetime
import random
from typing import List, Tuple
from bedrock_agentcore.memory import MemoryClient

class BasicMemoryExample:
    def __init__(self, region: str = "us-west-2"):
        """Initialize the memory client"""
        self.client = MemoryClient(region_name=region)
        self.memory = None
        # Track actors and sessions for hierarchical listing
        self.known_actors = set()
        self.known_sessions = {}
    
    def create_simple_memory(self) -> dict:
        """Create a basic memory instance without long-term strategies"""
        print("🔧 Creating basic memory instance...")
        
        # Check if memory already exists
        try:
            existing_memories = self.client.list_memories()
            for memory in existing_memories:
                if memory.get('name') == "BasicConversationMemory":
                    print(f"📝 Using existing memory with ID: {memory.get('id')}")
                    self.memory = memory
                    return self.memory
        except Exception as e:
            print(f"⚠️ Could not list existing memories: {e}")
        
        # Create new memory if none exists
        try:
            self.memory = self.client.create_memory(
                name="BasicConversationMemory",
                description="Simple memory for basic conversation tracking"
            )
            print(f"✅ Memory created with ID: {self.memory.get('id')}")
        except Exception as e:
            if "already exists" in str(e):
                print("📝 Memory already exists, attempting to retrieve it...")
                # Try to find and use the existing memory
                try:
                    existing_memories = self.client.list_memories()
                    print(f"📋 Found {len(existing_memories)} existing memories")
                    # Memory: [{'arn': 'arn:aws:bedrock-agentcore:us-west-2:<account_id>:memory/BasicConversationMemory-<memory_id>', 'id': '<memory_id>', 'status': 'ACTIVE', 'createdAt': datetime.datetime(2025, 7, 21, 9, 49, 28, 512000, tzinfo=tzlocal()), 'updatedAt': datetime.datetime(2025, 7, 21, 9, 49, 28, 512000, tzinfo=tzlocal()), 'memoryId': '<memory_id>'}]
                    # If we still haven't found it, just use the first memory if it exists
                    if existing_memories and len(existing_memories) > 0:
                        print(f"⚠️ Using first available memory: {existing_memories[0]}")
                        self.memory = existing_memories[0]
                        return self.memory
                        
                except Exception as list_error:
                    print(f"❌ Failed to retrieve existing memory: {list_error}")
                    raise e
                
                print("❌ Could not find any existing memories")
                raise Exception("Memory exists but could not be retrieved")
            else:
                print(f"❌ Failed to create memory: {e}")
                raise e
        
        return self.memory
    
    def store_conversation(self, actor_id: str, session_id: str, 
                          messages: List[Tuple[str, str]]) -> dict:
        """Store a conversation as events in memory"""
        print(f"💾 Storing conversation for actor {actor_id}, session {session_id}")
        
        # Track this actor and session
        self.known_actors.add(actor_id)
        if actor_id not in self.known_sessions:
            self.known_sessions[actor_id] = set()
        self.known_sessions[actor_id].add(session_id)
        
        event = self.client.create_event(
            memory_id=self.memory.get("id"),
            actor_id=actor_id,
            session_id=session_id,
            messages=messages
        )
        
        print(f"✅ Event stored with ID: {event.get('eventId')}")
        return event
    
    def retrieve_conversation_history(self, actor_id: str, session_id: str) -> list:
        """Retrieve conversation history from memory"""
        print(f"📖 Retrieving conversation history for {actor_id}/{session_id}")
        
        conversations = self.client.list_events(
            memory_id=self.memory.get("id"),
            actor_id=actor_id,
            session_id=session_id,
            max_results=50
        )
        
        print(f"✅ Retrieved {len(conversations)} events")
        # Print the complete conversation
        # print(f"📝 Conversation: {conversations}")
        return conversations
    
    def list_all_memories(self) -> list:
        """List all existing memory instances with hierarchical structure"""
        print("📋 Listing all memory instances with hierarchical structure...")
        
        memories = self.client.list_memories()
        
        print(f"✅ Found {len(memories)} memory instances:")
        print()

        for memory in memories:
            memory_id = memory.get('id') or memory.get('memoryId')
            print(f"📁 Memory: {memory_id}")
            print(f"   ├─ Status: {memory.get('status')}")
            print(f"   ├─ Created: {memory.get('createdAt')}")
            print(f"   └─ Actors and Sessions:")
            
            # List all known actors and their sessions
            if self.known_actors:
                for actor_id in sorted(self.known_actors):
                    print(f"      └─ 👤 Actor: {actor_id}")
                    
                    if actor_id in self.known_sessions:
                        sessions = sorted(self.known_sessions[actor_id])
                        for i, session_id in enumerate(sessions):
                            is_last_session = (i == len(sessions) - 1)
                            session_prefix = "         └─" if is_last_session else "         ├─"
                            print(f"{session_prefix} 🗂️  Session: {session_id}")
                            
                            # Get conversation details for this session
                            try:
                                conversations = self.client.list_events(
                                    memory_id=memory_id,
                                    actor_id=actor_id,
                                    session_id=session_id,
                                    max_results=50
                                )
                                
                                event_count = len(conversations)
                                event_prefix = "         │  " if not is_last_session else "            "
                                print(f"{event_prefix} 💬 Events: {event_count}")
                                
                                # Show hierarchy information for events (up to 3 events)
                                if conversations:
                                    events_to_show = min(3, len(conversations))
                                    for event_idx, event in enumerate(conversations[:events_to_show]):
                                        is_last_event = (event_idx == events_to_show - 1)
                                        event_id = event.get('eventId', 'Unknown')
                                        event_timestamp = event.get('eventTimestamp', 'Unknown')
                                        
                                        # Event header
                                        event_tree_prefix = "         │  " if not is_last_session else "            "
                                        event_branch = "└─" if is_last_event else "├─"
                                        print(f"{event_tree_prefix} {event_branch} 📅 Event: {event_id}")
                                        print(f"{event_tree_prefix} {'   ' if is_last_event else '│  '} 🕒 Time: {event_timestamp}")
                                        print(f"{event_tree_prefix} {'   ' if is_last_event else '│  '} 📝 Messages:")
                                        
                                        # Show messages from this event
                                        if 'payload' in event:
                                            payload = event['payload']
                                            
                                            # Extract messages from payload
                                            for j, item in enumerate(payload[:2]):  # Show first 2 messages
                                                if 'conversational' in item:
                                                    content = item['conversational']['content']['text']
                                                    role = item['conversational']['role']
                                                    truncated_msg = content[:80] + "..." if len(content) > 80 else content
                                                    msg_prefix = "         │  " if not is_last_session else "            "
                                                    msg_indent = "   " if is_last_event else "│  "
                                                    print(f"{msg_prefix} {msg_indent} • [{role}]: {truncated_msg}")
                                            
                                            if len(payload) > 2:
                                                remaining_messages = len(payload) - 2
                                                msg_prefix = "         │  " if not is_last_session else "            "
                                                msg_indent = "   " if is_last_event else "│  "
                                                print(f"{msg_prefix} {msg_indent} ... and {remaining_messages} more messages")
                                        
                                        # Add spacing between events (except for last event)
                                        if not is_last_event:
                                            spacing_prefix = "         │  " if not is_last_session else "            "
                                            print(f"{spacing_prefix} │")
                                    
                                    # Add vertical line after events but before "more events" summary
                                    if len(conversations) > 3:
                                        line_prefix = "         │  " if not is_last_session else "            "
                                        print(f"{line_prefix} │")
                                        remaining_events = len(conversations) - 3
                                        summary_prefix = "         │  " if not is_last_session else "            "
                                        print(f"{summary_prefix} └─ ... and {remaining_events} more events")
                                
                            except Exception as e:
                                error_prefix = "         │  " if not is_last_session else "            "
                                print(f"{error_prefix} ❌ Could not retrieve events: {str(e)[:50]}...")
                            
                            if not is_last_session:
                                print("         │")
            else:
                print("      └─ No tracked actors/sessions in current session")
            
            print()
        
        return memories

def demo_basic_usage():
    """Demonstrate basic AgentCore Memory usage"""
    print("🚀 Starting Basic AgentCore Memory Demo\n")
    
    # Initialize the example
    example = BasicMemoryExample()
    
    # Step 1: Create memory
    memory_info = example.create_simple_memory()
    print()
    
    # Step 2: Store a simple conversation, add randdom number to the message to allow we get different messages from each execution
    conversation = [
        ("Hello, I'm looking for help with my account" + str(random.randint(1, 1000000)), "USER"),
        ("I'd be happy to help you with your account. What specific issue are you experiencing?", "ASSISTANT"),
        ("I can't remember my password and the reset email isn't coming through", "USER"),
        ("Let me check your account settings. Can you confirm your email address?", "ASSISTANT"),
        ("Yes, it's john.doe@example.com", "USER"),
        ("I see the issue. Let me update your email settings and send a new reset link.", "ASSISTANT")
    ]
    
    event_info = example.store_conversation(
        actor_id="user_john_doe",
        session_id="password_reset_session_001",
        messages=conversation
    )
    print()
    
    # Step 3: Retrieve the conversation
    history = example.retrieve_conversation_history(
        actor_id="user_john_doe",
        session_id="password_reset_session_001"
    )
    print()
    
    # Step 4: Store another conversation for the same user
    followup_conversation = [
        ("Hi, I wanted to follow up on my password reset from " + datetime.now().strftime("%Y-%m-%d %H:%M:%S"), "USER"),
        ("Of course! I can see your previous session. Were you able to reset your password successfully?", "ASSISTANT"),
        ("Yes, thank you! Now I have a question about upgrading my plan", "USER"),
        ("I'd be happy to help with plan upgrades. What features are you interested in?", "ASSISTANT")
    ]
    
    example.store_conversation(
        actor_id="user_john_doe",
        session_id="plan_upgrade_session_002",
        messages=followup_conversation
    )
    print()
    
    # Step 4.1: Retrieve the followup conversation
    followup_history = example.retrieve_conversation_history(
        actor_id="user_john_doe",
        session_id="plan_upgrade_session_002"
    )
    print()

    # Step 5: List all memories
    all_memories = example.list_all_memories()
    print()
    
    print("🎉 Basic demo completed successfully!")
    return example

if __name__ == "__main__":
    try:
        # Run the basic demo
        demo = demo_basic_usage()
        
    except Exception as e:
        print(f"❌ Error running demo: {e}")