238 lines
8.6 KiB
Python
238 lines
8.6 KiB
Python
"""Tests for the admin-configurable prompt-prefix toggle."""
|
|
|
|
from __future__ import annotations
|
|
|
|
import json
|
|
from pathlib import Path
|
|
from types import SimpleNamespace
|
|
|
|
import pytest
|
|
from langchain_core.messages import AIMessage, HumanMessage
|
|
|
|
from app.gateway.routers import system_settings as ss_router
|
|
from app.gateway.services import _prepend_prompt_prefix, _prompt_prefix_opted_in
|
|
from deerflow.agents.lead_agent import prompt as prompt_module
|
|
from deerflow.config import system_settings as ss
|
|
|
|
|
|
@pytest.fixture
|
|
def tmp_runtime_home(monkeypatch: pytest.MonkeyPatch, tmp_path: Path) -> Path:
|
|
"""Redirect ``runtime_home()`` to a per-test tmp dir."""
|
|
|
|
def fake_home() -> Path:
|
|
return tmp_path
|
|
|
|
monkeypatch.setattr(ss, "runtime_home", fake_home)
|
|
return tmp_path
|
|
|
|
|
|
def test_load_returns_defaults_when_missing(tmp_runtime_home: Path) -> None:
|
|
settings = ss.load_system_settings()
|
|
assert settings.prompt_prefix.enabled is False
|
|
assert settings.prompt_prefix.prompt_prefix == ""
|
|
assert settings.prompt_prefix.ordinary_qa_markdown_format_enabled is False
|
|
assert settings.citation_display.reference_mode_enabled is False
|
|
assert settings.citation_display.cleanup_words == []
|
|
|
|
|
|
def test_save_then_load_roundtrips(tmp_runtime_home: Path) -> None:
|
|
settings = ss.SystemSettings(
|
|
prompt_prefix=ss.PromptPrefixSettings(
|
|
enabled=True,
|
|
switch_label="加入立场前缀",
|
|
prompt_prefix="你需要站在中方角度",
|
|
ordinary_qa_markdown_format_enabled=True,
|
|
)
|
|
)
|
|
ss.save_system_settings(settings)
|
|
|
|
file_path = tmp_runtime_home / "system_settings.json"
|
|
assert file_path.is_file()
|
|
saved = json.loads(file_path.read_text(encoding="utf-8"))
|
|
assert saved["prompt_prefix"]["prompt_prefix"] == "你需要站在中方角度"
|
|
|
|
reloaded = ss.load_system_settings()
|
|
assert reloaded.prompt_prefix.enabled is True
|
|
assert reloaded.prompt_prefix.switch_label == "加入立场前缀"
|
|
assert reloaded.prompt_prefix.prompt_prefix == "你需要站在中方角度"
|
|
assert reloaded.prompt_prefix.ordinary_qa_markdown_format_enabled is True
|
|
|
|
|
|
def test_runtime_markdown_format_switch_defaults_off_and_updates_without_restart(
|
|
tmp_runtime_home: Path,
|
|
) -> None:
|
|
assert prompt_module._ordinary_qa_markdown_format_enabled() is False
|
|
|
|
ss.save_system_settings(
|
|
ss.SystemSettings(
|
|
prompt_prefix=ss.PromptPrefixSettings(
|
|
ordinary_qa_markdown_format_enabled=True,
|
|
)
|
|
)
|
|
)
|
|
|
|
assert prompt_module._ordinary_qa_markdown_format_enabled() is True
|
|
|
|
|
|
def _patch_admin(monkeypatch: pytest.MonkeyPatch) -> None:
|
|
admin = SimpleNamespace(
|
|
id="admin1",
|
|
system_role="admin",
|
|
email="admin@example.com",
|
|
)
|
|
|
|
async def _optional(_request):
|
|
return admin
|
|
|
|
monkeypatch.setattr(ss_router, "get_optional_user_from_request", _optional)
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_prompt_settings_update_can_enable_markdown_format_without_changing_prefix(
|
|
tmp_runtime_home: Path,
|
|
monkeypatch: pytest.MonkeyPatch,
|
|
) -> None:
|
|
_patch_admin(monkeypatch)
|
|
current = ss.SystemSettings(
|
|
prompt_prefix=ss.PromptPrefixSettings(
|
|
enabled=True,
|
|
switch_label="加入立场前缀",
|
|
prompt_prefix="你需要站在中方角度",
|
|
)
|
|
)
|
|
ss.save_system_settings(current)
|
|
|
|
response = await ss_router.update_prompt_prefix(
|
|
ss_router.PromptPrefixUpdateRequest(ordinary_qa_markdown_format_enabled=True),
|
|
SimpleNamespace(),
|
|
)
|
|
|
|
assert response.ordinary_qa_markdown_format_enabled is True
|
|
assert response.enabled is True
|
|
assert response.switch_label == "加入立场前缀"
|
|
assert response.prompt_prefix == "你需要站在中方角度"
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_citation_display_update_preserves_unspecified_fields(
|
|
tmp_runtime_home: Path,
|
|
monkeypatch: pytest.MonkeyPatch,
|
|
) -> None:
|
|
_patch_admin(monkeypatch)
|
|
current = ss.SystemSettings(
|
|
citation_display=ss.CitationDisplaySettings(
|
|
cleanup_words=["广告", "免责声明"],
|
|
)
|
|
)
|
|
ss.save_system_settings(current)
|
|
|
|
response = await ss_router.update_citation_display(
|
|
ss_router.CitationDisplayUpdateRequest(reference_mode_enabled=True),
|
|
SimpleNamespace(),
|
|
)
|
|
|
|
assert response.reference_mode_enabled is True
|
|
assert response.cleanup_words == ["广告", "免责声明"]
|
|
|
|
payload = json.loads((tmp_runtime_home / "system_settings.json").read_text(encoding="utf-8"))
|
|
assert payload["citation_display"]["reference_mode_enabled"] is True
|
|
assert payload["citation_display"]["cleanup_words"] == ["广告", "免责声明"]
|
|
|
|
|
|
def test_prepend_prefix_tags_first_human_without_touching_content() -> None:
|
|
graph_input = {"messages": [HumanMessage(content="评价一下中美贸易")]}
|
|
out = _prepend_prompt_prefix(graph_input, "你需要站在中方角度")
|
|
msg = out["messages"][0]
|
|
assert isinstance(msg, HumanMessage)
|
|
# Content stays clean so the frontend never shows the prefix.
|
|
assert msg.content == "评价一下中美贸易"
|
|
# The prefix rides along in additional_kwargs for the middleware to inject.
|
|
assert msg.additional_kwargs["prompt_prefix"] == "你需要站在中方角度"
|
|
|
|
|
|
def test_prepend_prefix_preserves_block_content() -> None:
|
|
graph_input = {
|
|
"messages": [
|
|
HumanMessage(content=[{"type": "text", "text": "评价一下中美贸易"}])
|
|
]
|
|
}
|
|
out = _prepend_prompt_prefix(graph_input, "你需要站在中方角度")
|
|
msg = out["messages"][0]
|
|
assert msg.content == [{"type": "text", "text": "评价一下中美贸易"}]
|
|
assert msg.additional_kwargs["prompt_prefix"] == "你需要站在中方角度"
|
|
|
|
|
|
def test_prepend_prefix_only_touches_first_human() -> None:
|
|
graph_input = {
|
|
"messages": [
|
|
AIMessage(content="hi"),
|
|
HumanMessage(content="q1"),
|
|
HumanMessage(content="q2"),
|
|
]
|
|
}
|
|
out = _prepend_prompt_prefix(graph_input, "PREFIX")
|
|
# first human is index 1; index 2 must remain untouched
|
|
assert out["messages"][1].additional_kwargs.get("prompt_prefix") == "PREFIX"
|
|
assert out["messages"][1].content == "q1"
|
|
assert out["messages"][2].additional_kwargs.get("prompt_prefix") is None
|
|
assert out["messages"][2].content == "q2"
|
|
|
|
|
|
def test_prepend_prefix_no_messages_is_noop() -> None:
|
|
assert _prepend_prompt_prefix({"messages": []}, "X") == {"messages": []}
|
|
assert _prepend_prompt_prefix({}, "X") == {}
|
|
|
|
|
|
def test_prompt_prefix_middleware_injects_prefix_for_model_only() -> None:
|
|
"""The middleware prepends the prefix to the request copy, not to state."""
|
|
from deerflow.agents.middlewares.prompt_prefix_middleware import PromptPrefixMiddleware
|
|
|
|
mw = PromptPrefixMiddleware()
|
|
tagged = HumanMessage(
|
|
content="评价一下中美贸易",
|
|
additional_kwargs={"prompt_prefix": "你需要站在中方角度"},
|
|
)
|
|
patched = PromptPrefixMiddleware._apply([tagged])
|
|
assert patched is not None
|
|
assert patched[0].content == "你需要站在中方角度\n\n评价一下中美贸易"
|
|
# The original message in state is untouched.
|
|
assert tagged.content == "评价一下中美贸易"
|
|
assert mw is not None
|
|
|
|
|
|
def test_prompt_prefix_middleware_noop_without_tag() -> None:
|
|
from deerflow.agents.middlewares.prompt_prefix_middleware import PromptPrefixMiddleware
|
|
|
|
plain = HumanMessage(content="just a question")
|
|
assert PromptPrefixMiddleware._apply([plain]) is None
|
|
|
|
|
|
def test_prompt_prefix_middleware_block_content() -> None:
|
|
from deerflow.agents.middlewares.prompt_prefix_middleware import PromptPrefixMiddleware
|
|
|
|
tagged = HumanMessage(
|
|
content=[{"type": "text", "text": "评价一下中美贸易"}],
|
|
additional_kwargs={"prompt_prefix": "你需要站在中方角度"},
|
|
)
|
|
patched = PromptPrefixMiddleware._apply([tagged])
|
|
assert patched is not None
|
|
assert patched[0].content[0] == {"type": "text", "text": "你需要站在中方角度"}
|
|
assert patched[0].content[1] == {"type": "text", "text": "评价一下中美贸易"}
|
|
|
|
|
|
def test_opted_in_via_context() -> None:
|
|
body = SimpleNamespace(context={"prompt_prefix_enabled": True}, metadata=None)
|
|
assert _prompt_prefix_opted_in(body) is True
|
|
|
|
|
|
def test_opted_in_via_metadata() -> None:
|
|
body = SimpleNamespace(context=None, metadata={"prompt_prefix_enabled": True})
|
|
assert _prompt_prefix_opted_in(body) is True
|
|
|
|
|
|
def test_not_opted_in_by_default() -> None:
|
|
body = SimpleNamespace(context=None, metadata=None)
|
|
assert _prompt_prefix_opted_in(body) is False
|
|
body2 = SimpleNamespace(context={"mode": "ultra"}, metadata={})
|
|
assert _prompt_prefix_opted_in(body2) is False
|