from langgraph.graph import StateGraph, START, END

from .state import WorkflowState, get_state_schema
from .nodes import (
    wrap_init_memory,
    wrap_query_vdb_message,
    wrap_deepseek_chat,
    wrap_add_message,
)

def create_graph(init_args):
    builder = StateGraph(WorkflowState)

    # 添加节点
    builder.add_node("init_memory", wrap_init_memory)
    builder.add_node("wrap_query_vdb_message", wrap_query_vdb_message)
    builder.add_node("chatbot", wrap_deepseek_chat)
    builder.add_node("add_message", wrap_add_message)

    # 定义边
    builder.add_edge(START, "init_memory")
    builder.add_edge("init_memory", "wrap_query_vdb_message")
    builder.add_edge("wrap_query_vdb_message", "chatbot")
    builder.add_edge("chatbot", "add_message")
    builder.add_edge("add_message", END)

    return builder.compile()