litellm/tests/e2e/llm_translation/test_messages_e2e.py

"""Live e2e: POST /v1/messages (Anthropic Messages API) returns a real completion.

Registers an Anthropic deployment at runtime, drives the Messages endpoint through
the gateway, and asserts an assistant message with text came back, both
non-streaming and streamed. Migrated from
litellm-regression-tests/tests/test_inference_endpoints.py.
"""

from __future__ import annotations

from typing import Final

import pytest
from e2e_config import (
    STREAM_MIN_LEAD_SECONDS,
    provider_edge_base,
    provider_paces_stream,
    unique_marker,
)
from e2e_http import assert_client_error, require_successful_call, unwrap
from endpoints_client import EndpointsClient, MessagesResult
from lifecycle import ResourceManager
from models import (
    AnthropicAssistantTurn,
    AnthropicContentBlock,
    AnthropicCustomTool,
    AnthropicToolChoice,
    AnthropicToolResultBlock,
    AnthropicToolResultTurn,
    AnthropicMessagesBody,
    ChatMessage,
    JsonSchemaProperty,
    LiteLLMParamsBody,
    SpendLogRow,
    ToolInputSchema,
)
from pydantic import BaseModel, ConfigDict

pytestmark = [pytest.mark.e2e, pytest.mark.replayable]


class _OptionalMessagesBody(BaseModel):
    model: str | None = None
    messages: list[ChatMessage] | None = None
    max_tokens: int | None = None


class _MessagesEventDelta(BaseModel):
    text: str = ""


class _MessagesEventUsage(BaseModel):
    output_tokens: int | None = None


class _MessagesStreamEvent(BaseModel):
    """One Anthropic SSE event, keeping only what the stream's shape is asserted on.

    ``delta.text`` is populated on ``content_block_delta`` and absent on the
    ``message_delta`` that closes the turn, which is the event carrying ``usage``."""

    type: str
    delta: _MessagesEventDelta | None = None
    usage: _MessagesEventUsage | None = None


ANTHROPIC_BACKEND = "anthropic/claude-haiku-4-5"

WEATHER_TOOL = AnthropicCustomTool(
    name="get_weather",
    description="Get the current weather for a city.",
    input_schema=ToolInputSchema(
        properties={"city": JsonSchemaProperty(type="string")},
        required=["city"],
    ),
)


def _approx_equal(actual: float, expected: float) -> bool:
    """Within 1% or 1e-9 absolute - spend math, not exact float identity."""
    return abs(actual - expected) <= max(1e-9, abs(expected) * 1e-2)


def _anthropic_params() -> LiteLLMParamsBody:
    """The Anthropic deployment, wired through the record/replay edge when a fixture
    mode is active (LIT-5974). The mount base carries no ``/v1``: litellm's Anthropic
    handler appends ``/v1/messages`` to ``api_base`` itself, where the OpenAI handler
    appends only ``/chat/completions``."""
    base = provider_edge_base("anthropic")
    return LiteLLMParamsBody(
        model=ANTHROPIC_BACKEND, api_key="os.environ/ANTHROPIC_API_KEY", api_base=base
    )


