ai-video-fullstack/backend/services/brains/workflow_brain.py

"""Local graph-driven workflow assistant and its per-call state."""

from __future__ import annotations

from dataclasses import dataclass
from typing import Any

from loguru import logger
from models import AssistantConfig
from pipecat.adapters.schemas.function_schema import FunctionSchema
from pipecat.frames.frames import OutputTransportMessageUrgentFrame, TTSSpeakFrame
from pipecat.processors.aggregators.llm_context import LLMContext
from pipecat.processors.frame_processor import FrameProcessor

from services.brains.base import BaseBrain, BrainRuntime, BrainSpec
from services.workflow_engine import WorkflowEngine


@dataclass
class WorkflowState:
    current: str
    ended: bool = False
    turns_in_node: int = 0
    end_turn_id: str | None = None


class WorkflowBrain(BaseBrain):
    spec = BrainSpec(
        type="workflow",
        supported_runtime_modes=frozenset({"pipeline"}),
        owns_context=True,
    )
    _FALLBACK_AFTER_TURNS = 2

    def __init__(self, graph: dict[str, Any]):
        self._engine = WorkflowEngine(graph or {})
        if not self._engine.has_graph() or not self._engine.start_id:
            raise ValueError("WorkflowBrain 缺少有效的 startCall 节点")
        self._state = WorkflowState(current=self._engine.start_id)
        self._history: list[dict[str, str]] = []
        self._cfg: AssistantConfig | None = None
        self._runtime: BrainRuntime | None = None

    async def greeting(self, cfg: AssistantConfig) -> str:
        return self._engine.greeting() or cfg.greeting

    def system_prompt(self, cfg: AssistantConfig) -> str:
        return self._engine.system_prompt_for(self._state.current)

    def build_llm(self, cfg: AssistantConfig, context: LLMContext) -> FrameProcessor:
        from services.pipecat.service_factory import create_llm

        return create_llm(cfg)

    async def setup(self, cfg: AssistantConfig, runtime: BrainRuntime) -> None:
        self._cfg = cfg
        self._runtime = runtime
        for edge in self._engine.edges:
            if edge.get("target"):
                runtime.llm.register_function(
                    self._engine.edge_fn_name(edge),
                    self._make_transition_handler(edge),
                )
        self._apply_node(self._state.current)
        logger.info(
            f"工作流模式启用: 起始节点={self._engine.name(self._state.current)}"
        )

    async def on_connected(self) -> None:
        await self._emit_node_active(self._state.current)

    def record_user_message(self, content: str) -> None:
        if content:
            self._history.append({"role": "user", "content": content})

    async def on_assistant_text_start(self, turn_id: str) -> None:
        if self._state.ended and self._state.end_turn_id is None:
            self._state.end_turn_id = turn_id

    async def on_assistant_text_end(
        self,
        turn_id: str,
        content: str,
        interrupted: bool,
    ) -> None:
        if not content or interrupted:
            return
        self._history.append({"role": "assistant", "content": content})
        if turn_id == self._state.end_turn_id:
            runtime = self._require_runtime()
            runtime.call_end.begin("completed")
            runtime.call_end.arm_after_speech()
        elif not self._state.ended:
            self._state.turns_in_node += 1
            await self._fallback_route()

    def _apply_node(self, node_id: str) -> None:
        runtime = self._require_runtime()
        runtime.set_system_prompt(self._engine.system_prompt_for(node_id))
        if self._engine.is_end(node_id):
            runtime.set_tools([])
            return
        runtime.set_tools(
            [
                FunctionSchema(
                    name=self._engine.edge_fn_name(edge),
                    description=self._engine.edge_description(edge),
                    properties={},
                    required=[],
                )
                for edge in self._engine.outgoing(node_id)
            ]
        )

    async def _go_to_node(self, target: str) -> None:
        self._state.current = target
        self._state.turns_in_node = 0
        if self._engine.is_end(target):
            self._state.ended = True
        await self._emit_node_active(target)
        self._apply_node(target)

    async def _emit_node_active(self, node_id: str | None) -> None:
        if node_id:
            await self._require_runtime().queue_frame(
                OutputTransportMessageUrgentFrame(
                    message={"type": "node-active", "nodeId": node_id}
                )
            )

    async def _speak_transition(self, edge: dict | None) -> None:
        speech = self._engine.edge_transition_speech(edge)
        if speech:
            await self._require_runtime().queue_frame(
                TTSSpeakFrame(speech, append_to_context=False)
            )

    def _make_transition_handler(self, edge: dict):
        target = str(edge.get("target"))

        async def handler(params) -> None:
            logger.info(f"LLM 触发转移 → {self._engine.name(target)}")
            if not self._engine.is_end(target):
                await self._speak_transition(edge)
            await self._go_to_node(target)
            await params.result_callback({"status": "ok"})

        return handler

    async def _fallback_route(self) -> None:
        if self._state.ended:
            return
        if self._state.turns_in_node < self._FALLBACK_AFTER_TURNS:
            return
        if not self._engine.outgoing(self._state.current):
            return

        cfg = self._require_config()
        target = await self._engine.route(
            self._state.current,
            self._history,
            api_key=self._require(cfg.llm_api_key, "LLM apiKey"),
            base_url=self._require(cfg.llm_base_url, "LLM apiUrl"),
            model=self._require(cfg.model, "LLM modelId"),
        )
        if target and target != self._state.current:
            logger.info(f"文本兜底触发转移 → {self._engine.name(target)}")
            if not self._engine.is_end(target):
                await self._speak_transition(
                    self._engine.find_edge(self._state.current, target)
                )
            await self._go_to_node(target)

    def _require_runtime(self) -> BrainRuntime:
        if self._runtime is None:
            raise RuntimeError("WorkflowBrain 尚未绑定 pipeline runtime")
        return self._runtime

    def _require_config(self) -> AssistantConfig:
        if self._cfg is None:
            raise RuntimeError("WorkflowBrain 尚未初始化配置")
        return self._cfg

    @staticmethod
    def _require(value: str, label: str) -> str:
        if value:
            return value
        raise ValueError(f"缺少模型资源配置: {label}")