71 lines
3.0 KiB
Python
71 lines
3.0 KiB
Python
from types import SimpleNamespace
|
|
|
|
from deerflow.agents.lead_agent import prompt as prompt_module
|
|
from deerflow.config.skills_config import EsQueryRoutingConfig, SkillsConfig
|
|
|
|
|
|
def _app_config(*, enabled: bool, prompt: str = "自定义 ES 路由规则") -> SimpleNamespace:
|
|
return SimpleNamespace(
|
|
skills=SkillsConfig(
|
|
es_query_routing=EsQueryRoutingConfig(enabled=enabled, prompt=prompt),
|
|
)
|
|
)
|
|
|
|
|
|
def test_es_query_routing_config_is_disabled_by_default() -> None:
|
|
config = SkillsConfig()
|
|
|
|
assert config.es_query_routing.enabled is False
|
|
assert "es_query" in config.es_query_routing.prompt
|
|
|
|
|
|
def test_es_query_routing_section_honors_switch_prompt_and_allowlist() -> None:
|
|
enabled = _app_config(enabled=True, prompt=" 自定义 ES 路由规则 ")
|
|
|
|
assert prompt_module._build_es_query_routing_section(app_config=enabled) == (
|
|
'<skill_routing skill="es_query">\n自定义 ES 路由规则\n</skill_routing>'
|
|
)
|
|
assert (
|
|
prompt_module._build_es_query_routing_section(
|
|
app_config=enabled,
|
|
available_skills={"other_skill"},
|
|
)
|
|
== ""
|
|
)
|
|
assert "自定义 ES 路由规则" in prompt_module._build_es_query_routing_section(
|
|
app_config=enabled,
|
|
available_skills={"es_query"},
|
|
)
|
|
assert prompt_module._build_es_query_routing_section(app_config=_app_config(enabled=False)) == ""
|
|
assert prompt_module._build_es_query_routing_section(app_config=_app_config(enabled=True, prompt=" ")) == ""
|
|
|
|
|
|
def test_es_query_routing_is_injected_into_lead_and_allowed_custom_agent_prompts(monkeypatch) -> None:
|
|
config = _app_config(enabled=True)
|
|
monkeypatch.setattr(prompt_module, "_get_memory_context", lambda *args, **kwargs: "")
|
|
monkeypatch.setattr(prompt_module, "get_agent_soul", lambda *args, **kwargs: "")
|
|
monkeypatch.setattr(prompt_module, "get_skills_prompt_section", lambda *args, **kwargs: "<skill_system />")
|
|
monkeypatch.setattr(prompt_module, "get_deferred_tools_prompt_section", lambda *args, **kwargs: "")
|
|
monkeypatch.setattr(prompt_module, "_build_acp_section", lambda *args, **kwargs: "")
|
|
monkeypatch.setattr(prompt_module, "_build_custom_mounts_section", lambda *args, **kwargs: "")
|
|
monkeypatch.setattr(prompt_module, "_current_time_context", lambda: "")
|
|
|
|
lead_prompt = prompt_module.apply_prompt_template(app_config=config)
|
|
allowed_custom_prompt = prompt_module.apply_agent_only_prompt_template(
|
|
agent_id="agent-with-es",
|
|
agent_name="Agent with ES",
|
|
available_skills={"es_query"},
|
|
app_config=config,
|
|
)
|
|
denied_custom_prompt = prompt_module.apply_agent_only_prompt_template(
|
|
agent_id="agent-without-es",
|
|
agent_name="Agent without ES",
|
|
available_skills={"other_skill"},
|
|
app_config=config,
|
|
)
|
|
|
|
assert '<skill_routing skill="es_query">' in lead_prompt
|
|
assert "自定义 ES 路由规则" in lead_prompt
|
|
assert '<skill_routing skill="es_query">' in allowed_custom_prompt
|
|
assert '<skill_routing skill="es_query">' not in denied_custom_prompt
|