How to Stream Data from One Node to the Next

[panzhilin007]

How can we stream data from one node to the next? For example, in the following scenario: Node A calls an LLM and receives a streaming output, and Node B processes each token one by one based on Node A’s output.

I attempted the following approach, but the result suggests that the next node only starts executing after the previous node has fully completed. The previous node only yields the final piece of data, and Node B only starts execution once Node A has entirely finished.

Example:

import asyncio
from typing import AsyncGenerator, Annotated

from langchain_core.messages import HumanMessage, AnyMessage, AIMessage
from langgraph.graph import StateGraph, MessagesState, add_messages


class WorkflowState(MessagesState):
    chunks: Annotated[list[AnyMessage], add_messages]

async def producer_node(state: WorkflowState) -> AsyncGenerator[dict[str, list[AIMessage]], None]:
    chunk_list = []
    for message in state['messages']:
        chunk = AIMessage(content='chunk' + message.content[-1])
        chunk_list.append(chunk)
        yield {'chunks': [chunk]}


async def consumer_node(state: WorkflowState):
    if state["chunks"]:
        print(f"Processing: {state['chunks']}")


if __name__ == '__main__':
    builder = StateGraph(WorkflowState)
    builder.add_node("producer", producer_node)
    builder.add_node("consumer", consumer_node)
    builder.add_edge("producer", "consumer")
    builder.set_entry_point("producer")
    builder.set_finish_point("consumer")
    workflow = builder.compile()

    msg_state = WorkflowState(messages=[HumanMessage('message1'), HumanMessage('message2'), HumanMessage('message3')], chunks=[])


    async def run():
        async for update in workflow.astream(msg_state, stream_mode="updates"):
            if "producer" in update:
                print(f"Yielded chunk: {update['producer']}")


    asyncio.run(
        run()
    )

Result:

Yielded chunk: {'chunks': [AIMessage(content='chunk3', additional_kwargs={}, response_metadata={}, id='ec722354-5b4f-4fef-9598-8df91a74bb29')]}
Processing: [AIMessage(content='chunk3', additional_kwargs={}, response_metadata={}, id='ec722354-5b4f-4fef-9598-8df91a74bb29')]

[vvanglro]

Have the same needs.