litellm/tests/unit/llms/bedrock/embed/test_bedrock_embedding.py

import json
import asyncio
import os
from unittest.mock import Mock, patch
from unittest.mock import AsyncMock, MagicMock

import pytest
import httpx

import litellm
from litellm.llms.bedrock.embed.twelvelabs_marengo_transformation import TwelveLabsMarengoEmbeddingConfig
from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler, HTTPHandler
from litellm.llms.bedrock.embed.embedding import BedrockEmbedding
from tests.unit.llms.bedrock.event_loop_probe import EventLoopProbe

# Mock responses for different embedding models
titan_embedding_response = {"embedding": [0.1, 0.2, 0.3], "inputTextTokenCount": 10}

cohere_embedding_response = {"embeddings": [[0.1, 0.2, 0.3]], "inputTextTokenCount": 10}

twelvelabs_embedding_response = {
    "embedding": [0.1, 0.2, 0.3],
    "embeddingOption": "visual-text",
    "startSec": 0.0,
    "endSec": 1.0,
}

# Test data
test_input = "Hello world from litellm"
test_image_base64 = "data:image/png,test_image_base64_data"


@pytest.mark.parametrize(
    "model,input_type,embed_response",
    [
        ("bedrock/amazon.titan-embed-text-v1", "text", titan_embedding_response),
        ("bedrock/amazon.titan-embed-text-v2:0", "text", titan_embedding_response),
        ("bedrock/amazon.titan-embed-image-v1", "image", titan_embedding_response),
        ("bedrock/cohere.embed-english-v3", "text", cohere_embedding_response),
        ("bedrock/cohere.embed-multilingual-v3", "text", cohere_embedding_response),
        (
            "bedrock/twelvelabs.marengo-embed-2-7-v1:0",
            "text",
            twelvelabs_embedding_response,
        ),
        (
            "bedrock/twelvelabs.marengo-embed-2-7-v1:0",
            "image",
            twelvelabs_embedding_response,
        ),
    ],
)
def test_bedrock_embedding_with_api_key_bearer_token(model, input_type, embed_response):
    """Test embedding functionality with bearer token authentication"""
    client = HTTPHandler()
    test_api_key = "test-bearer-token-12345"

    with patch.object(client, "post") as mock_post:
        mock_response = Mock()
        mock_response.status_code = 200
        mock_response.text = json.dumps(embed_response)
        mock_response.json = lambda: json.loads(mock_response.text)
        mock_post.return_value = mock_response

        input_data = test_image_base64 if input_type == "image" else test_input

        # Add inputType parameter for TwelveLabs Marengo models
        kwargs = {
            "model": model,
            "input": input_data,
            "client": client,
            "aws_region_name": "us-east-1",
            "aws_bedrock_runtime_endpoint": "https://bedrock-runtime.us-east-1.amazonaws.com",
            "api_key": test_api_key,
        }

        # Add input_type parameter for TwelveLabs Marengo models (maps to inputType)
        if "twelvelabs.marengo-embed" in model:
            kwargs["input_type"] = input_type

        response = litellm.embedding(**kwargs)

        assert isinstance(response, litellm.EmbeddingResponse)
        assert isinstance(response.data[0]["embedding"], list)
        assert len(response.data[0]["embedding"]) == 3  # Based on mock response

        headers = mock_post.call_args.kwargs.get("headers", {})
        assert "Authorization" in headers
        assert headers["Authorization"] == f"Bearer {test_api_key}"


@pytest.mark.parametrize(
    "model,input_type,embed_response",
    [
        ("bedrock/amazon.titan-embed-text-v1", "text", titan_embedding_response),
    ],
)
def test_bedrock_embedding_with_env_variable_bearer_token(
    model, input_type, embed_response
):
    """Test embedding functionality with bearer token from environment variable"""
    client = HTTPHandler()
    test_api_key = "env-bearer-token-12345"

    with (
        patch.dict(os.environ, {"AWS_BEARER_TOKEN_BEDROCK": test_api_key}),
        patch.object(client, "post") as mock_post,
    ):

        mock_response = Mock()
        mock_response.status_code = 200
        mock_response.text = json.dumps(embed_response)
        mock_response.json = lambda: json.loads(mock_response.text)
        mock_post.return_value = mock_response

        response = litellm.embedding(
            model=model,
            input=test_input,
            client=client,
            aws_region_name="us-west-2",
            aws_bedrock_runtime_endpoint="https://bedrock-runtime.us-west-2.amazonaws.com",
        )

        assert isinstance(response, litellm.EmbeddingResponse)
        headers = mock_post.call_args.kwargs.get("headers", {})
        assert "Authorization" in headers
        assert headers["Authorization"] == f"Bearer {test_api_key}"


@pytest.mark.asyncio
async def test_async_bedrock_embedding_with_bearer_token():
    """Test async embedding functionality with bearer token authentication"""
    client = AsyncHTTPHandler()
    test_api_key = "async-bearer-token-12345"
    model = "bedrock/amazon.titan-embed-text-v1"

    with patch.object(client, "post") as mock_post:
        mock_response = Mock()
        mock_response.status_code = 200
        mock_response.text = json.dumps(titan_embedding_response)
        mock_response.json = Mock(return_value=titan_embedding_response)
        mock_post.return_value = mock_response

        response = await litellm.aembedding(
            model=model,
            input=test_input,
            client=client,
            aws_region_name="us-west-2",
            aws_bedrock_runtime_endpoint="https://bedrock-runtime.us-west-2.amazonaws.com",
            api_key=test_api_key,
        )

        assert isinstance(response, litellm.EmbeddingResponse)

        headers = mock_post.call_args.kwargs.get("headers", {})
        assert "Authorization" in headers
        assert headers["Authorization"] == f"Bearer {test_api_key}"


def test_bedrock_embedding_with_sigv4():
    """Test embedding falls back to SigV4 auth when no bearer token is provided"""
    model = "bedrock/amazon.titan-embed-text-v1"

    with patch(
        "litellm.llms.bedrock.embed.embedding.BedrockEmbedding.embeddings"
    ) as mock_bedrock_embed:
        mock_embedding_response = litellm.EmbeddingResponse()
        mock_embedding_response.data = [{"embedding": [0.1, 0.2, 0.3]}]
        mock_bedrock_embed.return_value = mock_embedding_response

        response = litellm.embedding(
            model=model,
            input=test_input,
            aws_region_name="us-west-2",
        )

        assert isinstance(response, litellm.EmbeddingResponse)
        mock_bedrock_embed.assert_called_once()


