Generative AI Google Brain

agents-agent-engine-memory-bank-get-started-with-memory-bank-custom-topics

4 min read

Customizing Memory Topics

Source notebook

Repo path: agents/agent_engine/memory_bank/get_started_with_memory_bank_custom_topics.ipynb · Open on GitHub · intermediate

Customizes Vertex AI Memory Bank topics for a financial advisor assistant and compares default vs custom extraction.

Summary

This notebook teaches how to configure Vertex AI Memory Bank with default managed topics, then replace generic extraction with domain-specific financial memory topics. It creates Agent Engine resources, records a financial advisor conversation in sessions, generates and retrieves memories, and compares default extraction against custom topics. It also introduces few-shot examples to show the extraction model the desired financial memory granularity.

Key code patterns

Initialize Vertex AI client

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

Sets the project and region used by all Agent Engine, session, and memory operations.

Configure Memory Bank models

default_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"
    ),
)

Uses an embedding model for similarity search and Gemini 2.5 Flash for memory extraction.

Define custom memory topics

custom_financial_topics = [
    MemoryTopic(
        custom_memory_topic=CustomMemoryTopic(
            label="investment_goals",
            description="Extract the client's specific financial objectives and goals."
        )
    )
]

Custom labels and detailed descriptions teach Memory Bank what domain-specific facts to extract.

Attach customization config

financial_customization = CustomizationConfig(
    memory_topics=custom_financial_topics
)
 
custom_memory_config = MemoryBankConfig(
    similarity_search_config=SimilaritySearchConfig(embedding_model=embedding_model),
    generation_config=GenerationConfig(model=generation_model),
    customization_configs=[financial_customization],
)

Adds custom financial topics to the Memory Bank configuration used by Agent Engine.

Create Agent Engine

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

Provisions an Agent Engine resource backed by the configured Memory Bank.

Create session and append events

session = client.agent_engines.sessions.create(
    name=custom_engine_name,
    user_id=client_id,
    config={"display_name": f"Custom topics session for {client_id}"},
)
 
client.agent_engines.sessions.events.append(
    name=session.response.name,
    author=client_id,
    invocation_id="0",
    timestamp=datetime.datetime.now(tz=datetime.timezone.utc),
    config={"content": {"role": "user", "parts": [{"text": message}]}}
)

Stores ordered conversation turns with required author, invocation_id, timestamp, and content fields.

Generate and retrieve memories

operation = client.agent_engines.memories.generate(
    name=custom_engine_name,
    vertex_session_source={"session": custom_session_name},
    config={"wait_for_completion": True},
)
 
results = client.agent_engines.memories.retrieve(
    name=custom_engine_name,
    scope={"user_id": client_id},
)

Runs memory extraction from a session and retrieves memories scoped to a user.

Few-shot extraction examples

GenerateMemoriesExample(
    conversation_source=ConversationSource(events=[...]),
    generated_memories=[
        ExampleGeneratedMemory(
            fact="Client has conservative risk tolerance - cannot accept major risks due to near retirement timeline"
        )
    ],
)

Shows the model the expected memory facts and level of granularity for financial conversations.

Models & APIs used

  • Models: text-embedding-005, gemini-2.5-flash
  • APIs / services: Vertex AI, Vertex AI API, Vertex AI Agent Engine, Vertex AI Memory Bank
  • SDKs / libraries: google-cloud-aiplatform, vertexai

When to use this

Use this pattern when a domain agent needs precise long-term memory extraction beyond Memory Bank’s default managed topics.

Gotchas & caveats

  • Vertex AI SDK version 1.111.0 or higher is required for all Memory Bank features.
  • Colab users may need to authenticate with google.colab.auth.authenticate_user().
  • A Google Cloud project must exist and the Vertex AI API must be enabled.
  • PROJECT_ID and LOCATION must be set, with the notebook defaulting LOCATION to us-central1.
  • Session events require author, invocation_id, timestamp, and content fields.
  • wait_for_completion=True blocks until memory generation finishes; the notebook notes production may use asynchronous generation.
  • The few-shot section uses Content and Part objects in examples, so those types must be available in the runtime.

Best practices

  • Create a baseline with default managed topics before adding custom topics.
  • Use the same conversation for default and custom engines to make the comparison apples-to-apples.
  • Keep custom topics focused on separate domain dimensions to reduce overlap and improve precision.
  • Write detailed topic descriptions with examples and exclusions for what belongs in other topics.
  • Use realistic few-shot conversations that show the desired memory extraction style.
  • Provide 2-5 few-shot examples per domain, according to the notebook’s stated guidance.
  • Use the same user_id across sessions when preserving client identity matters.
  • Retrieve full memory details with memories.get when displaying generated memories.

Graph View

Backlinks