NeoAnd/adk-python
1
0
Fork 0
mirror of https://github.com/EvolutionAPI/adk-python.git synced 2025-07-13 07:04:51 -06:00
Code Issues Actions Packages Projects Releases Wiki Activity

docs: Adds a sample agent to illustrate state usage via callbacks.

Browse Source 
PiperOrigin-RevId: 764981675
This commit is contained in:
Wei Sun (Jack) 2025-05-29 19:30:12 -07:00 committed by Copybara-Service
parent 4214c7eddd
commit 18fbe3cbfc
4 changed files with 251 additions and 0 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

66
contributing/samples/session_state_agent/README.md Normal file
View File

@ -0,0 +1,66 @@
# Sample Agent to demo session state persistence.
## Lifecycle of session state
After assigning a state using the context object (e.g.
`tool_context.state['log_query_var'] = 'log_query_var_value'`):
* The state is available for use in a later callback.
* Once the resulting event is processed by the runner and appneded in the
session, the state will be also persisted in the session.
This sample agent is for demonstrating the aforementioned behavior.
## Run the agent
Run below command:
```bash
$ adk run contributing/samples/session_state_agent --replay contributing/samples/session_state_agent/input.json
```
And you should see below output:
```bash
[user]: hello world!
===================== In before_agent_callback ==============================
** Asserting keys are cached in context: ['before_agent_callback_state_key'] pass ✅
** Asserting keys are already persisted in session: [] pass ✅
** Asserting keys are not persisted in session yet: ['before_agent_callback_state_key'] pass ✅
============================================================
===================== In before_model_callback ==============================
** Asserting keys are cached in context: ['before_agent_callback_state_key', 'before_model_callback_state_key'] pass ✅
** Asserting keys are already persisted in session: ['before_agent_callback_state_key'] pass ✅
** Asserting keys are not persisted in session yet: ['before_model_callback_state_key'] pass ✅
============================================================
===================== In after_model_callback ==============================
** Asserting keys are cached in context: ['before_agent_callback_state_key', 'before_model_callback_state_key', 'after_model_callback_state_key'] pass ✅
** Asserting keys are already persisted in session: ['before_agent_callback_state_key'] pass ✅
** Asserting keys are not persisted in session yet: ['before_model_callback_state_key', 'after_model_callback_state_key'] pass ✅
============================================================
[root_agent]: Hello! How can I help you verify something today?
===================== In after_agent_callback ==============================
** Asserting keys are cached in context: ['before_agent_callback_state_key', 'before_model_callback_state_key', 'after_model_callback_state_key', 'after_agent_callback_state_key'] pass ✅
** Asserting keys are already persisted in session: ['before_agent_callback_state_key', 'before_model_callback_state_key', 'after_model_callback_state_key'] pass ✅
** Asserting keys are not persisted in session yet: ['after_agent_callback_state_key'] pass ✅
============================================================
```
## Detailed Explanation
As rule of thumb, to read and write session state, user should assume the
state is available after writing via the context object
(`tool_context`, `callback_context` or `readonly_context`).
### Current Behavior
The current behavior of pesisting states are:
* for `before_agent_callback`: state delta will be persisted after all callbacks are processed.
* for `before_model_callback`: state delta will be persisted with the final LlmResponse,
aka. after `after_model_callback` is processed.
* for `after_model_callback`: state delta will be persisted together with the event of LlmResponse.
* for `after_agent_callback`: state delta will be persisted after all callbacks are processed.
**NOTE**: the current behavior is considered implementation detail and may be changed later. **DO NOT** rely on it.

1
contributing/samples/session_state_agent/__init__.py Normal file
View File

@ -0,0 +1 @@
from . import agent

180
contributing/samples/session_state_agent/agent.py Normal file
View File

