# Copyright 2025 Google LLC
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     https://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

%pip install --upgrade --quiet google-cloud-aiplatform>=1.123.0

import sys

if "google.colab" in sys.modules:
    from google.colab import auth

    auth.authenticate_user()

import os

import vertexai

# fmt: off
PROJECT_ID = "[your-project-id]"  # @param {type: "string", placeholder: "[your-project-id]", isTemplate: true}
LOCATION = "us-central1" # @param {type: "string", placeholder: "[your-project-id]", isTemplate: true}
# fmt: on

if not PROJECT_ID or PROJECT_ID == "[your-project-id]":
    PROJECT_ID = str(os.environ.get("GOOGLE_CLOUD_PROJECT"))
if not LOCATION:
    LOCATION = os.environ.get("GOOGLE_CLOUD_REGION")

client = vertexai.Client(project=PROJECT_ID, location=LOCATION)

import datetime
import os
import uuid
import warnings

warnings.filterwarnings("ignore")

# Import class-based types for Memory Bank
from vertexai import types

# Basic configuration types
MemoryBankConfig = types.ReasoningEngineContextSpecMemoryBankConfig
SimilaritySearchConfig = (
    types.ReasoningEngineContextSpecMemoryBankConfigSimilaritySearchConfig
)
GenerationConfig = types.ReasoningEngineContextSpecMemoryBankConfigGenerationConfig

# Advanced configuration types for governance
TtlConfig = types.ReasoningEngineContextSpecMemoryBankConfigTtlConfig
GranularTtlConfig = (
    types.ReasoningEngineContextSpecMemoryBankConfigTtlConfigGranularTtlConfig
)
CustomizationConfig = types.MemoryBankCustomizationConfig
MemoryTopic = types.MemoryBankCustomizationConfigMemoryTopic
ManagedMemoryTopic = types.MemoryBankCustomizationConfigMemoryTopicManagedMemoryTopic
CustomMemoryTopic = types.MemoryBankCustomizationConfigMemoryTopicCustomMemoryTopic
GenerateMemoriesExample = types.MemoryBankCustomizationConfigGenerateMemoriesExample
ConversationSource = (
    types.MemoryBankCustomizationConfigGenerateMemoriesExampleConversationSource
)
ConversationSourceEvent = (
    types.MemoryBankCustomizationConfigGenerateMemoriesExampleConversationSourceEvent
)
ExampleGeneratedMemory = (
    types.MemoryBankCustomizationConfigGenerateMemoriesExampleGeneratedMemory
)
ManagedTopicEnum = types.ManagedTopicEnum

print("✅ Libraries imported successfully!")

def display_generated_memories(operation, client, title="Generated Memories"):
    """Display memories from a generation operation with consistent formatting.

    Args:
        operation: The result from client.agent_engines.memories.generate()
        client: The Vertex AI client instance
        title: Title to display above the memories
    """
    if operation.response and operation.response.generated_memories:
        print(f"\n✅ {title}: {len(operation.response.generated_memories)}\n")

        for i, gen_memory in enumerate(operation.response.generated_memories, 1):
            # Skip deleted memories - we only want to show active ones
            if gen_memory.action != "DELETED" and gen_memory.memory:
                try:
                    # Fetch full memory details including topics and metadata
                    full_memory = client.agent_engines.memories.get(
                        name=gen_memory.memory.name
                    )
                    # Visual indicators: new vs updated memories
                    action_icon = "🆕" if gen_memory.action == "CREATED" else "🔄"
                    print(f"   {action_icon} {i}. {full_memory.fact}")
                except Exception as e:
                    print(f"   ⚠️ Could not retrieve memory: {e}")
    else:
        print(f"\n📭 No {title.lower()} found")


print("✅ Helper function defined!")

print("🧠 Creating basic Memory Bank configuration...\n")

# Create minimal Memory Bank configuration
basic_memory_config = MemoryBankConfig(
    # Embedding model for similarity search
    similarity_search_config=SimilaritySearchConfig(
        embedding_model=f"projects/{PROJECT_ID}/locations/{LOCATION}/publishers/google/models/text-embedding-005"
    ),
    # LLM for memory extraction
    generation_config=GenerationConfig(
        model=f"projects/{PROJECT_ID}/locations/{LOCATION}/publishers/google/models/gemini-2.5-flash"
    ),
)

