"Managing Agent State and Memory in Google ADK"

1. Short-Term Context: The Session and its State

In web development, a "session" typically tracks a user's activity on a site during a single visit. Google ADK uses a very similar concept. A Session object encapsulates a single conversation between a user and an agent, containing the history of all interactions (events) and a "scratchpad" for temporary data (state).

To get a clear understanding of this lifecycle—from the user's first message to how the agent uses session data—let's watch a short video.

As the video explained, the session.state is the key to maintaining context within a conversation. Let's dive deeper into its mechanics with the official documentation.

Hands-On: Giving Your Agent a Name to Remember

Let's apply this by making our agent remember the user's name.

First, let's update the agent's instructions to be more personalized. Open personal_assistant/instructions.py and modify code_review_instruction:

# personal_assistant/instructions.py

code_review_instruction = """
You are an expert AI assistant. Your purpose is to help {user_name} with their tasks.

Your capabilities include:
1.  **Code Review**: Provide constructive feedback on code snippets.
2.  **General Questions**: Answer questions using your tools.

Always address the user by their name, {user_name}, when appropriate.

You have access to tools to get currency exchange rates and to perform Google searches.
- Use 'get_fx_rate' for exchange rates.
- Use 'google_search_agent' for general knowledge questions.
"""

Notice the {user_name} placeholders. ADK will automatically substitute these with values from session.state.

But where does user_name come from? For now, we'll manually set it when we run our agent. The adk web command has flags to initialize the state.

Stop your adk web server if it's running, and restart it with the following command. Replace "YourName" with your actual name.

adk web --initial_state='{"user_name": "YourName"}'

Now, open your browser to http://localhost:8000, select your agent, and say "Hello". The agent should greet you by name, demonstrating it has access to the initial state.

This is useful for providing initial context, but what about information that emerges during the conversation? The documentation you read mentioned output_key as the simplest method for an agent to save its own output to the state.

Let's say we wanted the agent to remember the last greeting it generated. We could modify agent.py like this (you don't need to code this, just understand the principle):

# Example of using output_key
root_agent = Agent(
    name="code_reviewer_and_assistant",
    # ... other parameters
    output_key="last_response" # This line is new
)

With this change, after the agent generates a response, the Runner would automatically save that text into session.state['last_response']. This state variable would then be available for the next turn.

2. Long-Term Knowledge with MemoryService

session.state is powerful, but it's fundamentally ephemeral. It's designed for a single conversation. If you refresh the page, the session is gone. How can an agent remember that you prefer Python, or that your niece's birthday is next week?

This requires a long-term memory system. In ADK, this is handled by the MemoryService.

To understand how ADK approaches long-term memory, we'll watch a comprehensive video from Google's developer channel. It introduces the two main types of memory services and demonstrates how to use the more powerful, production-ready option: Vertex AI Memory Bank.

The Memory Workflow in Practice

The video covered a lot of ground. Let's distill it into the practical workflow you'll use:

1. Configure a MemoryService: For development, you can use InMemoryMemoryService. For production, you'd set up VertexAiMemoryBankService, which requires a Google Cloud project and an Agent Engine instance.
2. Save to Memory: After a conversation turn, you call memory_service.add_session_to_memory(). The service then extracts and persists meaningful facts. This is best done automatically using an after_agent_callback.
3. Retrieve from Memory: Before the next turn, the agent needs to retrieve relevant memories. This is done using a tool.
   • PreloadMemoryTool is the easiest. It runs automatically, fetches relevant memories based on the current query, and injects them into the agent's system instructions.
   • LoadMemoryTool gives the agent more control, allowing it to decide if and when to search its memory. This is more efficient as it avoids unnecessary lookups.

Hands-On: Giving Your Agent Long-Term Memory

Let's implement this. We will use the InMemoryMemoryService for simplicity, as it requires no cloud setup, but the code structure is identical for VertexAiMemoryBankService.

First, we need a callback function to save the session to memory. Create a new file, personal_assistant/callbacks.py, and add the following:

# personal_assistant/callbacks.py

import logging
from google.adk.agents import CallbackContext

async def save_to_memory_callback(context: CallbackContext):
    """
    An after-agent callback to automatically save the session to memory.
    """
    if not context.memory_service:
        logging.warning("MemoryService is not configured. Skipping save to memory.")
        return

    try:
        logging.info("Attempting to save session to memory...")
        await context.memory_service.add_session_to_memory(context.session)
        logging.info("Successfully added session to memory.")
    except Exception as e:
        logging.error(f"Failed to add session to memory: {e}")

Next, update personal_assistant/agent.py to use this callback and the PreloadMemoryTool.

# personal_assistant/agent.py

from google.adk.agents import Agent
from google.adk.tools import FunctionTool
from google.adk.tools.agent_tool import AgentTool
from google.adk.tools.preload_memory_tool import PreloadMemoryTool # Import PreloadMemoryTool

from .instructions import code_review_instruction
from .custom_functions import get_fx_rate
from .custom_agents import google_search_agent
from .callbacks import save_to_memory_callback # Import the callback

# The 'root_agent' variable is mandatory. ADK looks for it as the entry point.
root_agent = Agent(
    name="code_reviewer_and_assistant",
    description="A helpful assistant that reviews code, answers questions, and remembers past conversations.",
    instruction=code_review_instruction,
    tools=[
        FunctionTool(get_fx_rate),
        AgentTool(agent=google_search_agent),
        PreloadMemoryTool(), # Add the memory tool
    ],
    # Register the callback to run after the agent responds
    after_agent_callback=save_to_memory_callback,
)

Now, restart adk web (you can remove the --initial_state flag). Let's test the memory.

1. Turn 1 (Provide Information): "My favorite programming language is Python."
2. Turn 2 (Test Recall): "What is my favorite programming language?"

The agent should now correctly answer your question! It works because:

• After Turn 1, the save_to_memory_callback was triggered, saving the conversation to the in-memory service.
• Before Turn 2, the PreloadMemoryTool automatically searched the memory for context related to "favorite programming language," found the information from Turn 1, and added it to the agent's prompt.

Conclusion

Excellent work today. You've tackled one of the most important concepts in building truly useful agents: memory. Your agent is no longer a stateless machine but a learning companion that can maintain context and build knowledge over time.

Key Takeaways:

• State vs. Memory: session.state is a short-term scratchpad for a single conversation. A MemoryService provides a long-term, persistent knowledge base across conversations.
• Session Management: The Runner and SessionService handle the lifecycle of a conversation, while session.state holds dynamic data. You can access it in instructions with {key} syntax.
• Memory Workflow: The core memory loop is to save conversations (e.g., via an after_agent_callback) and retrieve relevant facts when needed (e.g., via PreloadMemoryTool or LoadMemoryTool).
• Service Tiers: ADK provides InMemory services for easy development and prototyping, and managed cloud services like VertexAiMemoryBankService for scalable, production-ready applications.

Now that we have built a single, competent agent with tools and memory, we can begin to explore more advanced architectural patterns. In our next lesson, we will look at multi-agent systems, where you'll learn how to make multiple specialized agents collaborate to solve complex problems that a single agent cannot handle alone.