amazon-bedrock-agentcore-samples/02-use-cases/SRE-agent/scripts/manage_memories.py

#!/usr/bin/env python3
"""
Helper script to manage memories in the SRE agent memory system.

Usage:
    uv run python scripts/manage_memories.py [ACTION] [OPTIONS]

Actions:
    list      List memories (default action)
    update    Update memories from configuration file (with duplicate removal)
    delete    Delete memories

Examples:
    uv run python scripts/manage_memories.py                           # List all memories
    uv run python scripts/manage_memories.py list --memory-type investigations  # List only investigations
    uv run python scripts/manage_memories.py update                    # Load user preferences from user_config.yaml (removes duplicates)
    uv run python scripts/manage_memories.py update --config-file custom.yaml  # Load from custom file
    uv run python scripts/manage_memories.py update --no-duplicate-check       # Skip duplicate removal check
    uv run python scripts/manage_memories.py delete --memory-id mem-123        # Delete specific memory resource
    uv run python scripts/manage_memories.py delete --memory-record-id mem-abc # Delete specific memory record
    uv run python scripts/manage_memories.py delete --all                      # Delete all memory resources
"""

import argparse
import json
import logging
import sys
from datetime import datetime, timezone
from pathlib import Path
from typing import Optional

import yaml

# Add the project root to path so we can import sre_agent
project_root = Path(__file__).parent.parent
sys.path.insert(0, str(project_root))

from bedrock_agentcore.memory import MemoryClient

from sre_agent.memory.client import SREMemoryClient
from sre_agent.memory.config import _load_memory_config
from sre_agent.memory.strategies import UserPreference, _save_user_preference

# Configure logging
logging.basicConfig(
    level=logging.INFO,
    format="%(asctime)s,p%(process)s,{%(filename)s:%(lineno)d},%(levelname)s,%(message)s",
)

logger = logging.getLogger(__name__)


def _read_memory_id() -> str:
    """Read memory ID from .memory_id file."""
    memory_id_file = Path(__file__).parent.parent / ".memory_id"

    if memory_id_file.exists():
        memory_id = memory_id_file.read_text().strip()
        logger.info(f"Found memory ID: {memory_id} in {memory_id_file}")
        return memory_id

    raise FileNotFoundError("Could not find .memory_id file in project root")


def _extract_actor_from_namespace(namespace: str, memory_type: str) -> str:
    """Extract actor ID from memory namespace."""
    if memory_type == "preferences":
        # Preferences namespace format: /sre/users/{user_id}/preferences
        prefix = "/sre/users/"
        suffix = "/preferences"
        if namespace.startswith(prefix) and namespace.endswith(suffix):
            # Extract user_id from between prefix and suffix
            user_id = namespace[len(prefix) : -len(suffix)]
            return user_id
    else:
        # Other namespace format: /sre/{memory_type}/{actor_id}
        prefix = f"/sre/{memory_type}/"
        if namespace.startswith(prefix):
            return namespace[len(prefix) :]
    return "unknown"


def _group_memories_by_actor(memories: list, memory_type: str) -> dict:
    """Group memories by actor ID extracted from namespaces."""
    actor_groups = {}

    for memory in memories:
        # Each memory can have multiple namespaces, use the first one
        if memory.get("namespaces") and len(memory["namespaces"]) > 0:
            namespace = memory["namespaces"][0]
            actor_id = _extract_actor_from_namespace(namespace, memory_type)

            if actor_id not in actor_groups:
                actor_groups[actor_id] = []
            actor_groups[actor_id].append(memory)

    return actor_groups