print("✅ Basic Memory Bank configuration created!")
print("   Embedding model: text-embedding-005")
print("   Generation model: gemini-2.5-flash")

print("\n🛠️ Creating Agent Engine with basic Memory Bank...\n")
print("⏳ Provisioning backend infrastructure...")

# Create the Agent Engine
agent_engine = client.agent_engines.create(
    config={"context_spec": {"memory_bank_config": basic_memory_config}}
)

agent_engine_name = agent_engine.api_resource.name

print("\n✅ Agent Engine created successfully!")
print(f"   Resource Name: {agent_engine_name}")

print("🆕 Creating your first memory manually\n")

# Generate a unique customer ID for this tutorial
customer_id = "customer_alex_" + str(uuid.uuid4())[:4]

# Create a memory manually
first_memory = client.agent_engines.memories.create(
    name=agent_engine_name,
    fact="Customer prefers to be contacted via email for all communications",
    scope={"user_id": customer_id},
)

print("✅ Memory created successfully!")

print("\n🔍 Retrieving all memories for the customer\n")

# Retrieve all memories for this customer
results = client.agent_engines.memories.retrieve(
    name=agent_engine_name, scope={"user_id": customer_id}
)

all_memories = list(results)

print(f"✅ Found {len(all_memories)} memory for customer {customer_id}:\n")

for i, retrieved_memory in enumerate(all_memories, 1):
    print(f"   {i}. {retrieved_memory.memory.fact}")

print("\n✏️ Updating the customer's communication preference\n")

first_memory = all_memories[0]

print(f"📋 Current preference: '{first_memory.memory.fact}'")
print(
    "\n💡 Customer calls to update: 'I prefer SMS for urgent issues, email for everything else'\n"
)

# Update the memory by generating from new information
# Memory Bank will recognize this updates the existing memory
update_operation = client.agent_engines.memories.generate(
    name=agent_engine_name,
    scope={"user_id": customer_id},
    direct_contents_source={
        "events": [
            {
                "content": {
                    "parts": [
                        {
                            "text": "Customer prefers SMS for urgent issues and email for all other communications"
                        }
                    ]
                }
            }
        ]
    },
    config={"wait_for_completion": True},
)

# Get the updated memory details
if update_operation.response and update_operation.response.generated_memories:
    updated_memory_ref = update_operation.response.generated_memories[0]
    updated_memory = client.agent_engines.memories.get(
        name=updated_memory_ref.memory.name
    )

    print("✅ Memory updated successfully!")
    print(f"   📝 New fact: {updated_memory.fact}")
    print(f"   🔄 Updated: {updated_memory.update_time}")
    print(f"   🔄 Action: {updated_memory_ref.action}")

    if (
        hasattr(updated_memory_ref, "previous_revision")
        and updated_memory_ref.previous_revision
    ):
        print(
            f"   📜 Previous revision preserved: {updated_memory_ref.previous_revision.split('/')[-1]}"
        )

print("\n📝 Adding more customer information\n")

# Create several more memories
additional_memories = [
    "Customer's account type is Enterprise Pro with priority support",
    "Customer reported issue with API authentication returning 401 errors",
    "Customer prefers morning appointments between 9 AM and 11 AM EST",
]

created_memories = []

for fact in additional_memories:
    memory = client.agent_engines.memories.create(
        name=agent_engine_name, fact=fact, scope={"user_id": customer_id}
    )
    created_memories.append(memory)
    print(f"   ✅ Created: {fact}")

print(f"\n✅ Created {len(created_memories)} additional memories")
print(f"   Total memories for {customer_id}: {len(created_memories) + 1}")

print("\n📚 Complete customer profile\n")

# Retrieve all memories
results = client.agent_engines.memories.retrieve(
    name=agent_engine_name, scope={"user_id": customer_id}
)

all_customer_memories = list(results)

print(f"✅ Found {len(all_customer_memories)} total memories for {customer_id}:\n")

for i, mem in enumerate(all_customer_memories, 1):
    print(f"   {i}. {mem.memory.fact}")

print("\n🗑️ Deleting a specific memory\n")

# Let's delete the API issue memory
memory_to_delete = all_customer_memories[2]

print("🔍 Memory to delete:")
print(f"   📝 {memory_to_delete.memory.fact}")
print(f"   🆔 ID: {memory_to_delete.memory.name.split('/')[-1]}")

print("\n⚠️ Confirming deletion...")

