from langgraph.graph import StateGraph, END
from nodes import *
from state import State

workflow = StateGraph(State)

# 添加节点
workflow.add_node("retrieve", retrieve_node)
workflow.add_node("summarize", summarize_node)
workflow.add_node("chatbot", call_model_node)

# 设置入口
workflow.set_entry_point("retrieve")

# 条件逻辑：检查消息数量
def should_summarize(state: State):
    if len(state["messages"]) > 10:
        return "summarize"
    return "chatbot"

workflow.add_conditional_edges(
    "retrieve", 
    should_summarize, 
    {
        "summarize": "summarize", 
        "chatbot": "chatbot"
    }
)
workflow.add_edge("summarize", "chatbot")
workflow.add_edge("chatbot", END)

# 编译时加入 Checkpointer (你可以使用你的 Postgres 实现来持久化 State)
app = workflow.compile() #checkpointer=postgres_checkpointer