diff --git a/tests/test_litellm/llms/bedrock/embed/test_bedrock_embedding.py b/tests/test_litellm/llms/bedrock/embed/test_bedrock_embedding.py index aec0b5fc6c..607f477d42 100644 --- a/tests/test_litellm/llms/bedrock/embed/test_bedrock_embedding.py +++ b/tests/test_litellm/llms/bedrock/embed/test_bedrock_embedding.py @@ -2,11 +2,12 @@ import json import os import sys from unittest.mock import Mock, patch + import pytest sys.path.insert(0, os.path.abspath("../../../../..")) # Adds the parent directory to the system path import litellm -from litellm.llms.custom_httpx.http_handler import HTTPHandler, AsyncHTTPHandler +from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler, HTTPHandler # Mock responses for different embedding models titan_embedding_response = { @@ -137,7 +138,7 @@ def test_bedrock_embedding_with_sigv4(): """Test embedding falls back to SigV4 auth when no bearer token is provided""" litellm.set_verbose = True 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]}] @@ -150,4 +151,85 @@ def test_bedrock_embedding_with_sigv4(): ) assert isinstance(response, litellm.EmbeddingResponse) - mock_bedrock_embed.assert_called_once() \ No newline at end of file + 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""" + litellm.set_verbose = True + 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)""" + litellm.set_verbose = True + 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"] \ No newline at end of file