class TestAnthropicMessages:
    def _register(
        self,
        endpoints_client: EndpointsClient,
        resources: ResourceManager,
        params: LiteLLMParamsBody | None = None,
    ) -> tuple[str, str]:
        model = f"e2e-messages-{unique_marker()}"
        model_id = endpoints_client.create_model(
            model, _anthropic_params() if params is None else params
        )
        resources.defer(lambda: endpoints_client.delete_model(model_id))
        return model, resources.key()

    @pytest.mark.covers("llm.messages.anthropic.basic.nonstream.works")
    def test_messages_returns_completion(
        self, endpoints_client: EndpointsClient, resources: ResourceManager
    ) -> None:
        model, key = self._register(endpoints_client, resources)

        result = endpoints_client.messages(key, model, "reply with one word")
        require_successful_call(result)
        parsed = MessagesResult.model_validate_json(result.body)
        assert parsed.role == "assistant", f"unexpected role: {result.body[:300]}"
        assert parsed.text.strip(), f"/v1/messages returned no text: {result.body[:300]}"

    @pytest.mark.covers("llm.messages.anthropic.basic.nonstream.cost_logged")
    def test_messages_logs_cost_matching_the_response_header(
        self, endpoints_client: EndpointsClient, resources: ResourceManager
    ) -> None:
        model = f"e2e-messages-cost-{unique_marker()}"
        model_id = endpoints_client.create_model(model, _anthropic_params())
        resources.defer(lambda: endpoints_client.delete_model(model_id))
        key = resources.key()

        result = endpoints_client.messages(key, model, f"reply with one word {unique_marker()}")
        require_successful_call(result)
        parsed = MessagesResult.model_validate_json(result.body)
        assert parsed.role == "assistant" and parsed.text.strip(), (
            f"/v1/messages returned no assistant text: {result.body[:300]}"
        )

        # The customer reads per-request cost off the response header (LIT-4076), so
        # it must be present and positive on /v1/messages, not only /chat/completions.
        header_cost = result.response_cost
        assert header_cost is not None and header_cost > 0, (
            "x-litellm-response-cost header missing or non-positive on /v1/messages; "
            f"headers={result.headers}"
        )

        # Correlate the spend row by the unique scoped key, not the Anthropic response
        # id: on /v1/messages the spend-log request_id is the proxy's own call id, which
        # need not equal the message body id, so an id-based poll can miss a correctly
        # logged row and time out. The key is fresh per test, so its only priced row is
        # this call.
        def _priced(rows: list[SpendLogRow]) -> bool:
            return any(r.spend is not None and r.spend > 0 for r in rows)

        rows = endpoints_client.proxy.poll_logs_for_key(key, predicate=_priced)
        priced = [r for r in rows if r.spend is not None and r.spend > 0]
        assert priced, (
            f"no priced /spend/logs row landed for key {key} within the poll window; got {rows}"
        )
        row = priced[0]
        assert (row.prompt_tokens or 0) > 0 and (row.completion_tokens or 0) > 0, (
            f"messages spend row missing token counts, so the cost is not real usage: {row}"
        )
        assert row.spend is not None and _approx_equal(row.spend, header_cost), (
            f"logged spend {row.spend} disagrees with the x-litellm-response-cost header {header_cost}; "
            "the customer bills against the header, so the two must match"
        )

    @pytest.mark.covers("llm.messages.anthropic.basic.stream.works")
    def test_messages_streams_completion(
        self, endpoints_client: EndpointsClient, resources: ResourceManager
    ) -> None:
        """Edge-wired like its non-streaming siblings, so record and replay both
        carry the streamed response.

        Asserts what the proxy controls: the event grammar (usage between the last
        content delta and ``message_stop``) and, on the clock, that the relay is
        incremental. How many deltas a reply is split into is the provider's choice, so
        the first content delta must instead reach the client well before
        ``message_stop``, which a buffered response cannot do. Replay serves chunks back
        to back, so only live and record runs judge the timing."""
        model, key = self._register(endpoints_client, resources)

        result = endpoints_client.proxy.messages_stream(
            key,
            AnthropicMessagesBody(
                model=model,
                max_tokens=800,
                stream=True,
                messages=[ChatMessage(role="user", content="Count from 1 to 200, one number per line.")],
            ),
        )
        require_successful_call(result)
        assert result.is_streaming, f"response was not streamed: {result.headers}"
        assert not result.stream_error, f"stream errored: {result.stream_error}"
        assert result.stream_events, "stream produced no SSE events"

        events = [
            _MessagesStreamEvent.model_validate_json(event) for event in result.stream_events
        ]
        types = [event.type for event in events]
        delta_positions = [
            index for index, event in enumerate(events) if event.type == "content_block_delta"
        ]
        assert delta_positions, f"stream carried no content deltas: {types}"
        text = "".join(
            event.delta.text
            for event in events
            if event.type == "content_block_delta" and event.delta is not None
        )
        assert text.strip(), f"content deltas assembled to no text: {result.stream_events[:5]}"

        usage_positions = [
            index
            for index, event in enumerate(events)
            if event.type == "message_delta" and event.usage is not None
        ]
        assert usage_positions, f"stream never reported usage: {types}"
        assert "message_stop" in types, f"stream never reached message_stop: {types}"
        stop_position = types.index("message_stop")
        assert delta_positions[-1] < usage_positions[0] < stop_position, (
            f"usage did not land between the last content delta and message_stop: {types}"
        )

        first_delta_at: Final = result.stream_event_arrivals[delta_positions[0]]
        stop_at: Final = result.stream_event_arrivals[stop_position]
        if provider_paces_stream():
            assert stop_at - first_delta_at >= STREAM_MIN_LEAD_SECONDS, (
                f"first content delta reached the client {first_delta_at:.2f}s after the request "
                f"and message_stop {stop_at:.2f}s after it; a relayed stream shows the first delta "
                f"at least {STREAM_MIN_LEAD_SECONDS}s before the end, so the response was buffered"
            )

    @pytest.mark.covers("llm.messages.anthropic.tool_use.nonstream.works")
    def test_messages_tool_use(
        self, endpoints_client: EndpointsClient, resources: ResourceManager
    ) -> None:
        model, key = self._register(endpoints_client, resources)

        response = unwrap(
            endpoints_client.proxy.messages(
                key,
                AnthropicMessagesBody(
                    model=model,
                    max_tokens=256,
                    tools=[WEATHER_TOOL],
                    messages=[
                        ChatMessage(role="user", content="What is the weather in Paris? Use the tool.")
                    ],
                ),
            )
        )
        assert response.content, f"no content blocks in response: {response}"
        assert any(block.type == "tool_use" for block in response.content), (
            f"model did not call the tool: {response}"
        )

    @pytest.mark.skip(reason="stage red: product gap, /v1/messages 500s (anthropic_messages TypeError) on missing messages instead of 400")
    @pytest.mark.covers("llm.messages.anthropic.input_validation.nonstream.works")
    def test_missing_messages_returns_error(
        self, endpoints_client: EndpointsClient, resources: ResourceManager
    ) -> None:
        model, key = self._register(endpoints_client, resources)
        result = endpoints_client.proxy.transport.send(
            "/v1/messages",
            headers=endpoints_client.proxy.transport.bearer(key),
            json=_OptionalMessagesBody(model=model, max_tokens=50),
        )
        assert_client_error(result, "messages missing messages")

    @pytest.mark.skip(reason="stage red: product gap, /v1/messages 500s (anthropic_messages TypeError) on missing max_tokens instead of 400")
    @pytest.mark.covers("llm.messages.anthropic.input_validation.nonstream.works")
    def test_missing_max_tokens_returns_error(
        self, endpoints_client: EndpointsClient, resources: ResourceManager
    ) -> None:
        model, key = self._register(endpoints_client, resources)
        result = endpoints_client.proxy.transport.send(
            "/v1/messages",
            headers=endpoints_client.proxy.transport.bearer(key),
            json=_OptionalMessagesBody(
                model=model, messages=[ChatMessage(role="user", content="hi")]
            ),
        )
        assert_client_error(result, "messages missing max_tokens")

    @pytest.mark.covers("llm.messages.anthropic.input_validation.nonstream.works")
    def test_missing_model_returns_error(
        self, endpoints_client: EndpointsClient, resources: ResourceManager
    ) -> None:
        _, key = self._register(endpoints_client, resources)
        result = endpoints_client.proxy.transport.send(
            "/v1/messages",
            headers=endpoints_client.proxy.transport.bearer(key),
            json=_OptionalMessagesBody(messages=[ChatMessage(role="user", content="hi")], max_tokens=50),
        )
        assert_client_error(result, "messages missing model")


