langchain源码分析：langgraph

from langchain_openai import ChatOpenAI
from langgraph.graph import StateGraph, START, END
from langgraph.graph.message import MessagesState
import sys
import os
base_dir = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
sys.path.insert(0, base_dir)
from init.openai import LLM
llm = LLM()

# 创建graph
graph_builder = StateGraph(MessagesState)

# 创建节点函数
def chatbot(state: MessagesState):
    return {"messages": [llm.invoke(state["messages"])]}

# 添加节点，命名为chatbot
graph_builder.add_node("chatbot", chatbot)

# 添加起始边
graph_builder.add_edge(START, "chatbot")

# 添加结束边
graph_builder.add_edge("chatbot", END)

# 代码编译为图
graph = graph_builder.compile()

# 生成图的流程图片，写到当前目录下
qa_async_png = graph.get_graph().draw_mermaid_png()
with open("chatbot.png", "wb") as f:
    f.write(qa_async_png)

# 流式调用大模型
for event in graph.stream({"messages": [{"role": "user", "content": "中国最好的大学是哪所？"}]}):
    for value in event.values():
        print("Assistant:", value["messages"][-1].content)

graph_builder = StateGraph(MessagesState)

graph_builder.add_node("chatbot", chatbot)

# 添加起始边
graph_builder.add_edge(START, "chatbot")
# 添加结束边
graph_builder.add_edge("chatbot", END)

graph = graph_builder.compile()

qa_async_png = graph.get_graph().draw_mermaid_png()

# 创建一个状态图
from langgraph.graph import StateGraph
from typing import TypedDict, List, Optional,Dict, Any
class PizzaState(TypedDict):
    order_id: str
    size: str
    flour_type: str
    sauce: str
    extra_cheese: bool
    toppings: List[str]
    dough_ready: Optional[bool]  # 后续步骤中添加的字段
    # 其他字段...
# 使用类型化状态创建图
pizza_workflow = StateGraph(PizzaState) # 不能用下面的原始Dict初始化，否则报错不能画图
# pizza_workflow = StateGraph(Dict[str, Any])  # 创建基于字典的状态图
# 接下来我们会向这个图中添加节点和边
def prepare_dough(state: Dict[str, Any]) -> Dict[str, Any]:
    """准备面团节点"""
    # 输入状态可能包含面粉类型、水量等信息
    print(f"正在准备 {state.get('flour_type', '普通')} 面粉的面团")

    # 处理后更新状态
    state["dough_ready"] = True
    state["dough_size"] = state.get("size", "中号")

    return state  # 返回更新后的状态
# 在图中添加这个节点
pizza_workflow.add_node("prepare_dough", prepare_dough)
def add_sauce(state: Dict[str, Any]) -> Dict[str, Any]:
    """添加酱料节点"""
    # 输入状态可能包含酱料类型、水量等信息
    print(f"正在添加 {state.get('sauce', '无')} 酱料")
    # 处理后更新状态
    state["sauce_added"] = True
    state["sauce_type"] = state.get("sauce", "无")
    return state  # 返回更新后的状态
# 在图中添加这个节点
pizza_workflow.add_node("add_sauce", add_sauce)
def add_cheese(state: Dict[str, Any]) -> Dict[str, Any]:
    """添加奶酪节点"""
    # 输入状态可能包含奶酪类型、水量等信息
    print(f"正在添加 {state.get('cheese_type', '无')} 奶酪")    
    # 处理后更新状态
    state["cheese_added"] = True
    state["cheese_type"] = state.get("cheese_type", "无")
    return state  # 返回更新后的状态
# 在图中添加这个节点
pizza_workflow.add_node("add_cheese", add_cheese)
# 标准边：定义固定的执行路径
pizza_workflow.add_edge("prepare_dough", "add_sauce")  # 准备面团后添加酱料
pizza_workflow.add_edge("add_sauce", "add_cheese")     # 添加酱料后撒奶酪
pizza_workflow.set_entry_point("prepare_dough")
def extra_cheese_node(state: Dict[str, Any]) -> Dict[str, Any]:
    """额外奶酪节点"""
    # 输入状态可能包含额外奶酪类型、水量等信息
    print(f"正在添加 {state.get('extra_cheese_type', '无')} 额外奶酪")    
    # 处理后更新状态
    state["extra_cheese_added"] = True
    state["extra_cheese_type"] = state.get("extra_cheese_type", "无")
    return state  # 返回更新后的状态
# 在图中添加额外奶酪节点
pizza_workflow.add_node("extra_cheese_node", extra_cheese_node)
def baking_node(state: Dict[str, Any]) -> Dict[str, Any]:
    """烘烤节点"""
    # 处理后更新状态
    state["baking_time"] = "12分钟"
    state["pizza_ready"] = True
    state["quality_check"] = "passed"
    return state  # 返回更新后的状态
    # 在图中添加烘烤节点
pizza_workflow.add_node("baking_node", baking_node)
# 条件边：根据状态决定下一步
def check_cheese_preference(state: Dict[str, Any]) -> str:
    """检查奶酪偏好，决定下一步"""
    if state.get("extra_cheese", False):
        return "add_extra_cheese"  # 如果客户要求额外奶酪，走这条路径
    return "regular_baking"        # 否则直接烘烤
# 添加条件边，根据函数返回值决定下一步
pizza_workflow.add_conditional_edges(
    "add_cheese",                 # 从这个节点出发
    check_cheese_preference,      # 使用这个函数判断
    {
        "add_extra_cheese": "extra_cheese_node",  # 返回值与下一个节点的映射
        "regular_baking": "baking_node"
    }
)
# 初始状态：客户订单信息
initial_state = {
    "order_id": "12345",
    "size": "大号",
    "flour_type": "全麦",
    "sauce": "番茄",
    "extra_cheese": True,
    "toppings": ["蘑菇", "橄榄", "青椒"]
}
# 编译工作流
compiled_workflow = pizza_workflow.compile()
# qa_async_png = compiled_workflow.get_graph().draw_mermaid_png()
graph = compiled_workflow.get_graph()
if hasattr(graph, 'draw_mermaid_png'):
    qa_async_png = graph.draw_mermaid_png()
with open("edg.png", "wb") as f:
    f.write(qa_async_png)
# 执行工作流，传入初始状态
final_state = compiled_workflow.invoke(initial_state)
# 最终状态会包含整个过程中添加的所有信息
print(final_state)
# 可能输出: {'order_id': '12345', 'size': '大号', ..., 'dough_ready': True, 
#           'sauce_added': True, 'cheese_type': '马苏里拉', 'baking_time': '12分钟',
#           'pizza_ready': True, 'quality_check': 'passed'}