207 lines
9.2 KiB
Python
207 lines
9.2 KiB
Python
"""将需要规划的需求整理为 ScreenPlan"""
|
||
|
||
import json
|
||
|
||
from langchain_core.messages import HumanMessage, SystemMessage
|
||
|
||
from app.clients.llm import get_llm
|
||
from app.processor.base import BaseNode
|
||
from app.processor.state import QueryState
|
||
from app.prompt.planner import PLANNER_SYSTEM_PROMPT
|
||
from app.schemas.screen_plan import ScreenPlan
|
||
|
||
# 规划输出未通过校验时，带着错误让模型重新规划的最大次数
|
||
# TODO 上游节点目前只有这里做了节点内重试，scope_node 和 requirement_node 尚未跟进
|
||
# TODO 将来可改为图级回路：规划校验失败 -> 规划版 reflection -> planner，
|
||
# 与动作阶段的重试统一由 route_failure 控制，届时可去掉这里的节点内重试
|
||
PLAN_MAX_REPAIRS = 1
|
||
|
||
|
||
class PlannerNode(BaseNode):
|
||
name = "planner_node"
|
||
|
||
async def process(self, state: QueryState) -> QueryState:
|
||
# 浅拷贝状态
|
||
updated_state: QueryState = state.copy()
|
||
# 获取已提取的需求和范围，作为规划内容及修改边界的依据
|
||
requirement = state["requirement"]
|
||
scope = state["change_scope"]
|
||
|
||
# 获取当前大屏和用户输入，供模型结合已有内容安排布局与组件
|
||
context = state["input_context"]
|
||
# 图保证组件目录已就绪后才进入规划
|
||
resources = state["component_resources"]
|
||
|
||
messages = [
|
||
# 系统消息提供稳定的规划规则和输出边界
|
||
SystemMessage(content=PLANNER_SYSTEM_PROMPT),
|
||
# 用户消息承载本次上下文、需求、范围及组件能力摘要
|
||
HumanMessage(
|
||
content=json.dumps(
|
||
{
|
||
"input_context": context.model_dump(mode="json"),
|
||
"requirement": requirement.model_dump(mode="json"),
|
||
"change_scope": scope.model_dump(mode="json"),
|
||
# 仅提供规划需要的组件摘要，省略完整默认配置
|
||
"component_resources": resources.to_planner_context(),
|
||
},
|
||
ensure_ascii=False,
|
||
)
|
||
),
|
||
]
|
||
llm_client = get_llm(ScreenPlan)
|
||
|
||
for attempt in range(PLAN_MAX_REPAIRS + 1):
|
||
plan = None
|
||
try:
|
||
plan = await llm_client.ainvoke(messages)
|
||
self._check(plan, state)
|
||
break
|
||
except ValueError as error:
|
||
# 解析失败、引用不存在的目标等可由模型自行纠正，超过次数才失败收尾
|
||
if attempt == PLAN_MAX_REPAIRS:
|
||
return self.fail_on_invalid_output(
|
||
state, error, "未能生成可用的大屏规划，请调整需求后重试"
|
||
)
|
||
self.logger.warning(f"规划未通过校验，重新规划: {error}")
|
||
feedback = f"上一次输出未通过校验：{error}\n请修正后重新输出完整规划。"
|
||
if plan is not None:
|
||
feedback = f"上一次输出：\n{plan.model_dump_json()}\n\n{feedback}"
|
||
messages = [*messages, HumanMessage(content=feedback)]
|
||
|
||
# 组件目录已在进入规划前完整获取，component 类资源需求不可能再有补充来源
|
||
# 保留会让动作阶段误判为“待获取资源”而失败，能力缺口应由 assumptions 表达
|
||
# TODO 接入数据源和 RAG 后，仅保留目录之外的资源需求
|
||
plan.resource_needs = [
|
||
need for need in plan.resource_needs if need.kind != "component"
|
||
]
|
||
|
||
# 合并需求约束和规划约束，按首次出现顺序去重，保留需求约束在前
|
||
plan.constraints = list(
|
||
dict.fromkeys([*requirement.constraints, *plan.constraints])
|
||
)
|
||
# 同样合并原始假设与规划补充的假设，避免重复记录或丢失前置条件
|
||
plan.assumptions = list(
|
||
dict.fromkeys([*requirement.assumptions, *plan.assumptions])
|
||
)
|
||
# 保存通过校验的规划，供后续动作生成使用
|
||
updated_state["screen_plan"] = plan
|
||
|
||
# 返回更新后的状态，由图继续处理规划对应的动作
|
||
return updated_state
|
||
|
||
@staticmethod
|
||
def _check(plan: ScreenPlan, state: QueryState) -> None:
|
||
"""校验规划的组件类型、目标引用和修改范围，不通过时抛出 ValueError"""
|
||
requirement = state["requirement"]
|
||
scope = state["change_scope"]
|
||
context = state["input_context"]
|
||
resources = state["component_resources"]
|
||
|
||
# 校验规划中的物料类型是否真实受支持
|
||
supported_types = {component.type for component in resources.components}
|
||
for component in plan.components:
|
||
# 如果新增组件没有指定类型，无法选择对应物料，立即报错
|
||
if component.disposition == "add" and component.component_type is None:
|
||
raise ValueError(f"新增规划组件未指定物料类型: {component.plan_id}")
|
||
# 新增或修改组件时，只要明确指定了类型，就必须在目录中存在
|
||
if (
|
||
component.disposition in ("add", "modify")
|
||
and component.component_type is not None
|
||
and component.component_type not in supported_types
|
||
):
|
||
raise ValueError(
|
||
f"规划使用了目录中不存在的组件类型: {component.component_type}"
|
||
)
|
||
|
||
# 收集当前大屏中真实存在的节点 ID，用于检查规划引用的已有目标
|
||
existing_ids = {node.id for node in context.screen.nodes}
|
||
# 汇总规划中明确引用的目标 ID，新增组件没有已有目标时不计入
|
||
planned_ids = {
|
||
component.target_id
|
||
for component in plan.components
|
||
if component.target_id is not None
|
||
}
|
||
# 如果规划引用了不存在的组件，则抛出异常并列出无效 ID
|
||
if unknown := planned_ids - existing_ids:
|
||
raise ValueError(f"规划引用了不存在的组件: {sorted(unknown)}")
|
||
|
||
# 汇总需求明确指定的目标 ID，避免规划遗漏用户提及的组件
|
||
required_ids = {
|
||
target for change in requirement.changes for target in change.target_ids
|
||
}
|
||
|
||
# 如果需求目标未在规划中出现，则停止处理并指出遗漏
|
||
if missing := required_ids - planned_ids:
|
||
raise ValueError(f"规划遗漏了需求中的组件: {sorted(missing)}")
|
||
# 如果是局部规划，还需保证对已有组件的修改和删除不超出指定范围
|
||
if scope.scope == "local":
|
||
# 只统计会改变已有组件的目标，保留和新增组件不纳入此项检查
|
||
changed_ids = {
|
||
component.target_id
|
||
for component in plan.components
|
||
if component.disposition in ("modify", "remove")
|
||
}
|
||
# 如果实际修改或删除的目标不在允许范围内，则拒绝扩大修改范围
|
||
if outside := changed_ids - set(scope.target_ids):
|
||
raise ValueError(f"局部规划修改了范围外的组件: {sorted(outside)}")
|
||
|
||
|
||
if __name__ == "__main__":
|
||
import asyncio
|
||
|
||
from pydantic import BaseModel
|
||
|
||
from app.processor.nodes.component_resource_node import ComponentResourceNode
|
||
from app.processor.nodes.input_node import InputNode
|
||
from app.schemas.change_scope import ChangeScope
|
||
from app.schemas.requirement import Requirement
|
||
|
||
async def main():
|
||
# 1. 创建 state，直接提供需求和范围，单独测试规划
|
||
state = InputNode().process(
|
||
{
|
||
"request": {
|
||
"prompt": "创建销售大屏，上方展示销售总额，下方展示月度趋势",
|
||
"screen": {"canvas": {"width": 1920, "height": 1080}, "nodes": []},
|
||
}
|
||
}
|
||
)
|
||
state["requirement"] = Requirement(
|
||
intent="create_screen",
|
||
goal="创建销售总览大屏",
|
||
needs_clarification=False,
|
||
changes=[
|
||
{
|
||
"target_scope": "screen",
|
||
"operation": "add",
|
||
"parameters": {"metrics": ["销售总额", "月度销售额"]},
|
||
}
|
||
],
|
||
constraints=["总额在上，趋势在下"],
|
||
)
|
||
state["change_scope"] = ChangeScope(
|
||
scope="screen",
|
||
needs_planning=True,
|
||
resource_needs=["销售总额和月度销售数据，用于指标与趋势展示"],
|
||
reason="需要安排区域和组件组合",
|
||
)
|
||
# 2. 执行节点并打印完整 state，仅调用 LLM
|
||
node = PlannerNode()
|
||
state = ComponentResourceNode().process(state)
|
||
final_state = await node.process(state)
|
||
print(
|
||
json.dumps(
|
||
final_state,
|
||
ensure_ascii=False,
|
||
indent=2,
|
||
default=lambda value: (
|
||
value.model_dump(mode="json")
|
||
if isinstance(value, BaseModel)
|
||
else str(value)
|
||
),
|
||
)
|
||
)
|
||
|
||
asyncio.run(main())
|