def test_bedrock_titan_v2_encoding_format_float():
    """Test amazon.titan-embed-text-v2:0 with encoding_format=float parameter"""
    client = HTTPHandler()
    test_api_key = "test-bearer-token-12345"
    model = "bedrock/amazon.titan-embed-text-v2:0"

    # Mock response with embeddingsByType for binary format (addressing issue #14680)
    titan_v2_response = {"embedding": [0.1, 0.2, 0.3], "inputTextTokenCount": 10}

    with patch.object(client, "post") as mock_post:
        mock_response = Mock()
        mock_response.status_code = 200
        mock_response.text = json.dumps(titan_v2_response)
        mock_response.json = lambda: json.loads(mock_response.text)
        mock_post.return_value = mock_response

        response = litellm.embedding(
            model=model,
            input=test_input,
            encoding_format="float",  # This should work but currently throws UnsupportedParamsError
            client=client,
            aws_region_name="us-east-1",
            aws_bedrock_runtime_endpoint="https://bedrock-runtime.us-east-1.amazonaws.com",
            api_key=test_api_key,
        )

        assert isinstance(response, litellm.EmbeddingResponse)
        assert isinstance(response.data[0]["embedding"], list)
        assert len(response.data[0]["embedding"]) == 3

        # Verify that the request contains embeddingTypes: ["float"] instead of encoding_format
        request_body = json.loads(mock_post.call_args.kwargs.get("data", "{}"))
        assert "embeddingTypes" in request_body
        assert request_body["embeddingTypes"] == ["float"]
        assert "encoding_format" not in request_body


def test_bedrock_titan_v2_encoding_format_base64():
    """Test amazon.titan-embed-text-v2:0 with encoding_format=base64 parameter (maps to binary)"""
    client = HTTPHandler()
    test_api_key = "test-bearer-token-12345"
    model = "bedrock/amazon.titan-embed-text-v2:0"

    # Mock response with embeddingsByType for binary format
    titan_v2_binary_response = {
        "embeddingsByType": {
            "binary": "YmluYXJ5X2VtYmVkZGluZ19kYXRh"  # base64 encoded binary data
        },
        "inputTextTokenCount": 10,
    }

    with patch.object(client, "post") as mock_post:
        mock_response = Mock()
        mock_response.status_code = 200
        mock_response.text = json.dumps(titan_v2_binary_response)
        mock_response.json = lambda: json.loads(mock_response.text)
        mock_post.return_value = mock_response

        response = litellm.embedding(
            model=model,
            input=test_input,
            encoding_format="base64",  # This should map to embeddingTypes: ["binary"]
            client=client,
            aws_region_name="us-east-1",
            aws_bedrock_runtime_endpoint="https://bedrock-runtime.us-east-1.amazonaws.com",
            api_key=test_api_key,
        )

        assert isinstance(response, litellm.EmbeddingResponse)

        # Verify that the request contains embeddingTypes: ["binary"] for base64 encoding
        request_body = json.loads(mock_post.call_args.kwargs.get("data", "{}"))
        assert "embeddingTypes" in request_body
        assert request_body["embeddingTypes"] == ["binary"]


def test_twelvelabs_input_type_parameter_mapping():
    """Test that input_type parameter is correctly mapped to inputType for TwelveLabs models"""
    client = HTTPHandler()
    test_api_key = "test-bearer-token-12345"
    model = "bedrock/twelvelabs.marengo-embed-2-7-v1:0"

    twelvelabs_response = {
        "data": [{"embedding": [0.1, 0.2, 0.3], "inputTextTokenCount": 10}]
    }

    with patch.object(client, "post") as mock_post:
        mock_response = Mock()
        mock_response.status_code = 200
        mock_response.text = json.dumps(twelvelabs_response)
        mock_response.json = lambda: json.loads(mock_response.text)
        mock_post.return_value = mock_response

        # Test with input_type parameter (new LiteLLM parameter)
        response = litellm.embedding(
            model=model,
            input=test_input,
            client=client,
            aws_region_name="us-east-1",
            aws_bedrock_runtime_endpoint="https://bedrock-runtime.us-east-1.amazonaws.com",
            api_key=test_api_key,
            input_type="text",  # New parameter that should map to inputType
        )

        assert isinstance(response, litellm.EmbeddingResponse)
        assert isinstance(response.data[0]["embedding"], list)
        assert len(response.data[0]["embedding"]) == 3

        # Verify that the request contains inputType (mapped from input_type)
        request_body = json.loads(mock_post.call_args.kwargs.get("data", "{}"))
        assert "inputType" in request_body
        assert request_body["inputType"] == "text"
        assert "input_type" not in request_body  # Should be mapped, not passed through


def test_twelvelabs_input_type_parameter_mapping_async_invoke():
    """Test that input_type parameter is correctly mapped to inputType for TwelveLabs async invoke models"""
    client = HTTPHandler()
    test_api_key = "test-bearer-token-12345"
    model = "bedrock/async_invoke/twelvelabs.marengo-embed-2-7-v1:0"

    async_invoke_response = {
        "invocationArn": "arn:aws:bedrock:us-east-1:123456789012:async-invoke/abc123def456"
    }

    with patch.object(client, "post") as mock_post:
        mock_response = Mock()
        mock_response.status_code = 200
        mock_response.text = json.dumps(async_invoke_response)
        mock_response.json = lambda: json.loads(mock_response.text)
        mock_post.return_value = mock_response

        # Test with input_type parameter for async invoke
        response = litellm.embedding(
            model=model,
            input=test_input,
            client=client,
            aws_region_name="us-east-1",
            aws_bedrock_runtime_endpoint="https://bedrock-runtime.us-east-1.amazonaws.com",
            api_key=test_api_key,
            output_s3_uri="s3://test-bucket/async-invoke-output/",
            input_type="text",  # New parameter that should map to inputType
        )

        assert isinstance(response, litellm.EmbeddingResponse)
        assert hasattr(response, "_hidden_params")
        assert response._hidden_params is not None
        assert hasattr(response._hidden_params, "_invocation_arn")

        # Verify that the request contains inputType in modelInput (mapped from input_type)
        request_body = json.loads(mock_post.call_args.kwargs.get("data", "{}"))
        assert "modelInput" in request_body
        assert "inputType" in request_body["modelInput"]
        assert request_body["modelInput"]["inputType"] == "text"
        assert "input_type" not in request_body  # Should be mapped, not passed through


