8000 docs: Adds a sample agent to illustrate state usage via `callbacks`. · google/adk-python@18fbe3c · GitHub
[go: up one dir, main page]
Skip to content 8000

Commit 18fbe3c

Browse files
Jacksunweicopybara-github
Jacksunwei
authored and
copybara-github
committed
docs: Adds a sample agent to illustrate state usage via callbacks.
PiperOrigin-RevId: 764981675
1 parent 4214c7e commit 18fbe3c

File tree

4 files changed

+251
-0
lines changed

4 files changed

+251
-0
lines changed

‎contributing/samples/session_state_agent/README.md

Lines changed: 66 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,66 @@
1+
# Sample Agent to demo session state persistence.
2+
3+
## Lifecycle of session state
4+
5+
After assigning a state using the context object (e.g.
6+
`tool_context.state['log_query_var'] = 'log_query_var_value'`):
7+
8+
* The state is available for use in a later callback.
9+
* Once the resulting event is processed by the runner and appneded in the
10+
session, the state will be also persisted in the session.
11+
12+
This sample agent is for demonstrating the aforementioned behavior.
13+
14+
## Run the agent
15+
16+
Run below command:
17+
18+
```bash
19+
$ adk run contributing/samples/session_state_agent --replay contributing/samples/session_state_agent/input.json
20+
```
21+
22+
And you should see below output:
23+
24+
```bash
25+
[user]: hello world!
26+
===================== In before_agent_callback ==============================
27+
** Asserting keys are cached in context: ['before_agent_callback_state_key'] pass ✅
28+
** Asserting keys are already persisted in session: [] pass ✅
29+
** Asserting keys are not persisted in session yet: ['before_agent_callback_state_key'] pass ✅
30+
============================================================
31+
===================== In before_model_callback ==============================
32+
** Asserting keys are cached in context: ['before_agent_callback_state_key', 'before_model_callback_state_key'] pass ✅
33+
** Asserting keys are already persisted in session: ['before_agent_callback_state_key'] pass ✅
34+
** Asserting keys are not persisted in session yet: ['before_model_callback_state_key'] pass ✅
35+
============================================================
36+
===================== In after_model_callback ==============================
37+
** Asserting keys are cached in context: ['before_agent_callback_state_key', 'before_model_callback_state_key', 'after_model_callback_state_key'] pass ✅
38+
** Asserting keys are already persisted in session: ['before_agent_callback_state_key'] pass ✅
39+
** Asserting keys are not persisted in session yet: ['before_model_callback_state_key', 'after_model_callback_state_key'] pass ✅
40+
============================================================
41+
[root_agent]: Hello! How can I help you verify something today?
42+
43+
===================== In after_agent_callback ==============================
44+
** 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 ✅
45+
** Asserting keys are already persisted in session: ['before_agent_callback_state_key', 'before_model_callback_state_key', 'after_model_callback_state_key'] pass ✅
46+
** Asserting keys are not persisted in session yet: ['after_agent_callback_state_key'] pass ✅
47+
============================================================
48+
```
49+
50+
## Detailed Explanation
51+
52+
As rule of thumb, to read and write session state, user should assume the
53+
state is available after writing via the context object
54+
(`tool_context`, `callback_context` or `readonly_context`).
55+
56+
### Current Behavior
57+
58+
The current behavior of pesisting states are:
59+
60+
* for `before_agent_callback`: state delta will be persisted after all callbacks are processed.
61+
* for `before_model_callback`: state delta will be persisted with the final LlmResponse,
62+
aka. after `after_model_callback` is processed.
63+
* for `after_model_callback`: state delta will be persisted together with the event of LlmResponse.
64+
* for `after_agent_callback`: state delta will be persisted after all callbacks are processed.
65+
66+
**NOTE**: the current behavior is considered implementation detail and may be changed later. **DO NOT** rely on it.

‎contributing/samples/session_state_agent/__init__.py

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1 @@
1+
from . import agent

‎contributing/samples/session_state_agent/agent.py

