Files
ai-video-fullstack/backend/tests/test_brains.py
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
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,
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()