def test_twelvelabs_missing_input_type_error():
    """Test that missing input_type parameter defaults to 'text' for TwelveLabs models"""
    client = HTTPHandler()
    test_api_key = "test-bearer-token-12345"

    # Test TwelveLabs model - should default to 'text' when input_type is missing
    twelvelabs_model = "bedrock/twelvelabs.marengo-embed-2-7-v1:0"
    twelvelabs_response = {
        "data": [{"embedding": [0.1, 0.2, 0.3], "inputTextTokenCount": 10}]
    }

    with patch.object(client, "post") as mock_post:
        mock_response = Mock()
        mock_response.status_code = 200
        mock_response.text = json.dumps(twelvelabs_response)
        mock_response.json = lambda: json.loads(mock_response.text)
        mock_post.return_value = mock_response

        # Test that missing input_type defaults to "text" for TwelveLabs
        response = litellm.embedding(
            model=twelvelabs_model,
            input=test_input,
            client=client,
            aws_region_name="us-east-1",
            aws_bedrock_runtime_endpoint="https://bedrock-runtime.us-east-1.amazonaws.com",
            api_key=test_api_key,
            # No input_type parameter - should default to "text"
        )

        # Verify the response is successful
        assert isinstance(response, litellm.EmbeddingResponse)

        # Verify that the request contains inputType: "text" by default
        request_body = json.loads(mock_post.call_args.kwargs.get("data", "{}"))
        assert "inputType" in request_body
        assert request_body["inputType"] == "text"

    # Test Amazon Titan model - should NOT throw error (input_type not required)
    titan_model = "bedrock/amazon.titan-embed-text-v1"
    titan_response = {"embedding": [0.1, 0.2, 0.3], "inputTextTokenCount": 10}

    with patch.object(client, "post") as mock_post:
        mock_response = Mock()
        mock_response.status_code = 200
        mock_response.text = json.dumps(titan_response)
        mock_response.json = lambda: json.loads(mock_response.text)
        mock_post.return_value = mock_response

        # Test that missing input_type does NOT throw an error for Amazon Titan
        response = litellm.embedding(
            model=titan_model,
            input=test_input,
            client=client,
            aws_region_name="us-east-1",
            aws_bedrock_runtime_endpoint="https://bedrock-runtime.us-east-1.amazonaws.com",
            api_key=test_api_key,
            # No input_type parameter - should work fine
        )

        # Should succeed without input_type
        assert isinstance(response, litellm.EmbeddingResponse)


@pytest.mark.parametrize(
    "model,embed_response",
    [
        ("bedrock/amazon.titan-embed-text-v1", titan_embedding_response),
        ("bedrock/amazon.titan-embed-text-v2:0", titan_embedding_response),
        ("bedrock/cohere.embed-english-v3", cohere_embedding_response),
    ],
)
def test_bedrock_embedding_header_forwarding(model, embed_response):
    """
    Test that custom headers are correctly forwarded to Bedrock embedding API calls.

    This test verifies the fix for the issue where headers configured via
    forward_client_headers_to_llm_api were not being passed to Bedrock embedding provider.

    Relevant Issue: https://github.com/BerriAI/litellm/pull/16042
    """
    client = HTTPHandler()
    test_api_key = "test-bearer-token-12345"

    # Headers that would be set by the proxy when forwarding client headers
    custom_headers = {
        "X-Custom-Header": "CustomValue",
        "X-BYOK-Token": "secret-token",
        "Extra-Header": "foobar",
    }

    with patch.object(client, "post") as mock_post:
        mock_response = Mock()
        mock_response.status_code = 200
        mock_response.text = json.dumps(embed_response)
        mock_response.json = lambda: json.loads(mock_response.text)
        mock_post.return_value = mock_response

        try:
            # Call embedding with custom headers via kwargs
            # This simulates what the proxy does when forward_client_headers_to_llm_api is set
            response = litellm.embedding(
                model=model,
                input=test_input,
                client=client,
                headers=custom_headers,  # This is how proxy passes forwarded headers
                aws_region_name="us-east-1",
                aws_bedrock_runtime_endpoint="https://bedrock-runtime.us-east-1.amazonaws.com",
                api_key=test_api_key,
            )

            assert isinstance(response, litellm.EmbeddingResponse)

            # Verify that the request was made
            assert mock_post.called, "HTTP client post should be called"

            # Get the actual call arguments
            call_kwargs = mock_post.call_args.kwargs
            headers = call_kwargs.get("headers", {})

            # Verify our custom headers are present in the request headers
            # Note: AWS SigV4 signing may modify header names to lowercase
            for header_key, header_value in custom_headers.items():
                header_found = (
                    header_key in headers
                    or header_key.lower() in headers
                    or any(k.lower() == header_key.lower() for k in headers.keys())
                )
                assert header_found, (
                    f"Header {header_key} should be in request headers. "
                    f"Found headers: {list(headers.keys())}"
                )

            print(f"✓ Test passed for {model}")
            print(f"  Headers correctly forwarded: {list(headers.keys())}")

        except Exception as e:
            pytest.fail(f"Failed to forward headers to {model}: {str(e)}")


