127 lines
4.4 KiB
Python
127 lines
4.4 KiB
Python
from __future__ import annotations
|
|
|
|
import unittest
|
|
from types import SimpleNamespace
|
|
from unittest.mock import AsyncMock, MagicMock, patch
|
|
|
|
from models import RuntimeModelResource
|
|
from services.vision import (
|
|
VISION_ANALYSIS_SYSTEM_PROMPT,
|
|
VISION_SYSTEM_HINT,
|
|
generate_vision_focus_prompt,
|
|
vision_analysis_system_prompt,
|
|
vision_system_hint,
|
|
vision_tool_description,
|
|
)
|
|
from services.workflow_engine import WorkflowEngine
|
|
|
|
|
|
class VisionFocusPromptTest(unittest.TestCase):
|
|
def test_focus_extends_but_does_not_replace_safety_prompts(self):
|
|
focus = "优先识别车辆数量和人员状态，只报告可确认的信息。"
|
|
|
|
system_hint = vision_system_hint(focus)
|
|
analysis_prompt = vision_analysis_system_prompt(focus)
|
|
|
|
self.assertIn(VISION_SYSTEM_HINT, system_hint)
|
|
self.assertIn(focus, system_hint)
|
|
self.assertIn(VISION_ANALYSIS_SYSTEM_PROMPT, analysis_prompt)
|
|
self.assertIn(focus, analysis_prompt)
|
|
self.assertIn("优先回答当前用户问题", analysis_prompt)
|
|
|
|
def test_empty_focus_preserves_previous_runtime_prompts(self):
|
|
self.assertEqual(vision_system_hint(), VISION_SYSTEM_HINT)
|
|
self.assertEqual(
|
|
vision_analysis_system_prompt(),
|
|
VISION_ANALYSIS_SYSTEM_PROMPT,
|
|
)
|
|
|
|
def test_tool_description_compacts_and_limits_focus(self):
|
|
description = vision_tool_description("车辆\n 数量 " * 300)
|
|
|
|
self.assertIn("当前业务重点", description)
|
|
self.assertNotIn("\n", description)
|
|
self.assertLess(len(description), 700)
|
|
|
|
def test_workflow_resolves_global_and_node_focus(self):
|
|
graph = {
|
|
"specVersion": 3,
|
|
"settings": {
|
|
"visionEnabled": True,
|
|
"visionFocusPrompt": "全局关注目标",
|
|
},
|
|
"nodes": [
|
|
{
|
|
"id": "global-agent",
|
|
"type": "agent",
|
|
"data": {"name": "全局 Agent", "prompt": "处理任务"},
|
|
},
|
|
{
|
|
"id": "custom-agent",
|
|
"type": "agent",
|
|
"data": {
|
|
"name": "独立 Agent",
|
|
"prompt": "检查车辆",
|
|
"inheritGlobalConfig": False,
|
|
"visionEnabled": True,
|
|
"visionFocusPrompt": "节点关注目标",
|
|
},
|
|
},
|
|
],
|
|
"edges": [],
|
|
}
|
|
engine = WorkflowEngine(graph)
|
|
|
|
self.assertEqual(
|
|
engine.agent_stage_config("global-agent").vision_focus_prompt,
|
|
"全局关注目标",
|
|
)
|
|
self.assertEqual(
|
|
engine.agent_stage_config("custom-agent").vision_focus_prompt,
|
|
"节点关注目标",
|
|
)
|
|
|
|
|
|
class VisionFocusGenerationTest(unittest.IsolatedAsyncioTestCase):
|
|
async def test_generates_editable_focus_with_selected_resource(self):
|
|
resource = RuntimeModelResource(
|
|
id="llm_1",
|
|
name="生成模型",
|
|
capability="LLM",
|
|
interface_type="openai-llm",
|
|
values={
|
|
"apiUrl": "https://llm.example/v1",
|
|
"modelId": "focus-model",
|
|
},
|
|
secrets={"apiKey": "secret"},
|
|
)
|
|
client = MagicMock()
|
|
client.chat.completions.create = AsyncMock(
|
|
return_value=SimpleNamespace(
|
|
choices=[
|
|
SimpleNamespace(
|
|
message=SimpleNamespace(
|
|
content="优先识别车辆数量、损伤位置和人员状态。"
|
|
)
|
|
)
|
|
]
|
|
)
|
|
)
|
|
client.close = AsyncMock()
|
|
|
|
with patch("services.vision.AsyncOpenAI", return_value=client):
|
|
result = await generate_vision_focus_prompt(
|
|
resource,
|
|
"你是交通事故处理助手，需要了解车辆和人员情况。",
|
|
)
|
|
|
|
self.assertEqual(result, "优先识别车辆数量、损伤位置和人员状态。")
|
|
request = client.chat.completions.create.await_args.kwargs
|
|
self.assertEqual(request["model"], "focus-model")
|
|
self.assertIn("交通事故处理助手", request["messages"][1]["content"])
|
|
client.close.assert_awaited_once()
|
|
|
|
|
|
if __name__ == "__main__":
|
|
unittest.main()
|