180 lines
6.6 KiB
Python
180 lines
6.6 KiB
Python
"""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()
|