def test_bedrock_embedding_extra_headers_and_headers_merge():
    """
    Test that both extra_headers and headers parameters are correctly merged for Bedrock embeddings.

    This ensures that headers from kwargs (forwarded by proxy) and extra_headers
    (passed explicitly) are both included in the final headers sent to the provider.
    """
    client = HTTPHandler()
    test_api_key = "test-bearer-token-12345"
    model = "bedrock/amazon.titan-embed-text-v1"

    # Headers from proxy (via kwargs["headers"])
    proxy_headers = {"X-Forwarded-Header": "ProxyValue"}

    # Explicit extra_headers
    explicit_headers = {"X-Explicit-Header": "ExplicitValue"}

    # Mock response
    embed_response = {"embedding": [0.1, 0.2, 0.3], "inputTextTokenCount": 10}

    with patch.object(client, "post") as mock_post:
        mock_response = Mock()
        mock_response.status_code = 200
        mock_response.text = json.dumps(embed_response)
        mock_response.json = lambda: json.loads(mock_response.text)
        mock_post.return_value = mock_response

        try:
            response = litellm.embedding(
                model=model,
                input=test_input,
                client=client,
                headers=proxy_headers,  # From proxy forwarding
                extra_headers=explicit_headers,  # Explicitly passed
                aws_region_name="us-east-1",
                aws_bedrock_runtime_endpoint="https://bedrock-runtime.us-east-1.amazonaws.com",
                api_key=test_api_key,
            )

            assert isinstance(response, litellm.EmbeddingResponse)

            call_kwargs = mock_post.call_args.kwargs
            headers = call_kwargs.get("headers", {})

            # Both sets of headers should be present
            # Note: AWS SigV4 signing may modify header names to lowercase
            proxy_header_found = any(
                k.lower() == "x-forwarded-header" for k in headers.keys()
            )
            assert proxy_header_found, (
                "Proxy forwarded header should be present. "
                f"Found headers: {list(headers.keys())}"
            )

            explicit_header_found = any(
                k.lower() == "x-explicit-header" for k in headers.keys()
            )
            assert explicit_header_found, (
                "Explicitly passed header should be present. "
                f"Found headers: {list(headers.keys())}"
            )

            print("✓ Both header sources correctly merged and forwarded")
            print(f"  Final headers: {list(headers.keys())}")

        except Exception as e:
            pytest.fail(f"Failed to merge and forward headers: {str(e)}")


def test_bedrock_cohere_v4_embedding_response_parsing():
    """
    Test parsing of Bedrock Cohere v4 embedding response which returns a dictionary of embeddings
    keyed by type (e.g. 'float', 'int8') instead of a direct list.
    """
    client = HTTPHandler()
    test_api_key = "test-bearer-token-12345"
    model = "bedrock/cohere.embed-v4:0"

    # Mock response for Cohere v4 with multiple embedding types
    cohere_v4_response = {
        "embeddings": {"float": [[0.1, 0.2, 0.3]], "int8": [[1, 2, 3]]},
        "response_type": "embeddings_by_type",
        "id": "test-id",
        "texts": ["test input"],
    }

    with patch.object(client, "post") as mock_post:
        mock_response = Mock()
        mock_response.status_code = 200
        mock_response.text = json.dumps(cohere_v4_response)
        mock_response.json = lambda: json.loads(mock_response.text)
        mock_post.return_value = mock_response

        response = litellm.embedding(
            model=model,
            input=["test input"],
            client=client,
            aws_region_name="us-east-1",
            aws_bedrock_runtime_endpoint="https://bedrock-runtime.us-east-1.amazonaws.com",
            api_key=test_api_key,
        )

        assert isinstance(response, litellm.EmbeddingResponse)

        # Verify we get two embedding objects back (one for float, one for int8)
        assert len(response.data) == 2

        # Check first embedding (float)
        assert response.data[0]["object"] == "embedding"
        assert response.data[0]["embedding"] == [0.1, 0.2, 0.3]
        assert response.data[0]["type"] == "float"

        # Check second embedding (int8)
        assert response.data[1]["object"] == "embedding"
        assert response.data[1]["embedding"] == [1, 2, 3]
        assert response.data[1]["type"] == "int8"


def test_bedrock_embedding_custom_headers_with_iam_role_and_custom_api_base():
    """
    Test that custom headers are correctly forwarded when using IAM role credentials
    (with session token) and a custom api_base.

    This test verifies the fix for the issue where custom headers were not being
    forwarded to Bedrock embeddings endpoint when using:
    - IAM role authentication (session tokens)
    - Custom api_base (proxy endpoint)

    The fix converts HeadersDict to regular dict before passing to httpx, ensuring
    headers are properly forwarded even with IAM roles and custom endpoints.

    Relevant Issue: Custom headers not forwarded with IAM roles + custom api_base
    """
    client = HTTPHandler()

    # Simulate IAM role credentials with session token
    aws_access_key_id = "AKIAIOSFODNN7EXAMPLE"
    aws_secret_access_key = "wJalrXUtnFEMI/K7MDENG/bPxRfiCYEXAMPLEKEY"
    aws_session_token = "AQoEXAMPLEH4aoAH0gNCAPyJxz4BlCFFxWNE1OPTgk5TthT+FvwqnKwRcOIfrRh3c/LTo6UDdyJwOOvEVPvLXCrrrUtdnniCEXAMPLE/IvU1dYUg2RVAJBanLiHb4IgRmpV3ZXrzoB348V+jZfXvYhEXAMPLEEXAMPLE"

    # Custom api_base (simulating a proxy endpoint)
    custom_api_base = "https://gateway.example.com/v1/bedrock-runtime/us-east-1"

    # Custom headers that need to be forwarded
    custom_headers = {
        "X-Custom-Header-1": "test-value-1",
        "X-Custom-Header-2": "test-value-2",
        "X-Forwarded-For": "192.168.1.1",
        "X-BYOK-Token": "secret-token-12345",
    }

    # Mock response
    embed_response = {"embedding": [0.1, 0.2, 0.3], "inputTextTokenCount": 10}

    with patch.object(client, "post") as mock_post:
        mock_response = Mock()
        mock_response.status_code = 200
        mock_response.text = json.dumps(embed_response)
        mock_response.json = lambda: json.loads(mock_response.text)
        mock_post.return_value = mock_response

        try:
            response = litellm.embedding(
                model="bedrock/amazon.titan-embed-text-v1",
                input=test_input,
                client=client,
                extra_headers=custom_headers,
                api_base=custom_api_base,
                aws_access_key_id=aws_access_key_id,
                aws_secret_access_key=aws_secret_access_key,
                aws_session_token=aws_session_token,  # IAM role session token
                aws_region_name="us-east-1",
            )

            assert isinstance(response, litellm.EmbeddingResponse)

            # Verify that the request was made
            assert mock_post.called, "HTTP client post should be called"

            # Get the actual call arguments
            call_kwargs = mock_post.call_args.kwargs
            headers = call_kwargs.get("headers", {})

            # Verify custom headers are present in the request
            # Note: HeadersDict should be converted to regular dict, so headers should be accessible
            for header_key, header_value in custom_headers.items():
                # Check if header exists (case-insensitive for HTTP headers)
                header_found = any(
                    k.lower() == header_key.lower() for k in headers.keys()
                )
                assert header_found, (
                    f"Custom header {header_key} should be in request headers. "
                    f"Found headers: {list(headers.keys())}"
                )

                # Verify the value matches
                header_value_found = None
                for k, v in headers.items():
                    if k.lower() == header_key.lower():
                        header_value_found = v
                        break

                assert header_value_found == header_value, (
                    f"Header {header_key} should have value {header_value}, "
                    f"but found {header_value_found}"
                )

            # Verify AWS signature headers are also present
            assert "Authorization" in headers, "AWS signature should be present"
            assert "X-Amz-Date" in headers, "AWS date header should be present"
            assert (
                "X-Amz-Security-Token" in headers
            ), "Session token header should be present"
            assert (
                headers["X-Amz-Security-Token"] == aws_session_token
            ), "Session token should match the provided token"

            # Verify the custom api_base was used
            called_url = call_kwargs.get("url", "")
            assert custom_api_base in str(called_url), (
                f"Custom api_base {custom_api_base} should be used. "
                f"Got URL: {called_url}"
            )

            print(
                "✓ Test passed: Custom headers forwarded with IAM role + custom api_base"
            )
            print(
                f"  Custom headers found: {[k for k in headers.keys() if k.lower().startswith('x-custom') or k.lower().startswith('x-forwarded')]}"
            )
            print(
                f"  AWS headers found: {[k for k in headers.keys() if k.lower().startswith('x-amz') or k.lower() == 'authorization']}"
            )

        except Exception as e:
            pytest.fail(
                f"Failed to forward headers with IAM role + custom api_base: {str(e)}"
            )


