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 schemas.agent.system_agent import AgentType, SystemAgentLLMConfig


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

    assert "English" in prompt
    assert "<!-- SAFETY_START -->" in prompt
    assert "<!-- OUTPUT_START -->" in prompt


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

    assert "简体中文" in prompt
    assert "<!-- SAFETY_START -->" 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(),
    )

    assert "focus_points" in prompt
    assert "段间用\\n\\n" 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
    assert prompt.count("<!-- OUTPUT_START -->") == 1


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

    assert "段间用\\n\\n" in prompt