def _list_memories_for_type(
    client: SREMemoryClient, memory_type: str, actor_id: Optional[str] = None
) -> None:
    """List all memories for a specific type."""
    print(f"\n=== {memory_type.upper()} MEMORIES ===")

    try:
        if actor_id and actor_id != "all":
            # List memories for specific actor
            memories = client.retrieve_memories(
                memory_type=memory_type,
                actor_id=actor_id,
                query="*",  # Wildcard query to get all memories
                max_results=100,
            )
            print(
                f"Found {len(memories)} {memory_type} memories for actor_id: {actor_id}"
            )

            for i, memory in enumerate(memories, 1):
                print(f"\n--- Memory {i} ---")
                print(json.dumps(memory, indent=2, default=str))

        else:
            # List memories across ALL actors using broader namespace
            print(f"Retrieving {memory_type} memories across ALL actors...")

            # Use different namespace patterns for different memory types
            if memory_type == "preferences":
                namespace = "/sre/users"  # For user preferences: /sre/users/{user_id}/preferences
            else:
                namespace = f"/sre/{memory_type}"  # For infrastructure/investigations: /sre/{type}/{actor_id}

            memories = client.client.retrieve_memories(
                memory_id=client.memory_id,
                namespace=namespace,
                query="*",
                actor_id=None,  # No actor restriction
                top_k=100,
            )
            print(f"Found {len(memories)} {memory_type} memories across all actors")

            # Group memories by actor
            actor_groups = _group_memories_by_actor(memories, memory_type)

            # Display grouped by actor
            for actor, actor_memories in sorted(actor_groups.items()):
                print(f"\n--- ACTOR: {actor} ({len(actor_memories)} memories) ---")

                for i, memory in enumerate(actor_memories, 1):
                    print(f"\n  Memory {i}:")
                    print(json.dumps(memory, indent=4, default=str))

    except Exception as e:
        logger.error(f"Failed to retrieve {memory_type} memories: {e}")
        print(f"Error retrieving {memory_type} memories: {e}")


def _list_all_memories() -> list:
    """List all memory resources."""
    try:
        memory_client = MemoryClient(region_name="us-east-1")
        memories = memory_client.list_memories(max_results=100)
        return memories
    except Exception as e:
        logger.error(f"Failed to list memory resources: {e}")
        return []


def _delete_memory(memory_id: str) -> bool:
    """Delete a specific memory resource."""
    try:
        memory_client = MemoryClient(region_name="us-east-1")
        logger.info(f"Deleting memory: {memory_id}")
        print(f"Deleting memory: {memory_id}...")

        result = memory_client.delete_memory_and_wait(
            memory_id=memory_id, max_wait=300, poll_interval=10
        )

        logger.info(f"Successfully deleted memory: {memory_id}")
        print(f"Successfully deleted memory: {memory_id}")
        return True

    except Exception as e:
        logger.error(f"Failed to delete memory {memory_id}: {e}")
        print(f"Error deleting memory {memory_id}: {e}")
        return False


def _delete_memory_record(memory_record_id: str) -> bool:
    """Delete a specific memory record by its ID."""
    try:
        # Read the memory resource ID from .memory_id file
        memory_id = _read_memory_id()

        # Use the SREMemoryClient for proper configuration
        client = SREMemoryClient(memory_name="sre_agent_memory")

        logger.info(
            f"Deleting memory record: {memory_record_id} from memory: {memory_id}"
        )
        print(f"Deleting memory record: {memory_record_id}...")

        # Use the underlying data plane client to delete the specific memory record
        result = client.client.gmdp_client.delete_memory_record(
            memoryId=memory_id, memoryRecordId=memory_record_id
        )

        logger.info(f"Successfully deleted memory record: {memory_record_id}")
        print(f"Successfully deleted memory record: {memory_record_id}")
        return True

    except Exception as e:
        logger.error(f"Failed to delete memory record {memory_record_id}: {e}")
        print(f"Error deleting memory record {memory_record_id}: {e}")
        return False


def _delete_all_memories() -> int:
    """Delete all memory resources."""
    memories = _list_all_memories()
    if not memories:
        print("No memories found to delete.")
        return 0

    print(f"Found {len(memories)} memory resources:")
    for memory in memories:
        memory_id = memory.get("id", "unknown")
        memory_name = memory.get("name", "unnamed")
        print(f"  - {memory_id} ({memory_name})")

    # Confirm deletion
    response = input(
        f"\nAre you sure you want to delete ALL {len(memories)} memory resources? (yes/no): "
    )
    if response.lower() not in ["yes", "y"]:
        print("Deletion cancelled.")
        return 0

    deleted_count = 0
    for memory in memories:
        memory_id = memory.get("id")
        if memory_id and _delete_memory(memory_id):
            deleted_count += 1

    print(f"\nDeleted {deleted_count} out of {len(memories)} memory resources.")
    return deleted_count


