84 lines
3.4 KiB
Python
84 lines
3.4 KiB
Python
"""输入 → 意图（是不是明确的问答）→ 问答节点或 screen agent；agent 追问时暂停等待用户"""
|
||
|
||
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.input_node import InputNode
|
||
from app.processor.nodes.intent_node import IntentNode
|
||
from app.processor.nodes.qa_node import QANode
|
||
from app.processor.nodes.screen_agent_node import ScreenAgentNode
|
||
from app.processor.nodes.wait_for_user_node import WaitForUserNode
|
||
from app.processor.state import QueryState, RunContext
|
||
|
||
|
||
def build_checkpoint_serde() -> JsonPlusSerializer:
|
||
# 允许恢复状态中使用的业务模型类型，使检查点读取后仍可按模型对象访问字段
|
||
# 内存与持久化检查点共用，否则持久化读回的状态无法还原为模型对象
|
||
return JsonPlusSerializer(
|
||
allowed_msgpack_modules=[
|
||
("app.schemas.input_context", "InputContext"),
|
||
("app.schemas.decision", "DecisionResponse"),
|
||
("app.schemas.actions.add", "AddAction"),
|
||
("app.schemas.actions.delete", "DeleteAction"),
|
||
("app.schemas.actions.move", "MoveAction"),
|
||
("app.schemas.actions.resize", "ResizeAction"),
|
||
("app.schemas.actions.update", "UpdateAction"),
|
||
("app.schemas.actions.update_canvas", "UpdateCanvasAction"),
|
||
("app.schemas.clarification", "Clarification"),
|
||
("app.schemas.clarification_reply", "ClarificationReply"),
|
||
]
|
||
)
|
||
|
||
|
||
def route_after_intent(state: QueryState) -> str:
|
||
return "qa_node" if state["intent"] == "qa" else "screen_agent_node"
|
||
|
||
|
||
def route_after_agent(state: QueryState) -> str:
|
||
return "wait_for_user_node" if state["status"] == "waiting_for_input" else END
|
||
|
||
|
||
def route_after_wait(state: QueryState) -> str:
|
||
# 用户回答后回到输入节点合并上下文，取消则结束
|
||
return "input_node" if state["status"] == "running" else END
|
||
|
||
|
||
def build_graph(checkpointer=None):
|
||
workflow = StateGraph(QueryState, context_schema=RunContext)
|
||
nodes = {
|
||
"input_node": InputNode(),
|
||
"intent_node": IntentNode(),
|
||
"qa_node": QANode(),
|
||
"screen_agent_node": ScreenAgentNode(),
|
||
"wait_for_user_node": WaitForUserNode(),
|
||
}
|
||
for node_name, node in nodes.items():
|
||
# 异步节点使用 BaseNode 的异步入口，保留统一的日志和进度封装
|
||
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(
|
||
"intent_node",
|
||
route_after_intent,
|
||
{"qa_node": "qa_node", "screen_agent_node": "screen_agent_node"},
|
||
)
|
||
workflow.add_edge("qa_node", END)
|
||
workflow.add_conditional_edges(
|
||
"screen_agent_node",
|
||
route_after_agent,
|
||
{"wait_for_user_node": "wait_for_user_node", END: END},
|
||
)
|
||
workflow.add_conditional_edges(
|
||
"wait_for_user_node", route_after_wait, {"input_node": "input_node", END: END}
|
||
)
|
||
|
||
# 调用方未提供检查点存储时使用进程内存，支持暂停恢复，进程退出后不保留
|
||
if checkpointer is None:
|
||
checkpointer = InMemorySaver(serde=build_checkpoint_serde())
|
||
return workflow.compile(checkpointer=checkpointer)
|