class _BridgeDelta(BaseModel):
    type: str | None = None
    partial_json: str | None = None
    stop_reason: str | None = None


class _BridgeEvent(BaseModel):
    type: str
    index: int | None = None
    content_block: AnthropicContentBlock | None = None
    delta: _BridgeDelta | None = None


class _ParcelInput(BaseModel):
    model_config = ConfigDict(extra="forbid", strict=True)
    parcel: str
    shelf: int


def _tool_from_stream(events: tuple[_BridgeEvent, ...]) -> AnthropicContentBlock:
    starts: Final = tuple(
        event
        for event in events
        if event.type == "content_block_start"
        and event.content_block is not None
        and event.content_block.type == "tool_use"
    )
    assert len(starts) == 1, "expected exactly one tool call"
    start: Final = starts[0]
    block: Final = start.content_block
    assert block is not None and block.id and start.index is not None
    fragments: Final = tuple(
        event
        for event in events
        if event.type == "content_block_delta" and event.delta is not None and event.delta.type == "input_json_delta"
    )
    assert fragments, "tool stream contained no argument fragments"
    assert all(event.index == start.index for event in fragments), "tool fragments changed index"
    positions: Final = tuple(i for i, event in enumerate(events) if event in fragments)
    stops: Final = tuple(
        i for i, event in enumerate(events) if event.type == "content_block_stop" and event.index == start.index
    )
    assert len(stops) == 1 and events.index(start) < positions[0] <= positions[-1] < stops[0]
    assert tuple(
        event.delta.stop_reason for event in events if event.type == "message_delta" and event.delta is not None
    ) == ("tool_use",)
    terminal_positions: Final = tuple(i for i, event in enumerate(events) if event.type == "message_delta")
    assert len(terminal_positions) == 1 and stops[0] < terminal_positions[0] < len(events) - 1
    assert tuple(i for i, event in enumerate(events) if event.type == "message_stop") == (len(events) - 1,), (
        "tool stream did not terminate exactly once"
    )
    arguments: Final = _ParcelInput.model_validate_json(
        "".join(event.delta.partial_json or "" for event in fragments if event.delta is not None)
    )
    return AnthropicContentBlock(type="tool_use", id=block.id, name=block.name, input=arguments.model_dump())