def _check_and_delete_existing_preference(
    client: SREMemoryClient, user_id: str, preference_type: str
) -> tuple[int, list]:
    """Check for existing preference events of the same type and delete them to prevent duplicates."""
    try:
        # List events for this user to find duplicate preferences
        # We use list_events because it shows individual preference events, not aggregated memories
        events = client.client.list_events(
            memory_id=client.memory_id,
            actor_id=user_id,
            session_id="preferences-default",
            max_results=100,  # Get more events to ensure we find all duplicates
            include_payload=True,
        )

        deleted_count = 0
        logger.debug(f"Found {len(events)} events for user {user_id}")

        # Track which events to delete (we'll batch delete them)
        events_to_delete = []

        for event in events:
            try:
                # Get the event payload (it's a list of message objects)
                payload = event.get("payload", [])

                for message_obj in payload:
                    # Extract the conversational content
                    conversational = message_obj.get("conversational", {})
                    role = conversational.get("role", "")
                    content_obj = conversational.get("content", {})
                    content = content_obj.get("text", "")

                    # We're looking for ASSISTANT messages with preference data
                    if role == "ASSISTANT" and content:
                        import json

                        try:
                            # Parse the content to check preference type
                            pref_data = json.loads(content)

                            # Check if this is a preference with the matching type
                            if pref_data.get("preference_type") == preference_type:
                                event_id = event.get("eventId")
                                event_time = event.get("eventTimestamp")
                                logger.info(
                                    f"Found existing {preference_type} preference event: {event_id} from {event_time}"
                                )
                                events_to_delete.append(event_id)

                        except json.JSONDecodeError:
                            # Not JSON or not a preference - skip
                            continue

            except Exception as e:
                logger.warning(f"Error checking event for duplicates: {e}")
                continue

        # Report on duplicate events found
        if events_to_delete:
            logger.info(
                f"Found {len(events_to_delete)} existing {preference_type} preference events for user {user_id}"
            )
            # Note: The Amazon Bedrock Agent Memory service doesn't support deleting individual events
            # Events are immutable and designed to accumulate over time
            # The memory strategies will aggregate all events, giving more weight to recent ones
            logger.info(
                "Note: Existing preference events cannot be deleted (events are immutable)"
            )
            logger.info(
                "New preference will be added and the memory strategy will aggregate all events"
            )
            deleted_count = 0  # We can't actually delete events

        if deleted_count > 0:
            logger.info(
                f"Deleted {deleted_count} existing {preference_type} preferences for user {user_id}"
            )

        return deleted_count, events_to_delete

    except Exception as e:
        logger.warning(
            f"Failed to check for existing preferences for user {user_id}, type {preference_type}: {e}"
        )
        return 0, []


def _load_user_preferences_from_yaml(yaml_file: Path) -> dict:
    """Load user preferences from YAML configuration file."""
    try:
        with open(yaml_file, "r") as f:
            config = yaml.safe_load(f)

        if not config or "users" not in config:
            raise ValueError("Invalid YAML format: missing 'users' section")

        logger.info(f"Loaded user preferences from {yaml_file}")
        return config["users"]

    except Exception as e:
        logger.error(f"Failed to load user preferences from {yaml_file}: {e}")
        raise


