Async streaming does not work

[corradolab]

Streaming does not work when invoked on a swarm.

In example

checkpointer = InMemorySaver()

swarm = create_swarm(
    agents=[agent1, agent2, agent3],
    default_active_agent="agent1"
)

graph = swarm.compile(checkpointer=checkpointer)

async def astream():
    async for chunk in graph.astream(
            {"messages": [{"role": "user", "content": "find the best customer"}]},
            {"configurable": {"thread_id": "1"}},
            stream_mode="messages"
        ):
        print(chunk)
        print("--------")
        await asyncio.sleep(0.5)

asyncio.run(astream())

This code works but the console print out a series of AIMessage

(AIMessage(content='', additional_kwargs={'tool_calls': [{'index': 0, 'id': 'call_UQqXoTvlKasYqTvmhz8qRyfY', 'function': {'arguments':...

which is not the expected beahviour.
If I use the same code to invoke directly an agent

async def astream():
    async for chunk in agent1.astream(  # <--- using agent instead of swarm here 
            {"messages": [{"role": "user", "content": "find the best customer"}]},
            {"configurable": {"thread_id": "1"}},
            stream_mode="messages"
        ):
        print(chunk)
        print("--------")
        await asyncio.sleep(0.5)

asyncio.run(astream())

I get a series of AIMessageChunck which is the expected beahviour

(AIMessageChunk(content='cco', additional_kwargs={}, response_metadata={}, id='run--13dcaa19-fcf9-4640-91ca-823e9b86eb27'), {'thread_, ...

[tingyiy]

Have you tried astream_events?