Add LangGraph streaming and error handling examples, demonstrating node updates, model token flow, retry policies, and error recovery mechanisms.
This commit is contained in:
Eric Wang
254
08_streaming.py
Normal file
254
08_streaming.py
Normal file
@@ -0,0 +1,254 @@
|
||||
"""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()
|
||||
Reference in New Issue
Block a user