Generative AI Google Brain

agents-gke-agents-with-memory-get-started-with-memory-for-adk-in-gke

3 min read

Get started with Sessions and Memory Bank for ADK agents in Google Kubernetes Engine

Source notebook

Repo path: agents/gke/agents_with_memory/get_started_with_memory_for_adk_in_gke.ipynb · Open on GitHub · advanced

Builds an ADK weather agent with Vertex AI Sessions, Memory Bank, and GKE deployment.

Summary

This notebook teaches how to add short-term and long-term memory to an ADK agent using Vertex AI Agent Engine Sessions and Memory Bank. It registers an Agent Engine, wires VertexAISessionService and VertexAiMemoryBankService into an ADK Runner, demonstrates memory carried across new sessions, then deploys the agent to Google Kubernetes Engine with ADK CLI.

Key code patterns

Vertex AI environment setup

os.environ["GOOGLE_GENAI_USE_VERTEXAI"] = "TRUE"
os.environ["GOOGLE_CLOUD_PROJECT"] = PROJECT_ID
os.environ["GOOGLE_CLOUD_LOCATION"] = LOCATION

Configures ADK and GenAI calls to use Vertex AI in the selected project and location.

Create Agent Engine memory backing

agent_engine = client.agent_engines.create(config={
    "display_name": AGENT_NAME,
    "context_spec": {"memory_bank_config": {"generation_config": {
        "model": f"projects/{PROJECT_ID}/locations/{LOCATION}/publishers/google/models/gemini-2.5-flash"
    }}}
})

Registers an Agent Engine resource so Sessions and Memory Bank can store data for the agent.

Attach managed session and memory services

session_service = VertexAiSessionService(
    project=PROJECT_ID, location=LOCATION, agent_engine_id=agent_engine_id)
memory_service = VertexAiMemoryBankService(
    project=PROJECT_ID, location=LOCATION, agent_engine_id=agent_engine_id)

Moves session state and long-term memories outside the agent runtime for scalable deployments.

Persist session after each turn

async def add_session_to_memory(callback_context):
    invocation_context = callback_context._invocation_context
    if invocation_context.memory_service:
        await invocation_context.memory_service.add_session_to_memory(
            invocation_context.session)

Lets Memory Bank extract and persist meaningful facts from completed agent turns.

Preload memory into agent context

root_agent = Agent(
    name="weather_agent",
    model="gemini-2.5-flash",
    tools=[get_weather, PreloadMemoryTool()],
    after_agent_callback=add_session_to_memory)

Uses ADK’s built-in memory tool to retrieve memories at the start of each turn.

Deploy ADK agent to GKE with memory

adk deploy gke \
  --session_service_uri=agentengine://{agent_engine_id} \
  --memory_service_uri=agentengine://{agent_engine_id} \
  --app_name {AGENT_NAME} --with_ui .

Passes Agent Engine URIs so the GKE-hosted agent uses managed sessions and memory.

Models & APIs used

  • Models: gemini-2.5-flash
  • APIs / services: Vertex AI, Vertex AI Agent Engine, Vertex AI Agent Engine Sessions, Vertex AI Agent Engine Memory Bank, Google Kubernetes Engine, Cloud Build, Artifact Registry
  • SDKs / libraries: google-adk, google-cloud-aiplatform, vertexai, google-genai

When to use this

Use this pattern when an ADK agent needs durable conversation state, cross-session memory, and deployment on GKE.

Gotchas & caveats

  • PROJECT_ID must be set to an existing Google Cloud project.
  • Manual gcloud authentication may be required depending on the Jupyter environment.
  • container.googleapis.com, cloudbuild.googleapis.com, artifactregistry.googleapis.com, and aiplatform.googleapis.com must be enabled.
  • kubectl and required GKE plugins must be installed before cluster access.
  • InMemorySessionService loses data when the runner shuts down and does not work reliably across scaled instances.
  • Memory Bank stores extracted meaningful memories, so not every conversation creates a memory.
  • GKE Standard clusters need Workload Identity enabled; Autopilot enables it by default.
  • Pods and external load balancer IP assignment can take a few minutes.
  • The cluster service account needs roles/aiplatform.user for Vertex AI access.

Best practices

  • Store production ADK session data outside the agent runtime.
  • Use VertexAISessionService for scalable managed session storage.
  • Use VertexAiMemoryBankService for persistent long-term memory.
  • Provide both a memory tool on the Agent and a memory service on the Runner.
  • Provide a session service on the Runner when using managed sessions.
  • Create new sessions in tests to verify Memory Bank carries knowledge across sessions.
  • Check pods are Running before opening the deployed agent UI.
  • Clean up Agent Engine and GKE resources to avoid unexpected costs.

Graph View

Backlinks