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

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()