ai-video-fullstack/backend/services/workflow_router.py

"""Small LLM router for Workflow conditional edges.

The router deliberately uses a separate, short completion.  Its only output is
a required function choice, so the current Agent cannot speak before the graph
has decided whether the user turn belongs to another node.
"""

from __future__ import annotations

import json
from collections.abc import Callable
from copy import deepcopy
from typing import Any

from loguru import logger
from models import AssistantConfig
from openai import AsyncOpenAI

from services.workflow.models import LLMRouteResult, RouteStatus


STAY_ON_CURRENT_NODE = "workflow_stay_on_current_node"
# Compatibility alias for callers saved before all source nodes supported LLM edges.
STAY_ON_CURRENT_AGENT = STAY_ON_CURRENT_NODE
MAX_ROUTING_HISTORY_ENTRIES = 20


def _routing_user_message(
    routing_input: str,
    current_user_message: dict[str, Any] | None,
) -> dict[str, Any]:
    """Combine routing metadata with the current text or multimodal turn."""
    if not current_user_message:
        return {"role": "user", "content": routing_input}

    content = current_user_message.get("content")
    if not isinstance(content, list):
        current_text = str(content or "").strip()
        suffix = f"\n\n[当前用户输入]\n{current_text}" if current_text else ""
        return {"role": "user", "content": f"{routing_input}{suffix}"}

    return {
        "role": "user",
        "content": [
            {
                "type": "text",
                "text": f"{routing_input}\n\n[当前用户输入如下]",
            },
            *deepcopy(content),
        ],
    }


class WorkflowLLMRouter:
    """Select one LLM edge without allowing the router to speak."""

    def __init__(self, cfg: AssistantConfig):
        self._cfg = cfg

    async def select_edge(
        self,
        *,
        node_name: str,
        node_prompt: str,
        edges: list[dict[str, Any]],
        history: list[dict[str, str]],
        variables: dict[str, Any],
        edge_name: Callable[[dict[str, Any]], str],
        edge_description: Callable[[dict[str, Any]], str],
        current_user_message: dict[str, Any] | None = None,
    ) -> LLMRouteResult:
        """Return a typed match, no-match or technical error."""
        if not edges:
            return LLMRouteResult(status=RouteStatus.NO_MATCH)

        names = {edge_name(edge) for edge in edges}
        stay_name = STAY_ON_CURRENT_NODE
        while stay_name in names:
            stay_name = f"_{stay_name}"

        tools = [
            {
                "type": "function",
                "function": {
                    "name": edge_name(edge),
                    "description": edge_description(edge),
                    "parameters": {"type": "object", "properties": {}},
                },
            }
            for edge in edges
        ]
        tools.append(
            {
                "type": "function",
                "function": {
                    "name": stay_name,
                    "description": "所有转移条件都不满足，留在当前节点。",
                    "parameters": {"type": "object", "properties": {}},
                },
            }
        )

        ordered_conditions = "\n".join(
            f"{index + 1}. {edge_description(edge)}"
            for index, edge in enumerate(edges)
        )
        router_prompt = (
            "你是工作流路由器，不是对话助手。收到一轮完整用户输入后，"
            "必须且只能调用一个提供的函数，禁止输出任何口头回复。\n"
            "按给出的顺序判断转移条件；选择第一个明确满足的转移函数。"
            "如果没有条件满足，调用留在当前节点的函数。\n\n"
            f"当前节点：{node_name}\n"
            f"当前节点任务：{node_prompt or '未配置'}\n"
            f"转移条件：\n{ordered_conditions}"
        )
        # WorkflowBrain records the current turn before routing. When the full
        # current message is supplied separately, keep only earlier history so
        # the text is not duplicated and the image remains attached to its turn.
        routing_history = (
            history[:-1] if current_user_message and history else history
        )
        recent_history = routing_history[-MAX_ROUTING_HISTORY_ENTRIES:]
        routing_input = json.dumps(
            {
                "conversation": recent_history,
                "session_variables": variables,
            },
            ensure_ascii=False,
            separators=(",", ":"),
        )
        extra_body = self._cfg.llm_values.get("extraBody")
        request_extra = (
            {"extra_body": extra_body} if isinstance(extra_body, dict) else {}
        )
        client = AsyncOpenAI(
            api_key=self._cfg.llm_api_key,
            base_url=self._cfg.llm_base_url,
            timeout=15.0,
        )
        try:
            response = await client.chat.completions.create(
                model=self._cfg.model,
                messages=[
                    {"role": "system", "content": router_prompt},
                    _routing_user_message(routing_input, current_user_message),
                ],
                tools=tools,
                tool_choice="required",
                temperature=0,
                **request_extra,
            )
            tool_calls = response.choices[0].message.tool_calls or []
            if not tool_calls:
                logger.warning("Workflow 路由 LLM 未返回函数调用，留在当前节点")
                return LLMRouteResult(
                    status=RouteStatus.ERROR,
                    error="路由模型没有返回函数调用",
                )
            selected = str(tool_calls[0].function.name or "")
            if selected == stay_name:
                return LLMRouteResult(status=RouteStatus.NO_MATCH)
            if selected not in names:
                logger.warning(f"Workflow 路由 LLM 返回未知函数:{selected}")
                return LLMRouteResult(
                    status=RouteStatus.ERROR,
                    error=f"路由模型返回未知函数:{selected}",
                )
            return LLMRouteResult(
                status=RouteStatus.MATCHED,
                function_name=selected,
            )
        except Exception as exc:  # noqa: BLE001 - routing failure must not end the call
            logger.warning(f"Workflow LLM 边判断失败，留在当前节点:{exc}")
            return LLMRouteResult(
                status=RouteStatus.ERROR,
                error=f"大模型路由失败:{type(exc).__name__}",
            )
        finally:
            await client.close()