From 18fbe3cbfc9f2af97e4b744ec0a7552331b1d8e3 Mon Sep 17 00:00:00 2001 From: "Wei Sun (Jack)" Date: Thu, 29 May 2025 19:30:12 -0700 Subject: [PATCH] docs: Adds a sample agent to illustrate state usage via `callbacks`. PiperOrigin-RevId: 764981675 --- .../samples/session_state_agent/README.md | 66 +++++++ .../samples/session_state_agent/__init__.py | 1 + .../samples/session_state_agent/agent.py | 180 ++++++++++++++++++ .../samples/session_state_agent/input.json | 4 + 4 files changed, 251 insertions(+) create mode 100644 contributing/samples/session_state_agent/README.md create mode 100644 contributing/samples/session_state_agent/__init__.py create mode 100644 contributing/samples/session_state_agent/agent.py create mode 100644 contributing/samples/session_state_agent/input.json diff --git a/contributing/samples/session_state_agent/README.md b/contributing/samples/session_state_agent/README.md new file mode 100644 index 0000000..bec0536 --- /dev/null +++ b/contributing/samples/session_state_agent/README.md @@ -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. diff --git a/contributing/samples/session_state_agent/__init__.py b/contributing/samples/session_state_agent/__init__.py new file mode 100644 index 0000000..02c597e --- /dev/null +++ b/contributing/samples/session_state_agent/__init__.py @@ -0,0 +1 @@ +from . import agent diff --git a/contributing/samples/session_state_agent/agent.py b/contributing/samples/session_state_agent/agent.py new file mode 100644 index 0000000..a4ed704 --- /dev/null +++ b/contributing/samples/session_state_agent/agent.py @@ -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, +) diff --git a/contributing/samples/session_state_agent/input.json b/contributing/samples/session_state_agent/input.json new file mode 100644 index 0000000..6f76f16 --- /dev/null +++ b/contributing/samples/session_state_agent/input.json @@ -0,0 +1,4 @@ +{ + "state": {}, + "queries": ["hello world!"] +}