Skip to main content

Command Palette

Search for a command to run...

State and Memory Management in Google ADK: A Practical Tutorial

Published
7 min read
State and Memory Management in Google ADK: A Practical Tutorial
J

Senior AI Engineer | Building & Telling Stories about AI/ML Systems | Software Engineer

Introduction

Building truly intelligent agents requires more than just processing individual requests—agents need to remember context, learn from interactions, and maintain continuity across conversations.

This is where Google ADK's state and memory management system becomes essential.

In this comprehensive tutorial, we'll explore how to leverage ADK's state management to create context-aware, personalized agents that can maintain information across interactions and provide meaningful, continuous experiences.

Understanding State vs. Conversation History

Before diving into implementation, it's crucial to understand the distinction between two types of memory in ADK:

  • Conversation History: The chronological sequence of messages exchanged between user and agent

  • State: A structured key-value store for persistent data that agents can read from and write to

Think of conversation history as a transcript, while state is like a notebook where agents jot down important facts, preferences, and contextual information.

The Four Scopes of State

ADK's state management system uses prefixes to determine data persistence scope:

1. Session State (No Prefix)

Data persists only for the current conversation session.

session.state["current_step"] = 3
session.state["last_query"] = "What's the weather in London?"

Use Cases:

  • Tracking progress in multi-step workflows

  • Storing temporary calculations

  • Maintaining conversation-specific context

2. User State (user: prefix)

Data persists across all sessions for a specific user.

session.state["user:preferred_language"] = "Spanish"
session.state["user:timezone"] = "America/New_York"
session.state["user:name"] = "Sarah"

Use Cases:

  • Personal preferences and settings

  • User profile information

  • Long-term learning about user behavior

3. Application State (app: prefix)

Data shared across all users and sessions.

session.state["app:version"] = "2.1.0"
session.state["app:feature_flags"] = {"new_ui": True}
session.state["app:maintenance_mode"] = False

Use Cases:

  • Global configuration settings

  • Feature flags

  • System-wide announcements

4. Temporary State (temp: prefix)

Data that exists only during current execution, not persisted.

session.state["temp:api_response"] = response_data
session.state["temp:intermediate_result"] = calculation

Use Cases:

  • Intermediate processing results

  • Temporary caching

  • Cross-tool data sharing within single execution

Setting Up Basic State Management

Let's start with a simple example of creating and managing sessions:

import asyncio
from google.adk.sessions import InMemorySessionService

# Initialize session service
session_service = InMemorySessionService()

# Application configuration
APP_NAME = "personal_assistant"
USER_ID = "user_12345"
SESSION_ID = "session_67890"

# Basic state operations
async def demonstrate_basic_state():
    # Create or retrieve session
    session = await session_service.create_session(
        app_name=APP_NAME,
        user_id=USER_ID,
        session_id=SESSION_ID
    )

    # Reading with default value
    user_name = session.state.get("user:name", "Guest")

    # Writing to different scopes
    session.state["user:name"] = "Alice"
    session.state["last_interaction"] = "greeting"
    session.state["app:total_users"] = session.state.get("app:total_users", 0) + 1

    user_name = session.state.get("user:name", "Guest")
    last_interaction = session.state.get("last_interaction", "none")
    total_users = session.state.get("app:total_users", 0)
    print(f"Hello, {user_name}!")
    print(f"Total app users: {total_users}")
    print(f"Last interaction: {last_interaction}")

# Run the async function
asyncio.run(demonstrate_basic_state())

Building State-Aware Tools

Tools are where agents interact with state most directly. Here's how to create tools that leverage state management:

from google.adk.tools.tool_context import ToolContext
from datetime import datetime
from typing import Dict, Any