@pytest.mark.asyncio
async def test_bedrock_embedding_custom_headers_with_iam_role_and_custom_api_base_async():
    """
    Test that custom headers are correctly forwarded in async mode when using IAM role
    credentials (with session token) and a custom api_base.

    This is the async version of the test above, verifying the fix works for both
    sync and async embedding calls.
    """
    client = AsyncHTTPHandler()

    # Simulate IAM role credentials with session token
    aws_access_key_id = "AKIAIOSFODNN7EXAMPLE"
    aws_secret_access_key = "wJalrXUtnFEMI/K7MDENG/bPxRfiCYEXAMPLEKEY"
    aws_session_token = "AQoEXAMPLEH4aoAH0gNCAPyJxz4BlCFFxWNE1OPTgk5TthT+FvwqnKwRcOIfrRh3c/LTo6UDdyJwOOvEVPvLXCrrrUtdnniCEXAMPLE/IvU1dYUg2RVAJBanLiHb4IgRmpV3ZXrzoB348V+jZfXvYhEXAMPLEEXAMPLE"

    # Custom api_base (simulating a proxy endpoint)
    custom_api_base = "https://gateway.example.com/v1/bedrock-runtime/us-west-2"

    # Custom headers that need to be forwarded
    custom_headers = {
        "X-Custom-Header-1": "test-value-1",
        "X-Custom-Header-2": "test-value-2",
        "X-Forwarded-For": "192.168.1.1",
        "X-BYOK-Token": "secret-token-12345",
    }

    # Mock response
    embed_response = {"embedding": [0.1, 0.2, 0.3], "inputTextTokenCount": 10}

    with patch.object(client, "post") as mock_post:
        mock_response = Mock()
        mock_response.status_code = 200
        mock_response.text = json.dumps(embed_response)
        mock_response.json = Mock(return_value=embed_response)
        mock_post.return_value = mock_response

        try:
            response = await litellm.aembedding(
                model="bedrock/amazon.titan-embed-text-v1",
                input=test_input,
                client=client,
                extra_headers=custom_headers,
                api_base=custom_api_base,
                aws_access_key_id=aws_access_key_id,
                aws_secret_access_key=aws_secret_access_key,
                aws_session_token=aws_session_token,  # IAM role session token
                aws_region_name="us-west-2",
            )

            assert isinstance(response, litellm.EmbeddingResponse)

            # Verify that the request was made
            assert mock_post.called, "HTTP client post should be called"

            # Get the actual call arguments
            call_kwargs = mock_post.call_args.kwargs
            headers = call_kwargs.get("headers", {})

            # Verify custom headers are present in the request
            for header_key, header_value in custom_headers.items():
                # Check if header exists (case-insensitive for HTTP headers)
                header_found = any(
                    k.lower() == header_key.lower() for k in headers.keys()
                )
                assert header_found, (
                    f"Custom header {header_key} should be in request headers. "
                    f"Found headers: {list(headers.keys())}"
                )

                # Verify the value matches
                header_value_found = None
                for k, v in headers.items():
                    if k.lower() == header_key.lower():
                        header_value_found = v
                        break

                assert header_value_found == header_value, (
                    f"Header {header_key} should have value {header_value}, "
                    f"but found {header_value_found}"
                )

            # Verify AWS signature headers are also present
            assert "Authorization" in headers, "AWS signature should be present"
            assert "X-Amz-Date" in headers, "AWS date header should be present"
            assert (
                "X-Amz-Security-Token" in headers
            ), "Session token header should be present"
            assert (
                headers["X-Amz-Security-Token"] == aws_session_token
            ), "Session token should match the provided token"

            # Verify the custom api_base was used
            called_url = call_kwargs.get("url", "")
            assert custom_api_base in str(called_url), (
                f"Custom api_base {custom_api_base} should be used. "
                f"Got URL: {called_url}"
            )

            print(
                "✓ Test passed (async): Custom headers forwarded with IAM role + custom api_base"
            )
            print(
                f"  Custom headers found: {[k for k in headers.keys() if k.lower().startswith('x-custom') or k.lower().startswith('x-forwarded')]}"
            )
            print(
                f"  AWS headers found: {[k for k in headers.keys() if k.lower().startswith('x-amz') or k.lower() == 'authorization']}"
            )

        except Exception as e:
            pytest.fail(
                f"Failed to forward headers with IAM role + custom api_base (async): {str(e)}"
            )


def test_titan_multimodal_embedding_image_cost_tracking():
    """Test that Titan multimodal embedding with image input populates image_count in Usage."""
    from litellm.llms.bedrock.embed.amazon_titan_multimodal_transformation import (
        AmazonTitanMultimodalEmbeddingG1Config,
    )

    config = AmazonTitanMultimodalEmbeddingG1Config()

    # Simulate response from AWS Bedrock
    response_list = [
        {
            "embedding": [0.1, 0.2, 0.3],
            "inputTextTokenCount": 0,
        }
    ]

    # Simulate batch_data with an image request (inputImage key set by _transform_request)
    batch_data = [{"inputImage": "/9j/4AAQSkZJRg=="}]

    result = config._transform_response(
        response_list=response_list,
        model="amazon.titan-embed-image-v1",
        batch_data=batch_data,
    )

    assert result.usage is not None
    assert result.usage.prompt_tokens_details is not None
    assert result.usage.prompt_tokens_details.image_count == 1