@ -0,0 +1,180 @@
# Copyright 2025 Google LLC
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
"""The agent to demo the session state lifecycle.
This agent illustrate how session state will be cached in context and persisted
in session state.
"""
import logging
from typing import Optional
from google.adk.agents.callback_context import CallbackContext
from google.adk.agents.llm_agent import Agent
from google.adk.models.llm_request import LlmRequest
from google.adk.models.llm_response import LlmResponse
from google.genai import types
logger = logging.getLogger('google_adk.' + __name__)
async def assert_session_values(
ctx: CallbackContext,
title: str,
*,
keys_in_ctx_session: Optional[list[str]] = None,
keys_in_service_session: Optional[list[str]] = None,
keys_not_in_service_session: Optional[list[str]] = None,
):
session_in_ctx = ctx._invocation_context.session
session_in_service = (
await ctx._invocation_context.session_service.get_session(
app_name=session_in_ctx.app_name,
user_id=session_in_ctx.user_id,
session_id=session_in_ctx.id,
)
)
assert session_in_service is not None
print(f'===================== {title} ==============================')
print(
f'** Asserting keys are cached in context: {keys_in_ctx_session}', end=' '
)
for key in keys_in_ctx_session or []:
assert key in session_in_ctx.state
print('\033[92mpass ✅\033[0m')
print(
'** Asserting keys are already persisted in session:'
f' {keys_in_service_session}',
end=' ',
)
for key in keys_in_service_session or []:
assert key in session_in_service.state
print('\033[92mpass ✅\033[0m')
print(
'** Asserting keys are not persisted in session yet:'
f' {keys_not_in_service_session}',
end=' ',
)
for key in keys_not_in_service_session or []:
assert key not in session_in_service.state
print('\033[92mpass ✅\033[0m')
print('============================================================')
async def before_agent_callback(
callback_context: CallbackContext,
) -> Optional[types.Content]:
if 'before_agent_callback_state_key' in callback_context.state:
return types.ModelContent('Sorry, I can only reply once.')
callback_context.state['before_agent_callback_state_key'] = (
'before_agent_callback_state_value'
)
await assert_session_values(
callback_context,
'In before_agent_callback',
keys_in_ctx_session=['before_agent_callback_state_key'],
keys_in_service_session=[],
keys_not_in_service_session=['before_agent_callback_state_key'],
)
async def before_model_callback(
callback_context: CallbackContext, llm_request: LlmRequest
):
callback_context.state['before_model_callback_state_key'] = (
'before_model_callback_state_value'
)
await assert_session_values(
callback_context,
'In before_model_callback',
keys_in_ctx_session=[
'before_agent_callback_state_key',
'before_model_callback_state_key',
],
keys_in_service_session=['before_agent_callback_state_key'],
keys_not_in_service_session=['before_model_callback_state_key'],
)
async def after_model_callback(
callback_context: CallbackContext, llm_response: LlmResponse
):
callback_context.state['after_model_callback_state_key'] = (
'after_model_callback_state_value'
)
await assert_session_values(
callback_context,
'In after_model_callback',
keys_in_ctx_session=[
'before_agent_callback_state_key',
'before_model_callback_state_key',
'after_model_callback_state_key',
],
keys_in_service_session=[
'before_agent_callback_state_key',
],
keys_not_in_service_session=[
'before_model_callback_state_key',
'after_model_callback_state_key',
],
)
async def after_agent_callback(callback_context: CallbackContext):
callback_context.state['after_agent_callback_state_key'] = (
'after_agent_callback_state_value'
)
await assert_session_values(
callback_context,
'In after_agent_callback',
keys_in_ctx_session=[
'before_agent_callback_state_key',
'before_model_callback_state_key',
'after_model_callback_state_key',
'after_agent_callback_state_key',
],
keys_in_service_session=[
'before_agent_callback_state_key',
'before_model_callback_state_key',
'after_model_callback_state_key',
],
keys_not_in_service_session=[
'after_agent_callback_state_key',
],
)
root_agent = Agent(
name='root_agent',
description='a verification agent.',
instruction=(
'Log all users query with `log_query` tool. Must always remind user you'
' cannot answer second query because your setup.'
),
model='gemini-2.0-flash-001',
before_agent_callback=before_agent_callback,
before_model_callback=before_model_callback,
after_model_callback=after_model_callback,
after_agent_callback=after_agent_callback,
)

4
contributing/samples/session_state_agent/input.json Normal file
View File

@ -0,0 +1,4 @@
{
"state": {},
"queries": ["hello world!"]
}