plyngent/tests/test_agent/test_messages_bridge.py

"""Tests for Anthropic Messages chat ↔ native conversion."""

from __future__ import annotations

from typing import TYPE_CHECKING

import msgspec
from msgspec import UNSET

from plyngent.agent.messages_bridge import (
    anthropic_response_to_assistant,
    anthropic_response_to_chat_completion,
    anthropic_stop_to_finish_reason,
    anthropic_usage_to_dict,
    chat_messages_to_anthropic,
    chat_param_to_anthropic_param,
    extended_thinking_rejected,
    provider_tools_to_anthropic_tools,
    tool_items_to_anthropic_tools,
)
from plyngent.lmproto.anthropic.model import (
    AnthropicAssistantMessage,
    AnthropicMessageResponse,
    AnthropicResponseText,
    AnthropicResponseToolUse,
    AnthropicServerToolUseContent,
    AnthropicThinkingContent,
    AnthropicToolDefinition,
    AnthropicToolResultContent,
    AnthropicToolUseContent,
    AnthropicUsage,
    AnthropicUserMessage,
    AnthropicWebSearchToolResultContent,
)
from plyngent.lmproto.openai_compatible.model import (
    AssistantChatMessage,
    AssistantFunctionTool,
    AssistantFunctionToolCall,
    ChatCompletionsParam,
    DeveloperChatMessage,
    ReasoningEffort,
    SystemChatMessage,
    ToolChatMessage,
    ToolFunction,
    ToolFunctionItem,
    UserChatMessage,
)

if TYPE_CHECKING:
    from plyngent.typedef import Unset


def test_tool_items_to_anthropic_tools() -> None:
    items = [
        ToolFunctionItem(
            function=ToolFunction(
                name="read_file",
                description="Read a file",
                parameters={"type": "object", "properties": {"path": {"type": "string"}}},
            )
        )
    ]
    tools = tool_items_to_anthropic_tools(items)
    assert len(tools) == 1
    assert tools[0].name == "read_file"
    assert tools[0].input_schema is not UNSET


def test_chat_messages_system_and_tool_round() -> None:
    messages = [
        SystemChatMessage(content="You are helpful."),
        UserChatMessage(content="hi"),
        AssistantChatMessage(
            content=None,
            tool_calls=[
                AssistantFunctionToolCall(
                    id="call_1",
                    function=AssistantFunctionTool(name="read_file", arguments='{"path":"a"}'),
                )
            ],
        ),
        ToolChatMessage(content="file body", tool_call_id="call_1"),
    ]
    system, items = chat_messages_to_anthropic(messages)
    assert system == "You are helpful."
    assert len(items) == 3
    assert isinstance(items[0], AnthropicUserMessage)
    assert items[0].content == "hi"
    assert isinstance(items[1], AnthropicAssistantMessage)
    assert isinstance(items[1].content, list)
    assert any(isinstance(b, AnthropicToolUseContent) and b.id == "call_1" for b in items[1].content)
    assert isinstance(items[2], AnthropicUserMessage)
    assert isinstance(items[2].content, list)
    assert isinstance(items[2].content[0], AnthropicToolResultContent)
    assert items[2].content[0].tool_use_id == "call_1"
    assert items[2].content[0].content == "file body"


def test_developer_becomes_user_prefix() -> None:
    system, items = chat_messages_to_anthropic(
        [
            SystemChatMessage(content="persona"),
            UserChatMessage(content="go"),
            DeveloperChatMessage(content="checkpoint"),
        ]
    )
    assert system == "persona"
    assert len(items) == 2
    assert isinstance(items[1], AnthropicUserMessage)
    assert items[1].content == "[developer] checkpoint"


def test_merge_consecutive_tool_results() -> None:
    messages = [
        UserChatMessage(content="go"),
        AssistantChatMessage(
            content=None,
            tool_calls=[
                AssistantFunctionToolCall(
                    id="c1",
                    function=AssistantFunctionTool(name="a", arguments="{}"),
                ),
                AssistantFunctionToolCall(
                    id="c2",
                    function=AssistantFunctionTool(name="b", arguments="{}"),
                ),
            ],
        ),
        ToolChatMessage(content="r1", tool_call_id="c1"),
        ToolChatMessage(content="r2", tool_call_id="c2"),
    ]
    _, items = chat_messages_to_anthropic(messages)
    assert len(items) == 3  # user, assistant, one merged tool-result user
    assert isinstance(items[2], AnthropicUserMessage)
    assert isinstance(items[2].content, list)
    assert len(items[2].content) == 2


