LangGraph流式输出与额外回调冲突问题求助

我需要实现模型输出流式传输，同时通过回调统计Token使用量，但目前所有添加Token统计回调的文档方法都会移除StreamMessagesHandler回调，导致无法实现流式输出。这是一个关键场景——既要流式输出内容，又要根据Token使用量向客户收费，因此需要实现添加新回调时不移除已有回调的功能。

使用版本：

langchain_core: 1.0.2
langchain: 1.0.3 
langgraph: 1.0.2

当前代码：

from dataclasses import dataclass

from langchain.chat_models import init_chat_model
from langchain_core.callbacks import UsageMetadataCallbackHandler
from langgraph.constants import START
from langgraph.graph import StateGraph

model_1 = init_chat_model(model="gpt-4o-mini")  # 需在环境变量中配置API_KEY，或直接传入

@dataclass
class MyState:
    topic: str

callback = UsageMetadataCallbackHandler()

model_with_callback = model_1.with_config(
    callbacks=[callback]
)

def call_model(state: MyState):
    """调用LLM生成关于指定主题的内容"""
    # 注意：即使使用.invoke而非.stream运行LLM，也会触发消息事件
    model_response = model_with_callback.invoke(
        [
            {"role": "user", "content": f"{state.topic}"}
        ]
    )
    return {"topic": model_response.content}


graph = (
    StateGraph(MyState)
    .add_node(call_model)
    .add_edge(START, "call_model")
    .compile()
)

stream_gen = graph.stream(
    {"topic": "法国的首都是什么？"},
    stream_mode="messages",
)
for message_chunk, metadata in stream_gen:
    if message_chunk.content:
        print(message_chunk.content, end="|", flush=True)
        print(callback.usage_metadata)

核心问题是with_config方法会直接覆盖原有回调列表，而非追加。要同时保留流式回调和Token统计回调，需获取模型默认回调列表后追加自定义回调，而非直接替换。

修改后的代码 #

from dataclasses import dataclass

from langchain.chat_models import init_chat_model
from langchain_core.callbacks import UsageMetadataCallbackHandler
from langgraph.constants import START
from langgraph.graph import StateGraph

model_1 = init_chat_model(model="gpt-4o-mini")  # 需在环境变量中配置API_KEY，或直接传入

@dataclass
class MyState:
    topic: str

# 创建Token统计回调
token_callback = UsageMetadataCallbackHandler()

# 保留原有回调，追加自定义统计回调
model_with_callback = model_1.with_config(
    callbacks=model_1.callbacks + [token_callback] if hasattr(model_1, 'callbacks') else [token_callback]
)

def call_model(state: MyState):
    """调用LLM生成关于指定主题的内容"""
    # 使用stream方法触发流式输出
    model_response = model_with_callback.stream(
        [
            {"role": "user", "content": f"{state.topic}"}
        ]
    )
    full_content = ""
    # 实时输出流式内容
    for chunk in model_response:
        full_content += chunk.content
        print(chunk.content, end="|", flush=True)
    return {"topic": full_content}


graph = (
    StateGraph(MyState)
    .add_node(call_model)
    .add_edge(START, "call_model")
    .compile()
)

stream_gen = graph.stream(
    {"topic": "法国的首都是什么？"},
    stream_mode="messages",
)
# 遍历流式输出，同步打印Token统计数据
for message_chunk, metadata in stream_gen:
    if message_chunk.content:
        print(token_callback.usage_metadata)

关键修改点 #

1. 回调追加而非覆盖：通过model_1.callbacks + [token_callback]保留模型默认的流式回调，避免被统计回调替换。
2. 改用stream方法：原代码使用invoke无法触发流式输出，改为stream方法并遍历返回的块，确保内容分块实时输出。
3. 同步统计Token：在流式输出过程中，可随时通过token_callback.usage_metadata获取最新的Token使用数据。

另一种简化方案 #

也可以直接在调用graph.stream时传入统计回调，无需修改模型的回调配置：

stream_gen = graph.stream(
    {"topic": "法国的首都是什么？"},
    stream_mode="messages",
    callbacks=[token_callback]
)

这种方式同样能在保留流式输出的同时，完成Token使用量统计。

内容的提问来源于stack exchange，提问作者Yason