litellm/tests/unit/llms/chat/test_converse_handler.py

import json
from collections.abc import AsyncIterator
from typing import Final
from unittest.mock import AsyncMock, MagicMock

import httpx
import pytest

import litellm
from litellm.llms.bedrock.chat import BedrockConverseLLM
from litellm.llms.bedrock.chat.converse_handler import make_sync_call
from litellm.llms.bedrock.common_utils import _get_all_bedrock_regions
from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler, HTTPHandler
from tests._support.stream_chunk_size import (
    LitellmParamsRecorder,
    keys_at_every_depth,
    record_litellm_params,
)


def test_encode_model_id_with_inference_profile():
    """
    Test instance profile is properly encoded when used as a model
    """
    test_model = "arn:aws:bedrock:us-east-1:12345678910:application-inference-profile/ujdtmcirjhevpi"
    expected_model = "arn%3Aaws%3Abedrock%3Aus-east-1%3A12345678910%3Aapplication-inference-profile%2Fujdtmcirjhevpi"
    bedrock_converse_llm = BedrockConverseLLM()
    returned_model = bedrock_converse_llm.encode_model_id(test_model)
    assert expected_model == returned_model


class TestBedrockRegionInModelPath:
    """
    Tests for region extraction from bedrock/{region}/{model} path format.

    When a user passes model="bedrock/ap-northeast-1/moonshotai.kimi-k2.5",
    get_llm_provider strips "bedrock/" and passes "ap-northeast-1/moonshotai.kimi-k2.5"
    to the converse handler. The handler must:
    1. Strip the region from modelId (so AWS gets "moonshotai.kimi-k2.5", not "ap-northeast-1%2Fmoonshotai.kimi-k2.5")
    2. Use the extracted region as aws_region_name for the API call
    """

    @pytest.mark.parametrize(
        "model,expected_model_id,expected_region",
        [
            # Region embedded in path — both modelId and region must be extracted
            (
                "ap-northeast-1/moonshotai.kimi-k2.5",
                "moonshotai.kimi-k2.5",
                "ap-northeast-1",
            ),
            (
                "us-east-1/moonshotai.kimi-k2.5",
                "moonshotai.kimi-k2.5",
                "us-east-1",
            ),
            (
                "us-west-2/anthropic.claude-haiku-4-5-20251001-v1:0",
                "anthropic.claude-haiku-4-5-20251001-v1%3A0",
                "us-west-2",
            ),
            # No region in path — modelId unchanged, no region injected
            (
                "moonshotai.kimi-k2.5",
                "moonshotai.kimi-k2.5",
                None,
            ),
            # Cross-region inference prefix (us., eu., ap.) — not a region path segment
            (
                "us.anthropic.claude-haiku-4-5-20251001-v1:0",
                "us.anthropic.claude-haiku-4-5-20251001-v1%3A0",
                None,
            ),
        ],
    )
    def test_region_and_model_id_extraction(
        self, model: str, expected_model_id: str, expected_region: str | None
    ) -> None:
        """
        Verify that completion() correctly extracts both modelId and aws_region_name
        from the bedrock/{region}/{model} path format.
        """
        bedrock_converse_llm = BedrockConverseLLM()
        optional_params: dict = {}

        # Simulate the modelId + region extraction logic from completion()
        _model_for_id = model
        _stripped = _model_for_id
        for rp in ["bedrock/converse/", "bedrock/", "converse/"]:
            if _stripped.startswith(rp):
                _stripped = _stripped[len(rp) :]
                break

        _region_from_model = None
        _potential_region = _stripped.split("/", 1)[0]
        if _potential_region in _get_all_bedrock_regions() and "/" in _stripped:
            _region_from_model = _potential_region
            _stripped = _stripped.split("/", 1)[1]
            _model_for_id = _stripped

        for _nova_prefix in ["nova-2/", "nova/"]:
            if _stripped.startswith(_nova_prefix):
                _model_for_id = _model_for_id.replace(_nova_prefix, "", 1)
                break

        model_id = bedrock_converse_llm.encode_model_id(model_id=_model_for_id)
        if _region_from_model is not None and "aws_region_name" not in optional_params:
            optional_params["aws_region_name"] = _region_from_model

        assert (
            model_id == expected_model_id
        ), f"modelId mismatch for {model!r}: got {model_id!r}, expected {expected_model_id!r}"
        assert (
            optional_params.get("aws_region_name") == expected_region
        ), (
            f"region mismatch for {model!r}: "
            f"got {optional_params.get('aws_region_name')!r}, expected {expected_region!r}"
        )

    def test_explicit_aws_region_name_not_overridden(self):
        """
        If aws_region_name is already set in optional_params, the region in the
        model path must NOT override it.
        """
        bedrock_converse_llm = BedrockConverseLLM()
        optional_params = {"aws_region_name": "eu-west-1"}
        model = "ap-northeast-1/moonshotai.kimi-k2.5"

        _model_for_id = model
        _stripped = model
        _region_from_model = None
        _potential_region = _stripped.split("/", 1)[0]
        if _potential_region in _get_all_bedrock_regions() and "/" in _stripped:
            _region_from_model = _potential_region
            _stripped = _stripped.split("/", 1)[1]
            _model_for_id = _stripped

        model_id = bedrock_converse_llm.encode_model_id(model_id=_model_for_id)
        if _region_from_model is not None and "aws_region_name" not in optional_params:
            optional_params["aws_region_name"] = _region_from_model

        # modelId is still correctly stripped
        assert model_id == "moonshotai.kimi-k2.5"
        # explicitly set region is preserved
        assert optional_params["aws_region_name"] == "eu-west-1"