def test_titan_multimodal_embedding_text_no_image_count():
    """Test that Titan multimodal embedding with text-only input does not set image_count."""
    from litellm.llms.bedrock.embed.amazon_titan_multimodal_transformation import (
        AmazonTitanMultimodalEmbeddingG1Config,
    )

    config = AmazonTitanMultimodalEmbeddingG1Config()

    response_list = [
        {
            "embedding": [0.1, 0.2, 0.3],
            "inputTextTokenCount": 5,
        }
    ]

    # Text-only request — no inputImage key
    batch_data = [{"inputText": "hello world"}]

    result = config._transform_response(
        response_list=response_list,
        model="amazon.titan-embed-image-v1",
        batch_data=batch_data,
    )

    assert result.usage is not None
    # prompt_tokens_details should be None for text-only (no image_count to report)
    assert result.usage.prompt_tokens_details is None


def test_titan_multimodal_embedding_backward_compat_no_batch_data():
    """Test that Titan transformer works without batch_data (backward compatibility)."""
    from litellm.llms.bedrock.embed.amazon_titan_multimodal_transformation import (
        AmazonTitanMultimodalEmbeddingG1Config,
    )

    config = AmazonTitanMultimodalEmbeddingG1Config()

    response_list = [
        {
            "embedding": [0.1, 0.2, 0.3],
            "inputTextTokenCount": 5,
        }
    ]

    # Call without batch_data — should not break
    result = config._transform_response(
        response_list=response_list,
        model="amazon.titan-embed-image-v1",
    )

    assert result.usage is not None
    assert result.usage.prompt_tokens == 5
    assert result.usage.prompt_tokens_details is None


def test_titan_image_embedding_cost_uses_per_image_rate():
    """
    End-to-end test: Titan image embedding with mocked AWS response
    should populate image_count for correct per-image cost calculation.
    """
    client = HTTPHandler()

    with patch.object(client, "post") as mock_post:
        mock_response = Mock()
        mock_response.status_code = 200
        embed_response = {
            "embedding": [0.1] * 1024,
            "inputTextTokenCount": 0,
        }
        mock_response.text = json.dumps(embed_response)
        mock_response.json = lambda: json.loads(mock_response.text)
        mock_post.return_value = mock_response

        response = litellm.embedding(
            model="bedrock/amazon.titan-embed-image-v1",
            input=["data:image/png;base64,iVBORw0KGgoAAAANSUhEUg=="],
            client=client,
            aws_access_key_id="fake",
            aws_secret_access_key="fake",
            aws_region_name="us-east-1",
        )

        assert isinstance(response, litellm.EmbeddingResponse)
        assert response.usage is not None
        assert response.usage.prompt_tokens_details is not None
        assert response.usage.prompt_tokens_details.image_count == 1


@pytest.mark.parametrize(
    "encoding_format,expected_embedding_types",
    [
        ("float", ["float"]),
        ("base64", ["float"]),
        (["float", "int8"], ["float", "int8"]),
    ],
)
def test_bedrock_cohere_embedding_types_wrapped_as_list(
    encoding_format, expected_embedding_types
):
    """
    Bedrock Cohere expects `embedding_types` as a JSON array, not a raw string.

    Regression test for: Bedrock returns
        Malformed input request: #/embedding_types: expected type: JSONArray, found: String
    when `encoding_format` is passed as a string.
    """
    client = HTTPHandler()
    model = "bedrock/cohere.embed-multilingual-v3"

    with patch.object(client, "post") as mock_post:
        mock_response = Mock()
        mock_response.status_code = 200
        mock_response.text = json.dumps(cohere_embedding_response)
        mock_response.json = lambda: json.loads(mock_response.text)
        mock_post.return_value = mock_response

        response = litellm.embedding(
            model=model,
            input=test_input,
            encoding_format=encoding_format,
            client=client,
            aws_region_name="us-east-1",
            aws_bedrock_runtime_endpoint="https://bedrock-runtime.us-east-1.amazonaws.com",
            api_key="test-bearer-token-12345",
        )

        assert isinstance(response, litellm.EmbeddingResponse)

        request_body = json.loads(mock_post.call_args.kwargs.get("data", "{}"))
        assert "embedding_types" in request_body
        assert request_body["embedding_types"] == expected_embedding_types
        assert isinstance(request_body["embedding_types"], list)


def test_load_credentials_assumes_role_with_external_id(monkeypatch):
    """A trust policy requiring sts:ExternalId must be satisfied by the deployment's aws_external_id."""
    import datetime

    import boto3
    from botocore.exceptions import ClientError

    from litellm.llms.bedrock.embed.embedding import BedrockEmbedding

    monkeypatch.delenv("AWS_EXTERNAL_ID", raising=False)

    class FakeSTSClient:
        def get_caller_identity(self):
            return {"Arn": "arn:aws:iam::111111111111:user/litellm-proxy-pod"}

        def assume_role(self, **params):
            if params.get("ExternalId") != "external-id-embed":
                raise ClientError(
                    {"Error": {"Code": "AccessDenied", "Message": "is not authorized to perform: sts:AssumeRole"}},
                    "AssumeRole",
                )
            return {
                "Credentials": {
                    "AccessKeyId": "ASIAEMBEDROLEKEY",
                    "SecretAccessKey": "assumed-secret",
                    "SessionToken": "assumed-session-token",
                    "Expiration": datetime.datetime.now(datetime.timezone.utc) + datetime.timedelta(minutes=30),
                }
            }

    optional_params = {
        "aws_access_key_id": "AKIAEMBEDCALLERKEY",
        "aws_secret_access_key": "pod-caller-secret",
        "aws_region_name": "us-east-1",
        "aws_role_name": "arn:aws:iam::999999999999:role/litellm-embed-role",
        "aws_session_name": "litellm-embed-session",
        "aws_external_id": "external-id-embed",
    }

    with patch.object(boto3, "client", return_value=FakeSTSClient()):
        credentials, aws_region_name = BedrockEmbedding()._load_credentials(optional_params)

    assert credentials.access_key == "ASIAEMBEDROLEKEY"
    assert credentials.token == "assumed-session-token"
    assert aws_region_name == "us-east-1"
    assert "aws_external_id" not in optional_params


