mirror of
https://github.com/tiennm99/litellm.git
synced 2026-08-06 06:24:16 +00:00
Add test coverage for Bedrock Titan V2 encoding_format parameter
- Test encoding_format='float' parameter mapping and response handling - Test encoding_format='base64' parameter mapping to binary format - Verify parameter transformation and response processing - Mock AWS API responses for both float and binary formats - Ensure OpenAI compatibility with new encoding_format support
This commit is contained in:
@@ -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()
|
||||
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"]
|
||||
Reference in New Issue
Block a user