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

from __future__ import annotations

import unittest
from types import SimpleNamespace

from models import AssistantConfig
from services.brains.base import RealtimeBrainRuntime
from services.node_specs import validate_graph
from services.vision import (
    VISION_SYSTEM_HINT,
    resolve_vision_question,
    vision_system_hint,
    vision_tool_schema_parts,
)
from services.workflow_engine import WorkflowEngine
from services.workflow.realtime import WorkflowRealtimeController


def workflow_settings(**overrides):
    settings = {
        "runtimeMode": "pipeline",
        "llmRoutingMode": "llm_router",
        "globalPrompt": "",
        "defaultRealtimeResourceId": "",
        "defaultLlmResourceId": "",
        "defaultAsrResourceId": "",
        "defaultTtsResourceId": "",
        "visionEnabled": False,
        "visionModelResourceId": "",
        "visionPromptMode": "runtime",
        "visionPrompt": "",
        "toolBindings": [],
        "builtInTools": {},
        "knowledgeBaseId": "",
        "knowledgeMode": "disabled",
        "knowledgeTopN": 5,
        "knowledgeScoreThreshold": 0.0,
        "enableInterrupt": True,
        "turnConfig": {
            "bargeIn": {"strategy": "vad"},
            "turnDetection": {"strategy": "smart_turn"},
        },
    }
    settings.update(overrides)
    return settings


class VisionPromptModeTest(unittest.TestCase):
    def test_fixed_mode_has_no_runtime_question_parameter(self):
        description, properties, required = vision_tool_schema_parts(
            "fixed",
            "识别车辆数量和人员状态",
        )

        self.assertIn("识别车辆数量和人员状态", description)
        self.assertEqual(properties, {})
        self.assertEqual(required, [])
        self.assertEqual(
            resolve_vision_question("fixed", " 识别车辆数量 ", {}),
            "识别车辆数量",
        )

    def test_runtime_mode_requires_model_generated_question(self):
        description, properties, required = vision_tool_schema_parts("runtime")

        self.assertNotIn("固定引导", description)
        self.assertIn("question", properties)
        self.assertEqual(required, ["question"])
        self.assertEqual(
            resolve_vision_question(
                "runtime",
                "不会使用这个配置",
                {"question": "用户手里拿着什么？"},
            ),
            "用户手里拿着什么？",
        )

    def test_guided_mode_exposes_guide_but_uses_runtime_question(self):
        guide = "只询问画面中可直接确认的车辆损伤"
        description, properties, required = vision_tool_schema_parts(
            "guided",
            guide,
        )

        self.assertIn(guide, description)
        self.assertIn(guide, properties["question"]["description"])
        self.assertEqual(required, ["question"])
        self.assertEqual(
            resolve_vision_question(
                "guided",
                guide,
                {"question": "右侧车门是否有明显凹陷？"},
            ),
            "右侧车门是否有明显凹陷？",
        )

    def test_mode_specific_system_hints(self):
        self.assertIn(VISION_SYSTEM_HINT, vision_system_hint("runtime"))
        self.assertIn("不需要生成", vision_system_hint("fixed", "固定问题"))
        self.assertIn("固定引导", vision_system_hint("guided", "业务引导"))

    def test_missing_required_config_or_question_is_rejected(self):
        with self.assertRaisesRegex(ValueError, "固定视觉提示词为空"):
            resolve_vision_question("fixed", "", {})
        with self.assertRaisesRegex(ValueError, "视觉问题为空"):
            resolve_vision_question("runtime", "", {})


