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

PiperOrigin-RevId: 764981675

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,
+)