def _stream_completion_with_spied_iter_bytes(model: str, stream_chunk_size: int | None = None) -> MagicMock:
    mock_response: Final = MagicMock()
    mock_response.status_code = 200
    mock_response.iter_bytes = MagicMock(return_value=iter([]))
    client: Final = HTTPHandler()
    client.post = MagicMock(return_value=mock_response)

    litellm.completion(
        model=model,
        messages=[{"role": "user", "content": "hi"}],
        stream=True,
        client=client,
        aws_access_key_id="fake",
        aws_secret_access_key="fake",
        aws_region_name="us-east-1",
        stream_chunk_size=stream_chunk_size,
    )
    return mock_response.iter_bytes


def test_make_sync_call_does_not_rechunk_stream_by_default():
    """Re-chunking the event stream into fixed 1024-byte blocks holds small
    early events in httpx's ByteChunker until 1024 bytes accumulate, delaying
    time-to-first-chunk by the whole generation when Bedrock trickles bytes
    (e.g. buffered tool-use streams)."""
    response = MagicMock()
    response.status_code = 200
    client = MagicMock()
    client.post = MagicMock(return_value=response)

    make_sync_call(
        client=client,
        api_base="https://bedrock-runtime.us-east-1.amazonaws.com/model/anthropic.claude-sonnet-4-6/converse-stream",
        headers={},
        data="{}",
        model="anthropic.claude-sonnet-4-6",
        messages=[],
        logging_obj=MagicMock(),
    )

    response.iter_bytes.assert_called_once_with(chunk_size=None)


def test_make_sync_call_honors_explicit_stream_chunk_size():
    response = MagicMock()
    response.status_code = 200
    client = MagicMock()
    client.post = MagicMock(return_value=response)

    make_sync_call(
        client=client,
        api_base="https://bedrock-runtime.us-east-1.amazonaws.com/model/anthropic.claude-sonnet-4-6/converse-stream",
        headers={},
        data="{}",
        model="anthropic.claude-sonnet-4-6",
        messages=[],
        logging_obj=MagicMock(),
        stream_chunk_size=2048,
    )

    response.iter_bytes.assert_called_once_with(chunk_size=2048)


def _converse_response_body() -> dict:
    return {
        "output": {"message": {"role": "assistant", "content": [{"text": "hi"}]}},
        "stopReason": "end_turn",
        "usage": {"inputTokens": 1, "outputTokens": 1, "totalTokens": 2},
    }


def test_converse_completion_forwards_bedrock_response_headers():
    mock_response = MagicMock()
    mock_response.status_code = 200
    mock_response.json = MagicMock(return_value=_converse_response_body())
    mock_response.text = json.dumps(_converse_response_body())
    mock_response.headers = httpx.Headers({"x-amzn-requestid": "req-123"})
    client = HTTPHandler()
    client.post = MagicMock(return_value=mock_response)

    response = litellm.completion(
        model="bedrock/converse/anthropic.claude-haiku-4-5-20251001-v1:0",
        messages=[{"role": "user", "content": "hi"}],
        client=client,
        aws_access_key_id="fake",
        aws_secret_access_key="fake",
        aws_region_name="us-east-1",
    )

    assert response._hidden_params["additional_headers"]["llm_provider-x-amzn-requestid"] == "req-123"


