"""LangGraph Streaming 示例：节点更新 + 模型 Token 流。 图结构： START -> model | 有 tool_calls v tools | +---------> model model -- 无 tool_calls --> END 本课重点不是工具本身，而是三种 stream_mode： 1. updates —— 每个节点执行完后，推送该节点的状态增量 2. messages —— 模型生成时，逐 token 推送 AIMessageChunk 3. 同时开两者 —— 一边打字机输出，一边看到节点何时完成 运行： uv run --with langgraph --with langchain-openai --with python-dotenv \ python 08_streaming.py """ import os from typing import Literal from dotenv import load_dotenv from langchain_core.tools import tool from langchain_openai import ChatOpenAI from langgraph.graph import END, START, MessagesState, StateGraph from langgraph.prebuilt import ToolNode load_dotenv() # ----------------------------------------------------------------------------- # 1. 工具与模型（沿用上一课的天气查询，方便观察节点切换） # ----------------------------------------------------------------------------- @tool def get_weather(city: str) -> str: """查询指定城市当前的天气。城市名称使用中文，例如北京、上海。""" weather_data = { "北京": "晴，25°C，北风 2 级", "上海": "多云，28°C，东南风 3 级", "广州": "雷阵雨，30°C，南风 2 级", "深圳": "阴，29°C，东风 2 级", } return weather_data.get(city, f"暂时没有 {city} 的天气数据") tools = [get_weather] model = ChatOpenAI( model=os.getenv("OPENAI_MODEL", "openai/gpt-5.4"), api_key=os.getenv("OPENAI_API_KEY"), base_url=os.getenv("OPENAI_BASE_URL", "https://api.qnaigc.com/v1"), temperature=0, max_retries=2, ) model_with_tools = model.bind_tools(tools) def call_model(state: MessagesState) -> dict: # 即使这里写的是 .invoke，stream_mode="messages" 仍能拿到 token 流。 # LangGraph 会自动把内部 LLM 调用的流式事件透传出来。 response = model_with_tools.invoke(state["messages"]) return {"messages": [response]} def route_after_model(state: MessagesState) -> Literal["tools", "__end__"]: last_message = state["messages"][-1] if last_message.tool_calls: return "tools" return END tool_node = ToolNode(tools) builder = StateGraph(MessagesState) builder.add_node("model", call_model) builder.add_node("tools", tool_node) builder.add_edge(START, "model") builder.add_conditional_edges( "model", route_after_model, {"tools": "tools", END: END}, ) builder.add_edge("tools", "model") graph = builder.compile() # ----------------------------------------------------------------------------- # 2. 演示 A：stream_mode="updates" —— 看节点级增量 # ----------------------------------------------------------------------------- def demo_updates(question: str) -> None: """每个节点跑完才出一包数据：{节点名: 该节点返回的状态增量}。""" print("\n--- stream_mode='updates' ---") for event in graph.stream( {"messages": [{"role": "user", "content": question}]}, stream_mode="updates", ): # event 形如：{"model": {"messages": [AIMessage(...)]}} # 或：{"tools": {"messages": [ToolMessage(...)]}} for node_name, update in event.items(): print(f"\n[{node_name}] 节点完成，增量：") for message in update.get("messages", []): message.pretty_print() # ----------------------------------------------------------------------------- # 3. 演示 B：stream_mode="messages" —— 看模型 Token 流 # ----------------------------------------------------------------------------- def demo_messages(question: str) -> None: """模型一边生成，一边吐 AIMessageChunk；适合做打字机效果。""" print("\n--- stream_mode='messages' ---") print("助手：", end="", flush=True) for chunk in graph.stream( {"messages": [{"role": "user", "content": question}]}, stream_mode="messages", version="v2", ): # v2 统一格式：{"type": "messages", "data": (message_chunk, metadata)} if chunk["type"] != "messages": continue message_chunk, metadata = chunk["data"] # 只打印 model 节点吐出的文本 token；工具调用阶段 content 常为空。 if metadata.get("langgraph_node") != "model": continue content = message_chunk.content if content: # content 可能是 str，也可能是 content blocks 列表，这里只处理文本。 if isinstance(content, str): print(content, end="", flush=True) elif isinstance(content, list): for block in content: if isinstance(block, dict) and block.get("type") == "text": print(block.get("text", ""), end="", flush=True) elif isinstance(block, str): print(block, end="", flush=True) print() # 换行 # ----------------------------------------------------------------------------- # 4. 演示 C：同时开 updates + messages —— 最实用的组合 # ----------------------------------------------------------------------------- def demo_both(question: str) -> None: """一边流式打印 token，一边在节点结束时打印节点名。""" print("\n--- stream_mode=['updates', 'messages'] ---") print("助手：", end="", flush=True) printed_tokens = False for chunk in graph.stream( {"messages": [{"role": "user", "content": question}]}, stream_mode=["updates", "messages"], version="v2", ): if chunk["type"] == "messages": message_chunk, metadata = chunk["data"] if metadata.get("langgraph_node") != "model": continue content = message_chunk.content if isinstance(content, str) and content: print(content, end="", flush=True) printed_tokens = True elif isinstance(content, list): for block in content: text = "" if isinstance(block, dict) and block.get("type") == "text": text = block.get("text", "") elif isinstance(block, str): text = block if text: print(text, end="", flush=True) printed_tokens = True elif chunk["type"] == "updates": # 节点结束后再换行并标注，避免和 token 混在同一行。 for node_name in chunk["data"]: if printed_tokens: print() printed_tokens = False print(f" ✓ 节点 [{node_name}] 完成") if node_name == "model": print("助手：", end="", flush=True) if printed_tokens: print() # ----------------------------------------------------------------------------- # 5. 入口：依次跑三种模式，对比输出差异 # ----------------------------------------------------------------------------- def main() -> None: print("=== LangGraph Streaming：节点更新 + Token 流 ===") print("建议先问：北京天气怎么样？ （会触发 tools -> model 两跳）") print("或问：用一句话介绍 LangGraph。 （只走 model）") print("输入 quit 退出。\n") print("命令：") print(" /u 只看 updates") print(" /m 只看 messages（token）") print(" /b 同时看两者（默认）") mode = "both" while True: question = input(f"\n用户[{mode}]：").strip() if question.lower() in {"quit", "exit", "q"}: break if not question: continue if question == "/u": mode = "updates" print("已切换到 updates") continue if question == "/m": mode = "messages" print("已切换到 messages") continue if question == "/b": mode = "both" print("已切换到 both") continue try: if mode == "updates": demo_updates(question) elif mode == "messages": demo_messages(question) else: demo_both(question) except Exception as exc: print(f"\n执行失败：{type(exc).__name__}: {exc}") if __name__ == "__main__": main()