my-v0-server/app/processor/nodes/planner_node.py

"""将需要规划的需求整理为 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


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"]

        if (
            requirement.needs_clarification
            or requirement.intent
            not in (
                "create_screen",
                "modify_screen",
            )
            or not scope.needs_planning
        ):
            raise ValueError("PlannerNode 只处理信息完整且需要规划的新建或修改需求")

        context = state["input_context"]
        resources = state.get("component_resources")
        if resources is not None and resources.status != "resolved":
            raise ValueError("组件目录尚未就绪，不能继续规划")

        result = await get_llm(ScreenPlan).ainvoke(
            [
                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()
                            if resources is not None
                            else None,
                        },
                        ensure_ascii=False,
                    )
                ),
            ]
        )

        plan = ScreenPlan.model_validate(result.model_dump())

        if resources is not None:
            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
                ):
                    if component.component_type not in supported_types:
                        raise ValueError(
                            f"规划使用了目录中不存在的组件类型: {component.component_type}"
                        )

        existing_ids = {node.id for node in context.screen.nodes}
        planned_ids = {
            component.target_id
            for component in plan.components
            if component.target_id is not None
        }
        if unknown := planned_ids - existing_ids:
            raise ValueError(f"规划引用了不存在的组件: {sorted(unknown)}")

        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)}")

        # 原始约束和假设不能在规划阶段被无声丢弃
        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


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())