docs(google-adk): rewrite integration guide to use official ADK MemoryService APIs (#5392)

Co-authored-by: Nishar <nishar@dayos.com>
Co-authored-by: Nishar Miya <miyannishar786@gmail.com>

docs/integrations/google-ai-adk.mdx

@@ -7,285 +7,338 @@ Integrate [**Mem0**](https://github.com/mem0ai/mem0) with [Google ADK (Agent Dev
 
 ## Overview
 
-1. Store and retrieve memories from Mem0 within Google ADK agents
+In this guide, we'll create a Google ADK agent that:
-2. Multi-agent workflows with shared memory across hierarchies
+1. Uses ADK's native `MemoryService` interface to connect Mem0
-3. Retrieve relevant memories from past conversations
+2. Automatically injects relevant memories using ADK's built-in `load_memory` tool
-4. Personalized responses based on user history
+3. Persists session history to Mem0 after each turn via an after-agent callback
+4. Shares memory seamlessly across multi-agent hierarchies
 
-## Prerequisites
+## Setup and Configuration
 
-Before setting up Mem0 with Google ADK, ensure you have:
+Install the necessary libraries:
 
-1. Installed the required packages:
 ```bash
 pip install google-adk mem0ai python-dotenv
 ```
 
-2. Valid API keys:
+Set up your API keys:
    - <a href="https://app.mem0.ai/dashboard/api-keys?utm_source=oss&utm_medium=integration-google-ai-adk" rel="nofollow">Mem0 API Key</a>
    - Google AI Studio API Key
 
-## Basic Integration Example
+<Note>Remember to get your API key from <a href="https://app.mem0.ai" rel="nofollow">Mem0 Platform</a> and set up a [Google AI Studio API Key](https://aistudio.google.com/apikey).</Note>
-
-The following example demonstrates how to create a Google ADK agent with Mem0 memory integration:
 
 ```python
 import os
-import asyncio
-from google.adk.agents import Agent
-from google.adk.runners import Runner
-from google.adk.sessions import InMemorySessionService
-from google.genai import types
-from mem0 import MemoryClient
 from dotenv import load_dotenv
 
 load_dotenv()
 
-# Set up environment variables
 # os.environ["GOOGLE_API_KEY"] = "your-google-api-key"
 # os.environ["MEM0_API_KEY"] = "your-mem0-api-key"
+```
 
-# Initialize Mem0 client
+## Implement Mem0MemoryService
-mem0 = MemoryClient()
 
-# Define memory function tools
+Create a custom `MemoryService` by implementing ADK's `BaseMemoryService`. Save the following as **`mem0_memory_service.py`**:
-def search_memory(query: str, user_id: str) -> dict:
-    """Search through past conversations and memories"""
-    # For Platform API, user_id goes in filters
-    filters = {"user_id": user_id}
-    memories = mem0.search(query, filters=filters)
-    if memories.get('results', []):
-        memory_list = memories['results']
-        memory_context = "\n".join([f"- {mem['memory']}" for mem in memory_list])
-        return {"status": "success", "memories": memory_context}
-    return {"status": "no_memories", "message": "No relevant memories found"}
 
-def save_memory(content: str, user_id: str) -> dict:
+```python
-    """Save important information to memory"""
+import asyncio
+import os
+from typing import Optional
+from typing_extensions import override
+
+from google.adk.memory.base_memory_service import BaseMemoryService, SearchMemoryResponse
+from google.adk.memory.memory_entry import MemoryEntry
+from google.adk.sessions import Session
+from google.genai.types import Content, Part
+from mem0 import MemoryClient
+
+
+class Mem0MemoryService(BaseMemoryService):
+    """MemoryService implementation backed by the Mem0 Platform."""
+
+    def __init__(self, api_key: Optional[str] = None):
+        super().__init__()
+        api_key = api_key or os.environ.get("MEM0_API_KEY")
+        self._client: Optional[MemoryClient] = MemoryClient(api_key=api_key) if api_key else None
+
+    @override
+    async def search_memory(
+        self, *, app_name: str, user_id: str, query: str
+    ) -> SearchMemoryResponse:
+        """Search for memories relevant to the current user and query."""
+        if not self._client:
+            return SearchMemoryResponse(memories=[])
+
+        try:
+            results = await asyncio.to_thread(
+                self._client.search,
+                query,
+                filters={"AND": [{"user_id": user_id}, {"app_id": app_name}]},
+                top_k=5,
+            )
+
+            entries = []
+            for mem in results.get("results", []):
+                text = mem.get("memory", "")
+                if not text:
+                    continue
+
+                raw_ts = mem.get("created_at") or mem.get("updated_at")
+                entries.append(
+                    MemoryEntry(
+                        content=Content(parts=[Part(text=text)]),
+                        author=mem.get("metadata", {}).get("author", "user"),
+                        timestamp=str(raw_ts) if raw_ts else None,
+                    )
+                )
+
+            return SearchMemoryResponse(memories=entries)
+
+        except Exception as e:
+            print(f"[Mem0MemoryService] search_memory error: {e}")
+            return SearchMemoryResponse(memories=[])
+
+    @override
+    async def add_session_to_memory(self, session: Session) -> None:
+        """Persist a completed ADK session into Mem0."""
+        if not self._client:
+            return
+
+        user_id = session.user_id
+        if not user_id:
+            return
+
+        app_name = getattr(session, "app_name", None)
+
+        try:
+            messages = []
+            for event in session.events:
+                if not (event.content and event.content.parts):
+                    continue
+                role = getattr(event.content, "role", None) or "user"
+                if role == "model":
+                    role = "assistant"
+                elif role not in ("user", "assistant"):
+                    continue
+                text_parts = [
+                    p.text for p in event.content.parts if hasattr(p, "text") and p.text
+                ]
+                if text_parts:
+                    messages.append({"role": role, "content": " ".join(text_parts)})
+
+            if messages:
+                metadata = {"app_id": app_name} if app_name else {}
+                await asyncio.to_thread(
+                    self._client.add, messages, user_id=user_id, metadata=metadata
+                )
+
+        except Exception as e:
+            print(f"[Mem0MemoryService] add_session_to_memory error: {e}")
+```
+
+## Add Auto-Save Callback
+
+This after-agent callback fires at the end of every turn and saves the session to Mem0. Save as **`memory_callbacks.py`**:
+
+```python
+async def save_session_to_memory(callback_context) -> None:
+    """Persist the completed session to Mem0 after each agent turn."""
     try:
-        result = mem0.add([{"role": "user", "content": content}], user_id=user_id)
+        await callback_context.add_session_to_memory()
-        return {"status": "success", "message": "Information saved to memory", "result": result}
+    except ValueError:
+        pass
     except Exception as e:
-        return {"status": "error", "message": f"Failed to save memory: {str(e)}"}
+        print(f"[save_session_to_memory] error: {e}")
+```
 
-# Create agent with memory capabilities
+## Basic Integration Example
-personal_assistant = Agent(
+
+The following example demonstrates creating an ADK agent with automatic Mem0 memory:
+
+```python
+import asyncio
+from google.adk.agents import LlmAgent
+from google.adk.runners import Runner
+from google.adk.sessions import InMemorySessionService
+from google.adk.tools import load_memory
+from google.genai.types import Content, Part
+
+from mem0_memory_service import Mem0MemoryService
+from memory_callbacks import save_session_to_memory
+
+memory_service = Mem0MemoryService()
+session_service = InMemorySessionService()
+
+agent = LlmAgent(
     name="personal_assistant",
     model="gemini-2.0-flash",
-    instruction="""You are a helpful personal assistant with memory capabilities.
+    instruction="""You are a helpful personal assistant.
-    Use the search_memory function to recall past conversations and user preferences.
+    Relevant memories from past conversations are provided to you automatically.
-    Use the save_memory function to store important information about the user.
+    Use them to personalize your responses.""",
-    Always personalize your responses based on available memory.""",
     description="A personal assistant that remembers user preferences and past interactions",
-    tools=[search_memory, save_memory]
+    tools=[load_memory],
+    after_agent_callback=save_session_to_memory,
 )
 
-async def chat_with_agent(user_input: str, user_id: str) -> str:
+runner = Runner(
-    """
+    agent=agent,
-    Handle user input with automatic memory integration.
+    session_service=session_service,
+    memory_service=memory_service,
+    app_name="memory_assistant",
+)
 
-    Args:
-        user_input: The user's message
-        user_id: Unique identifier for the user
 
-    Returns:
+async def chat(user_input: str, user_id: str) -> str:
-        The agent's response
-    """
-    # Set up session and runner
-    session_service = InMemorySessionService()
     session = await session_service.create_session(
         app_name="memory_assistant",
         user_id=user_id,
-        session_id=f"session_{user_id}"
     )
-    runner = Runner(agent=personal_assistant, app_name="memory_assistant", session_service=session_service)
+    content = Content(role="user", parts=[Part(text=user_input)])
-
+    async for event in runner.run_async(user_id=user_id, session_id=session.id, new_message=content):
-    # Create content and run agent
+        if event.is_final_response() and event.content and event.content.parts:
-    content = types.Content(role='user', parts=[types.Part(text=user_input)])
+            return event.content.parts[0].text
-    events = runner.run(user_id=user_id, session_id=session.id, new_message=content)
-
-    # Extract final response
-    for event in events:
-        if event.is_final_response():
-            response = event.content.parts[0].text
-
-            return response
-
     return "No response generated"
 
-# Example usage
+
 if __name__ == "__main__":
-    response = asyncio.run(chat_with_agent(
+    print(asyncio.run(chat(
         "I love Italian food and I'm planning a trip to Rome next month",
-        user_id="alice"
+        user_id="alice",
-    ))
+    )))
-    print(response)
+
+    print(asyncio.run(chat(
+        "Any food recommendations for my trip?",
+        user_id="alice",
+    )))
 ```
 
 ## Multi-Agent Hierarchy with Shared Memory
 
-Create specialized agents in a hierarchy that share memory:
+Because `memory_service` is passed to the `Runner`, every agent in the hierarchy shares the same memory automatically. Only the root coordinator needs the auto-save callback — ADK fires it once when the full turn completes:
 
 ```python
+import asyncio
+from google.adk.agents import LlmAgent
+from google.adk.runners import Runner
+from google.adk.sessions import InMemorySessionService
 from google.adk.tools.agent_tool import AgentTool
+from google.adk.tools import load_memory
+from google.genai.types import Content, Part
 
-# Travel specialist agent
+from mem0_memory_service import Mem0MemoryService
-travel_agent = Agent(
+from memory_callbacks import save_session_to_memory
+
+memory_service = Mem0MemoryService()
+session_service = InMemorySessionService()
+
+travel_agent = LlmAgent(
     name="travel_specialist",
     model="gemini-2.0-flash",
-    instruction="""You are a travel planning specialist. Use search_memory to
+    instruction="""You are a travel planning specialist.
-    understand the user's travel preferences and history before making recommendations.
+    Relevant memories about the user's travel preferences are provided automatically.
-    After providing advice, use save_memory to save travel-related information.""",
+    Use them to make personalized recommendations.""",
     description="Specialist in travel planning and recommendations",
-    tools=[search_memory, save_memory]
+    tools=[load_memory],
 )
 
-# Health advisor agent
+health_agent = LlmAgent(
-health_agent = Agent(
     name="health_advisor",
     model="gemini-2.0-flash",
-    instruction="""You are a health and wellness advisor. Use search_memory to
+    instruction="""You are a health and wellness advisor.
-    understand the user's health goals and dietary preferences.
+    Relevant memories about the user's health goals are provided automatically.
-    After providing advice, use save_memory to save health-related information.""",
+    Use them to give personalized advice.""",
     description="Specialist in health and wellness advice",
-    tools=[search_memory, save_memory]
+    tools=[load_memory],
 )
 
-# Coordinator agent that delegates to specialists
+coordinator = LlmAgent(
-coordinator_agent = Agent(
     name="coordinator",
     model="gemini-2.0-flash",
     instruction="""You are a coordinator that delegates requests to specialist agents.
-    For travel-related questions (trips, hotels, flights, destinations), delegate to the travel specialist.
+    For travel-related questions, delegate to the travel specialist.
-    For health-related questions (fitness, diet, wellness, exercise), delegate to the health advisor.
+    For health-related questions, delegate to the health advisor.
-    Use search_memory to understand the user before delegation.""",
+    Relevant memories about the user are provided automatically.""",
     description="Coordinates requests between specialist agents",
     tools=[
+        load_memory,
         AgentTool(agent=travel_agent, skip_summarization=False),
-        AgentTool(agent=health_agent, skip_summarization=False)
+        AgentTool(agent=health_agent, skip_summarization=False),
-    ]
+    ],
+    after_agent_callback=save_session_to_memory,
 )
 
-def chat_with_specialists(user_input: str, user_id: str) -> str:
+runner = Runner(
-    """
+    agent=coordinator,
-    Handle user input with specialist agent delegation and memory.
+    session_service=session_service,
+    memory_service=memory_service,
+    app_name="specialist_system",
+)
 
-    Args:
-        user_input: The user's message
-        user_id: Unique identifier for the user
 
-    Returns:
+async def chat_with_specialists(user_input: str, user_id: str) -> str:
-        The specialist agent's response
+    session = await session_service.create_session(
-    """
-    session_service = InMemorySessionService()
-    session = session_service.create_session(
         app_name="specialist_system",
         user_id=user_id,
-        session_id=f"session_{user_id}"
     )
-    runner = Runner(agent=coordinator_agent, app_name="specialist_system", session_service=session_service)
+    content = Content(role="user", parts=[Part(text=user_input)])
-
+    async for event in runner.run_async(user_id=user_id, session_id=session.id, new_message=content):
-    content = types.Content(role='user', parts=[types.Part(text=user_input)])
+        if event.is_final_response() and event.content and event.content.parts:
-    events = runner.run(user_id=user_id, session_id=session.id, new_message=content)
+            return event.content.parts[0].text
-
-    for event in events:
-        if event.is_final_response():
-            response = event.content.parts[0].text
-
-            # Store the conversation in shared memory
-            conversation = [
-                {"role": "user", "content": user_input},
-                {"role": "assistant", "content": response}
-            ]
-            mem0.add(conversation, user_id=user_id)
-
-            return response
-
     return "No response generated"
 
-# Example usage
-response = chat_with_specialists("Plan a healthy meal for my Italy trip", user_id="alice")
-print(response)
-```
-
-
-
-## Quick Start Chat Interface
-
-Simple interactive chat with memory and Google ADK:
-
-```python
-def interactive_chat():
-    """Interactive chat interface with memory and ADK"""
-    user_id = input("Enter your user ID: ") or "demo_user"
-    print(f"Chat started for user: {user_id}")
-    print("Type 'quit' to exit")
-    print("=" * 50)
-
-    while True:
-        user_input = input("\nYou: ")
-
-        if user_input.lower() == 'quit':
-            print("Goodbye! Your conversation has been saved to memory.")
-            break
-        else:
-            response = chat_with_specialists(user_input, user_id)
-            print(f"Assistant: {response}")
 
 if __name__ == "__main__":
-    interactive_chat()
+    response = asyncio.run(chat_with_specialists("Plan a healthy meal for my Italy trip", user_id="alice"))
+    print(response)
 ```
 
 ## Key Features
 
-### 1. Memory-Enhanced Function Tools
+1. **Automatic Memory Injection**: ADK's built-in `load_memory` tool searches Mem0 at the start of each turn and injects relevant memories directly into the agent context — no prompt instructions needed.
-- **Function Tools**: Standard Python functions that can search and save memories
+2. **Automatic Session Saving**: The `save_session_to_memory` callback persists every completed turn to Mem0 without any manual calls.
-- **Tool Context**: Access to session state and memory through function parameters
+3. **Native ADK Integration**: `Mem0MemoryService` implements ADK's `BaseMemoryService` and integrates via the `Runner` — works natively across the entire agent hierarchy.
-- **Structured Returns**: Dictionary-based returns with status indicators for better LLM understanding
+4. **User Scoping**: `user_id` is passed automatically from the ADK session context, ensuring memories are always scoped to the correct user.
-
+5. **Multi-Agent Support**: A single `Mem0MemoryService` instance shared through the `Runner` gives all agents — coordinators and specialists — access to the same user memory.
-### 2. Multi-Agent Memory Sharing
-- **Agent-as-a-Tool**: Specialists can be called as tools while maintaining shared memory
-- **Hierarchical Delegation**: Coordinator agents route to specialists based on context
-- **Memory Categories**: Store interactions with metadata for better organization
-
-### 3. Flexible Memory Operations
-- **Search Capabilities**: Retrieve relevant memories through conversation history
-- **User Segmentation**: Organize memories by user ID
-- **Memory Management**: Built-in tools for saving and retrieving information
 
 ## Configuration Options
 
-Customize memory behavior and agent setup:
+### Using Vertex AI
+
+To use Google Cloud Vertex AI instead of AI Studio, set the following environment variables before creating agents:
 
 ```python
-# Configure memory search with filters
+import os
-# For Platform API, all filters including user_id go in filters object
-memories = mem0.search(
-    query="travel preferences",
-    filters={
-        "AND": [
-            {"user_id": "alice"},
-            {"categories": {"contains": "travel"}}
-        ]
-    },
-    top_k=5
-)
-
-# Configure agent with custom model settings
-agent = Agent(
-    name="custom_agent",
-    model="gemini-2.0-flash",  # or use LiteLLM for other models
-    instruction="Custom agent behavior",
-    tools=[memory_tools],
-    # Additional ADK configurations
-)
-
-# Use Google Cloud Vertex AI instead of AI Studio
 os.environ["GOOGLE_GENAI_USE_VERTEXAI"] = "True"
 os.environ["GOOGLE_CLOUD_PROJECT"] = "your-project-id"
 os.environ["GOOGLE_CLOUD_LOCATION"] = "us-central1"
 ```
 
+### Advanced Memory Filtering
+
+You can customize how memories are searched by modifying `Mem0MemoryService.search_memory`. For example, to filter by category:
+
+```python
+results = await asyncio.to_thread(
+    self._client.search,
+    query,
+    filters={
+        "AND": [
+            {"user_id": user_id},
+            {"app_id": app_name},
+            {"categories": {"contains": "travel"}}
+        ]
+    },
+    top_k=10,
+)
+```
+
+<Note>`InMemorySessionService` stores sessions in memory and is intended for prototyping. For production, use a persistent session service and clean up sessions when they are no longer needed.</Note>
+
+## Conclusion
+
+By implementing `Mem0MemoryService` as an ADK `BaseMemoryService`, you get persistent, user-scoped memory across single agents and complex multi-agent hierarchies with minimal code. Memory injection and session saving happen automatically, keeping your agent prompts clean and your token usage efficient.
+
 <CardGroup cols={2}>
   <Card title="Healthcare Agent Cookbook" icon="heart-pulse" href="/cookbooks/integrations/healthcare-google-adk">
     Build HIPAA-compliant healthcare agents with Google ADK
@@ -294,4 +347,3 @@ os.environ["GOOGLE_CLOUD_LOCATION"] = "us-central1"
     Compare with OpenAI's agent framework
   </Card>
 </CardGroup>
-