Move callback samples to dedicated directory

Sample chained callback logs can we seen running asyncio_run.py

PiperOrigin-RevId: 757793406
This commit is contained in:
Selcuk Gun 2025-05-12 09:23:50 -07:00 committed by Copybara-Service
parent cdb4cac237
commit 812485fdfa
4 changed files with 356 additions and 87 deletions

View File

@ -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

View File

@ -0,0 +1,196 @@
# 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.planners import BuiltInPlanner
from google.adk.planners import PlanReActPlanner
from google.adk.tools.tool_context import ToolContext
from google.genai import types
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."
)
async def before_agent_callback(callback_context):
print('@before_agent_callback')
return None
async def after_agent_callback(callback_context):
print('@after_agent_callback')
return None
async def before_model_callback(callback_context, llm_request):
print('@before_model_callback')
return None
async def after_model_callback(callback_context, llm_response):
print('@after_model_callback')
return None
def after_agent_cb1(callback_context):
print('@after_agent_cb1')
def after_agent_cb2(callback_context):
print('@after_agent_cb2')
return types.Content(
parts=[
types.Part(
text='(stopped) after_agent_cb2',
),
],
)
def after_agent_cb3(callback_context):
print('@after_agent_cb3')
def before_agent_cb1(callback_context):
print('@before_agent_cb1')
def before_agent_cb2(callback_context):
print('@before_agent_cb2')
def before_agent_cb3(callback_context):
print('@before_agent_cb3')
def before_tool_cb1(tool, args, tool_context):
print('@before_tool_cb1')
def before_tool_cb2(tool, args, tool_context):
print('@before_tool_cb2')
def before_tool_cb3(tool, args, tool_context):
print('@before_tool_cb3')
def after_tool_cb1(tool, args, tool_context, tool_response):
print('@after_tool_cb1')
def after_tool_cb2(tool, args, tool_context, tool_response):
print('@after_tool_cb2')
return {'test': 'after_tool_cb2', 'response': tool_response}
def after_tool_cb3(tool, args, tool_context, tool_response):
print('@after_tool_cb3')
root_agent = Agent(
model='gemini-2.0-flash-exp',
name='data_processing_agent',
description=(
'hello world agent that can roll a dice of 8 sides and check prime'
' 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,
],
# planner=BuiltInPlanner(
# thinking_config=types.ThinkingConfig(
# include_thoughts=True,
# ),
# ),
generate_content_config=types.GenerateContentConfig(
safety_settings=[
types.SafetySetting( # avoid false alarm about rolling dice.
category=types.HarmCategory.HARM_CATEGORY_DANGEROUS_CONTENT,
threshold=types.HarmBlockThreshold.OFF,
),
]
),
before_agent_callback=[
before_agent_cb1,
before_agent_cb2,
before_agent_cb3,
],
after_agent_callback=[after_agent_cb1, after_agent_cb2, after_agent_cb3],
before_model_callback=before_model_callback,
after_model_callback=after_model_callback,
before_tool_callback=[before_tool_cb1, before_tool_cb2, before_tool_cb3],
after_tool_callback=[after_tool_cb1, after_tool_cb2, after_tool_cb3],
)

View File

@ -0,0 +1,145 @@
# 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.agents.run_config import RunConfig
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 = 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}')
async def run_prompt_bytes(session: Session, new_message: str):
content = types.Content(
role='user',
parts=[
types.Part.from_bytes(
data=str.encode(new_message), mime_type='text/plain'
)
],
)
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,
run_config=RunConfig(save_input_blobs_as_artifacts=True),
):
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?')
await run_prompt_bytes(session_11, 'Hi bytes')
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)
def main_sync():
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 = session_service.create_session(
app_name=app_name, user_id=user_id_1
)
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))
for event in runner.run(
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('------------------------------------')
run_prompt(session_11, 'Hi')
run_prompt(session_11, 'Roll a die with 100 sides.')
run_prompt(session_11, 'Roll a die again with 100 sides.')
run_prompt(session_11, 'What numbers did I got?')
end_time = time.time()
print('------------------------------------')
print('End time:', end_time)
print('Total time:', end_time - start_time)
if __name__ == '__main__':
print('--------------ASYNC--------------------')
asyncio.run(main())
print('--------------SYNC--------------------')
main_sync()

View File

@ -65,83 +65,6 @@ async def check_prime(nums: list[int]) -> str:
else f"{', '.join(str(num) for num in primes)} are prime numbers."
)
async def before_agent_callback(callback_context):
print('@before_agent_callback')
return None
async def after_agent_callback(callback_context):
print('@after_agent_callback')
return None
async def before_model_callback(callback_context, llm_request):
print('@before_model_callback')
return None
async def after_model_callback(callback_context, llm_response):
print('@after_model_callback')
return None
def after_agent_cb1(callback_context):
print('@after_agent_cb1')
def after_agent_cb2(callback_context):
print('@after_agent_cb2')
return types.Content(
parts=[
types.Part(
text='(stopped) after_agent_cb2',
),
],
)
def after_agent_cb3(callback_context):
print('@after_agent_cb3')
def before_agent_cb1(callback_context):
print('@before_agent_cb1')
def before_agent_cb2(callback_context):
print('@before_agent_cb2')
def before_agent_cb3(callback_context):
print('@before_agent_cb3')
def before_tool_cb1(tool, args, tool_context):
print('@before_tool_cb1')
def before_tool_cb2(tool, args, tool_context):
print('@before_tool_cb2')
def before_tool_cb3(tool, args, tool_context):
print('@before_tool_cb3')
def after_tool_cb1(tool, args, tool_context, tool_response):
print('@after_tool_cb1')
def after_tool_cb2(tool, args, tool_context, tool_response):
print('@after_tool_cb2')
return {'test': 'after_tool_cb2', 'response': tool_response}
def after_tool_cb3(tool, args, tool_context, tool_response):
print('@after_tool_cb3')
root_agent = Agent(
model='gemini-2.0-flash-exp',
name='data_processing_agent',
@ -183,14 +106,4 @@ root_agent = Agent(
),
]
),
before_agent_callback=[
before_agent_cb1,
before_agent_cb2,
before_agent_cb3,
],
after_agent_callback=[after_agent_cb1, after_agent_cb2, after_agent_cb3],
before_model_callback=before_model_callback,
after_model_callback=after_model_callback,
before_tool_callback=[before_tool_cb1, before_tool_cb2, before_tool_cb3],
after_tool_callback=[after_tool_cb1, after_tool_cb2, after_tool_cb3],
)