-
Notifications
You must be signed in to change notification settings - Fork 2.9k
Description
Environment:
- ADK Version: 1.25.0 (also reproduced in 1.24.1)
- Service: Agent Engine (Reasoning Engine)
- Memory Services:
VertexAiMemoryBankServiceandVertexAiSessionService - Environment: Gemini Enterprise
Description
I have deployed an agent using the Google ADK on Agent Engine. The agent is configured with VertexAiMemoryBankService and VertexAiSessionService.
The issue: When testing the agent directly through the Agent Engine Playground, the memory bank works as expected, and memories are correctly persisted. However, when the same deployed agent is accessed via Gemini Enterprise, the after_agent_callback (intended to save the session to memory) fails to generate/persist memories.
It appears that the invocation_context or the memory service state might not be correctly handled or triggered during the hand-off between Gemini Enterprise and the Reasoning Engine.
Code Snippets
1. Deployment Configuration: I am using factories to inject the services into AdkApp:
Python
def setup_memory_bank(env_vars: dict):
project = env_vars.get("GOOGLE_CLOUD_PROJECT")
location = env_vars.get("GOOGLE_CLOUD_LOCATION")
engine_id = env_vars.get('REASONING_ENGINE_ID')
return lambda: VertexAiMemoryBankService(
project=project,
location=location,
agent_engine_id=engine_id
)
def setup_vertexai_sessions(env_vars: dict):
project = env_vars.get("GOOGLE_CLOUD_PROJECT")
location = env_vars.get("GOOGLE_CLOUD_LOCATION")
engine_id = env_vars.get('REASONING_ENGINE_ID')
return lambda: VertexAiSessionService(
project=project,
location=location,
agent_engine_id=engine_id
)
# Deployment logic
adk_app = AdkApp(
agent=root_agent,
enable_tracing=True,
memory_service_builder=setup_memory_bank(env_vars),
session_service_builder=setup_vertexai_sessions(env_vars)
)
2. Memory Callback: This callback is registered as an after_agent_callback:
Python
async def add_session_to_memory(
callback_context: CallbackContext
) -> Optional[types.Content]:
"""Automatically save completed sessions to memory bank"""
if hasattr(callback_context, "_invocation_context"):
invocation_context = callback_context._invocation_context
if invocation_context.memory_service:
# This part seems to be bypassed or fails silently in Gemini Enterprise
await callback_context.add_events_to_memory(
events=invocation_context.session.events
)
3. Agent Initialization:
Python
root_agent = Agent(
model=ROOT_AGENT_MODEL,
name="my_agent",
tools=[PreloadMemoryTool(), question_analyser],
before_agent_callback=setup_before_agent_call,
after_agent_callback=add_session_to_memory,
generate_content_config=types.GenerateContentConfig(temperature=0.01),
)
Steps to Reproduce
- Deploy the agent to Agent Engine using ADK v1.25.0.
- Interact with the agent via Agent Engine Playground -> Result: Memory is saved.
- Register the agent and interact with it via Gemini Enterprise.
- Check the Memory Bank -> Result: No new memories are created/updated.
Expected Behavior
The after_agent_callback should consistently trigger add_events_to_memory regardless of whether the agent is called from the Playground or an external integration like Gemini Enterprise.
Additional Context
- Is there any specific requirement for Gemini Enterprise to pass down the session/invocation context so the ADK can identify the memory service?
- Tracing is enabled, but no explicit errors are shown regarding the memory write operation when coming from Gemini Enterprise.