def _parcel_result(tool: AnthropicContentBlock, result: AnthropicToolResultBlock) -> AnthropicToolResultTurn:
    assert tool.id and result.tool_use_id == tool.id, "tool result ID does not match the emitted call"
    return AnthropicToolResultTurn(content=[result])


def _request_tool(
    client: EndpointsClient, key: str, request: AnthropicMessagesBody, stream: bool
) -> AnthropicContentBlock:
    if stream:
        response: Final = client.proxy.messages_stream(key, request)
        require_successful_call(response)
        assert response.is_streaming and not response.stream_error
        return _tool_from_stream(tuple(_BridgeEvent.model_validate_json(event) for event in response.stream_events))
    response_body: Final = unwrap(client.proxy.messages(key, request))
    blocks: Final = tuple(block for block in response_body.content or () if block.type == "tool_use")
    assert len(blocks) == 1
    return blocks[0]


class TestOpenAIMessagesToolContinuation:
    @pytest.mark.parametrize("stream", [True, False], ids=["stream", "nonstream"])
    def test_required_tool_arguments_and_correlated_result(
        self, endpoints_client: EndpointsClient, resources: ResourceManager, stream: bool
    ) -> None:
        model: Final = f"e2e-bridge-tool-{unique_marker()}"
        base: Final = provider_edge_base("openai")
        model_id: Final = endpoints_client.create_model(
            model,
            LiteLLMParamsBody(
                model="openai/gpt-5.6", api_key="os.environ/OPENAI_API_KEY", api_base=f"{base}/v1" if base else None
            ),
        )
        resources.defer(lambda: endpoints_client.delete_model(model_id))
        key: Final = resources.key(models=[model])
        tool: Final = AnthropicCustomTool(
            name="locate_parcel",
            description="Look up the receipt for a parcel on a shelf. Return the receipt verbatim.",
            input_schema=ToolInputSchema(
                properties={"parcel": JsonSchemaProperty(type="string"), "shelf": JsonSchemaProperty(type="integer")},
                required=["parcel", "shelf"],
            ),
        )
        question: Final = ChatMessage(
            role="user",
            content="Call locate_parcel with parcel exactly amber-kite and shelf exactly 7. After the tool result, reply with only the receipt returned by the tool.",
        )
        request: Final = AnthropicMessagesBody(
            model=model,
            max_tokens=2048,
            messages=[question],
            tools=[tool],
            tool_choice=AnthropicToolChoice(type="tool", name=tool.name),
            stream=stream,
        )
        emitted: Final = _request_tool(endpoints_client, key, request, stream)
        assert emitted.id and emitted.name == "locate_parcel"
        assert emitted.input == {"parcel": "amber-kite", "shelf": 7}, "required tool arguments were lost or changed"
        receipt: Final = f"receipt-{unique_marker()}"
        result_turn: Final = _parcel_result(emitted, AnthropicToolResultBlock(tool_use_id=emitted.id, content=receipt))
        continuation: Final = unwrap(
            endpoints_client.proxy.messages(
                key,
                AnthropicMessagesBody(
                    model=model,
                    max_tokens=2048,
                    tools=[tool],
                    tool_choice=AnthropicToolChoice(type="none"),
                    messages=[question, AnthropicAssistantTurn(content=[emitted]), result_turn],
                ),
            )
        )
        answer: Final = "".join(block.text or "" for block in continuation.content or ())
        assert answer.strip() == receipt, "continuation did not consume the correlated tool result"
        assert all(block.type != "tool_use" for block in continuation.content or ())