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

import unittest

from models import AssistantConfig, RuntimeModelResource
from pipecat.frames.frames import LLMContextFrame
from pipecat.processors.aggregators.llm_context import LLMContext
from pipecat.processors.frame_processor import FrameDirection
from services.pipecat.pipeline import (
    KNOWLEDGE_CONTEXT_MARKER,
    KnowledgeRetrievalProcessor,
    UserTurnRoutingProcessor,
    _knowledge_tool_description,
    _workflow_vision_uses_main_llm,
)


class KnowledgeToolDescriptionTest(unittest.TestCase):
    def test_includes_bound_knowledge_scope(self):
        description = _knowledge_tool_description(
            AssistantConfig(
                knowledge_base_name="产品服务知识库",
                knowledge_base_description="产品价格、售后政策和退换货条件",
            )
        )

        self.assertIn("知识库名称：产品服务知识库", description)
        self.assertIn("资料适用范围：产品价格、售后政策和退换货条件", description)
        self.assertIn("与该范围无关的问题不要调用", description)

    def test_falls_back_when_metadata_is_empty(self):
        description = _knowledge_tool_description(AssistantConfig())

        self.assertEqual(
            description,
            "在当前助手绑定的知识库中检索与问题最相关的资料片段。",
        )

    def test_compacts_and_limits_description(self):
        description = _knowledge_tool_description(
            AssistantConfig(knowledge_base_description=("范围\n  内容 " * 200))
        )

        self.assertNotIn("\n  ", description)
        self.assertLess(len(description), 1000)

    def test_workflow_knowledge_uses_system_role(self):
        processor = KnowledgeRetrievalProcessor(None)
        messages = [
            {"role": "assistant", "content": "你好"},
            {
                "role": "developer",
                "content": f"{KNOWLEDGE_CONTEXT_MARKER}\n旧检索结果",
            },
        ]

        processor._set_context(
            messages,
            f"{KNOWLEDGE_CONTEXT_MARKER}\n新检索结果",
        )

        self.assertEqual(messages[0]["role"], "system")
        self.assertIn("新检索结果", messages[0]["content"])
        self.assertFalse(any(message["role"] == "developer" for message in messages))


class UserTurnRoutingProcessorTest(unittest.IsolatedAsyncioTestCase):
    async def test_routes_each_user_message_once_before_response_run(self):
        class FakeBrain:
            def __init__(self):
                self.turns = []

            async def on_user_turn_end(self, content, user_message=None):
                self.turns.append((content, user_message))
                return True

        brain = FakeBrain()
        processor = UserTurnRoutingProcessor(brain)
        forwarded = []

        async def push_frame(frame, direction):
            forwarded.append((frame, direction))

        processor.push_frame = push_frame
        context = LLMContext(messages=[{"role": "user", "content": "我叫李白"}])
        frame = LLMContextFrame(context)

        await processor.process_frame(frame, FrameDirection.DOWNSTREAM)
        self.assertEqual(
            brain.turns,
            [("我叫李白", {"role": "user", "content": "我叫李白"})],
        )
        self.assertEqual(forwarded, [])

        # A queued LLMRunFrame after the transition uses the same context. It
        # must reach the target Agent without invoking routing a second time.
        await processor.process_frame(frame, FrameDirection.DOWNSTREAM)
        self.assertEqual(
            brain.turns,
            [("我叫李白", {"role": "user", "content": "我叫李白"})],
        )
        self.assertEqual(forwarded, [(frame, FrameDirection.DOWNSTREAM)])

    async def test_routes_multimodal_user_message_by_its_text_part(self):
        class FakeBrain:
            def __init__(self):
                self.turns = []

            async def on_user_turn_end(self, content, user_message=None):
                self.turns.append((content, user_message))
                return False

        brain = FakeBrain()
        processor = UserTurnRoutingProcessor(brain)
        processor.push_frame = lambda *_args, **_kwargs: _async_none()
        context = LLMContext(
            messages=[
                {
                    "role": "user",
                    "content": [
                        {"type": "text", "text": "看看这张照片"},
                        {
                            "type": "image_url",
                            "image_url": {"url": "data:image/jpeg;base64,AA=="},
                        },
                    ],
                }
            ]
        )

        await processor.process_frame(
            LLMContextFrame(context),
            FrameDirection.DOWNSTREAM,
        )

        self.assertEqual(
            brain.turns,
            [
                (
                    "看看这张照片",
                    {
                        "role": "user",
                        "content": [
                            {"type": "text", "text": "看看这张照片"},
                            {
                                "type": "image_url",
                                "image_url": {
                                    "url": "data:image/jpeg;base64,AA=="
                                },
                            },
                        ],
                    },
                )
            ],
        )


class WorkflowVisionModeTest(unittest.TestCase):
    def test_uses_active_agent_llm_only_without_auxiliary_model(self):
        cfg = AssistantConfig(
            type="workflow",
            workflow_model_resources={
                "agent_llm": RuntimeModelResource(
                    id="agent_llm",
                    name="视觉 Agent",
                    capability="LLM",
                    interface_type="openai-llm",
                    support_image_input=True,
                )
            },
        )

        self.assertTrue(
            _workflow_vision_uses_main_llm(
                cfg,
                {
                    "enabled": True,
                    "llm_resource_id": "agent_llm",
                    "vision_model_resource_id": None,
                },
            )
        )
        self.assertFalse(
            _workflow_vision_uses_main_llm(
                cfg,
                {
                    "enabled": True,
                    "llm_resource_id": "agent_llm",
                    "vision_model_resource_id": "auxiliary_vision",
                },
            )
        )

    def test_rejects_a_non_visual_active_agent_llm(self):
        cfg = AssistantConfig(
            type="workflow",
            workflow_model_resources={
                "text_llm": RuntimeModelResource(
                    id="text_llm",
                    name="文本 Agent",
                    capability="LLM",
                    interface_type="openai-llm",
                    support_image_input=False,
                )
            },
        )

        with self.assertRaisesRegex(ValueError, "不支持图片输入"):
            _workflow_vision_uses_main_llm(
                cfg,
                {
                    "enabled": True,
                    "llm_resource_id": "text_llm",
                    "vision_model_resource_id": None,
                },
            )


async def _async_none():
    return None


if __name__ == "__main__":
    unittest.main()