# Delete the memory
client.agent_engines.memories.delete(name=memory_to_delete.memory.name)

print("\n✅ Memory deleted successfully!")

# Verify deletion
results = client.agent_engines.memories.retrieve(
    name=agent_engine_name, scope={"user_id": customer_id}
)

remaining_memories = list(results)

print(
    f"\n📊 Remaining memories: {len(remaining_memories)} (was {len(all_customer_memories)})"
)

print("\n💬 Creating customer support conversation\n")

# Create session for this customer
session = client.agent_engines.sessions.create(
    name=agent_engine_name,
    user_id=customer_id,
    config={"display_name": f"Support session for {customer_id}"},
)

session_name = session.response.name

print("✅ Session created successfully!")
print(f"   Session: {session_name}")
print(f"   Customer ID: {customer_id}")

# Define realistic customer support conversation
support_conversation = [
    {
        "role": "user",
        "message": "Hi, I'm having issues with my billing - I was charged twice for my enterprise subscription last month.",
    },
    {
        "role": "model",
        "message": "Hello! I apologize for the billing issue. Let me look into your enterprise subscription charges right away.",
    },
    {
        "role": "user",
        "message": "I'm on the Enterprise Pro plan at $299/month. The duplicate charge was on March 15th. My company is DataTech Solutions.",
    },
    {
        "role": "model",
        "message": "Thank you for those details. I see your Enterprise Pro subscription. I'll investigate the duplicate charge from March 15th.",
    },
    {
        "role": "user",
        "message": "Also, I prefer email communication over phone calls. I'm in EST timezone, so please don't call before 9 AM my time.",
    },
    {
        "role": "model",
        "message": "Noted! I've recorded your preference for email communication and your EST timezone availability after 9 AM.",
    },
    {
        "role": "user",
        "message": "I also opened ticket #SUP-2024-0847 last week about API authentication failures. Still not resolved.",
    },
    {
        "role": "model",
        "message": "I see ticket SUP-2024-0847 regarding API authentication. Let me prioritize getting that resolved for you.",
    },
    {
        "role": "user",
        "message": "The API returns 401 errors when using our production API key. We're using the REST API v2.1 endpoint.",
    },
    {
        "role": "model",
        "message": "Thank you for the technical details. I'll escalate the API authentication issue with our engineering team and get back to you within 24 hours via email.",
    },
]

print("\n⬆️ Adding conversation to session...\n")

invocation_id = 0

# Add each conversation turn to the session with timestamps
for turn in support_conversation:
    client.agent_engines.sessions.events.append(
        name=session_name,
        author=customer_id,
        invocation_id=str(invocation_id),
        timestamp=datetime.datetime.now(tz=datetime.timezone.utc),
        config={
            "content": {"role": turn["role"], "parts": [{"text": turn["message"]}]}
        },
    )

    invocation_id += 1
    # Display conversation with visual indicators
    icon = "👤" if turn["role"] == "user" else "🤖"
    print(f"{icon} {turn['message']}")

print("\n✅ Conversation added to session!")

print("\n🧠 Generating memories from conversation\n")

# Generate memories from the session
operation = client.agent_engines.memories.generate(
    name=agent_engine_name,
    vertex_session_source={"session": session_name},
    config={"wait_for_completion": True},
)

print("✅ Memory generation complete!")

# Display the generated memories
display_generated_memories(operation, client, "Automatically Extracted Memories")

print("\n📊 Complete Customer Profile (Manual + Auto-Generated)\n")

# Retrieve all memories for this customer
results = client.agent_engines.memories.retrieve(
    name=agent_engine_name, scope={"user_id": customer_id}
)

complete_profile = list(results)

print(f"✅ Total memories for {customer_id}: {len(complete_profile)}\n")

for i, mem in enumerate(complete_profile, 1):
    print(f"   {i}. {mem.memory.fact}")

print("\n⏱️ Configuring granular TTL for automatic data retention\n")

# Define granular TTL configuration with different retention periods
ttl_config = TtlConfig(
    granular_ttl_config=GranularTtlConfig(
        create_ttl="2592000s",  # 30 days  - Manual creates (temporary notes)
        generate_created_ttl="7776000s",  # 90 days  - New generated memories
        generate_updated_ttl="31536000s",  # 365 days - Updated/consolidated memories
    )
)

print("✅ Granular TTL configuration created!")

print("\n🔧 Updating Agent Engine with TTL governance\n")

