116 lines
3.8 KiB
Python
116 lines
3.8 KiB
Python
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
|