def test_anthropic_response_to_assistant_with_tools() -> None:
    response = AnthropicMessageResponse(
        id="msg_1",
        model="claude-test",
        stop_reason="tool_use",
        content=[
            AnthropicResponseText(text="calling"),
            AnthropicResponseToolUse(id="tu_1", name="add", input={"a": 1, "b": 2}),
        ],
        usage=AnthropicUsage(input_tokens=12, output_tokens=4),
    )
    assistant = anthropic_response_to_assistant(response)
    assert assistant.content == "calling"
    assert assistant.tool_calls is not UNSET
    assert len(assistant.tool_calls) == 1
    call = assistant.tool_calls[0]
    assert isinstance(call, AssistantFunctionToolCall)
    assert call.id == "tu_1"
    assert call.function.name == "add"
    assert '"a":1' in call.function.arguments or '"a": 1' in call.function.arguments


def test_anthropic_response_to_assistant_keeps_thinking_and_search_blocks() -> None:
    """A server-side search is not a registry call, but it is the round's evidence."""
    response = AnthropicMessageResponse(
        id="msg_search",
        model="deepseek-flash",
        stop_reason="end_turn",
        content=[
            AnthropicThinkingContent(thinking="Let me search.", signature="sig-1"),
            AnthropicServerToolUseContent(
                id="call_00_1",
                name="web_search",
                input={"query": "DeepSeek-V4"},
                caller={"type": "direct"},
            ),
            AnthropicWebSearchToolResultContent(
                tool_use_id="call_00_1",
                content=[{"type": "web_search_result", "url": "https://api-docs.deepseek.com"}],
            ),
            AnthropicResponseText(text="DeepSeek-V4 is out."),
        ],
        usage=AnthropicUsage(input_tokens=16592, output_tokens=1419),
    )
    assistant = anthropic_response_to_assistant(response)
    assert assistant.content == "DeepSeek-V4 is out."
    assert assistant.tool_calls is UNSET
    # The chain of thought is the round's reasoning text; the signature it must
    # go back with only exists on the block.
    assert assistant.reasoning_content == "Let me search."
    assert assistant.provider_blocks == [
        {"type": "thinking", "thinking": "Let me search.", "signature": "sig-1"},
        {
            "type": "server_tool_use",
            "id": "call_00_1",
            "name": "web_search",
            "input": {"query": "DeepSeek-V4"},
            "caller": {"type": "direct"},
        },
        {
            "type": "web_search_tool_result",
            "tool_use_id": "call_00_1",
            "content": [{"type": "web_search_result", "url": "https://api-docs.deepseek.com"}],
        },
    ]


def test_chat_param_to_anthropic_param_hands_provider_blocks_back() -> None:
    """The next request carries the round's thinking and search blocks first."""
    assistant = AssistantChatMessage(
        content="DeepSeek-V4 is out.",
        reasoning_content="Let me search.",
        provider_blocks=[
            {"type": "thinking", "thinking": "Let me search.", "signature": "sig-1"},
            {
                "type": "web_search_tool_result",
                "tool_use_id": "call_00_1",
                "content": [{"type": "web_search_result", "encrypted_content": "sd918tZXSvVp"}],
            },
        ],
        tool_calls=[
            AssistantFunctionToolCall(
                id="tu_1",
                function=AssistantFunctionTool(name="read_file", arguments='{"path": "a.py"}'),
            )
        ],
    )
    param = ChatCompletionsParam(model="deepseek-flash", messages=[assistant])
    built = chat_param_to_anthropic_param(param)
    # The wire shape: DeepSeek answers a trimmed ``encrypted_content`` with 422.
    data = msgspec.json.decode(msgspec.json.encode(built))
    assert data["messages"] == [
        {
            "role": "assistant",
            "content": [
                {"type": "thinking", "thinking": "Let me search.", "signature": "sig-1"},
                {
                    "type": "web_search_tool_result",
                    "tool_use_id": "call_00_1",
                    "content": [{"type": "web_search_result", "encrypted_content": "sd918tZXSvVp"}],
                },
                {"type": "text", "text": "DeepSeek-V4 is out."},
                {"type": "tool_use", "id": "tu_1", "name": "read_file", "input": {"path": "a.py"}},
            ],
        }
    ]