def _handle_update_action(args) -> None:
    """Handle update action to load user preferences from YAML."""
    try:
        # Default YAML file path (look in scripts directory first, then project root)
        yaml_file = Path(__file__).parent / "user_config.yaml"
        if not yaml_file.exists():
            yaml_file = Path(__file__).parent.parent / "user_config.yaml"

        # Use custom path if provided
        if hasattr(args, "config_file") and args.config_file:
            yaml_file = Path(args.config_file)

        if not yaml_file.exists():
            print(f"Error: Configuration file not found: {yaml_file}")
            sys.exit(1)

        print(f"Loading user preferences from: {yaml_file}")

        # Load user preferences from YAML
        users_config = _load_user_preferences_from_yaml(yaml_file)

        # Create memory client
        client = SREMemoryClient(memory_name="sre_agent_memory")

        total_added = 0
        total_deleted = 0
        total_users = len(users_config)

        if args.no_duplicate_check:
            print(
                f"Processing preferences for {total_users} users (duplicate check disabled)..."
            )
        else:
            print(
                f"Processing preferences for {total_users} users (removing duplicates)..."
            )

        # Process each user
        for user_id, user_config in users_config.items():
            if "preferences" not in user_config:
                print(f"Warning: No preferences found for user {user_id}")
                continue

            user_preferences = user_config["preferences"]
            print(
                f"\n--- Processing user: {user_id} ({len(user_preferences)} preferences) ---"
            )

            # Process each preference for this user
            for pref_data in user_preferences:
                try:
                    preference_type = pref_data["preference_type"]

                    # Check for and delete existing preferences of the same type (unless disabled)
                    deleted_count = 0
                    events_to_delete = []
                    if not args.no_duplicate_check:
                        deleted_count, events_to_delete = (
                            _check_and_delete_existing_preference(
                                client, user_id, preference_type
                            )
                        )

                        if deleted_count > 0:
                            # This should not happen anymore since we can't delete events
                            print(
                                f"  🗑️  Deleted {deleted_count} existing {preference_type} preferences for {user_id}"
                            )
                            total_deleted += deleted_count
                        elif events_to_delete:
                            # This is what will happen when duplicates are found
                            print(
                                f"  ℹ️  Found existing {preference_type} preference events for {user_id}"
                            )
                            print(
                                "     Note: Adding new preference (events accumulate over time)"
                            )

                    # Create UserPreference object
                    preference = UserPreference(
                        user_id=user_id,
                        preference_type=preference_type,
                        preference_value=pref_data["preference_value"],
                        context=pref_data.get(
                            "context",
                            f"Loaded from {yaml_file.name}. Do not add this memory to summary or semantic memory, only add it to user preferences long term memory.",
                        ),
                        timestamp=datetime.now(timezone.utc),
                    )

                    # Save preference using user_id as actor_id
                    success = _save_user_preference(
                        client,
                        user_id,  # Use user_id as actor_id for proper namespace
                        preference,
                    )

                    if success:
                        print(
                            f"  ✅ Added {preference.preference_type} preference for {user_id}"
                        )
                        total_added += 1
                    else:
                        print(
                            f"  ❌ Failed to add {preference.preference_type} preference for {user_id}"
                        )

                except Exception as e:
                    print(f"  ❌ Error processing preference for {user_id}: {e}")
                    logger.error(f"Error processing preference for {user_id}: {e}")

        print("\n=== SUMMARY ===")
        print(f"Successfully added {total_added} user preferences to memory")
        if total_deleted > 0:
            print(f"Deleted {total_deleted} duplicate preferences during update")
        print(f"Users processed: {total_users}")
        print(f"Memory ID: {client.memory_id}")

    except Exception as e:
        logger.error(f"Failed to update user preferences: {e}", exc_info=True)
        print(f"Error: {e}")
        sys.exit(1)


def _handle_delete_action(args) -> None:
    """Handle delete action."""
    if args.all:
        _delete_all_memories()
    elif args.memory_id:
        _delete_memory(args.memory_id)
    elif args.memory_record_id:
        _delete_memory_record(args.memory_record_id)
    else:
        print(
            "Error: Must specify either --memory-id, --memory-record-id, or --all for delete action"
        )
        sys.exit(1)


def _handle_list_action(args) -> None:
    """Handle list action."""
    try:
        # Read memory ID
        memory_id = _read_memory_id()
        print(f"Using memory ID: {memory_id}")

        # Load memory configuration
        memory_config = _load_memory_config()
        print(f"Memory configuration loaded from: {memory_config}")

        # Create memory client (memory_id is set internally during initialization)
        client = SREMemoryClient(memory_name="sre_agent_memory")

        # Verify the memory_id matches what we expect
        if client.memory_id != memory_id:
            logger.warning(
                f"Expected memory_id {memory_id}, but client initialized with {client.memory_id}"
            )
            logger.info(f"Using client memory_id: {client.memory_id}")

        # List memories
        if args.memory_type:
            # List specific memory type
            _list_memories_for_type(client, args.memory_type, args.actor_id)
        else:
            # List all memory types
            memory_types = ["preferences", "infrastructure", "investigations"]
            for memory_type in memory_types:
                _list_memories_for_type(client, memory_type, args.actor_id)

        print("\n=== SUMMARY ===")
        print(f"Memory ID: {memory_id}")
        if args.memory_type:
            print(f"Filtered by memory type: {args.memory_type}")
        if args.actor_id:
            print(f"Filtered by actor ID: {args.actor_id}")

    except Exception as e:
        logger.error(f"Failed to list memories: {e}", exc_info=True)
        print(f"Error: {e}")
        sys.exit(1)