# Create enhanced configuration with TTL
governance_memory_config = MemoryBankConfig(
    # Keep existing models
    similarity_search_config=SimilaritySearchConfig(
        embedding_model=f"projects/{PROJECT_ID}/locations/{LOCATION}/publishers/google/models/text-embedding-005"
    ),
    generation_config=GenerationConfig(
        model=f"projects/{PROJECT_ID}/locations/{LOCATION}/publishers/google/models/gemini-2.5-flash"
    ),
    # Add TTL governance
    ttl_config=ttl_config,
)

# Update the existing engine
updated_engine = client.agent_engines.update(
    name=agent_engine_name,
    config={"context_spec": {"memory_bank_config": governance_memory_config}},
)

print("✅ Agent Engine updated with TTL!")

print("\n🧪 Testing TTL with new memory creation\n")

# Create a new memory with manual creation (30-day TTL)
test_memory = client.agent_engines.memories.create(
    name=agent_engine_name,
    fact="Customer requested demo of new analytics dashboard feature",
    scope={"user_id": customer_id},
)

print("✅ Memory created with TTL!")

print("\n🔍 Inspecting memories with automatic topic labels\n")

# Get all memories for detailed inspection
pager = client.agent_engines.memories.list(name=agent_engine_name)
sample_memories = list(pager)[:5]  # First 5 memories

if sample_memories:
    print(f"Showing {len(sample_memories)} memories with full details:\n")

    for i, memory in enumerate(sample_memories, 1):
        # Get complete memory details
        full_memory = client.agent_engines.memories.get(name=memory.name)

        print(f"\n📌 Memory {i}:")
        print(f"   Fact: {full_memory.fact}")
        print(f"   Scope: {full_memory.scope}")
        print(f"   Created: {full_memory.create_time}")
        print(f"   Expires: {full_memory.expire_time}")

        # Show automatically attached topics
        if hasattr(full_memory, "topics") and full_memory.topics:
            topics = [
                topic.managed_memory_topic
                if hasattr(topic, "managed_memory_topic")
                else str(topic)
                for topic in full_memory.topics
            ]
            print(f"   Topics: {', '.join(str(t) for t in topics)}")
        else:
            print("   Topics: None")
else:
    print("No memories found")

def get_memories_by_topic(engine_name, topic_enum_value):
    """Retrieve memories filtered by a specific managed topic.

    Args:
        engine_name: The Agent Engine resource name
        topic_enum_value: The topic to filter by (e.g., 'USER_PERSONAL_INFO')

    Returns:
        List of memories matching the topic
    """
    # Get all memories
    all_memories_pager = client.agent_engines.memories.list(name=engine_name)
    filtered_memories = []

    # Filter by topic
    for memory in all_memories_pager:
        full_memory = client.agent_engines.memories.get(name=memory.name)
        topics = (
            full_memory.topics
            if hasattr(full_memory, "topics") and full_memory.topics
            else []
        )

        for topic in topics:
            if hasattr(topic, "managed_memory_topic"):
                topic_value = str(topic.managed_memory_topic)
                if topic_enum_value in topic_value:
                    filtered_memories.append(full_memory)
                    break

    return filtered_memories

print("\n📊 Demonstrating topic-based filtering\n")

# List of default managed topics to filter by
topics_to_filter = [
    ("USER_PERSONAL_INFO", "Personal Information: Name, company, account details"),
    ("USER_PREFERENCES", "Preferences: Communication channels, support preferences"),
    ("KEY_CONVERSATION_DETAILS", "Key Details: Ticket numbers, important outcomes"),
]

for topic_label, description in topics_to_filter:
    print(f"\n🔍 {description}")
    print(f"   Filtering by topic: {topic_label}\n")

    # Get memories for this topic
    topic_memories = get_memories_by_topic(agent_engine_name, topic_label)

    if topic_memories:
        print(f"   ✅ Found {len(topic_memories)} memory/memories:\n")
        for i, memory in enumerate(topic_memories, 1):
            print(f"      {i}. {memory.fact}")
            if memory.expire_time:
                print(f"         ⏱️ Expires: {memory.expire_time}")
    else:
        print(f"   📭 No memories found with topic '{topic_label}'")

    print()

print("\n📋 Viewing Complete Revision History\n")

# Get a memory with revision history
all_memories_pager = client.agent_engines.memories.list(name=agent_engine_name)
memory_with_revisions = None