def test_provider_tools_to_anthropic_tools_maps_the_search_intent() -> None:
    """The config default (``{"type": "web_search"}``) becomes the server tool."""
    assert provider_tools_to_anthropic_tools(None) == []
    assert provider_tools_to_anthropic_tools([]) == []
    assert provider_tools_to_anthropic_tools([{"type": "web_search"}]) == [
        {"type": "web_search_20250305", "name": "web_search"}
    ]
    # A cap (or any other key) survives the mapping.
    assert provider_tools_to_anthropic_tools([{"type": "web_search", "max_uses": 3}]) == [
        {"type": "web_search_20250305", "name": "web_search", "max_uses": 3}
    ]
    # An entry that already names the surface's tool passes through.
    assert provider_tools_to_anthropic_tools([{"type": "web_search_20260209", "name": "web_search"}]) == [
        {"type": "web_search_20260209", "name": "web_search"}
    ]
    # Entries without a type are dropped (nothing to declare on the wire).
    assert provider_tools_to_anthropic_tools([{"name": "web_search"}]) == []


def test_chat_param_to_anthropic_param_merges_provider_tools_after_local_ones() -> None:
    param = ChatCompletionsParam(
        model="deepseek-flash",
        messages=[UserChatMessage(content="hi")],
        tools=[
            ToolFunctionItem(
                function=ToolFunction(
                    name="read_file",
                    description="Read a file.",
                    parameters={"type": "object", "properties": {}},
                )
            )
        ],
    )
    built = chat_param_to_anthropic_param(param, provider_tools=[{"type": "web_search"}])
    tools = built.tools
    assert tools is not UNSET
    assert isinstance(tools[0], AnthropicToolDefinition)
    assert tools[0].name == "read_file"
    assert tools[1] == {"type": "web_search_20250305", "name": "web_search"}


def test_anthropic_response_to_chat_completion_usage() -> None:
    response = AnthropicMessageResponse(
        id="msg_1",
        model="claude-test",
        stop_reason="end_turn",
        content=[AnthropicResponseText(text="done")],
        usage=AnthropicUsage(input_tokens=9, output_tokens=3),
    )
    chat = anthropic_response_to_chat_completion(response)
    assert chat.choices[0].message.content == "done"
    assert chat.choices[0].finish_reason == "stop"
    assert chat.usage is not UNSET
    assert chat.usage["prompt_tokens"] == 9
    assert chat.usage["completion_tokens"] == 3


def test_stop_reason_mapping() -> None:
    assert anthropic_stop_to_finish_reason("end_turn", has_tool_calls=False) == "stop"
    assert anthropic_stop_to_finish_reason("tool_use", has_tool_calls=True) == "tool_calls"
    assert anthropic_stop_to_finish_reason("max_tokens", has_tool_calls=False) == "length"
    assert anthropic_stop_to_finish_reason("refusal", has_tool_calls=False) == "content_filter"
    assert anthropic_stop_to_finish_reason(None, has_tool_calls=True) == "tool_calls"


def test_usage_to_dict() -> None:
    assert anthropic_usage_to_dict(AnthropicUsage(input_tokens=1, output_tokens=2)) == {
        "prompt_tokens": 1,
        "completion_tokens": 2,
        "total_tokens": 3,
        "input_tokens": 1,
        "output_tokens": 2,
    }
    assert anthropic_usage_to_dict(AnthropicUsage()) is None


def test_chat_param_to_anthropic_param() -> None:
    param = ChatCompletionsParam(
        model="claude-test",
        messages=[
            SystemChatMessage(content="Be brief."),
            UserChatMessage(content="hi"),
        ],
        tools=[ToolFunctionItem(function=ToolFunction(name="f", description="d", parameters={"type": "object"}))],
        temperature=0.2,
        max_completion_tokens=1024,
    )
    body = chat_param_to_anthropic_param(param)
    assert body.model == "claude-test"
    assert body.max_tokens == 1024
    assert body.system == "Be brief."
    assert body.tools is not UNSET
    assert len(body.tools) == 1
    assert body.temperature == 0.2
    # No effort configured → no thinking block, provider default applies.
    assert body.thinking is UNSET


def _effort_param(effort: ReasoningEffort | Unset = UNSET) -> ChatCompletionsParam:
    return ChatCompletionsParam(
        model="claude-test",
        messages=[UserChatMessage(content="hi")],
        reasoning_effort=effort,
    )


