32aef14ddb
- Introduce RuntimeModelResource and RuntimeKnowledgeBase classes to manage workflow resources. - Update AssistantConfig to include workflow_model_resources and workflow_knowledge_bases for better integration. - Refactor validation and processing logic in routes and services to accommodate workflow types. - Implement dynamic variable support for workflow assistants and enhance graph normalization. - Add ToolExecutor for reusable tool execution across different assistant types. - Update various services to ensure compatibility with new workflow features and improve error handling.
517 lines
17 KiB
Python
517 lines
17 KiB
Python
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,
|
||
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
|
||
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_transition_and_end_are_owned_by_workflow_brain(self):
|
||
graph = {
|
||
"specVersion": 3,
|
||
"settings": {"globalPrompt": "全局规则"},
|
||
"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": "需求已收集",
|
||
},
|
||
}
|
||
],
|
||
}
|
||
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 = []
|
||
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)
|
||
|
||
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,
|
||
)
|
||
await brain.setup(cfg, runtime)
|
||
await brain.on_connected()
|
||
self.assertEqual(brain._manager.current_node, "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.assertEqual(
|
||
agent_config["context_strategy"].strategy.value,
|
||
"reset",
|
||
)
|
||
|
||
edge_function = next(
|
||
function
|
||
for function in brain._agent_config("agent")["functions"]
|
||
if function.name == "goto_finish"
|
||
)
|
||
_, terminal = await edge_function.handler({}, brain._manager)
|
||
self.assertEqual(terminal["name"], "end")
|
||
self.assertTrue(call_end.ending)
|
||
self.assertTrue(call_end.armed)
|
||
self.assertTrue(any(getattr(frame, "text", "") == "感谢来电" for frame in queued))
|
||
|
||
|
||
if __name__ == "__main__":
|
||
unittest.main()
|