[FEATURE] Async Initialization #330
Description
Problem Statement
I want to be able to make async calls during agent initialization. For example, it would be great to be able to pass in an mcp client as a tool, and have an agent automatically list and load the tools from that client:
agent = Agent(tools=[mcp_client]) # Call asyncio.run(mcp_client.list_tools())
agent("Hello!")However, this runs into a problem when the Agent is already running in an asyncio event_loop:
async def call_agent():
agent = Agent(tools=[mcp_client]) # RuntimeError: asyncio.run() cannot be called from a running event loopAdditionally, for the upcoming session manager, if I want to load an agent's session from a datastore asyncronously, we would get a similar error:
async def call_agent():
agent = Agent(session_manager=async_session_manager) # RuntimeError: asyncio.run() cannot be called from a running event loopI would like to maintain the existing devex of initializing an agent, but still allow for objects that require async calls to be passed into the agent's initializer.
async def call_agent():
agent = Agent(session_manager=async_session_manager) # This will succeedProposed Solution
We want to be able to initialize an agent regardless of the execution environment. In order to support this, we need to define the different cases where async resource initialization may affect agent initialization:
- Agent not already in asyncio event_loop, with async resources passed in
- Agent in asyncio event_loop, with async resources passed in
In order to support these cases, I propose the following changes:
1. Create new async initialize method on the Agent class
A new async initialize function will be added to the agent class. Any async resources that are passed into the agent initializer will be called as a part of the async initialize function. If the agent has async resurces passed in during initialization, and the agent is invoked before this initialize function is called, then an exception will be thrown.
class Agent:
...
async def initialize():
...
2. Call initialize in the agent class by default
Call the new initialize function during agent initialization. We can add a check to the Agent.init function to see if the execution environment is already in an asyncio event_loop to determine how this function is called:
class Agent:
def __init__(...):
...
try:
# Check if we are already in event loop
loop = asyncio.get_running_loop()
def wrap_async():
asyncio.run(self.initialize())
thread = threading.Thread(target=wrap_async)
thread.start()
thread.join()
except RuntimeException as e:
use asyncio if we are not in an event loop
asyncio.run(self.initialize())
...
3. (Optional) Add a auto_initialize=True parameter to toggle this behavior
Add a new parameter to the agent class called auto_initialize that will toggle if the logic described in the above point is executed. This gives an escape hatch for developers in case they want to run this initialize function themselves.
Use Case
- MCP clients being passed into the agent class, so that customers dont need to initialize them outside the class
- Session Manager, where we need to asyncronously load resources from a datastore
- Hooks, where we may want to call an async hook at init time.
Alternatives Solutions
Option 1: Class method for agent initialization
a = await Agent.async_create(mcp_client, session_manager)Add a classmethod to the Agent class that takes in and handles the initialization of async resources.
This comes with the advantage of a clearer separation of what code is sync vs async. The disadvantage is that if a hobbyist customer who may not be familiar with async python wants to use a mcp tool, they need to be able to understand and use this new class method.
Additional Context
- [FEATURE] Improved MCP DX to prevent common context manager scoping issue #198
- [FEATURE] Session Persistence #246
- [FEATURE] Implement general-purpose hook/callback system #231
Relevant stack overflow questions about asyncio.run: