from __future__ import annotations import unittest from types import SimpleNamespace from models import AssistantConfig, RuntimeTool from pipecat.frames.frames import ( LLMContextFrame, LLMFullResponseEndFrame, LLMFullResponseStartFrame, LLMTextFrame, ) 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 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", ) 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_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") class WorkflowBrainTests(unittest.IsolatedAsyncioTestCase): async def test_transition_and_end_are_owned_by_workflow_brain(self): graph = { "nodes": [ { "id": "start", "type": "startCall", "data": {"name": "开始", "prompt": "收集需求"}, }, { "id": "end", "type": "endCall", "data": {"name": "结束", "prompt": "礼貌结束"}, }, ], "edges": [ { "id": "finish", "source": "start", "target": "end", "data": {"condition": "需求已收集"}, } ], } brain = WorkflowBrain(graph) llm = FakeLLM() context = LLMContext(messages=[]) queued = [] prompts = [] visible_tools = [] call_end = FakeCallEnd() async def queue_frame(frame): queued.append(frame) runtime = BrainRuntime( context=context, llm=llm, queue_frame=queue_frame, set_system_prompt=prompts.append, set_tools=lambda tools: visible_tools.append(tools or []), call_end=call_end, ) await brain.setup(AssistantConfig(type="workflow", graph=graph), runtime) self.assertIn("goto_finish", llm.functions) self.assertIn("收集需求", prompts[-1]) self.assertEqual(visible_tools[-1][0].name, "goto_finish") params = FakeFunctionParams() await llm.functions["goto_finish"](params) self.assertEqual(params.result, {"status": "ok"}) self.assertIn("礼貌结束", prompts[-1]) self.assertEqual(visible_tools[-1], []) await brain.on_assistant_text_start("closing-turn") await brain.on_assistant_text_end( "closing-turn", "感谢来电，再见。", False, ) self.assertTrue(call_end.ending) self.assertTrue(call_end.armed) if __name__ == "__main__": unittest.main()