deerflow-code/offline-backend-20260512/backend/tests/test_prompt_prefix.py

"""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.services import _prepend_prompt_prefix, _prompt_prefix_opted_in
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 == ""


def test_save_then_load_roundtrips(tmp_runtime_home: Path) -> None:
    settings = ss.SystemSettings(
        prompt_prefix=ss.PromptPrefixSettings(
            enabled=True,
            switch_label="加入立场前缀",
            prompt_prefix="你需要站在中方角度",
        )
    )
    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 == "你需要站在中方角度"


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