# 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

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)

print("✅ Vertex AI client initialized!")
print(f"   Project: {PROJECT_ID}")
print(f"   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
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

    Returns:
        List of memory facts for further processing
    """
    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):
            if gen_memory.action != "DELETED" and gen_memory.memory:
                try:
                    # Retrieve full memory details including topics
                    full_memory = client.agent_engines.memories.get(
                        name=gen_memory.memory.name
                    )
                    action_icon = "🆕" if gen_memory.action == "CREATED" else "🔄"
                    print(f"   {action_icon} {i}. {full_memory.fact}")
                    memories.append(full_memory.fact)
                except Exception as e:
                    print(f"   ⚠️ Could not retrieve memory: {e}")
    else:
        print(f"\n📭 No {title.lower()} found")

    return memories


print("✅ Helper function defined successfully!")

print("🛠️ Creating Agent Engine with default managed topics...\n")
print("📋 This configuration uses ONLY the default topics:")
print("   - USER_PERSONAL_INFO")
print("   - USER_PREFERENCES")
print("   - KEY_CONVERSATION_DETAILS")
print("   - EXPLICIT_INSTRUCTIONS\n")

# Configure Memory Bank with default settings (no customization)
default_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 extracting memories from conversations
    generation_config=GenerationConfig(
        model=f"projects/{PROJECT_ID}/locations/{LOCATION}/publishers/google/models/gemini-2.5-flash"
    ),
    # Note: No customization_configs provided, so default managed topics are used
)

# Create the Agent Engine resource
# This provisions the backend infrastructure for Memory Bank
default_agent_engine = client.agent_engines.create(
    config={"context_spec": {"memory_bank_config": default_memory_config}}
)

default_engine_name = default_agent_engine.api_resource.name

print("✅ Agent Engine created with default configuration!")
print(f"   Resource Name: {default_engine_name}\n")

print("💬 Creating session for financial advisor client...\n")

# Generate unique client ID
client_id = "client_michael_" + str(uuid.uuid4())[:4]

# Create a session for this client
session = client.agent_engines.sessions.create(
    name=default_engine_name,
    user_id=client_id,
    config={"display_name": f"Financial consultation for {client_id}"},
)

session_name = session.response.name

print("✅ Session created successfully!")
print(f"   Session Name: {session_name}")
print(f"   Client ID: {client_id}\n")

print("Adding initial financial advisor conversation to session...\n")

# Realistic conversation between financial advisor and client
# This contains the structured financial data we want to extract
financial_conversation = [
    {
        "role": "user",
        "message": "Hi, I'm Michael Chen. I'm 42 years old and want to start planning for retirement.",
    },
    {
        "role": "model",
        "message": "Hello Michael! I'm glad you're thinking about retirement planning. Let's discuss your financial goals and current situation.",
    },
    {
        "role": "user",
        "message": "My main goal is to retire at 65 with a comfortable lifestyle. I also want to save for my daughter's college education - she's 10 now.",
    },
    {
        "role": "model",
        "message": "Excellent. Two important goals: retirement in 23 years and college savings in about 8 years. What's your current income and savings situation?",
    },
    {
        "role": "user",
        "message": "I earn $150,000 annually. I have about $200,000 in my 401k and $50,000 in savings. My house is worth $600,000 with a $300,000 mortgage.",
    },
    {
        "role": "model",
        "message": "Thank you for sharing. That gives us a good starting point. What's your comfort level with investment risk?",
    },
    {
        "role": "user",
        "message": "I'm pretty conservative. I don't want to risk losing my savings, but I understand I need some growth for retirement.",
    },
    {
        "role": "model",
        "message": "That's a balanced approach. We can structure a portfolio that's conservative but still provides growth potential. Let me note your risk tolerance as moderate-conservative.",
    },
    {
        "role": "user",
        "message": "Also, I'll need to access about $100,000 in 8 years for my daughter's college tuition. The rest can stay invested for retirement.",
    },
    {
        "role": "model",
        "message": "Perfect. We'll structure your investments with different time horizons - more liquid assets for the college fund and longer-term growth investments for retirement.",
    },
]

# Add each conversation turn to the session
# We track invocation_id to maintain conversation order
invocation_id = 0

for turn in financial_conversation:
    client.agent_engines.sessions.events.append(
        name=session_name,
        author=client_id,  # Required: who is speaking
        invocation_id=str(invocation_id),  # Required: conversation sequence
        timestamp=datetime.datetime.now(tz=datetime.timezone.utc),  # Required: when
        config={
            "content": {"role": turn["role"], "parts": [{"text": turn["message"]}]}
        },
    )

    invocation_id += 1
    icon = "👤" if turn["role"] == "user" else "🤖"
    print(f"{icon} {turn['message']}")

print("\n✅ Conversation added to session successfully!")
print("💡 Now let's see what memories are extracted with default topics...")

print("\n" + "=" * 80)
print("GENERATING MEMORIES WITH DEFAULT MANAGED TOPICS")
print("=" * 80 + "\n")

print("📋 This will extract information matching default topics:")
print("   - USER_PERSONAL_INFO (name, age, basic details)")
print("   - USER_PREFERENCES (general likes/dislikes)")
print("   - KEY_CONVERSATION_DETAILS (important outcomes)")
print("   - EXPLICIT_INSTRUCTIONS (direct remember/forget requests)\n")

print("⏳ Processing conversation...")

# Generate memories from the financial conversation
# This is a long-running operation that processes the entire session
default_operation = client.agent_engines.memories.generate(
    name=default_engine_name,
    vertex_session_source={"session": session_name},
    config={"wait_for_completion": True},  # Wait for completion (blocking call)
)

print("\n✅ Memory generation complete!")

# Display the generated memories using our helper function
default_memories = display_generated_memories(
    default_operation, client, "Memories Extracted with DEFAULT Managed Topics"
)

print("\n" + "=" * 80)
print("BASELINE RESULTS: Complete Memory List")
print("=" * 80 + "\n")

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

all_default_memories = list(results)

print(f"✅ Total memories captured with default topics: {len(all_default_memories)}\n")

if all_default_memories:
    print("📋 Complete list of extracted memories:")
    for i, memory in enumerate(all_default_memories, 1):
        print(f"   {i}. {memory.memory.fact}")
else:
    print("⚠️ No memories were extracted (this would indicate a problem)")

print("\n" + "=" * 80)
print("KEY OBSERVATION:")
print("=" * 80)
print("\nDefault topics provide a basic understanding but:")
print("  - Merge distinct financial concepts into generic categories")
print("  - Miss specific numerical details crucial for financial planning")
print("  - Don't distinguish between different investment timelines")
print("  - Lack structured categorization needed for domain expertise")
print("\nThis is exactly why financial services need CUSTOM TOPICS.")
print("Let's build them now...")

print("🎨 Defining custom financial advisor topics...\n")

custom_financial_topics = [
    # Topic 1: Investment Goals
    # What the client wants to achieve financially
    MemoryTopic(
        custom_memory_topic=CustomMemoryTopic(
            label="investment_goals",
            description="""Extract the client's specific financial objectives and goals. Include:

            - Primary goals (retirement planning, education funding, wealth accumulation, major purchases)
            - Specific targets with timelines (e.g., "retire at 65", "college fund needed in 8 years")
            - Target amounts or income requirements (e.g., "$80,000 annual retirement income")
            - Life milestones tied to goals (e.g., "daughter turns 18", "house purchase at 45")
            - Multiple goals should be captured separately (retirement vs. education vs. home purchase)

            Format: Extract as specific, actionable goals with quantitative details when available.
            Example: "Client's primary goal is retirement at age 65 with $80,000 annual income"
            Example: "Client needs $200,000 for daughter's college education in 8 years (daughter currently age 10)"

            Do NOT include risk tolerance or current financial situation here - those belong in other topics.""",
        )
    ),
    # Topic 2: Risk Tolerance
    # How much investment volatility the client can accept
    MemoryTopic(
        custom_memory_topic=CustomMemoryTopic(
            label="risk_tolerance",
            description="""Extract the client's comfort level with investment risk and volatility. Include:

            - Risk profile category (conservative, moderate-conservative, moderate, moderate-aggressive, aggressive)
            - Specific risk preferences and attitudes (e.g., "can't afford to lose principal", "wants growth", "balanced approach")
            - Emotional comfort with market fluctuations
            - Past experiences that influence risk appetite (e.g., "lost money in 2008", "comfortable with volatility")
            - Willingness to accept short-term losses for long-term gains
            - Specific risk constraints (e.g., "needs guaranteed income", "cannot risk college fund")

            Format: Capture both stated risk tolerance and behavioral/emotional indicators.
            Example: "Client has conservative risk tolerance - doesn't want to risk losing principal but understands need for some growth"
            Example: "Client is aggressive investor comfortable with high volatility given 20-year timeline"

            Do NOT include investment goals or financial capacity here - focus only on risk attitudes.""",
        )
    ),
    # Topic 3: Financial Capacity
    # Current financial situation and resources
    MemoryTopic(
        custom_memory_topic=CustomMemoryTopic(
            label="financial_capacity",
            description="""Extract the client's complete current financial situation and resources. Include:

            - Annual income (salary, bonuses, other income sources with specific amounts)
            - Current savings and investments (401k balance, IRA, brokerage accounts, emergency fund)
            - Regular contribution ability (monthly/annual amounts client can invest)
            - Real estate holdings (primary residence value, investment properties, mortgage balances)
            - Other assets (business ownership, vehicle value, inheritances expected)
            - Liabilities (mortgages, loans, credit card debt with amounts)
            - Net worth or equity positions when mentioned
            - Cash flow and saving capacity

            Format: Extract specific dollar amounts and account types. Calculate derived values when helpful.
            Example: "Client earns $150,000 annually"
            Example: "Client has $200,000 in 401k, $50,000 in savings (total liquid assets: $250,000)"
            Example: "Client owns home worth $600,000 with $300,000 mortgage (equity: $300,000)"

            Focus on current financial state, not goals or preferences.""",
        )
    ),
    # Topic 4: Account Preferences
    # Preferred account types and investment vehicles
    MemoryTopic(
        custom_memory_topic=CustomMemoryTopic(
            label="account_preferences",
            description="""Extract the client's preferences for account types and investment vehicles. Include:

            - Preferred account types (401k, traditional IRA, Roth IRA, 529 education plans, taxable brokerage)
            - Investment vehicle preferences (mutual funds, ETFs, individual stocks, bonds, alternatives)
            - Trading frequency and style (buy-and-hold, active trading, rebalancing preferences)
            - Tax considerations (preference for tax-advantaged accounts, tax-loss harvesting interest)
            - Employer benefits (401k match utilization, ESPP participation)
            - Diversification preferences (sector preferences, international exposure, concentration limits)
            - ESG or ethical investing preferences if mentioned

            Format: Capture both existing account usage and stated preferences for future investments.
            Example: "Client prefers Roth IRA for tax-free growth in retirement"
            Example: "Client wants diversified portfolio across asset classes"
            Example: "Client prefers index funds and ETFs over individual stock picking"

            Do NOT include dollar amounts or balances here - those belong in financial_capacity.""",
        )
    ),
    # Topic 5: Liquidity Timeline
    # When the client needs access to funds
    MemoryTopic(
        custom_memory_topic=CustomMemoryTopic(
            label="liquidity_timeline",
            description="""Extract when the client will need access to invested funds and liquidity requirements. Include:

            - Short-term needs (within 3 years): emergency fund, planned purchases, near-term expenses
            - Medium-term needs (3-10 years): home down payment, education funding, business investment
            - Long-term needs (10+ years): retirement, legacy planning, long-term wealth building
            - Specific amounts needed at specific times (e.g., "$100,000 in 8 years for college")
            - Flexibility in timing (hard deadlines vs. flexible goals)
            - Partial vs. full liquidity requirements
            - Rolling or phased access needs (e.g., "annual withdrawals starting at 65")

            Format: Be specific about amounts, timelines, and flexibility. Separate different time horizons clearly.
            Example: "Client needs $100,000 in 8 years for daughter's college tuition (hard deadline)"
            Example: "Client has 23-year timeline until retirement at 65 (flexible, long-term)"
            Example: "Client wants to maintain $50,000 emergency fund for immediate access"

            This helps structure portfolios with appropriate investment horizons and liquidity.""",
        )
    ),
]

print("✅ Custom financial topics defined successfully!\n")

print("⚙️ Creating customization configuration...\n")

# Package custom topics into a CustomizationConfig
# This object will be passed to the MemoryBankConfig
financial_customization = CustomizationConfig(memory_topics=custom_financial_topics)

print("✅ Customization configuration created!")

print("⚙️ Creating Memory Bank config with custom topics...\n")

# Create Memory Bank config with our custom topics
custom_memory_config = MemoryBankConfig(
    # Same embedding model for similarity search (text-embedding-005)
    similarity_search_config=SimilaritySearchConfig(
        embedding_model=f"projects/{PROJECT_ID}/locations/{LOCATION}/publishers/google/models/text-embedding-005"
    ),
    # Same LLM for memory extraction (gemini-2.5-flash)
    generation_config=GenerationConfig(
        model=f"projects/{PROJECT_ID}/locations/{LOCATION}/publishers/google/models/gemini-2.5-flash"
    ),
    # NEW: Add our custom financial topics
    # This is what makes the difference!
    customization_configs=[financial_customization],
)

print("✅ Memory Bank configuration created with custom topics!")

print("\n" + "=" * 80)
print("CREATING AGENT ENGINE WITH CUSTOM TOPICS")
print("=" * 80 + "\n")

print("🛠️ Creating new Agent Engine with custom financial topics...")

# Create new Agent Engine with custom financial topics
custom_agent_engine = client.agent_engines.create(
    config={"context_spec": {"memory_bank_config": custom_memory_config}}
)

custom_engine_name = custom_agent_engine.api_resource.name

print("\n✅ Agent Engine created with custom topics!")
print(f"   Resource Name: {custom_engine_name}\n")

print("💬 Creating session for custom topics engine...\n")

# Create new session for the custom engine
# Using the same client_id to maintain user identity
custom_session = client.agent_engines.sessions.create(
    name=custom_engine_name,
    user_id=client_id,  # Same client ID as baseline
    config={"display_name": f"Custom topics session for {client_id}"},
)

custom_session_name = custom_session.response.name

print("✅ Session created successfully!")
print(f"   Session Name: {custom_session_name}")
print(f"   Client ID: {client_id} (same as baseline)\n")

print("⬆️ Adding the same financial conversation to custom session...\n")

# Reset invocation counter for new session
invocation_id = 0

# Add the same conversation to the custom engine's session
for turn in financial_conversation:
    client.agent_engines.sessions.events.append(
        name=custom_session_name,
        author=client_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

print("✅ Conversation added to custom session successfully!")

print("\n" + "=" * 80)
print("GENERATING MEMORIES WITH CUSTOM FINANCIAL TOPICS")
print("=" * 80 + "\n")

print("⏳ Processing conversation...")

# Generate memories using custom topics
# This is where the magic happens!
custom_operation = client.agent_engines.memories.generate(
    name=custom_engine_name,
    vertex_session_source={"session": custom_session_name},
    config={"wait_for_completion": True},
)

print("\n✅ Memory generation complete!")

# Display memories extracted with custom topics
custom_memories = display_generated_memories(
    custom_operation, client, "Memories Extracted with CUSTOM Financial Topics"
)

print("\n" + "=" * 80 + "\n")
print("  SIDE-BY-SIDE COMPARISON: Default vs. Custom Topics")
print("\n" + "=" * 80 + "\n")

# Retrieve all memories from default engine
print("📊 BASELINE - Memories with DEFAULT Managed Topics:")
default_results = client.agent_engines.memories.retrieve(
    name=default_engine_name, scope={"user_id": client_id}
)
default_memories_list = list(default_results)

if default_memories_list:
    for i, memory in enumerate(default_memories_list, 1):
        print(f"   {i}. {memory.memory.fact}")
else:
    print("   (No memories extracted)")

print(f"\n✅ Total: {len(default_memories_list)} memories")

print("\n" + "-" * 80 + "\n")

# Retrieve all memories from custom engine
print("📊 IMPROVED - Memories with CUSTOM Financial Topics:")
custom_results = client.agent_engines.memories.retrieve(
    name=custom_engine_name, scope={"user_id": client_id}
)
custom_memories_list = list(custom_results)

if custom_memories_list:
    for i, memory in enumerate(custom_memories_list, 1):
        # Get full memory to access topics
        full_memory = client.agent_engines.memories.get(name=memory.memory.name)
        print(f"   {i}. {memory.memory.fact}")
else:
    print("   (No memories extracted)")

print(f"\n✅ Total: {len(custom_memories_list)} memories")

print("🎨 Creating few-shot examples for financial memory extraction...\n")

financial_few_shot_examples = [
    # Example 1: Conservative Retiree - Near-Term Needs
    # Demonstrates: near retirement, conservative risk, specific timeline
    GenerateMemoriesExample(
        conversation_source=ConversationSource(
            events=[
                ConversationSourceEvent(
                    content=Content(
                        role="user",
                        parts=[
                            Part(
                                text="I'm retiring in 3 years and need to move my 401k into something safer. I have about $400,000 saved. I can't take any major risks at this point."
                            )
                        ],
                    )
                ),
                ConversationSourceEvent(
                    content=Content(
                        role="model",
                        parts=[
                            Part(
                                text="I understand. With 3 years until retirement, capital preservation is key. Let's discuss safe investment options."
                            )
                        ],
                    )
                ),
            ]
        ),
        generated_memories=[
            ExampleGeneratedMemory(
                fact="Client's primary investment goal is retirement in 3 years; needs to transition 401k to safer investments"
            ),
            ExampleGeneratedMemory(
                fact="Client has conservative risk tolerance - cannot accept major risks due to near retirement timeline"
            ),
            ExampleGeneratedMemory(fact="Client has $400,000 in current 401k savings"),
            ExampleGeneratedMemory(
                fact="Client has 3-year timeline until retirement; needs fund access at that point (short-term liquidity requirement)"
            ),
        ],
    ),
    # Example 2: Aggressive Young Investor - Long-Term Growth
    # Demonstrates: long horizon, high risk tolerance, specific investment preferences
    GenerateMemoriesExample(
        conversation_source=ConversationSource(
            events=[
                ConversationSourceEvent(
                    content=Content(
                        role="user",
                        parts=[
                            Part(
                                text="I just got a $10,000 bonus and want to invest it aggressively. I'm 28 and won't need this money for at least 20 years. I'm thinking tech stocks or crypto."
                            )
                        ],
                    )
                ),
                ConversationSourceEvent(
                    content=Content(
                        role="model",
                        parts=[
                            Part(
                                text="Great! With a 20-year horizon at your age, you can handle volatility. Let's discuss growth-oriented options."
                            )
                        ],
                    )
                ),
            ]
        ),
        generated_memories=[
            ExampleGeneratedMemory(
                fact="Client's investment goal is long-term wealth building with 20+ year investment horizon"
            ),
            ExampleGeneratedMemory(
                fact="Client has aggressive risk tolerance - comfortable with high volatility given long timeline and young age (28)"
            ),
            ExampleGeneratedMemory(
                fact="Client has $10,000 from bonus available for investment; currently age 28"
            ),
            ExampleGeneratedMemory(
                fact="Client prefers growth-oriented investments: tech stocks and cryptocurrency"
            ),
            ExampleGeneratedMemory(
                fact="Client has 20+ year timeline with no liquidity needs until then (long-term investment horizon)"
            ),
        ],
    ),
    # Example 3: Balanced Investor - Structured Portfolio
    # Demonstrates: specific account type, asset allocation, regular contributions
    GenerateMemoriesExample(
        conversation_source=ConversationSource(
            events=[
                ConversationSourceEvent(
                    content=Content(
                        role="user",
                        parts=[
                            Part(
                                text="I want to open a Roth IRA and contribute the maximum each year. I'm comfortable with 70% stocks and 30% bonds. I can invest $6,500 annually."
                            )
                        ],
                    )
                ),
                ConversationSourceEvent(
                    content=Content(
                        role="model",
                        parts=[
                            Part(
                                text="Excellent choice! A Roth IRA with that 70/30 allocation is a solid retirement strategy."
                            )
                        ],
                    )
                ),
            ]
        ),
        generated_memories=[
            ExampleGeneratedMemory(
                fact="Client's investment goal is to maximize Roth IRA contributions for tax-advantaged retirement savings"
            ),
            ExampleGeneratedMemory(
                fact="Client has moderate to moderate-aggressive risk tolerance with 70% stocks / 30% bonds asset allocation preference"
            ),
            ExampleGeneratedMemory(
                fact="Client can contribute $6,500 annually (maximum Roth IRA contribution limit)"
            ),
            ExampleGeneratedMemory(
                fact="Client prefers Roth IRA account type with 70% stock / 30% bond allocation"
            ),
        ],
    ),
]

print("✅ Few-shot examples created successfully!\n")

print("\n⚙️ Updating customization configuration with few-shot examples...\n")

# Create advanced customization with both custom topics AND few-shot examples
advanced_financial_customization = CustomizationConfig(
    memory_topics=custom_financial_topics,  # Our 5 custom topics (WHAT to extract)
    generate_memories_examples=financial_few_shot_examples,  # Our 3 few-shot examples (HOW to extract)
)

print("✅ Advanced customization created!")

print("\n" + "=" * 80)
print("UPDATING AGENT ENGINE WITH FEW-SHOT EXAMPLES")
print("=" * 80 + "\n")

# Create Memory Bank config with custom topics AND few-shot examples
advanced_memory_config = MemoryBankConfig(
    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"
    ),
    customization_configs=[
        advanced_financial_customization
    ],  # Updated with few-shot examples
)

# Update the existing custom engine
print("⚙️ Updating Agent Engine configuration...")
updated_engine = client.agent_engines.update(
    name=custom_engine_name,
    config={"context_spec": {"memory_bank_config": advanced_memory_config}},
)

print("\n✅ Agent Engine updated with few-shot examples!")

print("\n💬 Creating new session to test few-shot examples...\n")

# Create new session for few-shot testing
# Still using the same client_id for continuity
fewshot_session = client.agent_engines.sessions.create(
    name=custom_engine_name,  # Using the updated engine with few-shot
    user_id=client_id,
    config={"display_name": f"Few-shot test session for {client_id}"},
)

fewshot_session_name = fewshot_session.response.name

print("✅ Session created successfully!")
print(f"   Session Name: {fewshot_session_name}")
print(f"   Engine: {custom_engine_name} (now with few-shot examples)")

print("\n⬆️ Adding same financial conversation to few-shot session...\n")

# Reset invocation counter
invocation_id = 0

# Add conversation to few-shot session
for turn in financial_conversation:
    client.agent_engines.sessions.events.append(
        name=fewshot_session_name,
        author=client_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

print("✅ Conversation added to few-shot session!")

print("\n" + "=" * 80)
print("GENERATING MEMORIES WITH CUSTOM TOPICS + FEW-SHOT EXAMPLES")
print("=" * 80 + "\n")

print("⏳ Processing conversation...")

# Generate memories with few-shot configuration
fewshot_operation = client.agent_engines.memories.generate(
    name=custom_engine_name,
    vertex_session_source={"session": fewshot_session_name},
    config={"wait_for_completion": True},
)

print("\n✅ Memory generation complete!")

# Display memories generated with few-shot examples
fewshot_memories = display_generated_memories(
    fewshot_operation,
    client,
    "Memories Extracted with CUSTOM TOPICS + FEW-SHOT EXAMPLES",
)

print("\n" + "=" * 80 + "\n")
print("  COMPREHENSIVE COMPARISON: Three Approaches to Memory Extraction")
print("\n" + "=" * 80 + "\n")

print("📊 APPROACH 1: Default Managed Topics (Baseline)")
print("\n")
print("⚙️ Configuration:")
print("   - 4 generic managed topics (USER_PERSONAL_INFO, USER_PREFERENCES, etc.)")
print("   - No domain customization")
print("   - No few-shot examples\n")

print("📋 Extracted Memories:")
default_results = client.agent_engines.memories.retrieve(
    name=default_engine_name, scope={"user_id": client_id}
)
default_list = list(default_results)

if default_list:
    for i, memory in enumerate(default_list, 1):
        print(f"   {i}. {memory.memory.fact}")
else:
    print("   (No memories extracted)")

print("\n" + "-" * 80 + "\n")

print("📊 APPROACH 2: Custom Financial Topics")
print("\n")
print("⚙️ Configuration:")
print("   - 5 domain-specific topics (investment_goals, risk_tolerance, etc.)")
print("   - Detailed topic descriptions")
print("   - No few-shot examples\n")

print("📋 Extracted Memories:")
custom_results = client.agent_engines.memories.retrieve(
    name=custom_engine_name, scope={"user_id": client_id}
)
custom_list = list(custom_results)

if custom_list:
    for i, memory in enumerate(custom_list, 1):
        full_memory = client.agent_engines.memories.get(name=memory.memory.name)
        print(f"   {i}. {memory.memory.fact}")
else:
    print("   (No memories extracted)")

print("\n" + "-" * 80 + "\n")

print("📊 APPROACH 3: Custom Topics + Few-Shot Examples (Optimized)")
print("\n")
print("⚙️ Configuration:")
print("   - 5 domain-specific topics (WHAT to extract)")
print("   - 3 few-shot examples (HOW to extract it)")
print("   - Detailed topic descriptions + extraction patterns\n")

print("📋 Extracted Memories:")
# The fewshot_memories were already extracted, display them
for i, fact in enumerate(fewshot_memories, 1):
    print(f"   {i}. {fact}")

def generate_personalized_financial_advice(client_id, query, engine_name):
    """Generate personalized financial advice using memories and Gemini."""
    from google import genai

    genai_client = genai.Client(vertexai=True, project=PROJECT_ID, location=LOCATION)

    print("\n🔍 Generating personalized advice for query:")
    print(f'   "{query}"\n')

    # Retrieve relevant memories using similarity search
    results = client.agent_engines.memories.retrieve(
        name=engine_name,
        scope={"user_id": client_id},
        similarity_search_params={
            "search_query": query,
            "top_k": 5,  # Get top 5 most relevant memories
        },
    )

    memories = list(results)

    print("📋 Client Context (from Memory Bank):\n")
    memory_context = []

    if memories:
        for i, mem in enumerate(memories, 1):
            print(f"   {i}. {mem.memory.fact}")
            memory_context.append(mem.memory.fact)
    else:
        print("⚠️ No relevant memories found.")
        print("\n💡 Without memory context, Gemini can only provide generic advice.")
        return None

    # Create prompt augmented with memory context
    prompt = f"""You are an expert financial advisor. Based on the following client information
from their previous conversations, provide personalized financial advice.

Client Information from Memory Bank:
{chr(10).join(f"- {fact}" for fact in memory_context)}

Client Question: {query}

Provide specific, actionable financial advice that:
1. Addresses their question directly
2. Considers their stated risk tolerance
3. Aligns with their investment goals and timeline
4. Takes into account their financial capacity
5. Respects their liquidity needs

Format your response with:
1. Direct Answer (personalized to their situation)
2. Rationale (based on their specific profile)
3. Specific Recommendations (with concrete numbers when appropriate)
4. Important Considerations
5. Next Steps"""

    # Generate personalized response with Gemini
    print("\n")
    print("💬 Personalized Financial Advice from Gemini:\n")

    response = genai_client.models.generate_content(
        model="gemini-2.5-flash",
        contents=prompt,
    )

    print(response.text)

    return response.text


print("✅ Function defined successfully!")

print("📊 TEST 1: Advice using DEFAULT MANAGED TOPICS")
print("-" * 80)
print(
    "❓ Query: 'Should I invest more aggressively to maximize returns for retirement?'\n"
)
print("📋 Memories available: From default managed topics (generic)\n")

advice_default = generate_personalized_financial_advice(
    client_id=client_id,
    query="Should I invest more aggressively to maximize returns for retirement?",
    engine_name=default_engine_name,
)

print("\n" + "=" * 80)
print("📊 TEST 2: Advice using CUSTOM TOPICS + FEW-SHOT EXAMPLES")
print("=" * 80 + "\n")

print(
    "❓ Query: 'Should I invest more aggressively to maximize returns for retirement?'\n"
)
print(
    "📋 Memories available: From custom topics + few-shot (structured, comprehensive)\n"
)

advice_custom = generate_personalized_financial_advice(
    client_id=client_id,
    query="Should I invest more aggressively to maximize returns for retirement?",
    engine_name=custom_engine_name,
)

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

delete_agent_engines = True  # Set to False if you want to keep the resources

if delete_agent_engines:
    try:
        # Delete default topics engine
        client.agent_engines.delete(name=default_engine_name, force=True)

        # Delete custom topics engine (includes few-shot configuration)
        client.agent_engines.delete(name=custom_engine_name, force=True)

        print("✅ All resources cleaned up successfully!")

    except Exception as e:
        print(f"\n⚠️ Error during cleanup: {e}")