def main():
    """Main function."""
    parser = argparse.ArgumentParser(
        description="Manage memories in the SRE agent memory system",
        formatter_class=argparse.RawDescriptionHelpFormatter,
        epilog="""
Examples:
    %(prog)s                                    # List all memory types grouped by actor
    %(prog)s list --memory-type investigations  # List only investigations grouped by actor
    %(prog)s list --actor-id sre-agent         # List memories for specific actor only
    %(prog)s update                            # Load user preferences from user_config.yaml (removes duplicates)
    %(prog)s update --config-file custom.yaml  # Load user preferences from custom YAML file
    %(prog)s update --no-duplicate-check       # Load preferences without removing duplicates
    %(prog)s delete --memory-id mem-123        # Delete specific memory resource
    %(prog)s delete --memory-record-id mem-abc # Delete specific memory record
    %(prog)s delete --all                      # Delete all memory resources (with confirmation)
        """,
    )

    # Add subcommands
    subparsers = parser.add_subparsers(dest="action", help="Action to perform")
    subparsers.required = False  # Make subcommands optional

    # List command (default)
    list_parser = subparsers.add_parser("list", help="List memories")
    list_parser.add_argument(
        "--memory-type",
        choices=["preferences", "infrastructure", "investigations"],
        help="Filter by memory type",
    )
    list_parser.add_argument(
        "--actor-id",
        help="Filter by actor ID. If not specified or 'all', shows memories grouped by actor (extracted from memory namespaces)",
    )

    # Update command
    update_parser = subparsers.add_parser(
        "update", help="Update memories from configuration file"
    )
    update_parser.add_argument(
        "--config-file",
        help="Path to YAML configuration file (default: user_config.yaml)",
    )
    update_parser.add_argument(
        "--no-duplicate-check",
        action="store_true",
        help="Skip checking for and removing duplicate preferences (default: duplicates are removed)",
    )

    # Delete command
    delete_parser = subparsers.add_parser("delete", help="Delete memories")
    delete_group = delete_parser.add_mutually_exclusive_group(required=True)
    delete_group.add_argument(
        "--memory-id", help="Memory ID to delete (deletes entire memory resource)"
    )
    delete_group.add_argument(
        "--memory-record-id",
        help="Memory record ID to delete (deletes individual memory record)",
    )
    delete_group.add_argument(
        "--all",
        action="store_true",
        help="Delete all memory resources (with confirmation prompt)",
    )

    # Global arguments (for backward compatibility when no subcommand is used)
    parser.add_argument(
        "--memory-type",
        choices=["preferences", "infrastructure", "investigations"],
        help="Filter by memory type (legacy, implies list action)",
    )
    parser.add_argument(
        "--actor-id", help="Filter by actor ID (legacy, implies list action)"
    )
    parser.add_argument(
        "--verbose", "-v", action="store_true", help="Enable verbose logging"
    )

    args = parser.parse_args()

    # Handle backward compatibility: if memory-type or actor-id is specified without subcommand, default to list
    if (
        not args.action
        and (hasattr(args, "memory_type") and args.memory_type)
        or (hasattr(args, "actor_id") and args.actor_id)
    ):
        args.action = "list"
    elif not args.action:
        args.action = "list"

    if args.verbose:
        logging.getLogger().setLevel(logging.DEBUG)

    # Handle actions
    if args.action == "list":
        _handle_list_action(args)
    elif args.action == "update":
        _handle_update_action(args)
    elif args.action == "delete":
        _handle_delete_action(args)
    else:
        print(f"Unknown action: {args.action}")
        parser.print_help()
        sys.exit(1)


if __name__ == "__main__":
    main()