def test_embedding_session_tags_sign_the_request_and_stay_out_of_the_body(monkeypatch):
    """The tagged STS session signs the InvokeModel call and the tags never reach the body (#34069)."""
    import datetime

    import boto3
    from botocore.exceptions import ClientError

    monkeypatch.delenv("AWS_BEARER_TOKEN_BEDROCK", raising=False)
    monkeypatch.delenv("AWS_WEB_IDENTITY_TOKEN_FILE", raising=False)
    monkeypatch.delenv("AWS_ROLE_ARN", raising=False)
    tags = [{"Key": "team", "Value": "genai"}]

    class FakeSTSClient:
        def get_caller_identity(self):
            return {"Arn": "arn:aws:iam::111111111111:user/litellm-proxy-pod"}

        def assume_role(self, **params):
            if list(params.get("Tags", ())) != tags:
                raise ClientError(
                    {"Error": {"Code": "AccessDenied", "Message": "is not authorized to perform: sts:TagSession"}},
                    "AssumeRole",
                )
            return {
                "Credentials": {
                    "AccessKeyId": "ASIAEMBEDTAGGED",
                    "SecretAccessKey": "assumed-secret",
                    "SessionToken": "assumed-session-token",
                    "Expiration": datetime.datetime.now(datetime.timezone.utc) + datetime.timedelta(minutes=30),
                }
            }

    client = HTTPHandler()
    with patch.object(boto3, "client", return_value=FakeSTSClient()), patch.object(client, "post") as mock_post:
        mock_response = Mock()
        mock_response.status_code = 200
        mock_response.text = json.dumps(titan_embedding_response)
        mock_response.json = lambda: json.loads(mock_response.text)
        mock_post.return_value = mock_response

        response = litellm.embedding(
            model="bedrock/amazon.titan-embed-text-v1",
            input=test_input,
            client=client,
            aws_region_name="us-east-1",
            aws_access_key_id="AKIAEMBEDCALLERKEY",
            aws_secret_access_key="pod-caller-secret",
            aws_role_name="arn:aws:iam::999999999999:role/litellm-embed-role",
            aws_session_name="litellm-embed-session",
            aws_session_tags=tags,
        )

    assert response.data[0]["embedding"] == titan_embedding_response["embedding"]
    sent = mock_post.call_args.kwargs
    assert "Credential=ASIAEMBEDTAGGED/" in sent["headers"]["Authorization"]
    assert "aws_session_tags" not in sent["data"]


def test_bedrock_embedding_bearer_token_never_runs_the_sigv4_credential_chain(monkeypatch):
    """The deployment's AWS profile does not exist, so resolving SigV4 credentials
    raises; a bearer-token deployment must still serve the request, since the
    bearer token alone signs it."""
    monkeypatch.setenv("AWS_BEARER_TOKEN_BEDROCK", "env-bearer-token-12345")
    client = HTTPHandler()

    with patch.object(client, "post") as mock_post:
        mock_response = Mock()
        mock_response.status_code = 200
        mock_response.text = json.dumps(titan_embedding_response)
        mock_response.json = lambda: json.loads(mock_response.text)
        mock_post.return_value = mock_response

        response = litellm.embedding(
            model="bedrock/amazon.titan-embed-text-v1",
            input=test_input,
            client=client,
            aws_region_name="us-west-2",
            aws_profile_name="litellm-no-such-aws-profile",
        )

    assert response.data[0]["embedding"] == titan_embedding_response["embedding"]
    assert mock_post.call_args.kwargs["headers"]["Authorization"] == "Bearer env-bearer-token-12345"


@pytest.mark.asyncio
async def test_async_single_func_embeddings_signs_off_the_event_loop(monkeypatch):
    """Regression for issue #40165: Titan, Nova, and TwelveLabs embeddings sign one SigV4 request per
    input, and botocore refreshes expiring credentials inside that signing with a blocking HTTP call,
    so each signing must run on a worker thread to keep the loop serving other requests."""
    monkeypatch.delenv("AWS_BEARER_TOKEN_BEDROCK", raising=False)
    probe = EventLoopProbe()
    client = MagicMock()
    client.__class__ = AsyncHTTPHandler
    client.post = AsyncMock(
        return_value=httpx.Response(
            200,
            json=titan_embedding_response,
            request=httpx.Request("POST", "https://bedrock-runtime.us-west-2.amazonaws.com"),
        )
    )

    release = asyncio.create_task(probe.release_refresh_from_the_loop())
    response = await BedrockEmbedding()._async_single_func_embeddings(
        client=client,
        timeout=None,
        batch_data=[{"inputText": test_input}],
        credentials=probe.credentials(),
        extra_headers=None,
        endpoint_url="https://bedrock-runtime.us-west-2.amazonaws.com/model/amazon.titan-embed-text-v1/invoke",
        aws_region_name="us-west-2",
        model="amazon.titan-embed-text-v1",
        logging_obj=MagicMock(),
        provider="amazon",
    )
    await release

    assert response.data[0]["embedding"] == titan_embedding_response["embedding"]
    assert "Authorization" in client.post.call_args.kwargs["headers"]
    assert probe.served_during_refresh is True
marengo_3_embedding_response = {"data": [{"embedding": [0.01 * i for i in range(512)]}]}
MARENGO_3_DUCK = "data:image/png;base64,ZHVjaw=="