class WorkflowVisionPromptModeTest(unittest.TestCase):
    def test_workflow_resolves_global_and_node_modes(self):
        graph = {
            "specVersion": 4,
            "settings": workflow_settings(
                **{
                    "visionEnabled": True,
                    "visionPromptMode": "fixed",
                    "visionPrompt": "全局固定问题",
                }
            ),
            "nodes": [
                {
                    "id": "global-agent",
                    "type": "agent",
                    "data": {"name": "全局 Agent", "prompt": "处理任务"},
                },
                {
                    "id": "custom-agent",
                    "type": "agent",
                    "data": {
                        "name": "独立 Agent",
                        "prompt": "检查车辆",
                        "inheritGlobalConfig": False,
                        "visionEnabled": True,
                        "visionPromptMode": "guided",
                        "visionPrompt": "节点引导",
                    },
                },
            ],
            "edges": [],
        }
        engine = WorkflowEngine(graph)

        global_stage = engine.agent_stage_config("global-agent")
        self.assertEqual(global_stage.vision_prompt_mode, "fixed")
        self.assertEqual(global_stage.vision_prompt, "全局固定问题")
        custom_stage = engine.agent_stage_config("custom-agent")
        self.assertEqual(custom_stage.vision_prompt_mode, "guided")
        self.assertEqual(custom_stage.vision_prompt, "节点引导")

    def test_fixed_or_guided_mode_requires_prompt_when_vision_is_enabled(self):
        graph = {
            "specVersion": 4,
            "settings": workflow_settings(
                **{
                    "visionEnabled": True,
                    "visionPromptMode": "fixed",
                    "visionPrompt": "",
                }
            ),
            "nodes": [
                {"id": "start", "type": "start", "data": {}},
                {
                    "id": "agent",
                    "type": "agent",
                    "data": {
                        "name": "Agent",
                        "prompt": "处理任务",
                        "inheritGlobalConfig": True,
                        "entryMode": "wait_user",
                    },
                }
            ],
            "edges": [
                {
                    "id": "begin",
                    "source": "start",
                    "target": "agent",
                    "data": {"mode": "always", "priority": 0},
                }
            ],
        }

        self.assertTrue(
            any("必须填写视觉提示词" in error for error in validate_graph(graph))
        )

    def test_invalid_mode_is_reported_instead_of_silently_rewritten(self):
        graph = {
            "specVersion": 4,
            "settings": workflow_settings(visionPromptMode="unexpected"),
            "nodes": [
                {"id": "start", "type": "start", "data": {}},
                {
                    "id": "agent",
                    "type": "agent",
                    "data": {
                        "name": "Agent",
                        "prompt": "处理任务",
                        "inheritGlobalConfig": True,
                        "entryMode": "wait_user",
                    },
                }
            ],
            "edges": [
                {
                    "id": "begin",
                    "source": "start",
                    "target": "agent",
                    "data": {"mode": "always", "priority": 0},
                }
            ],
        }

        self.assertTrue(
            any("视觉提示词模式无效" in error for error in validate_graph(graph))
        )

    def test_realtime_tool_contract_changes_with_mode(self):
        graph = {
            "specVersion": 4,
            "settings": workflow_settings(),
            "nodes": [
                {
                    "id": "fixed-agent",
                    "type": "agent",
                    "data": {
                        "name": "固定",
                        "prompt": "处理任务",
                        "inheritGlobalConfig": False,
                        "visionEnabled": True,
                        "visionPromptMode": "fixed",
                        "visionPrompt": "检查安全帽",
                    },
                },
                {
                    "id": "guided-agent",
                    "type": "agent",
                    "data": {
                        "name": "引导",
                        "prompt": "处理任务",
                        "inheritGlobalConfig": False,
                        "visionEnabled": True,
                        "visionPromptMode": "guided",
                        "visionPrompt": "只询问安全设备",
                    },
                },
            ],
            "edges": [],
        }
        engine = WorkflowEngine(graph)
        controller = WorkflowRealtimeController(
            cfg=AssistantConfig(type="workflow"),
            engine=engine,
            store=SimpleNamespace(),
            runtime=RealtimeBrainRuntime(
                realtime=SimpleNamespace(),
                queue_frame=lambda _frame: None,
                call_end=SimpleNamespace(ending=False),
            ),
        )

        fixed_tool, _fixed_handler = controller._vision_tool("fixed-agent", 0)
        guided_tool, _guided_handler = controller._vision_tool("guided-agent", 0)

        self.assertEqual(fixed_tool.properties, {})
        self.assertEqual(fixed_tool.required, ())
        self.assertIn("只询问安全设备", guided_tool.description)
        self.assertEqual(guided_tool.required, ("question",))


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