136 lines
4.6 KiB
Python
136 lines
4.6 KiB
Python
import unittest
|
|
|
|
from models import AssistantConfig
|
|
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,
|
|
)
|
|
|
|
|
|
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):
|
|
self.turns.append(content)
|
|
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, ["我叫李白"])
|
|
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, ["我叫李白"])
|
|
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):
|
|
self.turns.append(content)
|
|
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, ["看看这张照片"])
|
|
|
|
|
|
async def _async_none():
|
|
return None
|
|
|
|
|
|
if __name__ == "__main__":
|
|
unittest.main()
|