my-v0-server/app/processor/graph.py

"""输入解析、意图分流和需求理解流程"""

import inspect

from langgraph.checkpoint.memory import InMemorySaver
from langgraph.checkpoint.serde.jsonplus import JsonPlusSerializer
from langgraph.graph import END, StateGraph

from app.processor.nodes.clarify_node import ClarifyNode
from app.processor.nodes.input_node import InputNode
from app.processor.nodes.intent_node import IntentNode
from app.processor.nodes.requirement_node import RequirementNode
from app.processor.nodes.wait_for_user_node import WaitForUserNode
from app.processor.state import QueryState


def route_after_intent(state: QueryState) -> str:
    return state["intent"]


def route_after_requirement(state: QueryState) -> bool:
    return state["requirement"].needs_clarification


def route_after_wait(state: QueryState) -> str:
    return "input_node" if state["status"] == "running" else END


def build_graph(checkpointer=None):
    # 定义工作流
    workflow = StateGraph(QueryState)

    # 实例化节点
    nodes = {
        "input_node": InputNode(),
        "intent_node": IntentNode(),
        "requirement_node": RequirementNode(),
        "clarify_node": ClarifyNode(),
        "wait_for_user_node": WaitForUserNode(),
    }

    # 注册节点，异步节点使用 BaseNode 的异步入口
    for node_name, node in nodes.items():
        action = node._acall if inspect.iscoroutinefunction(node.process) else node
        workflow.add_node(node_name, action)

    # 定义入口
    workflow.set_entry_point("input_node")

    # 添加业务边
    workflow.add_edge("input_node", "intent_node")
    # 新建和修改进入需求理解，意图不明进入澄清
    workflow.add_conditional_edges(
        source="intent_node",
        path=route_after_intent,
        path_map={
            "create_screen": "requirement_node",
            "modify_screen": "requirement_node",
            "qa": END,
            "unclear": "clarify_node",
        },
    )
    workflow.add_conditional_edges(
        source="requirement_node",
        path=route_after_requirement,
        path_map={True: "clarify_node", False: END},
    )
    workflow.add_edge("clarify_node", "wait_for_user_node")
    workflow.add_conditional_edges(
        "wait_for_user_node",
        route_after_wait,
        {"input_node": "input_node", END: END},
    )
    # TODO 本地默认保存在内存，可传入持久化 checkpointer
    if checkpointer is None:
        serde = JsonPlusSerializer(
            allowed_msgpack_modules=[
                ("app.schemas.input_context", "InputContext"),
                ("app.schemas.decision", "DecisionResponse"),
                ("app.schemas.requirement", "Requirement"),
                ("app.schemas.clarification", "Clarification"),
                ("app.schemas.clarification_reply", "ClarificationReply"),
            ]
        )
        checkpointer = InMemorySaver(serde=serde)
    return workflow.compile(checkpointer=checkpointer)


# 创建全局图实例
query_app = build_graph()


if __name__ == "__main__":
    import asyncio
    import json
    from uuid import uuid4

    from langgraph.types import Command
    from pydantic import BaseModel

    from app.clients.decision import close_decision_engine

    async def main():
        # 1. 修改请求，测试完整流程
        state: QueryState = {
            "request": {
                "prompt": "帮我修改销售额统计这个图，但具体改哪里我还没说，请先问我",
                "screen": {
                    "canvas": {"width": 1920, "height": 1080},
                    "nodes": [
                        {
                            "id": "sales-chart",
                            "type": "bar-chart",
                            "name": "销售额统计",
                            "layout": {"x": 26, "y": 34.8, "width": 400, "height": 260},
                            "props": {"option": {"title": {"text": "销售额统计"}}},
                        }
                    ],
                    "dataSource": [],
                },
                "context": {"selected_node_ids": ["sales-chart"]},
            }
        }
        # 2. 执行图，程序退出时关闭共享连接
        config = {"configurable": {"thread_id": str(uuid4())}}
        try:
            final_state = await query_app.ainvoke(state, config=config)
            while final_state.get("__interrupt__"):
                pending = final_state["__interrupt__"][0].value
                print(pending["question"])
                options = pending.get("options", [])
                for index, option in enumerate(options, 1):
                    marker = "（推荐）" if option["recommended"] else ""
                    print(
                        f"{index}. {option['label']}{marker}：{option['description']}"
                    )
                if "error" in pending:
                    print(pending["error"])
                answer = input("输入编号或自由回答（/cancel 取消）：")
                reply = {"cancel": True} if answer == "/cancel" else {"answer": answer}
                if answer.isdigit() and 1 <= int(answer) <= len(options):
                    reply = {"option_id": options[int(answer) - 1]["id"]}
                    notes = input("补充说明（可留空）：").strip()
                    if notes:
                        reply["answer"] = notes
                final_state = await query_app.ainvoke(
                    Command(resume=reply), config=config
                )
        finally:
            await close_decision_engine()

        # 3. 打印完整 state
        print(
            json.dumps(
                {
                    key: value.model_dump(mode="json", exclude_unset=True)
                    if isinstance(value, BaseModel)
                    else value
                    for key, value in final_state.items()
                },
                ensure_ascii=False,
                indent=2,
            )
        )

    asyncio.run(main())