From 8eb1bdbc58dc709006988f5b6eec5fda25bd0c89 Mon Sep 17 00:00:00 2001 From: Shangjie Chen Date: Tue, 18 Nov 2025 10:25:47 -0800 Subject: [PATCH] chore: Add demo for rewind Co-authored-by: Shangjie Chen PiperOrigin-RevId: 833871446 --- contributing/samples/rewind_session/main.py | 166 ++++++++++++++++++++ 1 file changed, 166 insertions(+) create mode 100644 contributing/samples/rewind_session/main.py diff --git a/contributing/samples/rewind_session/main.py b/contributing/samples/rewind_session/main.py new file mode 100644 index 00000000..5856b90b --- /dev/null +++ b/contributing/samples/rewind_session/main.py @@ -0,0 +1,166 @@ +#!/usr/bin/env python3 +"""Simple test script for Rewind Session agent.""" + +# 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. + +import asyncio +import logging + +import agent +from google.adk.agents.run_config import RunConfig +from google.adk.cli.utils import logs +from google.adk.events.event import Event +from google.adk.runners import InMemoryRunner +from google.genai import types + +APP_NAME = "rewind_test_app" +USER_ID = "test_user" + +logs.setup_adk_logger(level=logging.ERROR) +logging.getLogger("google_genai.types").setLevel(logging.ERROR) + + +# ANSI color codes for terminal output +COLOR_RED = "\x1b[31m" +COLOR_BLUE = "\x1b[34m" +COLOR_YELLOW = "\x1b[33m" +COLOR_BOLD = "\x1b[1m" +RESET = "\x1b[0m" + + +def highlight(text: str) -> str: + """Adds color highlights to tool responses and agent text.""" + text = str(text) + return ( + text.replace("'red'", f"'{COLOR_RED}red{RESET}'") + .replace('"red"', f'"{COLOR_RED}red{RESET}"') + .replace("'blue'", f"'{COLOR_BLUE}blue{RESET}'") + .replace('"blue"', f'"{COLOR_BLUE}blue{RESET}"') + .replace("'version1'", f"'{COLOR_BOLD}{COLOR_YELLOW}version1{RESET}'") + .replace("'version2'", f"'{COLOR_BOLD}{COLOR_YELLOW}version2{RESET}'") + ) + + +async def call_agent_async( + runner: InMemoryRunner, user_id: str, session_id: str, prompt: str +) -> list[Event]: + """Helper function to call the agent and return events.""" + print(f"\nšŸ‘¤ User: {prompt}") + content = types.Content( + role="user", parts=[types.Part.from_text(text=prompt)] + ) + events = [] + try: + async for event in runner.run_async( + user_id=user_id, + session_id=session_id, + new_message=content, + run_config=RunConfig(), + ): + events.append(event) + if event.content and event.author and event.author != "user": + for part in event.content.parts: + if part.text: + print(f" šŸ¤– Agent: {highlight(part.text)}") + elif part.function_call: + print(f" šŸ› ļø Tool Call: {part.function_call.name}") + elif part.function_response: + print( + " šŸ“¦ Tool Response:" + f" {highlight(part.function_response.response)}" + ) + except Exception as e: + print(f"āŒ Error during agent call: {e}") + raise + return events + + +async def main(): + """Demonstrates session rewind.""" + print("šŸš€ Testing Rewind Session Feature") + print("=" * 50) + + runner = InMemoryRunner( + agent=agent.root_agent, + app_name=APP_NAME, + ) + + # Create a session + session = await runner.session_service.create_session( + app_name=APP_NAME, user_id=USER_ID + ) + print(f"Created session: {session.id}") + + # 1. Initial agent calls to set state and artifact + print("\n\n===== INITIALIZING STATE AND ARTIFACT =====") + await call_agent_async( + runner, USER_ID, session.id, "set state `color` to red" + ) + await call_agent_async( + runner, USER_ID, session.id, "save artifact file1 with content version1" + ) + + # 2. Check current state and artifact + print("\n\n===== STATE BEFORE UPDATE =====") + await call_agent_async( + runner, USER_ID, session.id, "what is the value of state `color`?" + ) + await call_agent_async(runner, USER_ID, session.id, "load artifact file1") + + # 3. Update state and artifact - THIS IS THE POINT WE WILL REWIND BEFORE + print("\n\n===== UPDATING STATE AND ARTIFACT =====") + events_update_state = await call_agent_async( + runner, USER_ID, session.id, "update state key color to blue" + ) + rewind_invocation_id = events_update_state[0].invocation_id + print(f"Will rewind before invocation: {rewind_invocation_id}") + + await call_agent_async( + runner, USER_ID, session.id, "save artifact file1 with content version2" + ) + + # 4. Check state and artifact after update + print("\n\n===== STATE AFTER UPDATE =====") + await call_agent_async( + runner, USER_ID, session.id, "what is the value of state key color?" + ) + await call_agent_async(runner, USER_ID, session.id, "load artifact file1") + + # 5. Perform rewind + print(f"\n\n===== REWINDING SESSION to before {rewind_invocation_id} =====") + await runner.rewind_async( + user_id=USER_ID, + session_id=session.id, + rewind_before_invocation_id=rewind_invocation_id, + ) + print("āœ… Rewind complete.") + + # 6. Check state and artifact after rewind + print("\n\n===== STATE AFTER REWIND =====") + await call_agent_async( + runner, USER_ID, session.id, "what is the value of state `color`?" + ) + await call_agent_async(runner, USER_ID, session.id, "load artifact file1") + + print("\n" + "=" * 50) + print("āœØ Rewind testing complete!") + print( + "šŸ”§ If rewind was successful, color should be 'red' and file1 content" + " should contain 'version1' in the final check." + ) + + +if __name__ == "__main__": + asyncio.run(main())