def test_chat_param_to_anthropic_param_thinking_from_effort() -> None:
    body = chat_param_to_anthropic_param(_effort_param("medium"))
    assert body.thinking is not UNSET
    assert body.thinking.type == "enabled"
    assert body.thinking.budget_tokens == 8192
    # A budget equal to max_tokens is invalid for Anthropic, so the default
    # output limit grows to leave room after it.
    assert body.max_tokens == 8192 + 1024


def test_chat_param_to_anthropic_param_thinking_raises_max_tokens() -> None:
    # Anthropic requires budget_tokens < max_tokens; the answer keeps room after
    # the reasoning budget instead of the budget being shrunk to fit.
    body = chat_param_to_anthropic_param(_effort_param("high"))
    assert body.thinking is not UNSET
    assert body.thinking.budget_tokens == 16384
    assert body.max_tokens == 16384 + 1024


def test_chat_param_to_anthropic_param_explicit_thinking_budget() -> None:
    body = chat_param_to_anthropic_param(_effort_param("minimal"), thinking_budget_tokens=20000)
    assert body.thinking is not UNSET
    assert body.thinking.budget_tokens == 20000
    assert body.max_tokens == 20000 + 1024


def test_chat_param_to_anthropic_param_thinking_budget_floor() -> None:
    body = chat_param_to_anthropic_param(_effort_param("minimal"), thinking_budget_tokens=100)
    assert body.thinking is not UNSET
    assert body.thinking.budget_tokens == 1024
    assert body.max_tokens == 8192


def test_chat_param_to_anthropic_param_thinking_disabled() -> None:
    body = chat_param_to_anthropic_param(_effort_param("none"))
    assert body.thinking is not UNSET
    assert body.thinking.type == "disabled"
    assert body.thinking.budget_tokens is UNSET


def test_chat_param_to_anthropic_param_adaptive_thinking_from_effort() -> None:
    body = chat_param_to_anthropic_param(_effort_param("high"), adaptive_thinking=True)
    assert body.thinking is not UNSET
    assert body.thinking.type == "adaptive"
    assert body.thinking.budget_tokens is UNSET
    assert body.output_config is not UNSET
    assert body.output_config.effort == "high"
    # Adaptive thinking spends no fixed budget, so max_tokens stays as asked.
    assert body.max_tokens == 8192


def test_chat_param_to_anthropic_param_adaptive_effort_levels() -> None:
    def effort_for(level: ReasoningEffort) -> str:
        body = chat_param_to_anthropic_param(_effort_param(level), adaptive_thinking=True)
        assert body.output_config is not UNSET
        return str(body.output_config.effort)

    assert effort_for("minimal") == "low"  # no such API level; the lowest one
    assert effort_for("xhigh") == "xhigh"


def test_chat_param_to_anthropic_param_adaptive_ignores_a_configured_budget() -> None:
    body = chat_param_to_anthropic_param(_effort_param("medium"), thinking_budget_tokens=20000, adaptive_thinking=True)
    assert body.thinking is not UNSET
    assert body.thinking.type == "adaptive"
    assert body.thinking.budget_tokens is UNSET
    assert body.max_tokens == 8192


def test_chat_param_to_anthropic_param_adaptive_disabled_carries_no_effort() -> None:
    body = chat_param_to_anthropic_param(_effort_param("none"), adaptive_thinking=True)
    assert body.thinking is not UNSET
    assert body.thinking.type == "disabled"
    assert body.output_config is UNSET


def test_chat_param_to_anthropic_param_adaptive_without_effort_sends_nothing() -> None:
    body = chat_param_to_anthropic_param(_effort_param(), adaptive_thinking=True)
    assert body.thinking is UNSET
    assert body.output_config is UNSET


def test_extended_thinking_rejected_matches_the_api_and_gateway_wording() -> None:
    documented = (
        'messages HTTP 400: {"error":{"message":"\\"thinking.type.enabled\\" is not supported for this '
        'model. Use \\"thinking.type.adaptive\\" and \\"output_config.effort\\" to control thinking behavior."}}'
    )
    gateway = (
        'messages HTTP 400: {"error":{"message":"claude-opus-5-5 requires adaptive thinking; omit thinking '
        'or use thinking.type=adaptive and output_config.effort"}}'
    )
    assert extended_thinking_rejected(documented)
    assert extended_thinking_rejected(gateway)
    assert not extended_thinking_rejected('messages HTTP 400: {"error":{"message":"max_tokens: required"}}')