# Find first memory for demonstration
for memory in all_memories_pager:
    full_memory = client.agent_engines.memories.get(name=memory.name)
    memory_with_revisions = full_memory
    break

if memory_with_revisions:
    print(f"📌 Memory Selected: {memory_with_revisions.fact}\n")
    print("📜 Revision History:")

    try:
        # List all revisions for this memory
        revisions_pager = client.agent_engines.memories.revisions.list(
            name=memory_with_revisions.name
        )

        revisions = list(revisions_pager)

        if revisions:
            for i, revision in enumerate(revisions, 1):
                print(f"\n📜 Revision {i}:")
                print(f"   🆔 Revision ID: {revision.name.split('/')[-1]}")
                print(f"   ⏱️ Timestamp: {revision.create_time}")
                print(f"   📝 Fact: {revision.fact}")
                if hasattr(revision, "expire_time") and revision.expire_time:
                    print(f"   🔚 Expires: {revision.expire_time}")
        else:
            print("\n   📭 No revision history available yet")
            print("   💡 Revisions accumulate as memories are updated or consolidated")

    except Exception as e:
        print(f"\n   ⚠️ Note: {e}")
else:
    print("📭 No memories available for revision inspection")

print("\n⏰ Time-Travel Query: Retrieving Specific Revision\n")

if memory_with_revisions and revisions:
    try:
        if revisions:
            # Get the first revision (earliest version)
            target_revision = revisions[0]
            revision_id = target_revision.name.split("/")[-1]

            print(f"🎯 Retrieving revision: {revision_id}")
            print(f"   ⏱️ Timestamp: {target_revision.create_time}\n")

            # Fetch the specific revision
            specific_revision = client.agent_engines.memories.revisions.get(
                name=target_revision.name
            )

            print("📋 Revision Details:")
            print(f"   📝 Memory Fact: {specific_revision.fact}")
            print(f"   🛠️ Created: {specific_revision.create_time}")
            if (
                hasattr(specific_revision, "expire_time")
                and specific_revision.expire_time
            ):
                print(f"   🔚 Expires: {specific_revision.expire_time}")
        else:
            print("📭 No revisions available for time-travel query")

    except Exception as e:
        print(f"⚠️ Could not retrieve revision: {e}")
else:
    print("📭 No memory available for time-travel demonstration")

print("\n🎯 Create initial memory with revision tracking\n")

# Create a new customer memory with revision labels for tracking
rollback_customer_id = "customer_sarah_" + str(uuid.uuid4())[:4]

initial_operation = client.agent_engines.memories.generate(
    name=agent_engine_name,
    scope={"user_id": rollback_customer_id},
    direct_contents_source={
        "events": [
            {
                "content": {
                    "parts": [
                        {
                            "text": "Customer prefers email communication only, no phone calls"
                        }
                    ]
                }
            }
        ]
    },
    config={
        "wait_for_completion": True,
        "revision_labels": {"data_source": "initial_preference", "verified": "true"},
    },
)

if initial_operation.response and initial_operation.response.generated_memories:
    rollback_memory = initial_operation.response.generated_memories[0].memory
    rollback_memory_name = rollback_memory.name

    # Get full memory details
    full_memory = client.agent_engines.memories.get(name=rollback_memory_name)

    print("✅ Initial memory created!")
    print(f"   📝 Fact: {full_memory.fact}")
    print(f"   🆔 Memory ID: {rollback_memory_name.split('/')[-1]}")
else:
    print("⚠️ Failed to create initial memory")

print("\n🎯 Update the memory (simulating an incorrect update)\n")

# Simulate an incorrect update - customer preference misinterpreted
update_operation = client.agent_engines.memories.generate(
    name=agent_engine_name,
    scope={"user_id": rollback_customer_id},
    direct_contents_source={
        "events": [
            {
                "content": {
                    "parts": [
                        {"text": "Customer prefers phone calls for all communications"}
                    ]
                }
            }
        ]
    },
    config={
        "wait_for_completion": True,
        "revision_labels": {
            "data_source": "updated_preference",
            "verified": "false",  # This update was not verified!
        },
    },
)

