Files
ai-video-fullstack/backend/tests/test_brains.py

1448 lines
51 KiB
Python
Raw Blame History

This file contains ambiguous Unicode characters
This file contains Unicode characters that might be confused with other characters. If you think that this is intentional, you can safely ignore this warning. Use the Escape button to reveal them.
from __future__ import annotations
import unittest
from types import SimpleNamespace
from unittest.mock import patch
from models import AssistantConfig, RuntimeTool
from pipecat.frames.frames import (
LLMContextFrame,
LLMFullResponseEndFrame,
LLMFullResponseStartFrame,
LLMMessagesAppendFrame,
LLMMessagesUpdateFrame,
LLMRunFrame,
LLMTextFrame,
OutputTransportMessageUrgentFrame,
TTSSpeakFrame,
)
from pipecat.processors.aggregators.llm_context import LLMContext
from pipecat.processors.frame_processor import FrameDirection
from schemas import AssistantUpsert, REALTIME_CAPABLE_TYPES
from services.brains import BrainRuntime, SPECS, build_brain
from services.brains.base import GREETING_CONTEXT_MARKER
from services.brains.dify_llm import (
DifyLLMService,
last_user_text,
normalize_api_base,
)
from services.brains.workflow_brain import WorkflowBrain
from services.runtime_variables import prepare_dynamic_config
from services.workflow.models import LLMRouteResult, RouteStatus
class FakeLLM:
def __init__(self):
self.functions = {}
def register_function(self, name, handler):
self.functions[name] = handler
class FakeCallEnd:
def __init__(self):
self.ending = False
self.reason = ""
self.armed = False
self.finished = False
self.response_started = False
self.waited_for_text: bool | None = None
self.tracked_speeches = 0
def begin(self, reason: str) -> None:
self.ending = True
self.reason = reason
def begin_response(self) -> None:
self.response_started = True
def arm_after_speech(self) -> None:
self.armed = True
def track_speech(self) -> None:
self.tracked_speeches += 1
async def arm_after_tracked_speech(self) -> None:
self.armed = True
if self.tracked_speeches == 0:
await self.finish()
async def finish_after_current_speech(self, *, has_text: bool) -> None:
self.waited_for_text = has_text
if has_text:
self.armed = True
else:
await self.finish()
async def finish(self) -> None:
self.finished = True
class FakeFunctionParams:
def __init__(self, arguments=None):
self.arguments = arguments or {}
self.result = None
self.properties = None
async def result_callback(self, result, properties=None):
self.result = result
self.properties = properties
class BrainRegistryTests(unittest.TestCase):
def test_capability_matrix(self):
self.assertEqual(
{
name: spec.supported_runtime_modes
for name, spec in SPECS.items()
},
{
"prompt": frozenset({"pipeline", "realtime"}),
"workflow": frozenset({"pipeline"}),
"dify": frozenset({"pipeline"}),
"fastgpt": frozenset({"pipeline"}),
},
)
self.assertEqual(
REALTIME_CAPABLE_TYPES,
{
name
for name, spec in SPECS.items()
if "realtime" in spec.supported_runtime_modes
},
)
def test_unknown_brain_does_not_fallback_to_prompt(self):
with self.assertRaisesRegex(ValueError, "尚未实现"):
build_brain(AssistantConfig(type="opencode"))
def test_workflow_realtime_is_rejected_at_schema_boundary(self):
with self.assertRaises(ValueError):
AssistantUpsert(
name="workflow",
type="workflow",
runtimeMode="realtime",
)
def test_prompt_realtime_keeps_dynamic_variable_definitions(self):
assistant = AssistantUpsert(
name="realtime prompt",
type="prompt",
runtimeMode="realtime",
dynamicVariableDefinitions={
"user_name": {
"type": "string",
"required": True,
"default": None,
}
},
)
self.assertIn("user_name", assistant.dynamic_variable_definitions)
def test_workflow_keeps_dynamic_variables_and_tool_bindings(self):
assistant = AssistantUpsert(
name="workflow",
type="workflow",
toolIds=["tool_a"],
dynamicVariableDefinitions={
"customer": {"type": "string", "required": False, "default": "王先生"}
},
graph={},
)
self.assertEqual(assistant.tool_ids, ["tool_a"])
self.assertIn("customer", assistant.dynamic_variable_definitions)
class DifyHelpersTests(unittest.TestCase):
def test_normalize_api_base(self):
self.assertEqual(
normalize_api_base("https://api.dify.ai"),
"https://api.dify.ai/v1",
)
self.assertEqual(
normalize_api_base("https://example.test/v1/chat-messages"),
"https://example.test/v1",
)
def test_last_user_text(self):
self.assertEqual(
last_user_text(
[
{"role": "user", "content": "first"},
{"role": "assistant", "content": "answer"},
{
"role": "user",
"content": [{"type": "text", "text": "latest"}],
},
]
),
"latest",
)
class DifyLLMServiceTests(unittest.IsolatedAsyncioTestCase):
async def test_streams_sdk_events_and_keeps_conversation_id(self):
class FakeDifyClient:
requests = []
async def achat_messages(self, request, **_kwargs):
self.requests.append(request)
async def events():
yield SimpleNamespace(
event="message",
answer="你好",
conversation_id="conversation-1",
)
yield SimpleNamespace(
event="message_end",
conversation_id="conversation-1",
)
return events()
client = FakeDifyClient()
service = DifyLLMService(
AssistantConfig(type="dify"),
client=client,
user_id="test-user",
)
frames = []
async def push_frame(frame, *_args, **_kwargs):
frames.append(frame)
service.push_frame = push_frame
context = LLMContext(messages=[{"role": "user", "content": "问题"}])
await service.process_frame(
LLMContextFrame(context),
FrameDirection.DOWNSTREAM,
)
self.assertIsInstance(frames[0], LLMFullResponseStartFrame)
self.assertIsInstance(frames[1], LLMTextFrame)
self.assertEqual(frames[1].text, "你好")
self.assertIsInstance(frames[-1], LLMFullResponseEndFrame)
self.assertEqual(service._conversation_id, "conversation-1")
context.add_message({"role": "user", "content": "追问"})
await service.process_frame(
LLMContextFrame(context),
FrameDirection.DOWNSTREAM,
)
self.assertEqual(client.requests[-1].conversation_id, "conversation-1")
class PromptBrainTests(unittest.IsolatedAsyncioTestCase):
async def test_realtime_prompt_brain_renders_dynamic_variables(self):
cfg = prepare_dynamic_config(
AssistantConfig(
type="prompt",
runtimeMode="realtime",
prompt="服务用户 {{user_name}}",
greeting="您好,{{user_name}}",
dynamic_variable_definitions={
"user_name": {"type": "string", "required": True}
},
),
{"user_name": "王先生"},
assistant_id="asst_realtime",
)
brain = build_brain(cfg)
self.assertEqual(brain.system_prompt(cfg), "服务用户 王先生")
self.assertEqual(await brain.greeting(cfg), "您好,王先生")
context = LLMContext(
messages=[{"role": "system", "content": brain.system_prompt(cfg)}]
)
brain.prepare_greeting_context("您好,王先生", context)
self.assertEqual(
[message["role"] for message in context.get_messages()],
["system", "system"],
)
self.assertNotEqual(context.get_messages()[0]["role"], "assistant")
self.assertEqual(
context.get_messages()[1]["content"],
f"{GREETING_CONTEXT_MARKER}\n您好,王先生",
)
async def test_end_call_tool_is_owned_by_prompt_brain(self):
brain = build_brain(
AssistantConfig(
type="prompt",
tools=[
RuntimeTool(
id="end-call",
name="结束通话",
function_name="end_call",
type="end_call",
definition={
"config": {
"message_type": "none",
"capture_reason": True,
}
},
)
],
)
)
llm = FakeLLM()
call_end = FakeCallEnd()
visible_tools = []
async def queue_frame(_frame):
pass
await brain.setup(
AssistantConfig(
type="prompt",
tools=[
RuntimeTool(
id="end-call",
name="结束通话",
function_name="end_call",
type="end_call",
definition={"config": {"capture_reason": True}},
)
],
),
BrainRuntime(
context=LLMContext(messages=[]),
llm=llm,
queue_frame=queue_frame,
set_system_prompt=lambda _prompt: None,
set_tools=lambda tools: visible_tools.extend(tools or []),
call_end=call_end,
),
)
self.assertEqual(visible_tools[0].name, "end_call")
params = FakeFunctionParams({"reason": "用户已完成咨询"})
await brain.on_assistant_text_start("closing-turn")
await llm.functions["end_call"](params)
self.assertEqual(call_end.reason, "用户已完成咨询")
self.assertFalse(call_end.finished)
await brain.on_assistant_text_end("closing-turn", "", False)
self.assertTrue(call_end.finished)
self.assertFalse(call_end.waited_for_text)
self.assertTrue(call_end.response_started)
self.assertEqual(params.result["action"], "ending_call")
async def test_end_call_waits_for_prompt_generated_closing_speech(self):
tool = RuntimeTool(
id="end-call",
name="结束通话",
function_name="end_call",
type="end_call",
definition={"config": {"message_type": "none"}},
)
cfg = AssistantConfig(type="prompt", tools=[tool])
brain = build_brain(cfg)
llm = FakeLLM()
call_end = FakeCallEnd()
await brain.setup(
cfg,
BrainRuntime(
context=LLMContext(messages=[]),
llm=llm,
queue_frame=lambda _frame: None,
set_system_prompt=lambda _prompt: None,
set_tools=lambda _tools: None,
call_end=call_end,
),
)
await brain.on_assistant_text_start("closing-turn")
await llm.functions["end_call"](FakeFunctionParams({}))
await brain.on_assistant_text_end(
"closing-turn",
"感谢您的来电。祝您生活愉快,再见。",
False,
)
self.assertFalse(call_end.finished)
self.assertTrue(call_end.armed)
self.assertTrue(call_end.waited_for_text)
async def test_http_tool_renders_secrets_and_updates_prompt_variable(self):
requests = []
class FakeResponse:
status_code = 200
content = b'{"order":{"status":"paid"}}'
def raise_for_status(self):
return None
def json(self):
return {"order": {"status": "paid"}}
class FakeClient:
def __init__(self, **_kwargs):
pass
async def __aenter__(self):
return self
async def __aexit__(self, *_args):
return None
async def request(self, method, url, **kwargs):
requests.append((method, url, kwargs))
return FakeResponse()
cfg = prepare_dynamic_config(
AssistantConfig(
type="prompt",
runtimeMode="pipeline",
prompt="订单状态:{{order_status}}",
dynamic_variable_definitions={
"order_status": {"type": "string", "default": "unknown"}
},
tools=[
RuntimeTool(
id="lookup",
name="查询订单",
function_name="lookup_order",
type="http",
description="查询订单状态",
definition={
"config": {
"method": "GET",
"url": "https://example.test/orders/{order_id}",
"headers": {"Authorization": "Bearer {{secret__token}}"},
"parameters": [
{
"name": "order_id",
"type": "string",
"location": "path",
"required": True,
},
{
"name": "Authorization",
"type": "string",
"location": "header",
"required": False,
},
],
"dynamic_variable_assignments": {
"order_status": "response.order.status"
},
}
},
secrets={"dynamic_variables": {"secret__token": "server-token"}},
)
],
),
{},
assistant_id="asst_1",
)
brain = build_brain(cfg)
llm = FakeLLM()
prompts = []
visible_tools = []
async def queue_frame(_frame):
pass
await brain.setup(
cfg,
BrainRuntime(
context=LLMContext(messages=[]),
llm=llm,
queue_frame=queue_frame,
set_system_prompt=prompts.append,
set_tools=lambda tools: visible_tools.extend(tools or []),
call_end=FakeCallEnd(),
),
)
params = FakeFunctionParams(
{"order_id": "A/1", "Authorization": "attacker-value"}
)
with patch("services.tool_executor.httpx.AsyncClient", FakeClient):
await llm.functions["lookup_order"](params)
self.assertEqual(requests[0][1], "https://example.test/orders/A%2F1")
self.assertEqual(
requests[0][2]["headers"]["Authorization"], "Bearer server-token"
)
self.assertEqual(params.result["updated_variables"], ["order_status"])
self.assertEqual(prompts[-1], "订单状态paid")
class WorkflowBrainTests(unittest.IsolatedAsyncioTestCase):
async def test_agent_vision_tool_is_scoped_to_effective_stage(self):
brain = WorkflowBrain(
{
"specVersion": 3,
"settings": {
"globalPrompt": "全局规则",
"defaultLlmResourceId": "llm_global",
"visionEnabled": True,
"visionModelResourceId": "vision_global",
},
"nodes": [
{"id": "start", "type": "start", "data": {}},
{
"id": "agent",
"type": "agent",
"data": {
"prompt": "观察用户需要展示的物品",
"inheritGlobalConfig": True,
},
},
],
"edges": [
{
"id": "begin",
"source": "start",
"target": "agent",
"data": {"mode": "always", "priority": 0},
}
],
}
)
scopes = []
vision_function = object()
async def queue_frame(_frame):
pass
brain._runtime = BrainRuntime(
context=LLMContext(messages=[]),
llm=FakeLLM(),
queue_frame=queue_frame,
set_system_prompt=lambda _prompt: None,
set_tools=lambda _tools: None,
call_end=FakeCallEnd(),
set_vision_scope=scopes.append,
vision_function=vision_function,
)
await brain._apply_agent_stage("agent")
inherited_config = brain._agent_config("agent")
self.assertIn(vision_function, inherited_config["functions"])
self.assertIn("fetch_user_image", inherited_config["role_message"])
self.assertEqual(
scopes[-1],
{
"enabled": True,
"vision_model_resource_id": "vision_global",
"llm_resource_id": "llm_global",
},
)
brain._engine.data("agent").update(
{
"inheritGlobalConfig": False,
"llmResourceId": "llm_agent",
"visionEnabled": False,
"visionModelResourceId": "",
}
)
await brain._apply_agent_stage("agent")
custom_config = brain._agent_config("agent")
self.assertNotIn(vision_function, custom_config["functions"])
self.assertNotIn("fetch_user_image", custom_config["role_message"])
self.assertFalse(scopes[-1]["enabled"])
async def test_initial_fixed_speech_waits_for_start_greeting_to_finish(self):
brain = WorkflowBrain(
{
"specVersion": 3,
"settings": {"globalPrompt": "全局规则"},
"nodes": [
{
"id": "start",
"type": "start",
"data": {"greeting": "欢迎使用"},
},
{
"id": "agent",
"type": "agent",
"data": {
"prompt": "收集用户信息",
"entryMode": "fixed_speech",
"entrySpeech": "请问您怎么称呼?",
},
},
],
"edges": [
{
"id": "begin",
"source": "start",
"target": "agent",
"data": {"mode": "always", "priority": 0},
}
],
}
)
queued = []
async def queue_frame(frame):
queued.append(frame)
class FakeManager:
def __init__(self):
self.current_node = None
async def initialize(self, config):
self.current_node = config["name"]
async def set_node_from_config(self, config):
self.current_node = config["name"]
for action in config.get("pre_actions", []):
await action["handler"](action, self)
brain._runtime = BrainRuntime(
context=LLMContext(messages=[]),
llm=FakeLLM(),
queue_frame=queue_frame,
set_system_prompt=lambda _prompt: None,
set_tools=lambda _tools: None,
call_end=FakeCallEnd(),
)
brain._manager = FakeManager()
await brain.on_connected(greeting_pending=True)
self.assertEqual(brain._manager.current_node, "start")
self.assertFalse(any(isinstance(frame, TTSSpeakFrame) for frame in queued))
await brain.on_greeting_finished()
self.assertEqual(brain._manager.current_node, "agent")
fixed_speech_frames = [
frame for frame in queued if isinstance(frame, TTSSpeakFrame)
]
self.assertEqual(len(fixed_speech_frames), 1)
self.assertEqual(fixed_speech_frames[0].text, "请问您怎么称呼?")
# Playback notifications may be duplicated by a transport reconnect;
# the initial entry behavior must still run only once.
await brain.on_greeting_finished()
self.assertEqual(
len([frame for frame in queued if isinstance(frame, TTSSpeakFrame)]),
1,
)
async def test_action_publishes_updated_session_variables(self):
tool = RuntimeTool(
id="lookup",
name="查询订单",
function_name="lookup_order",
type="http",
)
cfg = prepare_dynamic_config(
AssistantConfig(
type="workflow",
graph={
"specVersion": 3,
"settings": {},
"nodes": [
{"id": "start", "type": "start", "data": {}},
{
"id": "lookup_action",
"type": "action",
"data": {
"toolId": "lookup",
"resultAssignments": {
"order_status": "order.status"
},
},
},
],
"edges": [],
},
dynamic_variable_definitions={
"order_status": {"type": "string", "default": "pending"}
},
tools=[tool],
),
{},
assistant_id="asst_workflow_action",
)
brain = WorkflowBrain(cfg)
queued = []
async def queue_frame(frame):
queued.append(frame)
async def execute(_tool, _arguments, *, result_assignments=None):
self.assertEqual(result_assignments, {"order_status": "order.status"})
brain._store.assign("order_status", "paid")
return {
"status": "ok",
"updated_variables": ["order_status"],
}
brain._runtime = BrainRuntime(
context=LLMContext(messages=[]),
llm=FakeLLM(),
queue_frame=queue_frame,
set_system_prompt=lambda _prompt: None,
set_tools=lambda _tools: None,
call_end=FakeCallEnd(),
)
brain._tools.execute = execute
await brain._enter_action("lookup_action")
variable_events = [
frame.message
for frame in queued
if isinstance(frame, OutputTransportMessageUrgentFrame)
and frame.message.get("type") == "workflow-variables"
]
self.assertEqual(variable_events[-1]["reason"], "action")
self.assertEqual(variable_events[-1]["changed"], ["order_status"])
self.assertEqual(variable_events[-1]["variables"], {"order_status": "paid"})
async def test_nodes_without_outgoing_edges_remain_active(self):
queued = []
async def queue_frame(frame):
queued.append(frame)
runtime = BrainRuntime(
context=LLMContext(messages=[]),
llm=FakeLLM(),
queue_frame=queue_frame,
set_system_prompt=lambda _prompt: None,
set_tools=lambda _tools: None,
call_end=FakeCallEnd(),
)
class FakeManager:
def __init__(self, current_node=None):
self.current_node = current_node
async def initialize(self, config):
self.current_node = config["name"]
start_brain = WorkflowBrain(
{
"specVersion": 3,
"settings": {},
"nodes": [{"id": "start", "type": "start", "data": {}}],
"edges": [],
}
)
start_brain._runtime = runtime
start_brain._manager = FakeManager()
await start_brain.on_connected()
self.assertEqual(start_brain._manager.current_node, "start")
queued.clear()
await start_brain.on_client_ready()
ready_messages = [
frame.message
for frame in queued
if isinstance(frame, OutputTransportMessageUrgentFrame)
]
self.assertIn(
{"type": "node-active", "nodeId": "start"},
ready_messages,
)
self.assertTrue(
any(
message.get("type") == "workflow-variables"
and message.get("reason") == "client_ready"
and message.get("nodeId") == "start"
for message in ready_messages
)
)
agent_brain = WorkflowBrain(
{
"specVersion": 3,
"settings": {"globalPrompt": "全局规则"},
"nodes": [
{"id": "start", "type": "start", "data": {}},
{
"id": "agent",
"type": "agent",
"data": {"prompt": "持续回答"},
},
],
"edges": [
{
"id": "begin",
"source": "start",
"target": "agent",
"data": {"mode": "always", "priority": 0},
}
],
}
)
agent_brain._runtime = runtime
agent_brain._manager = FakeManager("agent")
queued.clear()
await agent_brain.on_client_ready()
self.assertTrue(
any(
isinstance(frame, OutputTransportMessageUrgentFrame)
and frame.message
== {"type": "node-active", "nodeId": "agent"}
for frame in queued
)
)
queued.clear()
handled = await agent_brain.on_user_turn_end("请继续回答")
self.assertTrue(handled)
self.assertEqual(agent_brain._manager.current_node, "agent")
self.assertTrue(any(isinstance(frame, LLMRunFrame) for frame in queued))
handoff_brain = WorkflowBrain(
{
"specVersion": 3,
"settings": {},
"nodes": [
{"id": "start", "type": "start", "data": {}},
{
"id": "handoff",
"type": "handoff",
"data": {"targetType": "human"},
},
],
"edges": [],
}
)
handoff_brain._runtime = runtime
handoff_config = await handoff_brain._resolve_path("handoff")
self.assertEqual(handoff_config["name"], "handoff")
self.assertTrue(
any(
isinstance(frame, OutputTransportMessageUrgentFrame)
and frame.message.get("type") == "handoff-requested"
for frame in queued
)
)
async def test_start_llm_conditions_wait_for_and_route_first_user_turn(self):
brain = WorkflowBrain(
{
"specVersion": 3,
"settings": {"globalPrompt": "全局规则"},
"nodes": [
{
"id": "start",
"type": "start",
"data": {"name": "Start", "greeting": "你想做什么?"},
},
{
"id": "eat",
"type": "agent",
"data": {
"name": "点饭",
"prompt": "帮助用户点饭",
"contextPolicy": "fresh",
"entryMode": "wait_user",
},
},
{"id": "drink", "type": "agent", "data": {}},
{"id": "run", "type": "agent", "data": {}},
],
"edges": [
{
"id": "eat",
"source": "start",
"target": "eat",
"data": {
"mode": "llm",
"priority": 10,
"condition": "用户想吃饭",
},
},
{
"id": "drink",
"source": "start",
"target": "drink",
"data": {
"mode": "llm",
"priority": 20,
"condition": "用户想喝水",
},
},
{
"id": "run",
"source": "start",
"target": "run",
"data": {
"mode": "llm",
"priority": 30,
"condition": "用户想跑步",
},
},
],
}
)
queued = []
async def queue_frame(frame):
queued.append(frame)
brain._runtime = BrainRuntime(
context=LLMContext(messages=[]),
llm=FakeLLM(),
queue_frame=queue_frame,
set_system_prompt=lambda _prompt: None,
set_tools=lambda _tools: None,
call_end=FakeCallEnd(),
)
class FakeManager:
def __init__(self):
self.current_node = None
self.config = None
async def initialize(self, config):
self.current_node = config["name"]
self.config = config
async def set_node_from_config(self, config):
self.current_node = config["name"]
self.config = config
class FakeRouter:
def __init__(self):
self.calls = 0
async def select_edge(self, **_kwargs):
self.calls += 1
return LLMRouteResult(
status=RouteStatus.MATCHED,
function_name="goto_eat",
)
manager = FakeManager()
router = FakeRouter()
brain._manager = manager
brain._router = router
await brain.on_connected()
self.assertEqual(manager.current_node, "start")
self.assertEqual(router.calls, 0)
handled = await brain.on_user_turn_end("我想吃饭")
self.assertTrue(handled)
self.assertEqual(router.calls, 1)
self.assertEqual(manager.current_node, "eat")
self.assertIn(
{"role": "user", "content": "我想吃饭"},
manager.config["task_messages"],
)
self.assertTrue(any(isinstance(frame, LLMRunFrame) for frame in queued))
self.assertIn("我想吃饭", brain._store.values["system__conversation_history"])
async def test_start_expression_condition_also_waits_for_user_turn(self):
brain = WorkflowBrain(
{
"specVersion": 3,
"settings": {},
"nodes": [
{"id": "start", "type": "start", "data": {}},
{"id": "agent", "type": "agent", "data": {}},
],
"edges": [
{
"id": "route",
"source": "start",
"target": "agent",
"data": {
"mode": "expression",
"priority": 10,
"expression": {
"combinator": "and",
"rules": [
{
"variable": "route",
"operator": "eq",
"value": "agent",
}
],
},
},
}
],
}
)
brain._store.values["route"] = "agent"
config = await brain._initial_node_config()
self.assertEqual(config["name"], "start")
self.assertEqual(brain._state.status.value, "waiting_user")
async def test_automatic_node_can_follow_llm_condition(self):
brain = WorkflowBrain(
{
"specVersion": 3,
"settings": {"globalPrompt": "全局规则"},
"nodes": [
{"id": "start", "type": "start", "data": {}},
{
"id": "handoff",
"type": "handoff",
"data": {"name": "人工转接", "targetType": "human"},
},
{
"id": "agent",
"type": "agent",
"data": {"name": "继续服务", "prompt": "继续处理"},
},
],
"edges": [
{
"id": "to-agent",
"source": "handoff",
"target": "agent",
"data": {
"mode": "llm",
"priority": 10,
"condition": "转接后仍需 AI 继续服务",
},
}
],
}
)
queued = []
async def queue_frame(frame):
queued.append(frame)
brain._runtime = BrainRuntime(
context=LLMContext(messages=[]),
llm=FakeLLM(),
queue_frame=queue_frame,
set_system_prompt=lambda _prompt: None,
set_tools=lambda _tools: None,
call_end=FakeCallEnd(),
)
class FakeRouter:
async def select_edge(self, **kwargs):
self.node_name = kwargs["node_name"]
return LLMRouteResult(
status=RouteStatus.MATCHED,
function_name="goto_to_agent",
)
router = FakeRouter()
brain._router = router
config = await brain._resolve_path("handoff")
self.assertEqual(config["name"], "agent")
self.assertEqual(router.node_name, "人工转接")
async def test_mixed_edge_conditions_follow_priority(self):
brain = WorkflowBrain(
{
"specVersion": 3,
"settings": {},
"nodes": [
{"id": "start", "type": "start", "data": {}},
{"id": "agent", "type": "agent", "data": {}},
{"id": "llm-target", "type": "end", "data": {}},
{"id": "expression-target", "type": "end", "data": {}},
{"id": "default-target", "type": "end", "data": {}},
],
"edges": [
{
"id": "llm",
"source": "agent",
"target": "llm-target",
"data": {
"mode": "llm",
"priority": 10,
"condition": "大模型条件成立",
},
},
{
"id": "expression",
"source": "agent",
"target": "expression-target",
"data": {
"mode": "expression",
"priority": 20,
"expression": {
"combinator": "and",
"rules": [
{
"variable": "route",
"operator": "eq",
"value": "expression",
}
],
},
},
},
{
"id": "default",
"source": "agent",
"target": "default-target",
"data": {"mode": "always", "priority": 30},
},
],
}
)
brain._store.values["route"] = "expression"
class FakeRouter:
def __init__(self):
self.result = LLMRouteResult(status=RouteStatus.NO_MATCH)
self.edge_ids = []
async def select_edge(self, **kwargs):
self.edge_ids = [edge["id"] for edge in kwargs["edges"]]
return self.result
router = FakeRouter()
brain._router = router
selected = await brain._select_edge("agent")
self.assertEqual(router.edge_ids, ["llm"])
self.assertEqual(selected["id"], "expression")
router.result = LLMRouteResult(
status=RouteStatus.MATCHED,
function_name="goto_llm",
)
selected = await brain._select_edge("agent")
self.assertEqual(selected["id"], "llm")
expression_edge = next(
edge for edge in brain._engine.edges if edge["id"] == "expression"
)
expression_edge["data"]["priority"] = 5
router.edge_ids = []
selected = await brain._select_edge("agent")
self.assertEqual(selected["id"], "expression")
self.assertEqual(router.edge_ids, [])
async def test_transition_and_end_are_owned_by_workflow_brain(self):
graph = {
"specVersion": 3,
"settings": {
"globalPrompt": "全局规则",
"defaultLlmResourceId": "llm_global",
"defaultAsrResourceId": "asr_global",
"defaultTtsResourceId": "tts_global",
"knowledgeBaseId": "kb_global",
"knowledgeMode": "automatic",
"enableInterrupt": False,
"turnConfig": {
"bargeIn": {"strategy": "transcription"},
"vad": {"confidence": 0.55},
},
},
"nodes": [
{
"id": "start",
"type": "start",
"data": {
"name": "Start",
"greeting": "欢迎,{{user_name}}",
},
},
{
"id": "agent",
"type": "agent",
"data": {
"name": "收集需求",
"prompt": "服务 {{user_name}}",
"contextPolicy": "fresh",
},
},
{
"id": "end",
"type": "end",
"data": {"name": "End", "message": "感谢来电", "scope": "session"},
},
],
"edges": [
{
"id": "begin",
"source": "start",
"target": "agent",
"data": {"mode": "always", "priority": 0},
},
{
"id": "finish",
"source": "agent",
"target": "end",
"data": {
"mode": "llm",
"priority": 10,
"condition": "需求已收集",
"transitionSpeech": "正在为你结束流程",
},
}
],
}
cfg = prepare_dynamic_config(
AssistantConfig(
type="workflow",
graph=graph,
dynamic_variable_definitions={
"user_name": {"type": "string", "required": True}
},
),
{"user_name": "王先生"},
assistant_id="asst_workflow",
)
brain = WorkflowBrain(cfg)
llm = FakeLLM()
context = LLMContext(messages=[])
queued = []
service_switches = []
knowledge_scopes = []
turn_configs = []
call_end = FakeCallEnd()
class FakeWorker:
def __init__(self):
self.frames = []
self.handlers = {}
def set_reached_downstream_filter(self, *_args):
pass
def event_handler(self, name):
def decorator(fn):
self.handlers[name] = fn
return fn
return decorator
async def queue_frame(self, frame):
self.frames.append(frame)
async def queue_frames(self, frames):
self.frames.extend(frames)
worker = FakeWorker()
pair = SimpleNamespace(
user=lambda: SimpleNamespace(_context=context),
assistant=lambda: SimpleNamespace(has_function_calls_in_progress=False),
)
async def queue_frame(frame):
queued.append(frame)
async def switch_services(llm_id, asr_id, tts_id):
service_switches.append((llm_id, asr_id, tts_id))
async def apply_turn_config(enable_interrupt, turn_config):
turn_configs.append((enable_interrupt, turn_config))
runtime = BrainRuntime(
context=context,
llm=llm,
queue_frame=queue_frame,
set_system_prompt=lambda _prompt: None,
set_tools=lambda _tools: None,
call_end=call_end,
worker=worker,
context_aggregator=pair,
switch_services=switch_services,
set_knowledge_scope=knowledge_scopes.append,
apply_turn_config=apply_turn_config,
)
await brain.setup(cfg, runtime)
greeting = await brain.greeting(cfg)
self.assertEqual(greeting, "欢迎,王先生")
greeting_message = {
"role": "system",
"content": f"{GREETING_CONTEXT_MARKER}\n欢迎,王先生",
}
brain.prepare_greeting_context(greeting, context)
self.assertEqual(context.get_messages(), [greeting_message])
await brain.on_connected()
self.assertEqual(brain._manager.current_node, "agent")
variable_events = [
frame.message
for frame in queued
if isinstance(frame, OutputTransportMessageUrgentFrame)
and frame.message.get("type") == "workflow-variables"
]
self.assertEqual(variable_events[0]["reason"], "initialized")
self.assertEqual(variable_events[0]["variables"], {"user_name": "王先生"})
self.assertNotIn("system__conversation_id", variable_events[0]["variables"])
self.assertEqual(
service_switches,
[("llm_global", "asr_global", "tts_global")],
)
self.assertEqual(knowledge_scopes[-1]["knowledge_base_id"], "kb_global")
self.assertEqual(turn_configs[-1][0], False)
self.assertEqual(turn_configs[-1][1]["vad"]["confidence"], 0.55)
brain._engine.data("agent").update(
{
"inheritGlobalConfig": False,
"llmResourceId": "llm_agent",
"asrResourceId": "asr_agent",
"ttsResourceId": "tts_agent",
"knowledgeBaseId": "kb_agent",
"knowledgeMode": "on_demand",
"enableInterrupt": True,
"turnConfig": {
"bargeIn": {"strategy": "vad"},
"turnDetection": {"strategy": "smart_turn"},
},
}
)
await brain._apply_agent_stage("agent")
self.assertEqual(
service_switches[-1],
("llm_agent", "asr_agent", "tts_agent"),
)
self.assertEqual(knowledge_scopes[-1]["knowledge_base_id"], "kb_agent")
self.assertEqual(turn_configs[-1][0], True)
self.assertEqual(
turn_configs[-1][1]["turnDetection"]["strategy"],
"smart_turn",
)
agent_config = brain._agent_config("agent")
self.assertIn("王先生", agent_config["role_message"])
self.assertIn("工作流路由已在用户一轮输入结束时完成", agent_config["role_message"])
self.assertEqual(agent_config["task_messages"], [greeting_message])
self.assertFalse(agent_config["respond_immediately"])
self.assertFalse(any(isinstance(frame, LLMRunFrame) for frame in worker.frames))
self.assertEqual(
agent_config["context_strategy"].strategy.value,
"reset",
)
brain._engine.data("agent")["entryMode"] = "generate"
generate_config = brain._agent_config("agent")
self.assertFalse(generate_config["respond_immediately"])
worker.frames.clear()
await brain._manager.set_node_from_config(generate_config)
self.assertFalse(any(isinstance(frame, LLMRunFrame) for frame in worker.frames))
await brain._after_node_activated(generate_config)
self.assertTrue(any(isinstance(frame, LLMRunFrame) for frame in queued))
brain._engine.data("agent").update(
{"entryMode": "fixed_speech", "entrySpeech": "您好,{{user_name}}"}
)
fixed_config = brain._agent_config("agent")
self.assertFalse(fixed_config["respond_immediately"])
self.assertNotIn("pre_actions", fixed_config)
self.assertEqual(
fixed_config["task_messages"],
[
greeting_message,
{"role": "assistant", "content": "您好,王先生"},
],
)
self.assertEqual(
brain._agent_config(
"agent",
[{"role": "assistant", "content": "正在进入下一阶段"}],
)["task_messages"],
[
greeting_message,
{"role": "assistant", "content": "正在进入下一阶段"},
{"role": "assistant", "content": "您好,王先生"},
],
)
worker.frames.clear()
queued.clear()
await brain._manager.set_node_from_config(fixed_config)
await brain._after_node_activated(fixed_config)
self.assertTrue(any(isinstance(frame, TTSSpeakFrame) for frame in queued))
self.assertFalse(any(isinstance(frame, LLMRunFrame) for frame in worker.frames))
context_updates = [
frame
for frame in worker.frames
if isinstance(frame, LLMMessagesUpdateFrame)
]
self.assertEqual(
context_updates[-1].messages,
[
greeting_message,
{"role": "assistant", "content": "您好,王先生"},
],
)
self.assertFalse(
any(
isinstance(frame, OutputTransportMessageUrgentFrame)
and frame.message.get("source") == "workflow-fixed-reply"
for frame in queued
)
)
await brain.on_client_ready()
fixed_reply_events = [
frame.message
for frame in queued
if isinstance(frame, OutputTransportMessageUrgentFrame)
and frame.message.get("source") == "workflow-fixed-reply"
]
self.assertEqual(fixed_reply_events[0]["content"], "您好,王先生")
self.assertEqual(fixed_reply_events[0]["nodeId"], "agent")
self.assertIn("您好,王先生", brain._store.values["system__conversation_history"])
self.assertFalse(
any(
function.name == "goto_finish"
for function in brain._agent_config("agent")["functions"]
)
)
await brain.on_assistant_text_end("old-turn", "需求已收集", False)
self.assertEqual(brain._manager.current_node, "agent")
class FakeRouter:
async def select_edge(self, **_kwargs):
return LLMRouteResult(
status=RouteStatus.MATCHED,
function_name="goto_finish",
)
brain._router = FakeRouter()
handled = await brain.on_user_turn_end("我的需求已经说完了")
self.assertTrue(handled)
self.assertEqual(brain._manager.current_node, "end")
self.assertIn("我的需求已经说完了", brain._store.values["system__conversation_history"])
self.assertTrue(call_end.ending)
self.assertTrue(call_end.armed)
self.assertTrue(any(getattr(frame, "text", "") == "感谢来电" for frame in queued))
transition_context_frames = [
frame
for frame in worker.frames
if isinstance(frame, LLMMessagesAppendFrame)
and frame.messages
== [{"role": "assistant", "content": "正在为你结束流程"}]
]
self.assertTrue(transition_context_frames)
transition_events = [
frame.message
for frame in queued
if isinstance(frame, OutputTransportMessageUrgentFrame)
and frame.message.get("source") == "workflow-edge-transition"
]
self.assertEqual(transition_events[0]["content"], "正在为你结束流程")
self.assertEqual(transition_events[0]["nodeId"], "end")
assistant_transcripts = [
frame.message.get("content")
for frame in queued
if isinstance(frame, OutputTransportMessageUrgentFrame)
and frame.message.get("type") == "transcript"
and frame.message.get("role") == "assistant"
]
self.assertEqual(
assistant_transcripts,
["您好,王先生", "正在为你结束流程", "感谢来电"],
)
self.assertIn(
"正在为你结束流程",
brain._store.values["system__conversation_history"],
)
self.assertIn(
"感谢来电",
brain._store.values["system__conversation_history"],
)
if __name__ == "__main__":
unittest.main()