def test_converse_streaming_forwards_bedrock_response_headers():
    mock_response = MagicMock()
    mock_response.status_code = 200
    mock_response.iter_bytes = MagicMock(return_value=iter([]))
    mock_response.headers = httpx.Headers({"x-amzn-requestid": "req-456"})
    client = HTTPHandler()
    client.post = MagicMock(return_value=mock_response)

    response = litellm.completion(
        model="bedrock/converse/anthropic.claude-haiku-4-5-20251001-v1:0",
        messages=[{"role": "user", "content": "hi"}],
        stream=True,
        client=client,
        aws_access_key_id="fake",
        aws_secret_access_key="fake",
        aws_region_name="us-east-1",
    )

    assert response._hidden_params["additional_headers"]["llm_provider-x-amzn-requestid"] == "req-456"


@pytest.mark.asyncio
async def test_async_converse_completion_forwards_bedrock_response_headers():
    mock_response = MagicMock()
    mock_response.status_code = 200
    mock_response.json = MagicMock(return_value=_converse_response_body())
    mock_response.text = json.dumps(_converse_response_body())
    mock_response.headers = httpx.Headers({"x-amzn-requestid": "req-abc"})
    client = AsyncHTTPHandler()
    client.post = AsyncMock(return_value=mock_response)

    response = await litellm.acompletion(
        model="bedrock/converse/anthropic.claude-haiku-4-5-20251001-v1:0",
        messages=[{"role": "user", "content": "hi"}],
        client=client,
        aws_access_key_id="fake",
        aws_secret_access_key="fake",
        aws_region_name="us-east-1",
    )

    assert response._hidden_params["additional_headers"]["llm_provider-x-amzn-requestid"] == "req-abc"


@pytest.mark.asyncio
async def test_async_converse_streaming_forwards_bedrock_response_headers():
    async def _no_bytes(chunk_size: int | None = None) -> AsyncIterator[bytes]:
        return
        yield b""

    mock_response = MagicMock()
    mock_response.status_code = 200
    mock_response.aiter_bytes = _no_bytes
    mock_response.headers = httpx.Headers({"x-amzn-requestid": "req-def"})
    client = AsyncHTTPHandler()
    client.post = AsyncMock(return_value=mock_response)

    response = await litellm.acompletion(
        model="bedrock/converse/anthropic.claude-haiku-4-5-20251001-v1:0",
        messages=[{"role": "user", "content": "hi"}],
        stream=True,
        client=client,
        aws_access_key_id="fake",
        aws_secret_access_key="fake",
        aws_region_name="us-east-1",
    )

    assert response._hidden_params["additional_headers"]["llm_provider-x-amzn-requestid"] == "req-def"


def test_completion_plumbs_stream_chunk_size_through_converse() -> None:
    iter_bytes_spy = _stream_completion_with_spied_iter_bytes(
        model="bedrock/anthropic.claude-3-5-sonnet-20240620-v1:0"
    )
    iter_bytes_spy.assert_called_once_with(chunk_size=None)

    iter_bytes_spy = _stream_completion_with_spied_iter_bytes(
        model="bedrock/anthropic.claude-3-5-sonnet-20240620-v1:0",
        stream_chunk_size=2048,
    )
    iter_bytes_spy.assert_called_once_with(chunk_size=2048)


def _stream_converse_completion_with_spied_client(
    monkeypatch: pytest.MonkeyPatch, stream_chunk_size: int | None = None
) -> tuple[MagicMock, MagicMock, LitellmParamsRecorder]:
    recorder: Final = record_litellm_params(monkeypatch)
    mock_response: Final = MagicMock()
    mock_response.status_code = 200
    mock_response.iter_bytes = MagicMock(return_value=iter([]))
    client: Final = HTTPHandler()
    client.post = MagicMock(return_value=mock_response)

    litellm.completion(
        model="bedrock/converse/anthropic.claude-haiku-4-5-20251001-v1:0",
        messages=[{"role": "user", "content": "hi"}],
        stream=True,
        client=client,
        aws_access_key_id="fake",
        aws_secret_access_key="fake",
        aws_region_name="us-east-1",
        stream_chunk_size=stream_chunk_size,
    )
    return mock_response.iter_bytes, client.post, recorder