@pytest.mark.parametrize(
    "model,kwargs,expected_body,expected_usage_details",
    [
        (
            "bedrock/us.twelvelabs.marengo-embed-3-0-v1:0",
            {"input_type": "text"},
            {"inputType": "text", "text": {"inputText": "a duck on water"}},
            {"query_count": 1},
        ),
        (
            "bedrock/twelvelabs.marengo-embed-3-0-v1:0",
            {"input_type": "text"},
            {"inputType": "text", "text": {"inputText": "a duck on water"}},
            {"query_count": 1},
        ),
        (
            "bedrock/us.twelvelabs.marengo-embed-3-0-v1:0",
            {"input_type": "text_image", "media_source": MARENGO_3_DUCK},
            {
                "inputType": "text_image",
                "text_image": {"inputText": "a duck on water", "mediaSource": {"base64String": "ZHVjaw=="}},
            },
            {"query_count": 1, "image_count": 1},
        ),
        (
            "bedrock/us.twelvelabs.marengo-embed-3-0-v1:0",
            {"input_type": "multi_input", "media_sources": {"bird": MARENGO_3_DUCK}},
            {
                "inputType": "multi_input",
                "multi_input": {
                    "inputText": "a duck on water",
                    "mediaSources": [{"name": "bird", "mediaType": "image", "base64String": "ZHVjaw=="}],
                },
            },
            {"query_count": 1, "image_count": 1},
        ),
    ],
)
def test_marengo_3_embedding_sends_the_nested_payload_and_parses_512_dims(
    model, kwargs, expected_body, expected_usage_details
):
    client = HTTPHandler()

    with patch.object(client, "post") as mock_post:
        mock_response = Mock()
        mock_response.status_code = 200
        mock_response.text = json.dumps(marengo_3_embedding_response)
        mock_response.json = lambda: json.loads(mock_response.text)
        mock_post.return_value = mock_response

        response = litellm.embedding(
            model=model,
            input="a duck on water",
            client=client,
            aws_region_name="us-east-1",
            api_key="test-bearer-token-12345",
            **kwargs,
        )

    assert json.loads(mock_post.call_args.kwargs["data"]) == expected_body
    assert mock_post.call_args.kwargs["url"].endswith(f"/model/{model.removeprefix('bedrock/').replace(':', '%3A')}/invoke")
    assert len(response.data[0]["embedding"]) == 512
    assert response.data[0]["embedding"][:2] == [0.0, 0.01]
    assert response.usage.prompt_tokens == 0
    assert response.usage.total_tokens == 0
    assert response.usage.prompt_tokens_details.model_dump(exclude_none=True) == expected_usage_details


def test_marengo_3_image_embedding_sends_the_media_under_the_image_key():
    client = HTTPHandler()

    with patch.object(client, "post") as mock_post:
        mock_response = Mock()
        mock_response.status_code = 200
        mock_response.text = json.dumps(marengo_3_embedding_response)
        mock_response.json = lambda: json.loads(mock_response.text)
        mock_post.return_value = mock_response

        response = litellm.embedding(
            model="bedrock/us.twelvelabs.marengo-embed-3-0-v1:0",
            input=MARENGO_3_DUCK,
            client=client,
            aws_region_name="us-east-1",
            api_key="test-bearer-token-12345",
            input_type="image",
        )

    assert json.loads(mock_post.call_args.kwargs["data"]) == {
        "inputType": "image",
        "image": {"mediaSource": {"base64String": "ZHVjaw=="}},
    }
    assert len(response.data[0]["embedding"]) == 512
    assert response.data[0]["embedding"][:2] == [0.0, 0.01]
    assert response.usage.prompt_tokens == 0
    assert response.usage.prompt_tokens_details.model_dump(exclude_none=True) == {"image_count": 1}


def test_marengo_2_7_embedding_keeps_the_flat_payload():
    client = HTTPHandler()

    with patch.object(client, "post") as mock_post:
        mock_response = Mock()
        mock_response.status_code = 200
        mock_response.text = json.dumps(twelvelabs_embedding_response)
        mock_response.json = lambda: json.loads(mock_response.text)
        mock_post.return_value = mock_response

        response = litellm.embedding(
            model="bedrock/us.twelvelabs.marengo-embed-2-7-v1:0",
            input="a duck on water",
            client=client,
            aws_region_name="us-east-1",
            api_key="test-bearer-token-12345",
            input_type="text",
        )

    assert json.loads(mock_post.call_args.kwargs["data"]) == {
        "inputType": "text",
        "inputText": "a duck on water",
        "textTruncate": "end",
    }
    assert response.data[0]["embedding"] == [0.1, 0.2, 0.3]
    assert response.usage.prompt_tokens == 0
    assert response.usage.prompt_tokens_details.model_dump(exclude_none=True) == {"query_count": 1}


def test_marengo_usage_counts_text_requests_and_images_across_a_batch():
    duck = {"mediaType": "image", "base64String": "ZHVjaw=="}
    response = TwelveLabsMarengoEmbeddingConfig()._transform_response(
        response_list=[marengo_3_embedding_response, marengo_3_embedding_response, marengo_3_embedding_response],
        model="us.twelvelabs.marengo-embed-3-0-v1:0",
        batch_data=[
            {"inputType": "text", "text": {"inputText": "a duck"}},
            {"inputType": "image", "image": {"mediaSource": {"base64String": "ZHVjaw=="}}},
            {"inputType": "multi_input", "multi_input": {"mediaSources": [{"name": "a", **duck}, {"name": "b", **duck}]}},
        ],
    )

    assert [item["index"] for item in response.data] == [0, 1, 2]
    assert response.usage.prompt_tokens == 0
    assert response.usage.total_tokens == 0
    assert response.usage.prompt_tokens_details.model_dump(exclude_none=True) == {"query_count": 1, "image_count": 3}


def test_marengo_usage_without_request_data_bills_nothing():
    response = TwelveLabsMarengoEmbeddingConfig()._transform_response(
        response_list=[marengo_3_embedding_response], model="us.twelvelabs.marengo-embed-3-0-v1:0"
    )

    assert len(response.data[0]["embedding"]) == 512
    assert response.usage.prompt_tokens == 0
    assert response.usage.prompt_tokens_details is None


def test_marengo_response_items_without_an_embedding_are_skipped():
    response = TwelveLabsMarengoEmbeddingConfig()._transform_response(
        response_list=[{"data": [{"embeddingOption": "visual-text", "startSec": 0.0}, {"embedding": [0.1, 0.2, 0.3]}]}],
        model="us.twelvelabs.marengo-embed-3-0-v1:0",
    )

    assert [item["embedding"] for item in response.data] == [[0.1, 0.2, 0.3]]
    assert response.data[0]["index"] == 0


def test_marengo_3_text_image_without_media_source_is_a_bad_request():
    with pytest.raises(litellm.BadRequestError, match=r"text_image.*media_source"):
        litellm.embedding(
            model="bedrock/us.twelvelabs.marengo-embed-3-0-v1:0",
            input="a duck on water",
            aws_region_name="us-east-1",
            api_key="test-bearer-token-12345",
            input_type="text_image",
        )