-
Notifications
You must be signed in to change notification settings - Fork 3.1k
Description
🔴 Required Information
Is your feature request related to a specific problem?
The problem I am trying to solve is with call backs not being asynchronous. In the src>google>adk>flows>>llm_flows>base_llm_flow.py the before_model_callback methods are handled by async def _handle_before_model_callback. Inside this method there exists code that iterates over the methods passed in.
if not agent.canonical_before_model_callbacks:
return
for callback in agent.canonical_before_model_callbacks:
callback_response = callback(
callback_context=callback_context, llm_request=llm_request
)
if inspect.isawaitable(callback_response):
callback_response = await callback_response
if callback_response:
return callback_response
This for loop has an await which makes the for loop waits on each iteration for the individual call backs to complete. This forces the methods being passed in to run in a sequential order even when using run_async and using async def callbacks.
Describe the Solution You'd Like
I have already implemented a patch that wraps the calls in the for loop inside asyncio . This allows the guardrails to run in a true asynchronous way, and reduces the amount of time taken to run guardrails. This is different from the approach taken by https://docs.vijil.ai/tutorials/protect-agents/adk , in which they run an asyncio set of tasks inside one guardrail on their own set of guardrails., a task for each guardrail The approach I'm describing runs all of the individual guardrails sent in asynchronously keeping inline with current patterns.
What I would like to see is a flag added to allow this behaviour to be switched on and off.
Impact on your work
This impacts the length of time it takes for an agent to validate using guardrails and lengthens the process time for requests significantly.
Willingness to contribute
Yes
🟡 Recommended Information
Describe Alternatives You've Considered
I've looked at https://docs.vijil.ai/tutorials/protect-agents/adk and also creating a custom agent.
Proposed API / Implementation
I cannot share the exact patch I have written , but here is some Ai generated code
callbacks = agent.canonical_before_model_callbacks
if not callbacks:
return
async def run_one(callback):
result = callback(
callback_context=callback_context,
llm_request=llm_request
)
if inspect.isawaitable(result):
return await result
else:
# Run sync callback in a thread
return await asyncio.to_thread(lambda: result)
tasks = [asyncio.create_task(run_one(cb)) for cb in callbacks]
As soon as one returns a non‑empty value, cancel the rest
for task in asyncio.as_completed(tasks):
result = await task
if result:
for t in tasks:
if not t.done():
t.cancel()
return result
return None
I would recommend using FIRST_COMPLETED with some logic to test if the guardrail passed or blocked.
Additional Context
If the async approach can't be implemented, is there a reason as to why. Do async call backs pose any kind of risk or other unknown considerations, like bugs etc. I have tested my own patch and it seems to work as expected.