How to stream specific nodes？

[zzhaoh]

I built a multi-agent system using only nodes and edges. Here's an example:

async def chatbot(state: MessagesState) -> dict:
    chat_prompt = """some prompt to chat"""
    response = await llm.ainvoke([SystemMessage(content=chat_prompt)] + state["messages"])
    return {"messages": [{"content": response.content}]}

async def generate_document(state: MessagesState):
    generate_prompt = """some prompt to generate document"""
    response = await llm.ainvoke([SystemMessage(content=generate_prompt)] + state["messages"])
    # Do something else, like saving the document to a database

async def intent_router(state: MessagesState):
    prompt = """some prompt to classify user input into one of two intents"""
    intent = await llm.invoke([SystemMessage(content=prompt)] + state["messages"])
    response = intent.strip().lower()
    if response not in {"generate", "chatbot"}:
        response = "chatbot"
    return {"intent": response}

graph_builder = StateGraph(MessagesState)
graph_builder.add_node("chatbot", chatbot)
graph_builder.add_node("generate", generate_document)
graph_builder.add_node("intent_router", intent_router)
graph_builder.add_edge(START, "intent_router")
graph_builder.add_conditional_edges(
    "intent_router",
    path=lambda state: state["intent"]
)
graph = graph_builder.compile()

Then I used astream() to observe the token output of the LLM. However, it streamed all tokens — even from the intent_router node, which I don't want. Here's what I tried:

async for chunk in graph.astream(
    input=state,
    config=config,
    stream_mode="messages",
):
    print(chunk)

So I changed the stream_mode to "custom", and had to modify the code of chatbot and generate_document like this:

async def chatbot(state: MessagesState) -> dict:
    chat_prompt = """some prompt to chat"""
    full_content = ""
    stream = get_stream_writer()
    async for chunk in llm.astream([SystemMessage(content=chat_prompt)] + state["messages"]):
        if chunk.content:
            full_content += chunk.content
            stream(chunk.content)

    return {"messages": [{"content": full_content}]}

It works well, but it’s a bit inconvenient.

❓My Questions:
Is there a better way to achieve token-level streaming only from certain nodes, like chatbot and generate, while skipping nodes like intent_router?

How can I avoid this issue when using create_react_agent() and create_supervisor()?

Thanks!

[feng-1985]

did you solve it ?

[zzhaoh]

Not yet, it seems like I have to build my own react_agent so I can use custom stream_mode.