from langgraph.graph import StateGraph, START, END
from typing import TypedDict

# ============================================================
# 第 1 步：定义 State —— 描述整个工作流在运行时的数据结构
# ============================================================
class CounterState(TypedDict):
    """工作流的状态：一个简单的计数器"""
    count: int           # 当前计数

# ============================================================
# 第 2 步：写节点函数 —— 每个节点做一件事
# ============================================================
def step_one(state: CounterState) -> dict:
    """第一个处理节点：count + 1"""
    new_count = state["count"] + 1
    return {
        "count": new_count
    }


def step_two(state: CounterState) -> dict:
    """第二个处理节点：count * 2"""
    new_count = state["count"] * 2
    return {
        "count": new_count
    }


# ============================================================
# 第 3 步：创建 StateGraph 对象
# ============================================================
builder = StateGraph(CounterState)

# ============================================================
# 第 4 步：注册节点，连接边
# ============================================================
builder.add_node("step_one", step_one)       # 把函数注册为节点
builder.add_node("step_two", step_two)

builder.add_edge(START, "step_one")          # START → step_one
builder.add_edge("step_one", "step_two")     # step_one → step_two
builder.add_edge("step_two", END)            # step_two → END

# ============================================================
# 第 5 步：编译图（编译后才可执行）
# ============================================================
graph = builder.compile()

# ============================================================
# 第 6 步：执行
# ============================================================
result = graph.invoke({"count": 1, "log": []})
print(result)
# 输出: {'count': 4}

#显示图结构
from IPython.display import Image,display
display(Image(graph.get_graph().draw_mermaid_png()))