def test_completion_stream_chunk_size_reaches_iter_bytes_but_not_converse_body(
    monkeypatch: pytest.MonkeyPatch,
) -> None:
    iter_bytes_spy, post_spy, recorder = _stream_converse_completion_with_spied_client(
        monkeypatch, stream_chunk_size=64
    )

    iter_bytes_spy.assert_called_once_with(chunk_size=64)
    data: Final = post_spy.call_args.kwargs["data"]
    assert "stream_chunk_size" not in keys_at_every_depth(json.loads(data)), data
    assert len(recorder.seen) == 1
    assert recorder.seen[0]["stream_chunk_size"] == 64


def test_completion_without_stream_chunk_size_uses_default_chunking(monkeypatch: pytest.MonkeyPatch) -> None:
    iter_bytes_spy, _, recorder = _stream_converse_completion_with_spied_client(monkeypatch)

    iter_bytes_spy.assert_called_once_with(chunk_size=None)
    assert len(recorder.seen) == 1
    assert recorder.seen[0]["stream_chunk_size"] is None


async def _astream_converse_completion_with_spied_client(
    monkeypatch: pytest.MonkeyPatch, stream_chunk_size: int | None = None
) -> tuple[MagicMock, AsyncMock, LitellmParamsRecorder]:
    async def _no_bytes(chunk_size: int | None = None) -> AsyncIterator[bytes]:
        return
        yield b""

    mock_response: Final = MagicMock()
    mock_response.status_code = 200
    recorder: Final = record_litellm_params(monkeypatch)
    mock_response.aiter_bytes = MagicMock(return_value=_no_bytes())
    aiter_bytes_spy: Final = mock_response.aiter_bytes
    client: Final = AsyncHTTPHandler()
    client.post = AsyncMock(return_value=mock_response)

    await litellm.acompletion(
        model="bedrock/converse/anthropic.claude-haiku-4-5-20251001-v1:0",
        messages=[{"role": "user", "content": "hi"}],
        stream=True,
        client=client,
        aws_access_key_id="fake",
        aws_secret_access_key="fake",
        aws_region_name="us-east-1",
        stream_chunk_size=stream_chunk_size,
    )
    return aiter_bytes_spy, client.post, recorder


@pytest.mark.asyncio
async def test_acompletion_stream_chunk_size_reaches_aiter_bytes_but_not_converse_body(
    monkeypatch: pytest.MonkeyPatch,
) -> None:
    aiter_bytes_spy, post_spy, recorder = await _astream_converse_completion_with_spied_client(
        monkeypatch, stream_chunk_size=64
    )

    aiter_bytes_spy.assert_called_once_with(chunk_size=64)
    data: Final = post_spy.call_args.kwargs["data"]
    assert "stream_chunk_size" not in keys_at_every_depth(json.loads(data)), data
    assert len(recorder.seen) == 1
    assert recorder.seen[0]["stream_chunk_size"] == 64


@pytest.mark.asyncio
async def test_acompletion_without_stream_chunk_size_uses_default_chunking(
    monkeypatch: pytest.MonkeyPatch,
) -> None:
    aiter_bytes_spy, _, recorder = await _astream_converse_completion_with_spied_client(monkeypatch)

    aiter_bytes_spy.assert_called_once_with(chunk_size=None)
    assert len(recorder.seen) == 1
    assert recorder.seen[0]["stream_chunk_size"] is None


@pytest.mark.parametrize("stream_chunk_size,expected_chunk_size", [(64, 64), (None, None)])
def test_router_deployment_stream_chunk_size_reaches_iter_bytes(
    monkeypatch: pytest.MonkeyPatch, stream_chunk_size: int | None, expected_chunk_size: int | None
) -> None:
    recorder: Final = record_litellm_params(monkeypatch)
    mock_response: Final = MagicMock()
    mock_response.status_code = 200
    mock_response.iter_bytes = MagicMock(return_value=iter([]))
    client: Final = HTTPHandler()
    client.post = MagicMock(return_value=mock_response)
    deployment_params: Final = {
        "model": "bedrock/converse/anthropic.claude-haiku-4-5-20251001-v1:0",
        "aws_access_key_id": "fake",
        "aws_secret_access_key": "fake",
        "aws_region_name": "us-east-1",
    }
    router: Final = litellm.Router(
        model_list=[
            {
                "model_name": "converse-chunked",
                "litellm_params": deployment_params
                | ({} if stream_chunk_size is None else {"stream_chunk_size": stream_chunk_size}),
            }
        ]
    )

    router.completion(
        model="converse-chunked",
        messages=[{"role": "user", "content": "hi"}],
        stream=True,
        client=client,
    )

    mock_response.iter_bytes.assert_called_once_with(chunk_size=expected_chunk_size)
    data: Final = client.post.call_args.kwargs["data"]
    assert "stream_chunk_size" not in keys_at_every_depth(json.loads(data)), data
    assert len(recorder.seen) == 1
    assert recorder.seen[0]["stream_chunk_size"] == stream_chunk_size


