litellm/tests/llm_translation/test_bedrock_nova_embedding.py
yuneng-jiang 6a0d03914c
test: drop the cwd-relative sys.path.insert calls from the test suite (#37802)
* test: drop the cwd-relative sys.path.insert calls from the test suite

TQ003 stands at 1,077 across 1,058 files, and 1,015 of them are the same shape:
sys.path.insert(0, os.path.abspath("../..")) and its deeper siblings. The
argument resolves against the working directory rather than the file, so from
the repo root, where every job runs pytest, it inserts the directory two levels
above the checkout. It has never pointed at litellm. The package is installed
into the environment anyway, which is what actually makes the import work, and
what the rule's message has said all along.

Removing them leaves 1,634 imports of sys and os with no remaining reference,
and those go too, except where another test module imports the name back out of
the file. The rest of TQ003 is 62 call sites that resolve against __file__ or a
variable, which are a different question and are left alone.

Collection is identical either way: 45,871 tests and the same 51 pre-existing
collection errors before and after, and ruff reports no new undefined name.

* test: drop the duplicate imports the sys.path sweep exposed to F811

* test(pre-call-utils): restore the os import the new bedrock tests need
2026-08-22 09:25:58 -07:00

677 lines
23 KiB
Python

"""
Test suite for Amazon Nova Multimodal Embeddings integration with LiteLLM.
Tests cover:
- Synchronous text embeddings
- Synchronous image embeddings
- Synchronous video/audio embeddings
- Asynchronous embeddings with segmentation
- Different embedding purposes and dimensions
- Error handling
"""
import json
from unittest.mock import MagicMock, Mock, patch
import pytest
import litellm
from litellm.llms.bedrock.embed.amazon_nova_transformation import (
AmazonNovaEmbeddingConfig,
)
class TestNovaTransformationRequest:
"""Test request transformation for Nova embeddings."""
def test_text_embedding_sync_request(self):
"""Test synchronous text embedding request transformation."""
config = AmazonNovaEmbeddingConfig()
inference_params = {
"embeddingPurpose": "GENERIC_INDEX",
"embedding_dimension": 1024,
"truncation_mode": "END",
}
request = config._transform_request(
input="Hello, world!",
inference_params=inference_params,
async_invoke_route=False,
)
assert request["schemaVersion"] == "nova-multimodal-embed-v1"
assert request["taskType"] == "SINGLE_EMBEDDING"
assert "singleEmbeddingParams" in request
params = request["singleEmbeddingParams"]
assert params["embeddingPurpose"] == "GENERIC_INDEX"
assert params["embeddingDimension"] == 1024
assert params["text"]["truncationMode"] == "END"
assert params["text"]["value"] == "Hello, world!"
def test_text_embedding_async_request(self):
"""Test asynchronous text embedding request transformation."""
config = AmazonNovaEmbeddingConfig()
inference_params = {
"embeddingPurpose": "TEXT_RETRIEVAL",
"embeddingDimension": 3072,
"text": {
"value": "Long text content...",
"segmentationConfig": {"maxLengthChars": 10000},
},
"output_s3_uri": "s3://my-bucket/output/",
}
request = config._transform_request(
input="Long text content...",
inference_params=inference_params,
async_invoke_route=True,
model_id="amazon.nova-2-multimodal-embeddings-v1:0",
output_s3_uri="s3://my-bucket/output/",
)
assert "modelId" in request
assert "modelInput" in request
assert "outputDataConfig" in request
model_input = request["modelInput"]
assert model_input["taskType"] == "SEGMENTED_EMBEDDING"
assert "segmentedEmbeddingParams" in model_input
params = model_input["segmentedEmbeddingParams"]
assert params["embeddingPurpose"] == "TEXT_RETRIEVAL"
assert params["embeddingDimension"] == 3072
assert params["text"]["segmentationConfig"]["maxLengthChars"] == 10000
def test_image_embedding_request(self):
"""Test image embedding request transformation."""
config = AmazonNovaEmbeddingConfig()
# Mock base64 image data
image_data = "iVBORw0KGgoAAAANSUhEUgAAAAEAAAABCAYAAAAfFcSJAAAADUlEQVR42mNk+M9QDwADhgGAWjR9awAAAABJRU5ErkJggg=="
inference_params = {
"embeddingPurpose": "IMAGE_RETRIEVAL",
"embeddingDimension": 1024,
"image": {
"format": "png",
"source": {"bytes": image_data},
"detailLevel": "STANDARD_IMAGE",
},
}
request = config._transform_request(
input=image_data,
inference_params=inference_params,
async_invoke_route=False,
)
params = request["singleEmbeddingParams"]
assert params["embeddingPurpose"] == "IMAGE_RETRIEVAL"
assert params["embeddingDimension"] == 1024
assert params["image"]["format"] == "png"
assert params["image"]["detailLevel"] == "STANDARD_IMAGE"
assert "source" in params["image"]
assert "bytes" in params["image"]["source"]
def test_video_embedding_request(self):
"""Test video embedding request transformation."""
config = AmazonNovaEmbeddingConfig()
inference_params = {
"embeddingPurpose": "VIDEO_RETRIEVAL",
"embeddingDimension": 3072,
"video": {
"format": "mp4",
"source": {"s3Location": {"uri": "s3://my-bucket/video.mp4"}},
"embeddingMode": "AUDIO_VIDEO_COMBINED",
},
}
request = config._transform_request(
input="s3://my-bucket/video.mp4",
inference_params=inference_params,
async_invoke_route=False,
)
params = request["singleEmbeddingParams"]
assert params["embeddingPurpose"] == "VIDEO_RETRIEVAL"
assert params["embeddingDimension"] == 3072
assert params["video"]["format"] == "mp4"
assert params["video"]["embeddingMode"] == "AUDIO_VIDEO_COMBINED"
assert (
params["video"]["source"]["s3Location"]["uri"] == "s3://my-bucket/video.mp4"
)
def test_audio_embedding_request(self):
"""Test audio embedding request transformation."""
config = AmazonNovaEmbeddingConfig()
inference_params = {
"embeddingPurpose": "AUDIO_RETRIEVAL",
"embeddingDimension": 1024,
"audio": {
"format": "mp3",
"source": {"s3Location": {"uri": "s3://my-bucket/audio.mp3"}},
},
}
request = config._transform_request(
input="s3://my-bucket/audio.mp3",
inference_params=inference_params,
async_invoke_route=False,
)
params = request["singleEmbeddingParams"]
assert params["embeddingPurpose"] == "AUDIO_RETRIEVAL"
assert params["embeddingDimension"] == 1024
assert params["audio"]["format"] == "mp3"
assert (
params["audio"]["source"]["s3Location"]["uri"] == "s3://my-bucket/audio.mp3"
)
def test_async_invoke_requires_output_s3_uri(self):
"""Test that async invoke requires output_s3_uri."""
config = AmazonNovaEmbeddingConfig()
inference_params = {
"embedding_purpose": "GENERIC_INDEX",
}
with pytest.raises(ValueError, match="output_s3_uri is required"):
config._transform_request(
input="Test text",
inference_params=inference_params,
async_invoke_route=True,
model_id="amazon.nova-2-multimodal-embeddings-v1:0",
output_s3_uri=None,
)
def test_default_embedding_purpose(self):
"""Test default embedding purpose is GENERIC_INDEX."""
config = AmazonNovaEmbeddingConfig()
request = config._transform_request(
input="Test text",
inference_params={},
async_invoke_route=False,
)
params = request["singleEmbeddingParams"]
assert params["embeddingPurpose"] == "GENERIC_INDEX"
def test_default_embedding_dimension(self):
"""Test default embedding dimension is 3072."""
config = AmazonNovaEmbeddingConfig()
request = config._transform_request(
input="Test text",
inference_params={},
async_invoke_route=False,
)
params = request["singleEmbeddingParams"]
assert params["embeddingDimension"] == 3072
def test_data_url_image_parsing(self):
"""Test that data URL images are properly parsed and transformed."""
config = AmazonNovaEmbeddingConfig()
# Test with JPEG image data URL
jpeg_data_url = "data:image/jpeg;base64,/9j/4AAQSkZJRgABAQAASABIAAD"
request = config._transform_request(
input=jpeg_data_url,
inference_params={"dimensions": 1024},
async_invoke_route=False,
)
params = request["singleEmbeddingParams"]
assert "image" in params
assert params["image"]["format"] == "jpeg"
assert "source" in params["image"]
assert params["image"]["source"]["bytes"] == "/9j/4AAQSkZJRgABAQAASABIAAD"
assert params["embeddingDimension"] == 1024
assert params["embeddingPurpose"] == "GENERIC_INDEX"
def test_data_url_png_image_parsing(self):
"""Test that data URL PNG images are properly parsed."""
config = AmazonNovaEmbeddingConfig()
# Test with PNG image data URL
png_data_url = (
"data:image/png;base64,iVBORw0KGgoAAAANSUhEUgAAAAEAAAABCAYAAAAfFcSJ"
)
request = config._transform_request(
input=png_data_url,
inference_params={},
async_invoke_route=False,
)
params = request["singleEmbeddingParams"]
assert "image" in params
assert params["image"]["format"] == "png"
assert (
params["image"]["source"]["bytes"]
== "iVBORw0KGgoAAAANSUhEUgAAAAEAAAABCAYAAAAfFcSJ"
)
def test_data_url_jpg_format_conversion(self):
"""Test that jpg format is converted to jpeg."""
config = AmazonNovaEmbeddingConfig()
# Test with jpg (should be converted to jpeg)
jpg_data_url = "data:image/jpg;base64,/9j/4AAQSkZJRg"
request = config._transform_request(
input=jpg_data_url,
inference_params={},
async_invoke_route=False,
)
params = request["singleEmbeddingParams"]
assert (
params["image"]["format"] == "jpeg"
) # Should be converted from jpg to jpeg
def test_data_url_video_parsing(self):
"""Test that data URL videos are properly parsed."""
config = AmazonNovaEmbeddingConfig()
video_data_url = "data:video/mp4;base64,AAAAIGZ0eXBpc29t"
request = config._transform_request(
input=video_data_url,
inference_params={},
async_invoke_route=False,
)
params = request["singleEmbeddingParams"]
assert "video" in params
assert params["video"]["format"] == "mp4"
assert params["video"]["source"]["bytes"] == "AAAAIGZ0eXBpc29t"
def test_data_url_audio_parsing(self):
"""Test that data URL audio files are properly parsed."""
config = AmazonNovaEmbeddingConfig()
audio_data_url = "data:audio/mp3;base64,SUQzBAAAAAAAI1RTU0UAAAA"
request = config._transform_request(
input=audio_data_url,
inference_params={},
async_invoke_route=False,
)
params = request["singleEmbeddingParams"]
assert "audio" in params
assert params["audio"]["format"] == "mp3"
assert params["audio"]["source"]["bytes"] == "SUQzBAAAAAAAI1RTU0UAAAA"
class TestNovaTransformationResponse:
"""Test response transformation for Nova embeddings."""
def test_text_embedding_response(self):
"""Test text embedding response transformation."""
config = AmazonNovaEmbeddingConfig()
response_list = [
{
"embeddings": [
{
"embeddingType": "TEXT",
"embedding": [0.1, 0.2, 0.3, 0.4, 0.5],
}
]
}
]
result = config._transform_response(
response_list, model="amazon.nova-2-multimodal-embeddings-v1:0"
)
assert result.model == "amazon.nova-2-multimodal-embeddings-v1:0"
assert len(result.data) == 1
assert result.data[0].embedding == [0.1, 0.2, 0.3, 0.4, 0.5]
assert result.data[0].index == 0
assert result.data[0].object == "embedding"
assert result.usage.total_tokens > 0
def test_multiple_embeddings_response(self):
"""Test response with multiple embeddings."""
config = AmazonNovaEmbeddingConfig()
response_list = [
{
"embeddings": [
{
"embeddingType": "TEXT",
"embedding": [0.1, 0.2, 0.3],
}
]
},
{
"embeddings": [
{
"embeddingType": "TEXT",
"embedding": [0.4, 0.5, 0.6],
}
]
},
]
result = config._transform_response(
response_list, model="amazon.nova-2-multimodal-embeddings-v1:0"
)
assert len(result.data) == 2
assert result.data[0].embedding == [0.1, 0.2, 0.3]
assert result.data[1].embedding == [0.4, 0.5, 0.6]
assert result.data[0].index == 0
assert result.data[1].index == 1
def test_video_embedding_response_separate_mode(self):
"""Test video embedding response with separate audio/video."""
config = AmazonNovaEmbeddingConfig()
response_list = [
{
"embeddings": [
{
"embeddingType": "VIDEO",
"embedding": [0.1, 0.2, 0.3],
},
{
"embeddingType": "AUDIO",
"embedding": [0.4, 0.5, 0.6],
},
]
}
]
result = config._transform_response(
response_list, model="amazon.nova-2-multimodal-embeddings-v1:0"
)
assert len(result.data) == 2
assert result.data[0].embedding == [0.1, 0.2, 0.3]
assert result.data[1].embedding == [0.4, 0.5, 0.6]
def test_image_embedding_response_with_image_count(self):
"""Test that Nova image embedding response populates image_count for cost tracking."""
config = AmazonNovaEmbeddingConfig()
response_list = [
{
"embeddings": [
{
"embeddingType": "IMAGE",
"embedding": [0.1, 0.2, 0.3],
}
]
}
]
# Simulate batch_data with image in singleEmbeddingParams
batch_data = [
{
"schemaVersion": "nova-multimodal-embed-v1",
"taskType": "SINGLE_EMBEDDING",
"singleEmbeddingParams": {
"embeddingPurpose": "GENERIC_INDEX",
"embeddingDimension": 3072,
"image": {
"format": "jpeg",
"source": {"bytes": "/9j/4AAQSkZJRg=="},
},
},
}
]
result = config._transform_response(
response_list=response_list,
model="amazon.nova-2-multimodal-embeddings-v1:0",
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_text_embedding_response_no_image_count(self):
"""Test that Nova text embedding response does not set image_count."""
config = AmazonNovaEmbeddingConfig()
response_list = [
{
"embeddings": [
{
"embeddingType": "TEXT",
"embedding": [0.1, 0.2, 0.3],
"truncatedCharLength": 20,
}
]
}
]
batch_data = [
{
"schemaVersion": "nova-multimodal-embed-v1",
"taskType": "SINGLE_EMBEDDING",
"singleEmbeddingParams": {
"embeddingPurpose": "GENERIC_INDEX",
"embeddingDimension": 3072,
"text": {"value": "hello world", "truncationMode": "END"},
},
}
]
result = config._transform_response(
response_list=response_list,
model="amazon.nova-2-multimodal-embeddings-v1:0",
batch_data=batch_data,
)
assert result.usage is not None
assert result.usage.prompt_tokens_details is None
def test_nova_embedding_backward_compat_no_batch_data(self):
"""Test that Nova transformer works without batch_data (backward compatibility)."""
config = AmazonNovaEmbeddingConfig()
response_list = [
{
"embeddings": [
{
"embeddingType": "TEXT",
"embedding": [0.1, 0.2, 0.3, 0.4, 0.5],
}
]
}
]
# Call without batch_data — should not break
result = config._transform_response(
response_list=response_list,
model="amazon.nova-2-multimodal-embeddings-v1:0",
)
assert result.usage is not None
assert result.usage.total_tokens > 0
assert result.usage.prompt_tokens_details is None
def test_async_invoke_response(self):
"""Test async invoke response transformation."""
config = AmazonNovaEmbeddingConfig()
response = {
"invocationArn": "arn:aws:bedrock:us-east-1:123456789012:async-invoke/abc123"
}
result = config._transform_async_invoke_response(
response, model="amazon.nova-2-multimodal-embeddings-v1:0"
)
assert result.model == "amazon.nova-2-multimodal-embeddings-v1:0"
assert len(result.data) == 1
assert result.data[0].embedding == [] # Empty for async jobs
assert result.usage.total_tokens == 0
assert hasattr(result, "_hidden_params")
assert hasattr(result._hidden_params, "_invocation_arn")
assert (
result._hidden_params._invocation_arn
== "arn:aws:bedrock:us-east-1:123456789012:async-invoke/abc123"
)
class TestNovaEmbeddingIntegration:
"""Integration tests for Nova embeddings through LiteLLM."""
@pytest.mark.skip(reason="Requires AWS credentials and actual API calls")
def test_sync_text_embedding_e2e(self):
"""End-to-end test for synchronous text embedding."""
response = litellm.embedding(
model="bedrock/amazon.nova-2-multimodal-embeddings-v1:0",
input=["Hello, world!"],
aws_region_name="us-east-1",
)
assert response is not None
assert len(response.data) == 1
assert len(response.data[0].embedding) > 0
@pytest.mark.skip(reason="Requires AWS credentials and actual API calls")
def test_async_text_embedding_e2e(self):
"""End-to-end test for asynchronous text embedding."""
response = litellm.embedding(
model="bedrock/async_invoke/amazon.nova-2-multimodal-embeddings-v1:0",
input=["Long text content for segmentation..."],
aws_region_name="us-east-1",
output_s3_uri="s3://my-bucket/output/",
segmentation_config={"maxLengthChars": 10000},
)
assert response is not None
assert hasattr(response, "_hidden_params")
assert hasattr(response._hidden_params, "_invocation_arn")
@pytest.mark.skip(reason="Requires AWS credentials and actual API calls")
def test_image_embedding_e2e(self):
"""End-to-end test for image embedding."""
response = litellm.embedding(
model="bedrock/amazon.nova-2-multimodal-embeddings-v1:0",
input=["s3://my-bucket/image.png"],
aws_region_name="us-east-1",
input_type="image",
format="png",
embedding_purpose="IMAGE_RETRIEVAL",
)
assert response is not None
assert len(response.data) == 1
@pytest.mark.skip(reason="Requires AWS credentials and actual API calls")
def test_video_embedding_e2e(self):
"""End-to-end test for video embedding."""
response = litellm.embedding(
model="bedrock/amazon.nova-2-multimodal-embeddings-v1:0",
input=["s3://my-bucket/video.mp4"],
aws_region_name="us-east-1",
input_type="video",
format="mp4",
embedding_mode="AUDIO_VIDEO_COMBINED",
embedding_purpose="VIDEO_RETRIEVAL",
)
assert response is not None
assert len(response.data) == 1
@pytest.mark.skip(reason="Requires AWS credentials and actual API calls")
def test_different_dimensions(self):
"""Test different embedding dimensions."""
for dimension in [256, 384, 1024, 3072]:
response = litellm.embedding(
model="bedrock/amazon.nova-2-multimodal-embeddings-v1:0",
input=["Test text"],
aws_region_name="us-east-1",
dimensions=dimension,
)
assert response is not None
assert len(response.data[0].embedding) == dimension
@pytest.mark.skip(reason="Requires AWS credentials and actual API calls")
def test_different_embedding_purposes(self):
"""Test different embedding purposes."""
purposes = [
"GENERIC_INDEX",
"GENERIC_RETRIEVAL",
"TEXT_RETRIEVAL",
"CLASSIFICATION",
"CLUSTERING",
]
for purpose in purposes:
response = litellm.embedding(
model="bedrock/amazon.nova-2-multimodal-embeddings-v1:0",
input=["Test text"],
aws_region_name="us-east-1",
embedding_purpose=purpose,
)
assert response is not None
assert len(response.data) == 1
class TestNovaProviderDetection:
"""Test provider detection for Nova models."""
def test_nova_provider_detection(self):
"""Test that Nova provider is correctly detected."""
from litellm.llms.bedrock.base_aws_llm import BaseAWSLLM
provider = BaseAWSLLM.get_bedrock_embedding_provider(
"amazon.nova-2-multimodal-embeddings-v1:0"
)
# Should detect "amazon" as provider since "nova" is in the model name
# but the provider detection looks at the first part before the dot
assert provider in ["amazon", "nova"]
def test_nova_in_model_name(self):
"""Test that models with 'nova' in the name are detected."""
from litellm.llms.bedrock.base_aws_llm import BaseAWSLLM
# Test various Nova model name formats
test_models = [
"amazon.nova-2-multimodal-embeddings-v1:0",
"us.amazon.nova-2-multimodal-embeddings-v1:0",
]
for model in test_models:
provider = BaseAWSLLM.get_bedrock_embedding_provider(model)
assert provider is not None
if __name__ == "__main__":
# Run basic transformation tests
print("Running Nova Embedding Transformation Tests...")
test_request = TestNovaTransformationRequest()
test_request.test_text_embedding_sync_request()
test_request.test_text_embedding_async_request()
test_request.test_image_embedding_request()
test_request.test_video_embedding_request()
test_request.test_audio_embedding_request()
test_response = TestNovaTransformationResponse()
test_response.test_text_embedding_response()
test_response.test_multiple_embeddings_response()
test_response.test_async_invoke_response()
print("All transformation tests passed!")