diff --git a/docs/my-website/docs/embedding/supported_embedding.md b/docs/my-website/docs/embedding/supported_embedding.md index e63d9403665..0e8252b409b 100644 --- a/docs/my-website/docs/embedding/supported_embedding.md +++ b/docs/my-website/docs/embedding/supported_embedding.md @@ -263,6 +263,8 @@ print(response) | Model Name | Function Call | |----------------------|---------------------------------------------| +| Amazon Nova Multimodal Embeddings | `embedding(model="bedrock/amazon.nova-2-multimodal-embeddings-v1:0", input=input)` | [Nova Docs](../providers/bedrock_embedding#amazon-nova-multimodal-embeddings) | +| Amazon Nova (Async) | `embedding(model="bedrock/async_invoke/amazon.nova-2-multimodal-embeddings-v1:0", input=input, input_type="text", output_s3_uri="s3://bucket/")` | [Nova Async Docs](../providers/bedrock_embedding#asynchronous-embeddings-with-segmentation) | | Titan Embeddings - G1 | `embedding(model="amazon.titan-embed-text-v1", input=input)` | | Cohere Embeddings - English | `embedding(model="cohere.embed-english-v3", input=input)` | | Cohere Embeddings - Multilingual | `embedding(model="cohere.embed-multilingual-v3", input=input)` | diff --git a/docs/my-website/docs/providers/bedrock_embedding.md b/docs/my-website/docs/providers/bedrock_embedding.md index 76c9606533e..e2e7c0dcedd 100644 --- a/docs/my-website/docs/providers/bedrock_embedding.md +++ b/docs/my-website/docs/providers/bedrock_embedding.md @@ -4,7 +4,8 @@ | Provider | LiteLLM Route | AWS Documentation | Cost Tracking | |----------|---------------|-------------------|---------------| -| Amazon Titan | `bedrock/amazon.*` | [Amazon Titan Embeddings](https://docs.aws.amazon.com/bedrock/latest/userguide/titan-embedding-models.html) | ✅ | +| Amazon Titan | `bedrock/amazon.titan-*` | [Amazon Titan Embeddings](https://docs.aws.amazon.com/bedrock/latest/userguide/titan-embedding-models.html) | ✅ | +| Amazon Nova | `bedrock/amazon.nova-*` | [Amazon Nova Embeddings](https://docs.aws.amazon.com/bedrock/latest/userguide/nova-embed.html) | ✅ | | Cohere | `bedrock/cohere.*` | [Cohere Embeddings](https://docs.aws.amazon.com/bedrock/latest/userguide/model-parameters-cohere-embed.html) | ✅ | | TwelveLabs | `bedrock/us.twelvelabs.*` | [TwelveLabs](https://docs.aws.amazon.com/bedrock/latest/userguide/model-parameters-twelvelabs.html) | ✅ | @@ -16,6 +17,7 @@ LiteLLM supports AWS Bedrock's async-invoke feature for embedding models that re | Provider | Async Invoke Route | Use Case | |----------|-------------------|----------| +| Amazon Nova | `bedrock/async_invoke/amazon.nova-2-multimodal-embeddings-v1:0` | Multimodal embeddings with segmentation for long text, video, and audio | | TwelveLabs Marengo | `bedrock/async_invoke/us.twelvelabs.marengo-embed-2-7-v1:0` | Video, audio, image, and text embeddings | ### Required Parameters @@ -116,7 +118,7 @@ def check_async_job_status(invocation_arn, aws_region_name="us-east-1"): """Check the status of an async invoke job using LiteLLM batch API""" try: response = retrieve_batch( - batch_id=invocation_arn, + batch_id=invocation_arn, # Pass the invocation ARN here custom_llm_provider="bedrock", aws_region_name=aws_region_name ) @@ -128,11 +130,47 @@ def check_async_job_status(invocation_arn, aws_region_name="us-east-1"): # Check status status = check_async_job_status(invocation_arn, "us-east-1") if status: - print(f"Job Status: {status.status}") - print(f"Output Location: {status.output_file_id}") + print(f"Job Status: {status.status}") # "in_progress", "completed", or "failed" + print(f"Output Location: {status.metadata['output_file_id']}") # S3 URI where results are stored ``` -**Note:** The actual embedding results are stored in S3. The `output_file_id` from the batch status can be used to locate the results file in your S3 bucket. +#### Polling Until Complete + +Here's a complete example of polling for job completion: + +```python +def wait_for_async_job(invocation_arn, aws_region_name="us-east-1", max_wait=3600): + """Poll job status until completion""" + start_time = time.time() + + while True: + status = retrieve_batch( + batch_id=invocation_arn, + custom_llm_provider="bedrock", + aws_region_name=aws_region_name, + ) + + if status.status == "completed": + print("✅ Job completed!") + return status + elif status.status == "failed": + error_msg = status.metadata.get('failure_message', 'Unknown error') + raise Exception(f"❌ Job failed: {error_msg}") + else: + elapsed = time.time() - start_time + if elapsed > max_wait: + raise TimeoutError(f"Job timed out after {max_wait} seconds") + + print(f"⏳ Job still processing... (elapsed: {elapsed:.0f}s)") + time.sleep(10) # Wait 10 seconds before checking again + +# Wait for completion +completed_status = wait_for_async_job(invocation_arn) +output_s3_uri = completed_status.metadata['output_file_id'] +print(f"Results available at: {output_s3_uri}") +``` + +**Note:** The actual embedding results are stored in S3. When the job is completed, download the results from the S3 location specified in `status.metadata['output_file_id']`. The results will be in JSON/JSONL format containing the embedding vectors. ### Error Handling @@ -179,7 +217,7 @@ except Exception as e: ### Limitations -- Async-invoke is currently only supported for TwelveLabs Marengo models +- Async-invoke is supported for TwelveLabs Marengo and Amazon Nova models - Results are stored in S3 and must be retrieved separately using the output file ID - Job status checking requires using LiteLLM's `retrieve_batch()` function - No built-in polling mechanism in LiteLLM (must implement your own status checking loop) @@ -259,6 +297,7 @@ print(response) | Model Name | Usage | Supported Additional OpenAI params | |----------------------|---------------------------------------------|-----| +| **Amazon Nova Multimodal Embeddings** | `embedding(model="bedrock/amazon.nova-2-multimodal-embeddings-v1:0", input=input)` | Supports multimodal input (text, image, video, audio), multiple purposes, dimensions (256, 384, 1024, 3072) | | Titan Embeddings V2 | `embedding(model="bedrock/amazon.titan-embed-text-v2:0", input=input)` | [here](https://github.com/BerriAI/litellm/blob/f5905e100068e7a4d61441d7453d7cf5609c2121/litellm/llms/bedrock/embed/amazon_titan_v2_transformation.py#L59) | | Titan Embeddings - V1 | `embedding(model="bedrock/amazon.titan-embed-text-v1", input=input)` | [here](https://github.com/BerriAI/litellm/blob/f5905e100068e7a4d61441d7453d7cf5609c2121/litellm/llms/bedrock/embed/amazon_titan_g1_transformation.py#L53) | Titan Multimodal Embeddings | `embedding(model="bedrock/amazon.titan-embed-image-v1", input=input)` | [here](https://github.com/BerriAI/litellm/blob/f5905e100068e7a4d61441d7453d7cf5609c2121/litellm/llms/bedrock/embed/amazon_titan_multimodal_transformation.py#L28) | diff --git a/litellm/__init__.py b/litellm/__init__.py index dfe35e33253..778d972c6eb 100644 --- a/litellm/__init__.py +++ b/litellm/__init__.py @@ -1257,6 +1257,9 @@ from .llms.bedrock.embed.cohere_transformation import BedrockCohereEmbeddingConf from .llms.bedrock.embed.twelvelabs_marengo_transformation import ( TwelveLabsMarengoEmbeddingConfig, ) +from .llms.bedrock.embed.amazon_nova_transformation import ( + AmazonNovaEmbeddingConfig, +) from .llms.openai.openai import OpenAIConfig, MistralEmbeddingConfig from .llms.openai.image_variations.transformation import OpenAIImageVariationConfig from .llms.deepinfra.chat.transformation import DeepInfraConfig diff --git a/litellm/batches/main.py b/litellm/batches/main.py index 995c45b925a..57a9857dd6a 100644 --- a/litellm/batches/main.py +++ b/litellm/batches/main.py @@ -1045,28 +1045,44 @@ def _handle_async_invoke_status( # Transform response to a LiteLLMBatch object from litellm.types.utils import LiteLLMBatch + # Normalize status to lowercase (AWS returns 'Completed', 'Failed', etc.) + aws_status_raw = status_response.get("status", "") + aws_status_lower = aws_status_raw.lower() + # Map AWS status values to LiteLLM expected values + status_mapping = { + "completed": "completed", + "failed": "failed", + "inprogress": "in_progress", + "in_progress": "in_progress", + } + normalized_status = status_mapping.get(aws_status_lower, aws_status_lower) + + # Get output S3 URI safely + output_s3_uri = "" + try: + output_s3_uri = status_response["outputDataConfig"]["s3OutputDataConfig"]["s3Uri"] + except (KeyError, TypeError): + pass + + # Use BedrockBatchesConfig's timestamp parsing method (expects raw AWS status string) + from litellm.llms.bedrock.batches.transformation import BedrockBatchesConfig + created_at, in_progress_at, completed_at, failed_at, _, _ = BedrockBatchesConfig()._parse_timestamps_and_status(status_response, aws_status_raw) result = LiteLLMBatch( id=status_response["invocationArn"], object="batch", - status=status_response["status"], - created_at=status_response["submitTime"], - in_progress_at=status_response["lastModifiedTime"], - completed_at=status_response.get("endTime"), - failed_at=( - status_response.get("endTime") - if status_response["status"] == "failed" - else None - ), + status=normalized_status, + created_at=created_at, + in_progress_at=in_progress_at, + completed_at=completed_at, + failed_at=failed_at, request_counts=BatchRequestCounts( total=1, - completed=1 if status_response["status"] == "completed" else 0, - failed=1 if status_response["status"] == "failed" else 0, + completed=1 if normalized_status == "completed" else 0, + failed=1 if normalized_status == "failed" else 0, ), metadata=dict( **{ - "output_file_id": status_response["outputDataConfig"][ - "s3OutputDataConfig" - ]["s3Uri"], + "output_file_id": output_s3_uri, "failure_message": status_response.get("failureMessage") or "", "model_arn": status_response["modelArn"], } diff --git a/litellm/constants.py b/litellm/constants.py index f1126b19090..b5826b43bfe 100644 --- a/litellm/constants.py +++ b/litellm/constants.py @@ -862,6 +862,7 @@ BEDROCK_EMBEDDING_PROVIDERS_LITERAL = Literal[ "cohere", "amazon", "twelvelabs", + "nova", ] BEDROCK_CONVERSE_MODELS = [ @@ -922,6 +923,7 @@ cohere_embedding_models: set = set( bedrock_embedding_models: set = set( [ "amazon.titan-embed-text-v1", + "amazon.nova-2-multimodal-embeddings-v1:0", "cohere.embed-english-v3", "cohere.embed-multilingual-v3", "cohere.embed-v4:0", diff --git a/litellm/llms/bedrock/base_aws_llm.py b/litellm/llms/bedrock/base_aws_llm.py index 72e270428ac..ed658c793af 100644 --- a/litellm/llms/bedrock/base_aws_llm.py +++ b/litellm/llms/bedrock/base_aws_llm.py @@ -387,9 +387,16 @@ class BaseAWSLLM: Handles scenarios like: 1. model=cohere.embed-english-v3:0 -> Returns `cohere` 2. model=amazon.titan-embed-text-v1 -> Returns `amazon` - 3. model=us.twelvelabs.marengo-embed-2-7-v1:0 -> Returns `twelvelabs` - 4. model=twelvelabs.marengo-embed-2-7-v1:0 -> Returns `twelvelabs` + 3. model=amazon.nova-2-multimodal-embeddings-v1:0 -> Returns `nova` + 4. model=us.twelvelabs.marengo-embed-2-7-v1:0 -> Returns `twelvelabs` + 5. model=twelvelabs.marengo-embed-2-7-v1:0 -> Returns `twelvelabs` """ + # Special case: Check for "nova" in model name first (before "amazon") + # This handles amazon.nova-* models + if "nova" in model.lower(): + if "nova" in get_args(BEDROCK_EMBEDDING_PROVIDERS_LITERAL): + return cast(BEDROCK_EMBEDDING_PROVIDERS_LITERAL, "nova") + # Handle regional models like us.twelvelabs.marengo-embed-2-7-v1:0 if "." in model: parts = model.split(".") diff --git a/litellm/llms/bedrock/embed/amazon_nova_transformation.py b/litellm/llms/bedrock/embed/amazon_nova_transformation.py new file mode 100644 index 00000000000..ada49d0ff21 --- /dev/null +++ b/litellm/llms/bedrock/embed/amazon_nova_transformation.py @@ -0,0 +1,260 @@ +""" +Transformation logic from OpenAI /v1/embeddings format to Bedrock Amazon Nova /invoke and /async-invoke format. + +Why separate file? Make it easy to see how transformation works + +Supports: +- Synchronous embeddings (SINGLE_EMBEDDING) +- Asynchronous embeddings with segmentation (SEGMENTED_EMBEDDING) +- Multimodal inputs: text, image, video, audio +- Multiple embedding purposes and dimensions + +Docs - https://docs.aws.amazon.com/bedrock/latest/userguide/nova-embed.html +""" + +from typing import List, Optional + +from litellm.types.utils import Embedding, EmbeddingResponse, Usage + + +class AmazonNovaEmbeddingConfig: + """ + Reference: https://docs.aws.amazon.com/bedrock/latest/userguide/nova-embed.html + + Amazon Nova Multimodal Embeddings supports: + - Text, image, video, and audio inputs + - Synchronous (InvokeModel) and asynchronous (StartAsyncInvoke) APIs + - Multiple embedding purposes and dimensions + """ + + def __init__(self) -> None: + pass + + def get_supported_openai_params(self) -> List[str]: + return [ + "dimensions", + ] + + def map_openai_params( + self, non_default_params: dict, optional_params: dict + ) -> dict: + """Map OpenAI-style parameters to Nova parameters.""" + for k, v in non_default_params.items(): + if k == "dimensions": + # Map OpenAI dimensions to Nova embedding_dimension + optional_params["embedding_dimension"] = v + elif k in self.get_supported_openai_params(): + optional_params[k] = v + return optional_params + + def _transform_request( + self, + input: str, + inference_params: dict, + async_invoke_route: bool = False, + model_id: Optional[str] = None, + output_s3_uri: Optional[str] = None, + ) -> dict: + """ + Transform OpenAI-style input to Nova format. + + Only handles OpenAI params (dimensions). All other Nova-specific params + should be passed via inference_params and will be passed through as-is. + + Args: + input: The input text or media reference + inference_params: Additional parameters (will be passed through) + async_invoke_route: Whether this is for async invoke + model_id: Model ID (for async invoke) + output_s3_uri: S3 URI for output (for async invoke) + + Returns: + dict: Nova embedding request + """ + # Determine task type + task_type = "SEGMENTED_EMBEDDING" if async_invoke_route else "SINGLE_EMBEDDING" + + # Build the base request structure + request: dict = { + "schemaVersion": "nova-multimodal-embed-v1", + "taskType": task_type, + } + + # Start with inference_params (user-provided params) + embedding_params = inference_params.copy() + + embedding_params.pop("output_s3_uri", None) + + # Map OpenAI dimensions to embeddingDimension if provided + if "dimensions" in embedding_params: + embedding_params["embeddingDimension"] = embedding_params.pop("dimensions") + elif "embedding_dimension" in embedding_params: + embedding_params["embeddingDimension"] = embedding_params.pop("embedding_dimension") + + # Add required embeddingPurpose if not provided (required by Nova API) + if "embeddingPurpose" not in embedding_params: + embedding_params["embeddingPurpose"] = "GENERIC_INDEX" + + # Add required embeddingDimension if not provided (required by Nova API) + if "embeddingDimension" not in embedding_params: + embedding_params["embeddingDimension"] = 3072 + + # For text input, add basic text structure if user hasn't provided text/image/video/audio + if "text" not in embedding_params and "image" not in embedding_params and "video" not in embedding_params and "audio" not in embedding_params: + # Default to text if no modality specified + if input.startswith("s3://"): + embedding_params["text"] = { + "source": {"s3Location": {"uri": input}}, + "truncationMode": "END" # Required by Nova API + } + else: + embedding_params["text"] = { + "value": input, + "truncationMode": "END" # Required by Nova API + } + + # Set the embedding params in the request + if task_type == "SINGLE_EMBEDDING": + request["singleEmbeddingParams"] = embedding_params + else: + request["segmentedEmbeddingParams"] = embedding_params + + # For async invoke, wrap in the async invoke format + if async_invoke_route and model_id: + return self._wrap_async_invoke_request( + model_input=request, + model_id=model_id, + output_s3_uri=output_s3_uri, + ) + + return request + + def _wrap_async_invoke_request( + self, + model_input: dict, + model_id: str, + output_s3_uri: Optional[str] = None, + ) -> dict: + """ + Wrap the transformed request in the AWS Bedrock async invoke format. + + Args: + model_input: The transformed Nova embedding request + model_id: The model identifier (without async_invoke prefix) + output_s3_uri: S3 URI for output data config + + Returns: + dict: The wrapped async invoke request + """ + import urllib.parse + + # Clean the model ID + unquoted_model_id = urllib.parse.unquote(model_id) + if unquoted_model_id.startswith("async_invoke/"): + unquoted_model_id = unquoted_model_id.replace("async_invoke/", "") + + # Validate that the S3 URI is not empty + if not output_s3_uri or output_s3_uri.strip() == "": + raise ValueError("output_s3_uri is required for async invoke requests") + + return { + "modelId": unquoted_model_id, + "modelInput": model_input, + "outputDataConfig": { + "s3OutputDataConfig": { + "s3Uri": output_s3_uri + } + }, + } + + def _transform_response( + self, response_list: List[dict], model: str + ) -> EmbeddingResponse: + """ + Transform Nova response to OpenAI format. + + Nova response format: + { + "embeddings": [ + { + "embeddingType": "TEXT" | "IMAGE" | "VIDEO" | "AUDIO" | "AUDIO_VIDEO_COMBINED", + "embedding": [0.1, 0.2, ...], + "truncatedCharLength": 100 # Optional, only for text + } + ] + } + """ + embeddings: List[Embedding] = [] + total_tokens = 0 + + for response in response_list: + # Nova response has an "embeddings" array + if "embeddings" in response and isinstance(response["embeddings"], list): + for item in response["embeddings"]: + if "embedding" in item: + embedding = Embedding( + embedding=item["embedding"], + index=len(embeddings), + object="embedding", + ) + embeddings.append(embedding) + + # Estimate token count + # For text, use truncatedCharLength if available + if "truncatedCharLength" in item: + total_tokens += item["truncatedCharLength"] // 4 + else: + # Rough estimate based on embedding dimension + total_tokens += len(item["embedding"]) // 4 + elif "embedding" in response: + # Direct embedding response (fallback) + embedding = Embedding( + embedding=response["embedding"], + index=len(embeddings), + object="embedding", + ) + embeddings.append(embedding) + total_tokens += len(response["embedding"]) // 4 + + usage = Usage(prompt_tokens=total_tokens, total_tokens=total_tokens) + + return EmbeddingResponse(data=embeddings, model=model, usage=usage) + + def _transform_async_invoke_response( + self, response: dict, model: str + ) -> EmbeddingResponse: + """ + Transform async invoke response (invocation ARN) to OpenAI format. + + AWS async invoke returns: + { + "invocationArn": "arn:aws:bedrock:us-east-1:123456789012:async-invoke/abc123" + } + + We transform this to a job-like embedding response with the ARN in hidden params. + """ + invocation_arn = response.get("invocationArn", "") + + # Create a placeholder embedding object for the job + embedding = Embedding( + embedding=[], # Empty embedding for async jobs + index=0, + object="embedding", + ) + + # Create usage object (empty for async jobs) + usage = Usage(prompt_tokens=0, total_tokens=0) + + # Create hidden params with job ID + from litellm.types.llms.base import HiddenParams + + hidden_params = HiddenParams() + setattr(hidden_params, "_invocation_arn", invocation_arn) + + return EmbeddingResponse( + data=[embedding], + model=model, + usage=usage, + hidden_params=hidden_params, + ) + diff --git a/litellm/llms/bedrock/embed/embedding.py b/litellm/llms/bedrock/embed/embedding.py index fea29935975..7152d7ce15c 100644 --- a/litellm/llms/bedrock/embed/embedding.py +++ b/litellm/llms/bedrock/embed/embedding.py @@ -27,6 +27,7 @@ from litellm.types.utils import EmbeddingResponse, LlmProviders from ..base_aws_llm import BaseAWSLLM from ..common_utils import BedrockError +from .amazon_nova_transformation import AmazonNovaEmbeddingConfig from .amazon_titan_g1_transformation import AmazonTitanG1Config from .amazon_titan_multimodal_transformation import ( AmazonTitanMultimodalEmbeddingG1Config, @@ -175,6 +176,12 @@ class BedrockEmbedding(BaseAWSLLM): response=response_list[0], model=model ) ) + elif provider == "nova": + returned_response = ( + AmazonNovaEmbeddingConfig()._transform_async_invoke_response( + response=response_list[0], model=model + ) + ) else: # For other providers, create a generic async response invocation_arn = response_list[0].get("invocationArn", "") @@ -222,6 +229,10 @@ class BedrockEmbedding(BaseAWSLLM): response_list=response_list, model=model ) ) + elif provider == "nova": + returned_response = AmazonNovaEmbeddingConfig()._transform_response( + response_list=response_list, model=model + ) ########################################################## # Validate returned response @@ -366,7 +377,7 @@ class BedrockEmbedding(BaseAWSLLM): is_async_invoke=is_async_invoke, ) - def embeddings( + def embeddings( # noqa: PLR0915 self, model: str, input: List[str], @@ -467,6 +478,17 @@ class BedrockEmbedding(BaseAWSLLM): ) ) batch_data.append(twelvelabs_request) + elif provider == "nova": + batch_data = [] + for i in input: + nova_request = AmazonNovaEmbeddingConfig()._transform_request( + input=i, + inference_params=inference_params, + async_invoke_route=has_async_invoke, + model_id=modelId, + output_s3_uri=inference_params.get("output_s3_uri"), + ) + batch_data.append(nova_request) ### SET RUNTIME ENDPOINT ### endpoint_url, proxy_endpoint_url = self.get_runtime_endpoint( @@ -581,22 +603,39 @@ class BedrockEmbedding(BaseAWSLLM): aws_region_name=aws_region_name, ) - # Construct the status check URL - status_url = f"{endpoint_url}/async-invoke/{invocation_arn}" - # Prepare headers + from urllib.parse import quote + + # Encode the ARN for use in URL path + encoded_arn = quote(invocation_arn, safe="") + status_url = f"{endpoint_url.rstrip('/')}/async-invoke/{encoded_arn}" + + # Prepare headers for GET request headers = {"Content-Type": "application/json"} - # Get AWS signed headers - prepped = self.get_request_headers( # type: ignore - credentials=credentials, - aws_region_name=aws_region_name, - extra_headers=None, - endpoint_url=status_url, - data="", # GET request, no body + # Use AWSRequest directly for GET requests (get_request_headers hardcodes POST) + try: + from botocore.auth import SigV4Auth + from botocore.awsrequest import AWSRequest + except ImportError: + raise ImportError( + "Missing boto3 to call bedrock. Run 'pip install boto3'." + ) + + # Create AWSRequest with GET method and encoded URL + request = AWSRequest( + method="GET", + url=status_url, + data=None, # GET request, no body headers=headers, - api_key=None, ) + + # Sign the request - SigV4Auth will create canonical string from request URL + sigv4 = SigV4Auth(credentials, "bedrock", aws_region_name) + sigv4.add_auth(request) + + # Prepare the request + prepped = request.prepare() # LOGGING if logging_obj is not None: diff --git a/litellm/types/llms/bedrock.py b/litellm/types/llms/bedrock.py index 330308e179c..3696f679640 100644 --- a/litellm/types/llms/bedrock.py +++ b/litellm/types/llms/bedrock.py @@ -427,6 +427,133 @@ class TwelveLabsAsyncInvokeStatusResponse(TypedDict): failureMessage: Optional[str] +# Amazon Nova Multimodal Embeddings types +NOVA_EMBEDDING_PURPOSES = Literal[ + "GENERIC_INDEX", + "GENERIC_RETRIEVAL", + "TEXT_RETRIEVAL", + "IMAGE_RETRIEVAL", + "VIDEO_RETRIEVAL", + "DOCUMENT_RETRIEVAL", + "AUDIO_RETRIEVAL", + "CLASSIFICATION", + "CLUSTERING", +] + +NOVA_EMBEDDING_DIMENSIONS = Literal[256, 384, 1024, 3072] + +NOVA_TRUNCATION_MODES = Literal["START", "END", "NONE"] + +NOVA_DETAIL_LEVELS = Literal["STANDARD_IMAGE", "DOCUMENT_IMAGE"] + +NOVA_EMBEDDING_MODES = Literal["AUDIO_VIDEO_COMBINED", "AUDIO_VIDEO_SEPARATE"] + +NOVA_EMBEDDING_TYPES = Literal[ + "TEXT", "IMAGE", "VIDEO", "AUDIO", "AUDIO_VIDEO_COMBINED" +] + + +class NovaSourceS3Location(TypedDict): + uri: str + + +class NovaSourceObject(TypedDict, total=False): + bytes: str # base64 encoded + s3Location: NovaSourceS3Location + + +class NovaTextParams(TypedDict, total=False): + truncationMode: NOVA_TRUNCATION_MODES + value: str + source: NovaSourceObject + + +class NovaImageParams(TypedDict, total=False): + format: str # png, jpeg, gif, webp + source: Required[NovaSourceObject] + detailLevel: NOVA_DETAIL_LEVELS + + +class NovaVideoParams(TypedDict, total=False): + format: str # mp4, mov, mkv, webm, flv, mpeg, mpg, wmv, 3gp + source: Required[NovaSourceObject] + embeddingMode: Required[NOVA_EMBEDDING_MODES] + + +class NovaAudioParams(TypedDict, total=False): + format: str # mp3, wav, ogg + source: Required[NovaSourceObject] + + +class NovaTextSegmentationConfig(TypedDict, total=False): + maxLengthChars: int # 800-50,000, default 32,000 + + +class NovaMediaSegmentationConfig(TypedDict, total=False): + durationSeconds: int # 1-30, default 5 + + +class NovaTextParamsWithSegmentation(NovaTextParams, total=False): + segmentationConfig: NovaTextSegmentationConfig + + +class NovaVideoParamsWithSegmentation(NovaVideoParams, total=False): + segmentationConfig: NovaMediaSegmentationConfig + + +class NovaAudioParamsWithSegmentation(NovaAudioParams, total=False): + segmentationConfig: NovaMediaSegmentationConfig + + +class NovaSingleEmbeddingParams(TypedDict, total=False): + embeddingPurpose: Required[NOVA_EMBEDDING_PURPOSES] + embeddingDimension: NOVA_EMBEDDING_DIMENSIONS + text: NovaTextParams + image: NovaImageParams + video: NovaVideoParams + audio: NovaAudioParams + + +class NovaSegmentedEmbeddingParams(TypedDict, total=False): + embeddingPurpose: Required[NOVA_EMBEDDING_PURPOSES] + embeddingDimension: NOVA_EMBEDDING_DIMENSIONS + text: NovaTextParamsWithSegmentation + image: NovaImageParams + video: NovaVideoParamsWithSegmentation + audio: NovaAudioParamsWithSegmentation + + +class NovaEmbeddingRequest(TypedDict, total=False): + schemaVersion: str # "nova-multimodal-embed-v1" + taskType: Literal["SINGLE_EMBEDDING", "SEGMENTED_EMBEDDING"] + singleEmbeddingParams: NovaSingleEmbeddingParams + segmentedEmbeddingParams: NovaSegmentedEmbeddingParams + + +class NovaEmbeddingItem(TypedDict, total=False): + embeddingType: NOVA_EMBEDDING_TYPES + embedding: Required[List[float]] + truncatedCharLength: int # Only for text + + +class NovaEmbeddingResponse(TypedDict): + embeddings: List[NovaEmbeddingItem] + + +class NovaS3OutputDataConfig(TypedDict): + s3Uri: str + + +class NovaOutputDataConfig(TypedDict): + s3OutputDataConfig: NovaS3OutputDataConfig + + +class NovaAsyncInvokeRequest(TypedDict): + modelId: str + modelInput: NovaEmbeddingRequest + outputDataConfig: NovaOutputDataConfig + + AmazonEmbeddingRequest = Union[ AmazonTitanMultimodalEmbeddingRequest, AmazonTitanV2EmbeddingRequest, diff --git a/litellm/utils.py b/litellm/utils.py index b3b114071b7..f551cc96736 100644 --- a/litellm/utils.py +++ b/litellm/utils.py @@ -2827,6 +2827,8 @@ def get_optional_params_embeddings( # noqa: PLR0915 object = litellm.BedrockCohereEmbeddingConfig() elif "twelvelabs" in model or "marengo" in model: object = litellm.TwelveLabsMarengoEmbeddingConfig() + elif "nova" in model.lower(): + object = litellm.AmazonNovaEmbeddingConfig() else: # unmapped model supported_params = [] _check_valid_arg(supported_params=supported_params) diff --git a/tests/llm_translation/test_bedrock_nova_embedding.py b/tests/llm_translation/test_bedrock_nova_embedding.py new file mode 100644 index 00000000000..8cc77b3c3cb --- /dev/null +++ b/tests/llm_translation/test_bedrock_nova_embedding.py @@ -0,0 +1,469 @@ +""" +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 +import os +import sys +from unittest.mock import MagicMock, Mock, patch + +import pytest + +sys.path.insert( + 0, os.path.abspath("../..") +) # Adds the parent directory to the system path + +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 + + +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_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!") +