Skip to content

Asynchronous model call backs. #4788

@critico

Description

@critico

🔴 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.

Metadata

Metadata

Assignees

No one assigned

    Labels

    core[Component] This issue is related to the core interface and implementation

    Projects

    No projects

    Milestone

    No milestone

    Relationships

    None yet

    Development

    No branches or pull requests

    Issue actions