diff --git a/litellm/llms/bedrock/embed/amazon_titan_v2_transformation.py b/litellm/llms/bedrock/embed/amazon_titan_v2_transformation.py index 8056e9e9b2c..ff748b58e8e 100644 --- a/litellm/llms/bedrock/embed/amazon_titan_v2_transformation.py +++ b/litellm/llms/bedrock/embed/amazon_titan_v2_transformation.py @@ -10,7 +10,7 @@ Docs - https://docs.aws.amazon.com/bedrock/latest/userguide/model-parameters-tit """ import types -from typing import List, Optional +from typing import List, Optional, Union from litellm.types.llms.bedrock import ( AmazonTitanV2EmbeddingRequest, @@ -30,9 +30,7 @@ class AmazonTitanV2Config: normalize: Optional[bool] = None dimensions: Optional[int] = None - def __init__( - self, normalize: Optional[bool] = None, dimensions: Optional[int] = None - ) -> None: + def __init__(self, normalize: Optional[bool] = None, dimensions: Optional[int] = None) -> None: locals_ = locals().copy() for key, value in locals_.items(): if key != "self" and value is not None: @@ -57,32 +55,56 @@ class AmazonTitanV2Config: } def get_supported_openai_params(self) -> List[str]: - return ["dimensions"] + return ["dimensions", "encoding_format"] - def map_openai_params( - self, non_default_params: dict, optional_params: dict - ) -> dict: + def map_openai_params(self, non_default_params: dict, optional_params: dict) -> dict: for k, v in non_default_params.items(): if k == "dimensions": optional_params["dimensions"] = v + elif k == "encoding_format": + # Map OpenAI encoding_format to AWS embeddingTypes + if v == "float": + optional_params["embeddingTypes"] = ["float"] + elif v == "base64": + # base64 maps to binary format in AWS + optional_params["embeddingTypes"] = ["binary"] + else: + # For any other encoding format, default to float + optional_params["embeddingTypes"] = ["float"] return optional_params - def _transform_request( - self, input: str, inference_params: dict - ) -> AmazonTitanV2EmbeddingRequest: + def _transform_request(self, input: str, inference_params: dict) -> AmazonTitanV2EmbeddingRequest: return AmazonTitanV2EmbeddingRequest(inputText=input, **inference_params) # type: ignore - def _transform_response( - self, response_list: List[dict], model: str - ) -> EmbeddingResponse: + def _transform_response(self, response_list: List[dict], model: str) -> EmbeddingResponse: total_prompt_tokens = 0 transformed_responses: List[Embedding] = [] for index, response in enumerate(response_list): _parsed_response = AmazonTitanV2EmbeddingResponse(**response) # type: ignore + + # According to AWS docs, embeddingsByType is always present + # If binary was requested (encoding_format="base64"), use binary data + # Otherwise, use float data from embeddingsByType or fallback to embedding field + embedding_data: Union[List[float], List[int]] + + if ("embeddingsByType" in _parsed_response and + "binary" in _parsed_response["embeddingsByType"]): + # Use binary data if available (for encoding_format="base64") + embedding_data = _parsed_response["embeddingsByType"]["binary"] + elif ("embeddingsByType" in _parsed_response and + "float" in _parsed_response["embeddingsByType"]): + # Use float data from embeddingsByType + embedding_data = _parsed_response["embeddingsByType"]["float"] + elif "embedding" in _parsed_response: + # Fallback to legacy embedding field + embedding_data = _parsed_response["embedding"] + else: + raise ValueError(f"No embedding data found in response: {response}") + transformed_responses.append( Embedding( - embedding=_parsed_response["embedding"], + embedding=embedding_data, index=index, object="embedding", )