def save_user_preference(
    preference_type: str, 
    value: str, 
    tool_context: ToolContext
) -> Dict[str, Any]:
    """Save a user preference with timestamp.

    Args:
        preference_type: Type of preference (e.g., 'cuisine', 'music_genre')
        value: The preference value
        tool_context: Automatically injected by ADK

    Returns:
        dict: Operation status and details
    """
    # Store preference with user scope
    preference_key = f"user:preference_{preference_type}"
    timestamp_key = f"user:preference_{preference_type}_updated"

    # Save the preference and when it was set
    tool_context.state[preference_key] = value
    tool_context.state[timestamp_key] = datetime.now().isoformat()

    return {
        "status": "success",
        "message": f"Saved {preference_type} preference: {value}",
        "updated_at": tool_context.state[timestamp_key]
    }

Creating State-Aware Agents

Now let's build agents that intelligently use state information:

from google.adk.agents import Agent
from google.adk.tools.tool_context import ToolContext
from datetime import datetime
from typing import Dict, Any, List

def save_user_preference(
    preference_type: str, 
    value: str, 
    tool_context: ToolContext
) -> Dict[str, Any]:
    """Save a user preference with timestamp.

    Args:
        preference_type: Type of preference (e.g., 'cuisine', 'music_genre')
        value: The preference value
        tool_context: Automatically injected by ADK

    Returns:
        dict: Operation status and details
    """
    # Store preference with user scope
    preference_key = f"user:preference_{preference_type}"
    timestamp_key = f"user:preference_{preference_type}_updated"

    # Save the preference and when it was set
    tool_context.state[preference_key] = value
    tool_context.state[timestamp_key] = datetime.now().isoformat()

    return {
        "status": "success",
        "message": f"Saved {preference_type} preference: {value}",
        "updated_at": tool_context.state[timestamp_key]
    }

def get_user_profile(tool_context: ToolContext) -> Dict[str, Any]:
    """Retrieve comprehensive user profile information.

    Args:
        tool_context: Automatically injected by ADK

    Returns:
        dict: User profile data including preferences and history
    """
    # Get user name
    user_name = tool_context.state.get("user:name", "Guest")

    # Collect all user preferences
    preferences = {}

    # Check for common preference types
    common_preferences = [
        "cuisine", "music_genre", "favorite_color", "language", "outdoor_activity",
        "timezone", "notification_preference", "theme", "accessibility", "reading_genre"
    ]

    for pref_type in common_preferences:
        pref_key = f"user:preference_{pref_type}"
        if pref_key in tool_context.state:
            preferences[pref_type] = tool_context.state[pref_key]
            # Get timestamp if available
            timestamp_key = f"{pref_key}_updated"
            if timestamp_key in tool_context.state:
                preferences[f"{pref_type}_updated"] = tool_context.state[timestamp_key]

    # Get conversation history
    last_interaction = tool_context.state.get("last_interaction", "none")
    total_interactions = tool_context.state.get("user:total_interactions", 0)

    return {
        "user_name": user_name,
        "preferences": preferences,
        "last_interaction": last_interaction,
        "total_interactions": total_interactions,
        "profile_retrieved_at": datetime.now().isoformat()
    }

def track_conversation_flow(
    flow_type: str,
    step: str,
    data: str,
    tool_context: ToolContext
) -> Dict[str, Any]:
    """Track multi-step conversation flows for better context.

    Args:
        flow_type: Type of flow (e.g., 'booking', 'planning', 'troubleshooting')
        step: Current step in the flow
        data: Relevant data for this step
        tool_context: Automatically injected by ADK

    Returns:
        dict: Flow tracking status and current state
    """
    # Store flow information
    flow_key = f"user:flow_{flow_type}"
    step_key = f"user:flow_{flow_type}_step"
    data_key = f"user:flow_{flow_type}_data"
    timestamp_key = f"user:flow_{flow_type}_updated"

    tool_context.state[flow_key] = flow_type
    tool_context.state[step_key] = step
    tool_context.state[data_key] = data
    tool_context.state[timestamp_key] = datetime.now().isoformat()

    return {
        "status": "success",
        "message": f"Tracked {flow_type} flow: {step}",
        "current_step": step,
        "flow_data": data,
        "updated_at": tool_context.state[timestamp_key]
    }

