dify/api/tests/unit_tests/services/agent/test_prompt_mentions.py
yyh 9b432d0d29
fix: enforce tool provider type contracts (#39380)
Co-authored-by: autofix-ci[bot] <114827586+autofix-ci[bot]@users.noreply.github.com>
2026-07-22 01:30:55 +00:00

296 lines
13 KiB
Python

"""Unit tests for the prompt mention contract (ENG-616).
Token form: ``[§<kind>:<id>[:<label>]§]``. Mentions are pointers into the Agent
config lists; expansion replaces them with canonical names and the scrub pass
guarantees no mention-shaped marker survives to the model.
"""
from __future__ import annotations
from urllib.parse import quote
import pytest
from models.agent_config_entities import AgentSoulConfig, WorkflowNodeJobConfig, WorkflowPreviousNodeOutputRef
from services.agent.prompt_mentions import (
MAX_MENTION_REF_ID_LENGTH,
NODE_JOB_PROMPT_ALLOWED_KINDS,
SOUL_PROMPT_ALLOWED_KINDS,
MentionKind,
build_node_job_mention_resolver,
build_soul_mention_resolver,
expand_prompt_mentions,
extract_workflow_node_output_selectors,
extract_workflow_variable_selectors,
normalize_previous_node_output_selector,
parse_prompt_mentions,
scrub_mention_markers,
workflow_previous_node_output_refs_from_selectors,
)
# ── parse ─────────────────────────────────────────────────────────────────────
def test_parse_extracts_kind_id_and_optional_label():
prompt = "Use [§skill:tender-analyzer%2FSKILL.md:tender-analyzer§] then ask [§human:c-1§]."
mentions = parse_prompt_mentions(prompt)
assert [(m.kind, m.ref_id, m.label) for m in mentions] == [
(MentionKind.SKILL, "tender-analyzer%2FSKILL.md", "tender-analyzer"),
(MentionKind.HUMAN, "c-1", None),
]
assert prompt[mentions[0].start : mentions[0].end] == mentions[0].raw
def test_parse_supports_ids_with_slash_and_dot():
mentions = parse_prompt_mentions("[§tool:langgenius/tavily/tavily_search:tavily§] [§node_output:node-1.tenders§]")
assert mentions[0].ref_id == "langgenius/tavily/tavily_search"
assert mentions[1].ref_id == "node-1.tenders"
def test_parse_supports_legacy_bare_output_tokens():
mentions = parse_prompt_mentions("Use §output:summary:summary§")
assert [(mention.kind, mention.ref_id, mention.label) for mention in mentions] == [
(MentionKind.OUTPUT, "summary", "summary")
]
def test_parse_ignores_legacy_template_forms_and_unknown_kinds():
prompt = "{{var}} {{#context#}} {{#sys.query#}} [§bogus_kind:x§]"
assert parse_prompt_mentions(prompt) == []
def test_parse_skips_oversized_id_or_label():
long_id = "x" * (MAX_MENTION_REF_ID_LENGTH + 1)
assert parse_prompt_mentions(f"[§skill:{long_id}§]") == []
def test_parse_accepts_long_unicode_encoded_drive_key_within_drive_limit():
encoded_drive_key = quote("" * 512)
mentions = parse_prompt_mentions(f"[§skill:{encoded_drive_key}:Long Skill§]")
assert [(mention.kind, mention.ref_id) for mention in mentions] == [(MentionKind.SKILL, encoded_drive_key)]
# ── expand + scrub ────────────────────────────────────────────────────────────
def test_expand_uses_resolver_and_degrades_unresolved_to_label_then_id():
prompt = "A [§skill:s1:Skill One§] B [§human:h1:EMAIL · DAVE§] C [§knowledge:k1§]"
def resolver(mention):
return "resolved-skill" if mention.kind == MentionKind.SKILL else None
expanded = expand_prompt_mentions(prompt, resolver)
assert expanded == "A resolved-skill B EMAIL · DAVE C k1"
assert "" not in expanded
def test_expand_scrubs_unknown_kind_tokens_but_keeps_legacy_forms():
prompt = "x [§wat:id-1:Label§] y {{#context#}} z {{#node.var#}}"
expanded = expand_prompt_mentions(prompt, lambda m: None)
# unknown mention-shaped token degraded to its label; legacy forms untouched
assert expanded == "x Label y {{#context#}} z {{#node.var#}}"
def test_scrub_degrades_colon_tokens_without_label_to_id_part():
assert scrub_mention_markers("see [§weird_kind:some-id§]") == "see some-id"
def test_expand_empty_prompt_is_noop():
assert expand_prompt_mentions("", lambda m: "x") == ""
def test_extract_workflow_variable_selectors_supports_frontend_agent_task_format():
selectors = extract_workflow_variable_selectors(
"Read {{#node-1.output#}}, compare {{#start.question#}}, leave {{#context#}} alone."
)
assert selectors == [("node-1", "output"), ("start", "question")]
def test_extract_workflow_node_output_selectors_supports_current_frontend_markers_only():
selectors = extract_workflow_node_output_selectors(
"Read {{#node-1.output#}}, {{#sys.query#}}, [§node_output:legacy-node.report:PREV/report§], "
"and {{#node-1.output#}} again."
)
assert selectors == [("node-1", "output")]
def test_workflow_previous_node_output_refs_from_selectors_materializes_refs():
refs = workflow_previous_node_output_refs_from_selectors([("node-1", "output"), ("node-2", "report", "url")])
assert refs == [
WorkflowPreviousNodeOutputRef(selector=["node-1", "output"], node_id="node-1", output="output"),
WorkflowPreviousNodeOutputRef(
selector=["node-2", "report", "url"],
node_id="node-2",
output="report",
),
]
# ── soul resolver ─────────────────────────────────────────────────────────────
@pytest.fixture
def soul() -> AgentSoulConfig:
return AgentSoulConfig.model_validate(
{
"tools": {
"dify_tools": [
{
"plugin_id": "langgenius/tavily",
"provider": "tavily",
"provider_type": "plugin",
"tool_name": "tavily_search",
"credential_type": "unauthorized",
},
],
"cli_tools": [{"id": "ct-1", "name": "ffmpeg"}],
},
"knowledge": {
"sets": [
{
"id": "kb-1",
"name": "产品手册",
"datasets": [{"id": "ds-1", "name": "产品手册"}],
"query": {"mode": "generated_query"},
"retrieval": {"mode": "multiple", "top_k": 4},
}
]
},
"human": {"contacts": [{"id": "c-1", "name": "David Hayes", "channel": "email"}]},
}
)
def test_soul_resolver_resolves_each_kind(soul: AgentSoulConfig):
resolver = build_soul_mention_resolver(soul)
prompt = (
"Use [§tool:tavily/tavily_search:tavily§], run [§cli_tool:ct-1:ffmpeg§], "
"ground in [§knowledge:kb-1§], ask [§human:c-1§]."
)
expanded = expand_prompt_mentions(prompt, resolver)
assert expanded == ("Use tavily_search, run ffmpeg, ground in 产品手册, ask EMAIL · David Hayes.")
def test_soul_resolver_unknown_ids_degrade(soul: AgentSoulConfig):
expanded = expand_prompt_mentions("[§knowledge:missing:旧产品手册§]", build_soul_mention_resolver(soul))
assert expanded == "旧产品手册"
def test_soul_resolver_cli_tool_resolves_by_id_and_keeps_name_alias(soul: AgentSoulConfig):
resolver = build_soul_mention_resolver(soul)
# id is the contract; the name alias keeps tokens minted before ids existed working
assert expand_prompt_mentions("[§cli_tool:ct-1§]", resolver) == "ffmpeg"
assert expand_prompt_mentions("[§cli_tool:ffmpeg§]", resolver) == "ffmpeg"
# a rename only changes the expansion, never breaks the id reference
soul.tools.cli_tools[0].name = "ffmpeg-v7"
assert expand_prompt_mentions("[§cli_tool:ct-1§]", build_soul_mention_resolver(soul)) == "ffmpeg-v7"
@pytest.fixture
def soul_with_provider_entry(soul: AgentSoulConfig) -> AgentSoulConfig:
# provider-level entry (tool_name omitted) = all tools of the provider
soul.tools.dify_tools.append(
soul.tools.dify_tools[0].model_copy(
update={"plugin_id": "langgenius/duckduckgo", "provider": "duckduckgo", "tool_name": None}
)
)
return soul
def test_soul_resolver_provider_all_tools_mention(soul_with_provider_entry: AgentSoulConfig):
resolver = build_soul_mention_resolver(soul_with_provider_entry)
# [§tool:<provider>/*§] = all tools of that provider
assert expand_prompt_mentions("Use [§tool:duckduckgo/*:DuckDuckGo 全部§].", resolver) == (
"Use all duckduckgo tools."
)
# plugin-prefixed alias of the same provider
assert expand_prompt_mentions("[§tool:langgenius/duckduckgo/duckduckgo/*§]", resolver) == "all duckduckgo tools"
# without a provider-level entry the mention dangles -> degrades to label
bare = build_soul_mention_resolver(AgentSoulConfig.model_validate({}))
assert expand_prompt_mentions("[§tool:duckduckgo/*:DuckDuckGo 全部§]", bare) == "DuckDuckGo 全部"
def test_soul_resolver_single_tool_resolves_via_provider_level_entry(soul_with_provider_entry: AgentSoulConfig):
# one tool offered through the provider-level ("all") entry still resolves
resolver = build_soul_mention_resolver(soul_with_provider_entry)
assert expand_prompt_mentions("[§tool:duckduckgo/ddg_search§]", resolver) == "ddg_search"
# ── node-job resolver ─────────────────────────────────────────────────────────
@pytest.fixture
def node_job() -> WorkflowNodeJobConfig:
return WorkflowNodeJobConfig.model_validate(
{
"workflow_prompt": "",
"previous_node_output_refs": [{"selector": ["start-1", "tenders"], "name": "START/tenders"}],
# declared output names are JSON-schema-friendly identifiers (no dots)
"declared_outputs": [{"name": "qna_report", "type": "file"}],
"human_contacts": [{"id": "c-1", "name": "David Hayes", "channel": "email"}],
}
)
def test_node_job_resolver_resolves_each_kind(node_job: WorkflowNodeJobConfig):
resolver = build_node_job_mention_resolver(node_job)
prompt = "Read [§node_output:start-1.tenders§] and produce [§output:qna_report§]; if unsure contact [§human:c-1§]."
expanded = expand_prompt_mentions(prompt, resolver)
assert expanded == (
"Read START/tenders and produce qna_report (file output; create the file locally, run "
"`dify-agent file upload <path>`, then set final_output.qna_report to a `tool_file` mapping "
"using the returned `reference`; if replying to the user in natural language, use the returned "
"`download_url`; do not call final_output before upload succeeds, and do not use the local path, "
"filename, URL, or a synthesized dify-file-ref as the reference); "
"if unsure contact EMAIL · David Hayes."
)
def test_node_job_resolver_accepts_legacy_reversed_output_token(node_job: WorkflowNodeJobConfig):
resolver = build_node_job_mention_resolver(node_job)
expanded = expand_prompt_mentions("[§qna_report:qna_report:output§]", resolver)
assert "final_output.qna_report" in expanded
def test_node_job_resolver_matches_ref_by_node_id_and_output_fields():
node_job = WorkflowNodeJobConfig.model_validate(
{"previous_node_output_refs": [{"node_id": "n-2", "output": "text"}]}
)
expanded = expand_prompt_mentions("[§node_output:n-2.text:LLM/text§]", build_node_job_mention_resolver(node_job))
# ref has no display name -> degrade to the mention label
assert expanded == "LLM/text"
def test_normalize_previous_node_output_selector_returns_canonical_selector():
assert normalize_previous_node_output_selector(
WorkflowPreviousNodeOutputRef(selector=["node-1", "report", "url"])
) == ("node-1", "report", "url")
assert normalize_previous_node_output_selector(WorkflowPreviousNodeOutputRef(node_id="node-2", output="text")) == (
"node-2",
"text",
)
assert normalize_previous_node_output_selector(WorkflowPreviousNodeOutputRef(selector=["node-3", 1])) is None
# ── allowlists ────────────────────────────────────────────────────────────────
def test_per_surface_allowlists_match_design():
assert {
MentionKind.SKILL,
MentionKind.FILE,
MentionKind.TOOL,
MentionKind.CLI_TOOL,
MentionKind.KNOWLEDGE,
MentionKind.HUMAN,
} == SOUL_PROMPT_ALLOWED_KINDS
assert {MentionKind.NODE_OUTPUT, MentionKind.OUTPUT, MentionKind.HUMAN} == NODE_JOB_PROMPT_ALLOWED_KINDS