feat: History Management Sample

Merge #891 This creates a sample relating to discussion #826 - how to manage context windows. COPYBARA_INTEGRATE_REVIEW=#891 from calvingiles:history-management-sample 2827817 PiperOrigin-RevId: 784920438

Author: calvingiles

contributing/samples/history_management/__init__.py

@@ -0,0 +1,15 @@
+# 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.
+
+from . import agent

contributing/samples/history_management/agent.py

@@ -0,0 +1,110 @@
+# 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 random
+
+from google.adk import Agent
+from google.adk.agents.callback_context import CallbackContext
+from google.adk.models import LlmRequest
+from google.adk.tools.tool_context import ToolContext
+
+
+def roll_die(sides: int, tool_context: ToolContext) -> int:
+ """Roll a die and return the rolled result.
+
+ Args:
+ sides: The integer number of sides the die has.
+
+ Returns:
+ An integer of the result of rolling the die.
+ """
+ result = random.randint(1, sides)
+ if not 'rolls' in tool_context.state:
+ tool_context.state['rolls'] = []
+
+ tool_context.state['rolls'] = tool_context.state['rolls'] + [result]
+ return result
+
+
+async def check_prime(nums: list[int]) -> str:
+ """Check if a given list of numbers are prime.
+
+ Args:
+ nums: The list of numbers to check.
+
+ Returns:
+ A str indicating which number is prime.
+ """
+ primes = set()
+ for number in nums:
+ number = int(number)
+ if number <= 1:
+ continue
+ is_prime = True
+ for i in range(2, int(number**0.5) + 1):
+ if number % i == 0:
+ is_prime = False
+ break
+ if is_prime:
+ primes.add(number)
+ return (
+ 'No prime numbers found.'
+ if not primes
+ else f"{', '.join(str(num) for num in primes)} are prime numbers."
+ )
+
+
+def create_slice_history_callback(n_recent_turns):
+ async def before_model_callback(callback_context: CallbackContext, llm_request: LlmRequest):
+ if n_recent_turns < 1:
+ return
+
+ user_indexes = [i for i, content in enumerate(llm_request.contents) if content.role == "user"]
+
+ if n_recent_turns > len(user_indexes):
+ return
+
+ suffix_idx = user_indexes[-n_recent_turns]
+ llm_request.contents = llm_request.contents[suffix_idx:]
+
+ return before_model_callback
+
+
+root_agent = Agent(
+ model='gemini-2.0-flash',
+ name='short_history_agent',
+ description=(
+ 'an agent that maintains only the last turn in its context window.'
+ ' numbers.'
+ ),
+ instruction="""
+ You roll dice and answer questions about the outcome of the dice rolls.
+ You can roll dice of different sizes.
+ You can use multiple tools in parallel by calling functions in parallel(in one request and in one round).
+ It is ok to discuss previous dice roles, and comment on the dice rolls.
+ When you are asked to roll a die, you must call the roll_die tool with the number of sides. Be sure to pass in an integer. Do not pass in a string.
+ You should never roll a die on your own.
+ When checking prime numbers, call the check_prime tool with a list of integers. Be sure to pass in a list of integers. You should never pass in a string.
+ You should not check prime numbers before calling the tool.
+ When you are asked to roll a die and check prime numbers, you should always make the following two function calls:
+ 1. You should first call the roll_die tool to get a roll. Wait for the function response before calling the check_prime tool.
+ 2. After you get the function response from roll_die tool, you should call the check_prime tool with the roll_die result.
+ 2.1 If user asks you to check primes based on previous rolls, make sure you include the previous rolls in the list.
+ 3. When you respond, you must include the roll_die result from step 1.
+ You should always perform the previous 3 steps when asking for a roll and checking prime numbers.
+ You should not rely on the previous history on prime results.
+ """,
+ tools=[roll_die, check_prime],
+ before_model_callback=create_slice_history_callback(n_recent_turns=2),
+)

contributing/samples/history_management/main.py

@@ -0,0 +1,80 @@
+# 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 time
+import warnings
+
+import agent
+from dotenv import load_dotenv
+from google.adk import Runner
+from google.adk.artifacts import InMemoryArtifactService
+from google.adk.cli.utils import logs
+from google.adk.sessions import InMemorySessionService
+from google.adk.sessions import Session
+from google.genai import types
+
+load_dotenv(override=True)
+warnings.filterwarnings('ignore', category=UserWarning)
+logs.log_to_tmp_folder()
+
+
+async def main():
+ app_name = 'my_app'
+ user_id_1 = 'user1'
+ session_service = InMemorySessionService()
+ artifact_service = InMemoryArtifactService()
+ runner = Runner(
+ app_name=app_name,
+ agent=agent.root_agent,
+ artifact_service=artifact_service,
+ session_service=session_service,
+ )
+ session_11 = await session_service.create_session(
+ app_name=app_name, user_id=user_id_1
+ )
+
+ async def run_prompt(session: Session, new_message: str):
+ content = types.Content(
+ role='user', parts=[types.Part.from_text(text=new_message)]
+ )
+ print('** User says:', content.model_dump(exclude_none=True))
+ async for event in runner.run_async(
+ user_id=user_id_1,
+ session_id=session.id,
+ new_message=content,
+ ):
+ if event.content.parts and event.content.parts[0].text:
+ print(f'** {event.author}: {event.content.parts[0].text}')
+
+ start_time = time.time()
+ print('Start time:', start_time)
+ print('------------------------------------')
+ await run_prompt(session_11, 'Hi')
+ await run_prompt(session_11, 'Roll a die with 100 sides')
+ await run_prompt(session_11, 'Roll a die again with 100 sides.')
+ await run_prompt(session_11, 'What numbers did I got?')
+ print(
+ await artifact_service.list_artifact_keys(
+ app_name=app_name, user_id=user_id_1, session_id=session_11.id
+ )
+ )
+ end_time = time.time()
+ print('------------------------------------')
+ print('End time:', end_time)
+ print('Total time:', end_time - start_time)
+
+
+if __name__ == '__main__':
+ asyncio.run(main())