if update_operation.response and update_operation.response.generated_memories:
    updated_ref = update_operation.response.generated_memories[0]

    # Get the updated memory
    updated_memory = client.agent_engines.memories.get(name=updated_ref.memory.name)

    print("✅ Memory updated (incorrectly)!")
    print(f"   📝 New fact: {updated_memory.fact}")
    print(f"   🔄 Action: {updated_ref.action}")

    # Save the previous revision for rollback
    if hasattr(updated_ref, "previous_revision") and updated_ref.previous_revision:
        previous_revision_name = updated_ref.previous_revision
        previous_revision_id = previous_revision_name.split("/")[-1]
        print(f"   📜 Previous revision ID: {previous_revision_id}")
else:
    print("⚠️ Failed to update memory")

print("\n📜 Viewing revision history before rollback\n")

# List all revisions to see the history
revisions_before = list(
    client.agent_engines.memories.revisions.list(name=rollback_memory_name)
)

print(f"✅ Found {len(revisions_before)} revisions:\n")

for i, rev in enumerate(revisions_before, 1):
    revision_id = rev.name.split("/")[-1]
    labels = rev.labels if hasattr(rev, "labels") and rev.labels else {}

    print(f"📌 Revision {i}:")
    print(f"   🆔 ID: {revision_id}")
    print(f"   📝 Fact: {rev.fact}")
    print(f"   ⏱️ Created: {rev.create_time}")
    print(f"   🏷️ Labels: {labels}")

print("\nRollback to the previous (correct) revision\n")

# Rollback to the previous revision using the revision ID we saved
rollback_operation = client.agent_engines.memories.rollback(
    name=rollback_memory_name, target_revision_id=previous_revision_id
)

print("✅ Rollback operation completed!")
print(f"   🎯 Target revision: {previous_revision_id}")

# Verify the rollback by getting the current memory state
restored_memory = client.agent_engines.memories.get(name=rollback_memory_name)

print("\n📋 Memory after rollback:")
print(f"   📝 Current fact: {restored_memory.fact}")
print(f"   ⏱️ Last updated: {restored_memory.update_time}")

print("\n📊 Comparing revision history after rollback\n")

# List revisions after rollback
revisions_after = list(
    client.agent_engines.memories.revisions.list(name=rollback_memory_name)
)

print(f"✅ Revision count after rollback: {len(revisions_after)}\n")

for i, rev in enumerate(revisions_after, 1):
    revision_id = rev.name.split("/")[-1]
    labels = rev.labels if hasattr(rev, "labels") and rev.labels else {}

    print(f"📌 Revision {i} (after rollback):")
    print(f"   🆔 ID: {revision_id}")
    print(f"   📝 Fact: {rev.fact}")
    print(f"   🏷️ Labels: {labels}")

    # Highlight if this is the verified original
    if labels.get("verified") == "true":
        print("   ✅ This is the verified original revision")
    print()

print("\n🔍 Filtering revisions by labels\n")

# Example 1: Get all verified revisions
print("Example 1: Filter for verified revisions only\n")

try:
    verified_revisions = list(
        client.agent_engines.memories.revisions.list(
            name=rollback_memory_name, config={"filter": 'labels.verified="true"'}
        )
    )

    print(f"✅ Found {len(verified_revisions)} verified revision(s):\n")

    for rev in verified_revisions:
        print(f"   📝 {rev.fact}")
        print(f"   ⏱️ Created: {rev.create_time}")
        print(f"   🏷️ Labels: {rev.labels}")

except Exception as e:
    print(f"⚠️ Filter operation: {e}")

# Example 2: Get revisions from a specific data source
print("\nExample 2: Filter by data source\n")

try:
    source_revisions = list(
        client.agent_engines.memories.revisions.list(
            name=rollback_memory_name,
            config={"filter": 'labels.data_source="initial_preference"'},
        )
    )

    print(
        f"✅ Found {len(source_revisions)} revision(s) from 'initial_preference' source:\n"
    )

    for rev in source_revisions:
        print(f"   📝 {rev.fact}")

except Exception as e:
    print(f"⚠️ Filter operation: {e}")

print("\n🧹 Cleaning up resources...\n")

delete_agent_engine = True

if delete_agent_engine:
    print("⏳ Deleting Agent Engine (this will remove all memories and sessions)...\n")

    # Delete the agent engine
    # The force=True parameter ensures all contained resources are also deleted
    client.agent_engines.delete(name=agent_engine_name, force=True)

    print("✅ Agent Engine deleted successfully!")
else:
    print("⚠️ Cleanup skipped - Agent Engine and all data retained")
    print(f"   📋 Resource Name: {agent_engine_name}")