eryao/backend/tests/unit/test_agentscope_prompts.py

from __future__ import annotations

from core.agentscope.prompts.agent_prompt import build_agent_prompt
from core.agentscope.prompts.system_prompt import build_system_prompt
from core.agentscope.prompts.user_prompt import build_follow_up_user_prompt
from schemas.agent.system_agent import AgentType, SystemAgentLLMConfig


def test_system_prompt_safety_has_refusal_rules_en() -> None:
    prompt = build_system_prompt(
        agent_type=AgentType.WORKER,
        language="en-US",
        llm_config=SystemAgentLLMConfig(),
    )

    assert "<!-- SAFETY_START -->" in prompt
    assert "-> REFUSE" in prompt
    assert "Tarot" in prompt
    assert "Ba Zi" in prompt


def test_system_prompt_safety_has_refusal_rules_zh() -> None:
    prompt = build_system_prompt(
        agent_type=AgentType.WORKER,
        language="zh-CN",
        llm_config=SystemAgentLLMConfig(),
    )

    assert "<!-- SAFETY_START -->" in prompt
    assert "必须拒绝" in prompt
    assert "塔罗" in prompt
    assert "八字" in prompt


def test_system_prompt_no_language_constraint_in_system() -> None:
    prompt = build_system_prompt(
        agent_type=AgentType.WORKER,
        language="en-US",
        llm_config=SystemAgentLLMConfig(),
    )

    assert "<!-- OUTPUT_START -->" not in prompt
    assert "MUST respond in" not in prompt


def test_system_prompt_safety_restricts_to_divination() -> None:
    prompt = build_system_prompt(
        agent_type=AgentType.WORKER,
        language="zh-CN",
        llm_config=SystemAgentLLMConfig(),
    )

    assert "六爻" in prompt or "解卦" in prompt


def test_system_prompt_does_not_contain_env_section() -> None:
    prompt = build_system_prompt(
        agent_type=AgentType.WORKER,
        language="zh-CN",
        llm_config=SystemAgentLLMConfig(),
    )

    assert "USER_CONTEXT_JSON" not in prompt
    assert "Runtime Context" not in prompt
    assert "<!-- ENV_START -->" not in prompt


def test_agent_prompt_keeps_only_identity_and_domain_flow() -> None:
    prompt = build_agent_prompt(
        agent_type=AgentType.WORKER,
        llm_config=SystemAgentLLMConfig(),
        language="zh-CN",
    )

    assert "focus_points" in prompt
    assert "断卦要点" in prompt
    assert "[role_playing]" in prompt
    assert "[output_json_rules]" in prompt


def test_system_prompt_sections_are_not_duplicated() -> None:
    prompt = build_system_prompt(
        agent_type=AgentType.WORKER,
        language="zh-CN",
        llm_config=SystemAgentLLMConfig(),
    )

    assert prompt.count("<!-- SAFETY_START -->") == 1
    assert prompt.count("<!-- AGENT_START -->") == 1


def test_system_prompt_requires_paragraph_breaks_for_answer() -> None:
    prompt = build_agent_prompt(
        agent_type=AgentType.WORKER,
        llm_config=SystemAgentLLMConfig(),
        language="zh-CN",
    )

    assert "具体解析" in prompt


def test_user_prompt_has_language_constraint_en() -> None:
    prompt = build_follow_up_user_prompt(question="test question", language="en-US")

    assert "CRITICAL: YOUR ENTIRE RESPONSE MUST BE IN ENGLISH" in prompt
    assert "═══════════════════════════════════════════════════════════════════════════════" in prompt
    assert "[SCOPE CHECK - REFUSE IF:]" in prompt


def test_user_prompt_has_language_constraint_zh() -> None:
    prompt = build_follow_up_user_prompt(question="test question", language="zh-CN")

    assert "关键：必须全程使用简体中文回答" in prompt
    assert "═══════════════════════════════════════════════════════════════════════════════" in prompt
    assert "【范围检查" in prompt