def update_user_interaction(
    interaction_type: str,
    details: str,
    tool_context: ToolContext
) -> Dict[str, Any]:
    """Update user interaction history and preferences.

    Args:
        interaction_type: Type of interaction (e.g., 'question', 'request', 'feedback')
        details: Details about the interaction
        tool_context: Automatically injected by ADK

    Returns:
        dict: Interaction update status
    """
    # Update interaction count
    current_count = tool_context.state.get("user:total_interactions", 0)
    tool_context.state["user:total_interactions"] = current_count + 1

    # Store interaction details
    interaction_key = f"user:interaction_{current_count + 1}"
    tool_context.state[interaction_key] = {
        "type": interaction_type,
        "details": details,
        "timestamp": datetime.now().isoformat()
    }

    # Update last interaction
    tool_context.state["last_interaction"] = f"{interaction_type}: {details}"

    return {
        "status": "success",
        "message": f"Updated interaction: {interaction_type}",
        "interaction_count": current_count + 1,
        "timestamp": datetime.now().isoformat()
    }

# Tools list
state_tools = [
    save_user_preference, 
    get_user_profile, 
    track_conversation_flow,
    update_user_interaction
]

# Personal assistant with state awareness
root_agent = Agent(
    name="personal_assistant",
    model="gemini-2.0-flash-exp",
    instruction="""
    You are a highly personalized assistant that remembers user preferences and context.

    STARTUP BEHAVIOR:
    - Always check user state at the beginning of each interaction
    - If user:name exists, greet them by name
    - If this is a returning user, reference relevant previous preferences
    - Check for any ongoing conversation flows and offer to continue them

    STATE USAGE GUIDELINES:
    - Use save_user_preference tool when users express preferences
    - Use get_user_profile tool to understand user background before making recommendations
    - Use track_conversation_flow for multi-step processes (booking, planning, troubleshooting)
    - Use update_user_interaction to track user engagement and build context

    PERSONALIZATION:
    - Tailor responses based on user:preferences
    - Reference previous interactions when relevant
    - Maintain consistency with established user relationships
    - Learn from user feedback and adjust recommendations accordingly

    MEMORY MANAGEMENT:
    - Store important decisions and outcomes
    - Remember user goals and aspirations
    - Track what works well for each user
    - Maintain conversation context across sessions

    CONVERSATION FLOW:
    - For new users, focus on learning their preferences
    - For returning users, reference their history and preferences
    - Proactively suggest improvements based on past interactions
    - Handle multi-step processes with clear progress tracking
    """,
    tools=state_tools,
    output_key="last_assistant_response"
)

See the whole iteration with the agent.

See the events in the session.

And the current state of the agent.

Conclusion

State and memory management in Google ADK enables the creation of truly intelligent, context-aware agents.

By leveraging the four scopes of state (session, user, application, and temporary), implementing robust session services, and following best practices, you can build agents that:

  • Maintain continuity across conversations

  • Personalize experiences based on user preferences

  • Track complex workflows and processes

  • Learn and adapt from user interactions

  • Scale efficiently in production environments

The key to successful state management is thoughtful design of your data model, appropriate use of scoping, and robust error handling.

Start simple with session-scoped state, then expand to user and application scopes as your agent's sophistication grows.

Remember that state is not just about storage—it's about creating meaningful, continuous relationships between users and your AI agents.

PS:

If you like this article, share it with others ♻️

Would help a lot ❤️

And feel free to follow me for more content like this.

More from this blog

Juan C Olamendy

122 posts

🤖 Talk about AI/ML · AI-preneur 🛠️ Build AI tools 🚀 Share my journey 𓀙 🔗 pixela.io · 🛍️ shoppingbot.ai