def test_converse_stream_rejects_non_int_stream_chunk_size_before_calling_bedrock(monkeypatch: pytest.MonkeyPatch):
    record_litellm_params(monkeypatch)
    client = HTTPHandler()
    client.post = MagicMock()

    with pytest.raises(litellm.BadRequestError):
        litellm.completion(
            model="bedrock/converse/anthropic.claude-haiku-4-5-20251001-v1:0",
            messages=[{"role": "user", "content": "hi"}],
            stream=True,
            client=client,
            aws_access_key_id="fake",
            aws_secret_access_key="fake",
            aws_region_name="us-east-1",
            stream_chunk_size="sixty-four",
        )

    client.post.assert_not_called()


def test_converse_non_stream_ignores_invalid_stream_chunk_size():
    mock_response = MagicMock()
    mock_response.status_code = 200
    mock_response.json = MagicMock(return_value=_converse_response_body())
    mock_response.text = json.dumps(_converse_response_body())
    mock_response.headers = httpx.Headers()
    client = HTTPHandler()
    client.post = MagicMock(return_value=mock_response)

    response = litellm.completion(
        model="bedrock/converse/anthropic.claude-haiku-4-5-20251001-v1:0",
        messages=[{"role": "user", "content": "hi"}],
        client=client,
        aws_access_key_id="fake",
        aws_secret_access_key="fake",
        aws_region_name="us-east-1",
        stream_chunk_size="64",
    )

    assert response.choices[0].message.content == "hi"
    client.post.assert_called_once()


def _bedrock_error_response(status_code: int, request_id: str) -> httpx.Response:
    return httpx.Response(
        status_code=status_code,
        headers={
            "x-amzn-RequestId": request_id,
            "x-amzn-ErrorType": "InternalServerException",
        },
        text=json.dumps({"message": "Amazon Bedrock is unable to process your request."}),
        request=httpx.Request("POST", "https://bedrock-runtime.us-east-1.amazonaws.com/"),
    )


def test_converse_completion_error_forwards_bedrock_response_headers():
    error_response = _bedrock_error_response(500, "req-err-123")
    client = HTTPHandler()
    client.post = MagicMock(
        side_effect=httpx.HTTPStatusError(
            "server error",
            request=error_response.request,
            response=error_response,
        )
    )

    with pytest.raises(litellm.ServiceUnavailableError) as exc_info:
        litellm.completion(
            model="bedrock/converse/anthropic.claude-haiku-4-5-20251001-v1:0",
            messages=[{"role": "user", "content": "hi"}],
            client=client,
            aws_access_key_id="fake",
            aws_secret_access_key="fake",
            aws_region_name="us-east-1",
        )

    assert exc_info.value.response.headers["x-amzn-requestid"] == "req-err-123"


@pytest.mark.asyncio
async def test_async_converse_completion_error_forwards_bedrock_response_headers():
    error_response = _bedrock_error_response(500, "req-err-456")
    client = AsyncHTTPHandler()
    client.post = AsyncMock(
        side_effect=httpx.HTTPStatusError(
            "server error",
            request=error_response.request,
            response=error_response,
        )
    )

    with pytest.raises(litellm.ServiceUnavailableError) as exc_info:
        await litellm.acompletion(
            model="bedrock/converse/anthropic.claude-haiku-4-5-20251001-v1:0",
            messages=[{"role": "user", "content": "hi"}],
            client=client,
            aws_access_key_id="fake",
            aws_secret_access_key="fake",
            aws_region_name="us-east-1",
        )

    assert exc_info.value.response.headers["x-amzn-requestid"] == "req-err-456"