Lines changed: 180 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,180 @@
1+
# Copyright 2025 Google LLC
2+
#
3+
# Licensed under the Apache License, Version 2.0 (the "License");
4+
# you may not use this file except in compliance with the License.
5+
# You may obtain a copy of the License at
6+
#
7+
# http://www.apache.org/licenses/LICENSE-2.0
8+
#
9+
# Unless required by applicable law or agreed to in writing, software
10+
# distributed under the License is distributed on an "AS IS" BASIS,
11+
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12+
# See the License for the specific language governing permissions and
13+
# limitations under the License.
14+
15+
"""The agent to demo the session state lifecycle.
16+
17+
This agent illustrate how session state will be cached in context and persisted
18+
in session state.
19+
"""
20+
21+
22+
import logging
23+
from typing import Optional
24+
25+
from google.adk.agents.callback_context import CallbackContext
26+
from google.adk.agents.llm_agent import Agent
27+
from google.adk.models.llm_request import LlmRequest
28+
from google.adk.models.llm_response import LlmResponse
29+
from google.genai import types
30+
31+
logger = logging.getLogger('google_adk.' + __name__)
32+
33+
34+
async def assert_session_values(
35+
ctx: CallbackContext,
36+
title: str,
37+
*,
38+
keys_in_ctx_session: Optional[list[str]] = None,
39+
keys_in_service_session: Optional[list[str]] = None,
40+
keys_not_in_service_session: Optional[list[str]] = None,
41+
):
42+
session_in_ctx = ctx._invocation_context.session
43+
session_in_service = (
44+
await ctx._invocation_context.session_service.get_session(
45+
app_name=session_in_ctx.app_name,
46+
user_id=session_in_ctx.user_id,
47+
session_id=session_in_ctx.id,
48+
)
49+
)
50+
assert session_in_service is not None
51+
52+
print(f'===================== {title} ==============================')
53+
print(
54+
f'** Asserting keys are cached in context: {keys_in_ctx_session}', end=' '
55+
)
56+
for key in keys_in_ctx_session or []:
57+
assert key in session_in_ctx.state
58+
print('\033[92mpass ✅\033[0m')
59+
60+
print(
61+
'** Asserting keys are already persisted in session:'
62+
f' {keys_in_service_session}',
63+
end=' ',
64+
)
65+
for key in keys_in_service_session or []:
66+
assert key in session_in_service.state
67+
print('\033[92mpass ✅\033[0m')
68+
69+
print(
70+
'** Asserting keys are not persisted in session yet:'
71+
f' {keys_not_in_service_session}',
72+
end=' ',
73+
)
74+
for key in keys_not_in_service_session or []:
75+
assert key not in session_in_service.state
76+
print('\033[92mpass ✅\033[0m')
77+
print('============================================================')
78+
79+
80+
async def before_agent_callback(
81+
callback_context: CallbackContext,
82+
) -> Optional[types.Content]:
83+
if 'before_agent_callback_state_key' in callback_context.state:
84+
return types.ModelContent('Sorry, I can only reply once.')
85+
86+
callback_context.state['before_agent_callback_state_key'] = (
87+
'before_agent_callback_state_value'
88+
)
89+
90+
await assert_session_values(
91+
callback_context,
92+
'In before_agent_callback',
93+
keys_in_ctx_session=['before_agent_callback_state_key'],
94+
keys_in_service_session=[],
95+
keys_not_in_service_session=['before_agent_callback_state_key'],
96+
)
97+
98+
99+
async def before_model_callback(
100+
callback_context: CallbackContext, llm_request: LlmRequest
101+
):
102+
callback_context.state['before_model_callback_state_key'] = (
103+
'before_model_callback_state_value'
104+
)
105+
106+
await assert_session_values(
107+
callback_context,
108+
'In before_model_callback',
109+
keys_in_ctx_session=[
110+
'before_agent_callback_state_key',
111+
'before_model_callback_state_key',
112+
],
113+
keys_in_service_session=['before_agent_callback_state_key'],
114+
keys_not_in_service_session=['before_model_callback_state_key'],
115+
)
116+
117+
118+
async def after_model_callback(
119+
callback_context: CallbackContext, llm_response: LlmResponse
120+
):
121+
callback_context.state['after_model_callback_state_key'] = (
122+
'after_model_callback_state_value'
123+
)
124+
125+
await assert_session_values(
126+
callback_context,
127+
'In after_model_callback',
128+
keys_in_ctx_session=[
129+
'before_agent_callback_state_key',
130+
'before_model_callback_state_key',
131+
'after_model_callback_state_key',
132+
],
133+
keys_in_service_session=[
134+
'before_agent_callback_state_key',
135+
],
136+
keys_not_in_service_session=[
137+
'before_model_callback_state_key',
138+
'after_model_callback_state_key',
139+
],
140+
)
141+
142+
143+
async def after_agent_callback(callback_context: CallbackContext):
144+
callback_context.state['after_agent_callback_state_key'] = (
145+
'after_agent_callback_state_value'
146+
)
147+
148+
await assert_session_values(
149+
callback_context,
150+
'In after_agent_callback',
151+
keys_in_ctx_session=[
152+
'before_agent_callback_state_key',
153+
'before_model_callback_state_key',
154+
'after_model_callback_state_key',
155+
'after_agent_callback_state_key',
156+
],
157+
keys_in_service_session=[
158+
'before_agent_callback_state_key',
159+
'before_model_callback_state_key',
160+
'after_model_callback_state_key',
161+
],
162+
keys_not_in_service_session=[
163+
'after_agent_callback_state_key',
164+
],
165+
)
166+
167+
168+
root_agent = Agent(
169+
name='root_agent',
170+
description='a verification agent.',
171+
instruction=(
172+
'Log all users query with `log_query` tool. Must always remind user you'
173+
' cannot answer second query because your setup.'
174+
),
175+
model='gemini-2.0-flash-001',
176+
before_agent_callback=before_agent_callback,
177+
before_model_callback=before_model_callback,
178+
after_model_callback=after_model_callback,
179+
after_agent_callback=after_agent_callback,
180+
)

‎contributing/samples/session_state_agent/input.json

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,4 @@
1+
{
2+
"state": {},
3+
"queries": ["hello world!"]
4+
}

0 commit comments

Comments
 (0)
0