Files
ai-video-fullstack/backend/tests/test_brains.py
T
Xin Wang 72856bf3a7 Enhance workflow routing and agent configuration management
- Introduce WorkflowLLMRouter for pre-response LLM routing, allowing agents to determine the appropriate function to call based on user input.
- Implement UserTurnRoutingProcessor to manage user turns before reaching the LLM, ensuring proper routing and handling of user messages.
- Refactor WorkflowBrain to integrate new routing logic and enhance agent stage configuration, including entry modes and resource management.
- Update service factory to support dynamic LLM resource configuration based on workflow settings.
- Add tests for new routing functionality and ensure proper handling of user messages in various scenarios.
2026-07-14 09:36:28 +08:00

707 lines
24 KiB
Python

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,
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.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
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
def begin(self, reason: str) -> None:
self.ending = True
self.reason = reason
def arm_after_speech(self) -> None:
self.armed = True
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), "您好,王先生")
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 llm.functions["end_call"](params)
self.assertEqual(call_end.reason, "用户已完成咨询")
self.assertTrue(call_end.finished)
self.assertEqual(params.result["action"], "ending_call")
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_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")
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()
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_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",
},
"nodes": [
{
"id": "start",
"type": "start",
"data": {"name": "Start"},
},
{
"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 = []
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))
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,
)
await brain.setup(cfg, runtime)
await brain.on_connected()
self.assertEqual(brain._manager.current_node, "agent")
self.assertEqual(
service_switches,
[("llm_global", "asr_global", "tts_global")],
)
self.assertEqual(knowledge_scopes[-1]["knowledge_base_id"], "kb_global")
brain._engine.data("agent").update(
{
"inheritGlobalConfig": False,
"llmResourceId": "llm_agent",
"asrResourceId": "asr_agent",
"ttsResourceId": "tts_agent",
"knowledgeBaseId": "kb_agent",
"knowledgeMode": "on_demand",
}
)
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")
agent_config = brain._agent_config("agent")
self.assertIn("王先生", agent_config["role_message"])
self.assertIn("工作流路由已在用户一轮输入结束时完成", agent_config["role_message"])
self.assertEqual(agent_config["task_messages"], [])
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.assertTrue(generate_config["respond_immediately"])
worker.frames.clear()
await brain._manager.set_node_from_config(generate_config)
self.assertTrue(any(isinstance(frame, LLMRunFrame) for frame in worker.frames))
brain._engine.data("agent").update(
{"entryMode": "fixed_speech", "entrySpeech": "您好,{{user_name}}"}
)
fixed_config = brain._agent_config("agent")
self.assertFalse(fixed_config["respond_immediately"])
self.assertEqual(
fixed_config["pre_actions"][0]["type"],
"workflow_fixed_speech",
)
self.assertEqual(fixed_config["pre_actions"][0]["text"], "您好,王先生")
self.assertEqual(
fixed_config["task_messages"],
[{"role": "assistant", "content": "您好,王先生"}],
)
worker.frames.clear()
queued.clear()
await brain._manager.set_node_from_config(fixed_config)
self.assertTrue(any(isinstance(frame, TTSSpeakFrame) for frame in queued))
self.assertFalse(any(isinstance(frame, LLMRunFrame) for frame in worker.frames))
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 "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))
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()