zhenxun_bot/zhenxun/services/ai/flow/workflow/engine.py

import asyncio
from collections.abc import AsyncIterator
import contextlib
from typing import TYPE_CHECKING, Any
import uuid

from nonebot.params import Depends
from pydantic import BaseModel

if TYPE_CHECKING:
    from .nodes import NodeSource

from zhenxun.services.ai.core.exceptions import ControlFlowExit, ToolRetryError
from zhenxun.services.ai.core.messages import PromptInput, UsageInfo
from zhenxun.services.ai.core.stream_events import EventBus
from zhenxun.services.ai.flow.base import BaseRunnable, BaseRuntimeConfig
from zhenxun.services.ai.run.blackboard import BlackboardManager
from zhenxun.services.ai.run.context import RunContext
from zhenxun.services.ai.run.models import (
    AgentRunEnd,
    AgentRunError,
    AgentRunResult,
    StreamedRunResult,
)
from zhenxun.services.ai.tools.core.tool import FunctionTool
from zhenxun.services.ai.utils.logger import log_flow as logger
from zhenxun.utils.message import MessageUtils

from .nodes import Steps
from .types import (
    StepInput,
    StepOutput,
    WorkflowRunResult,
)


class Workflow(BaseRunnable[WorkflowRunResult]):
    """
    工作流顶层容器 (The Workflow Facade)。
    继承自 BaseRunnable，支持被作为节点嵌套在 Team 或 其他工作流中。
    """

    def __init__(
        self,
        name: str,
        steps: list["NodeSource"],
        description: str = "",
        blackboard: type[BaseModel] | BaseModel | None = None,
    ):
        """
        静态图元工作流容器初始化。

        参数:
            name: 工作流的名称标识。
            steps: 工作流的节点列表（按列表顺序构成串行或嵌套结构）。
            description: 工作流的说明描述，用于被 Agent 调用时理解其功能。
            blackboard: (可选) 结构化黑板。可传入 Schema 类型类，或直接传入带有初始数据的 Schema 实例对象。
        """  # noqa: E501
        self.name = name
        self.description = description
        self.id = uuid.uuid4().hex

        self.root_steps = Steps(steps=steps, name=f"{self.name}_Root")
        self.runtime_config = BaseRuntimeConfig(stateless=True)
        self.persona = None

        self.blackboard_schema = None
        self.initial_blackboard_state = None
        if blackboard is not None:
            if isinstance(blackboard, type) and issubclass(blackboard, BaseModel):
                self.blackboard_schema = blackboard
            elif isinstance(blackboard, BaseModel):
                self.blackboard_schema = type(blackboard)
                self.initial_blackboard_state = blackboard
            else:
                raise ValueError(
                    "blackboard 参数必须是 Pydantic BaseModel 的子类(类型)或其实例"
                )

    def _build_result(
        self,
        initial_input: StepInput,
        safe_context: RunContext,
        final_output: StepOutput,
    ) -> WorkflowRunResult:
        """根据执行链上的全量输出构建最终的工作流执行结果对象"""
        flat_outputs = {}

        def _extract(out: StepOutput):
            flat_outputs[out.step_name] = out
            if out.steps:
                for o in out.steps:
                    _extract(o)

        if final_output:
            _extract(final_output)

        status = "completed" if final_output and final_output.success else "error"

        return WorkflowRunResult(
            workflow_id=self.id,
            workflow_name=self.name,
            status=status,
            original_input=initial_input.input,
            state=safe_context.state,
            step_outputs=flat_outputs,
            last_step_content=final_output.content if final_output else None,
            final_output=final_output,
        )

    def bind(self, **kwargs: Any) -> Any:
        """DI 注入语法糖"""

        async def _dependency() -> "Workflow":
            return self

        return Depends(_dependency)

    async def reply(
        self,
        prompt: PromptInput | None = None,
        reply_to: bool = False,
        *,
        context: RunContext | None = None,
        **kwargs: Any,
    ) -> WorkflowRunResult:
        """
        工作流交互执行语法糖，隐式提取上下文并自动将最终流水线产出发送回复给用户。

        参数:
            prompt: 传入工作流入口根节点的初始参数或指令。
            reply_to: 是否将结果作为回复消息发送 (at用户或引用原消息)。
            context: 显式传入的会话与运行上下文。
            kwargs: 追加的工作流附带参数 (additional_data)。

        返回:
            WorkflowRunResult: 包含执行状态、断点快照、各节点产出的全量工作流结果对象。
        """
        ctx = context or RunContext()
        bot = ctx.get_bot()
        event = ctx.get_event()

        res = await self.run(prompt=prompt, context=ctx, **kwargs)

        if bot and event:
            if res.status == "completed" and res.final_output:
                msg = (
                    str(res.final_output.content)
                    if res.final_output.content
                    else "执行完毕"
                )
                await MessageUtils.build_message(msg).send(reply_to=reply_to)
            elif res.status == "error":
                err_msg = res.final_output.error if res.final_output else "未知异常"
                await MessageUtils.build_message(
                    f"❌ 工作流执行发生错误: {err_msg}"
                ).send(reply_to=reply_to)

        return res

    async def run(
        self,
        prompt: PromptInput | None = None,
        *,
        context: RunContext | None = None,
        **kwargs: Any,
    ) -> WorkflowRunResult:
        """
        工作流单次运行阻塞核心入口，遍历所有图元节点直至终止。

        参数:
            prompt: 传入工作流入口根节点的初始参数或指令。
            context: 显式传入的会话与运行上下文。
            kwargs: 追加的工作流附带参数 (additional_data)。

        返回:
            WorkflowRunResult: 包含执行状态、断点快照、各节点产出的全量工作流结果对象。
        """
        session_id = (
            context.session_id if context and context.session_id else f"wf_{self.id}"
        )
        safe_context = context or RunContext(session_id=session_id)

        if self.blackboard_schema and not safe_context.session.blackboard:
            safe_context.session.blackboard = BlackboardManager(
                schema=self.blackboard_schema,
                initial_state=self.initial_blackboard_state,
            )

        logger.debug(f"🏭 **工作流 [{self.name}] 启动**")

        initial_input = StepInput(input=prompt)
        if kwargs:
            initial_input.additional_data.update(kwargs)

        try:
            final_output = await self.root_steps.aexecute(initial_input, safe_context)

            logger.debug(f"🏭 **工作流 [{self.name}] 运行结束**")

            return self._build_result(initial_input, safe_context, final_output)

        except BaseException as e:
            if isinstance(e, ControlFlowExit):
                logger.debug(f"⏭️ 工作流执行被业务控制流安全中止: {e}")
                dummy_output = StepOutput(content=str(e), success=False)
                return self._build_result(initial_input, safe_context, dummy_output)

            raise e

    @contextlib.asynccontextmanager
    async def run_stream(
        self,
        prompt: PromptInput | None = None,
        *,
        context: RunContext | None = None,
        **kwargs: Any,
    ) -> AsyncIterator["StreamedRunResult[Any]"]:
        """对齐 BaseRunnable 接口的流式上下文管理器"""
        event_bus = EventBus()
        if context:
            context.run.event_bus = event_bus

        async def _execution_task():
            try:
                async for event in self._internal_stream(prompt, context, **kwargs):
                    await event_bus.emit(event)
            except BaseException as e:
                await event_bus.emit(AgentRunError(error=e))
            finally:
                await event_bus.end()

        task = asyncio.create_task(_execution_task())
        try:
            yield StreamedRunResult[Any](event_bus)
        finally:
            if not task.done():
                task.cancel()

    async def _internal_stream(
        self,
        prompt: PromptInput | None = None,
        context: RunContext | None = None,
        **kwargs: Any,
    ) -> AsyncIterator[Any]:
        """流式执行工作流节点树的内部实现"""
        session_id = (
            context.session_id if context and context.session_id else f"wf_{self.id}"
        )
        safe_context = context or RunContext(session_id=session_id)

        if self.blackboard_schema and not safe_context.session.blackboard:
            safe_context.session.blackboard = BlackboardManager(
                schema=self.blackboard_schema,
                initial_state=self.initial_blackboard_state,
            )

        logger.debug(f"🏭 **工作流 [{self.name}] 启动**")

        initial_input = StepInput(input=prompt)
        if kwargs:
            initial_input.additional_data.update(kwargs)

        try:
            final_output = None
            async for event in self.root_steps.aexecute_stream(
                initial_input, safe_context
            ):
                if isinstance(event, StepOutput):
                    final_output = event
                else:
                    yield event

            if final_output:
                logger.debug(f"🏭 **工作流 [{self.name}] 运行结束**")

                wf_result = self._build_result(
                    initial_input, safe_context, final_output
                )
                agent_res = AgentRunResult(
                    output=wf_result.last_step_content,
                    structured_data=wf_result,
                    usage=UsageInfo(),
                )
                yield AgentRunEnd(result=agent_res)
        except Exception:
            pass

    def as_tool(self, tool_name: str | None = None) -> FunctionTool:
        """将工作流封装并导出为可供 Agent 直接调用的 FunctionTool 实例"""

        async def _execute_workflow_tool(prompt: str, context: RunContext) -> str:
            run_result = await self.run(prompt=prompt, context=context)
            output = run_result.final_output

            if output and output.success:
                return (
                    f"工作流 [{self.name}] 执行完毕。最终流水线产出:\n{output.content}"
                )

            raise ToolRetryError(
                f"工作流执行失败: {output.error if output else 'unknown'}，"
                "请尝试换种方式处理。"
            )

        final_tool_name = tool_name or f"trigger_workflow_{self.id}"

        tool_desc = (
            f"触发执行专属流水线: {self.name}。\n"
            f"描述: {self.description}\n"
            f"注意：如果你认为该工作流能完全解决用户的问题，请立刻调用此工具，"
            f"并将用户的诉求提炼后作为 prompt 传入。"
        )

        return FunctionTool(
            func=_execute_workflow_tool, name=final_tool_name, description=tool_desc
        )