scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-09-06 08:16:43 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 78

Merge pull request #14687 from timelfrink/fix/issue-14685-bedrock-titan-v2-encoding-format

Browse source 
Fix: Bedrock Titan V2 encoding_format parameter support
This commit is contained in:
Krish Dholakia 2025-09-18 17:54:42 -07:00 • committed by GitHub
commit 9238e1d3bf
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
4 changed files with 187 additions and 23 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

53
docs/my-website/docs/providers/bedrock.md
View file

@ -1822,6 +1822,59 @@ Here's an example of using a bedrock model with LiteLLM. For a complete list, re
| Mixtral 8x7B Instruct | `completion(model='bedrock/mistral.mixtral-8x7b-instruct-v0:1', messages=messages)` | `os.environ['AWS_ACCESS_KEY_ID']`, `os.environ['AWS_SECRET_ACCESS_KEY']`, `os.environ['AWS_REGION_NAME']` |
## Bedrock Embedding
### API keys
This can be set as env variables or passed as **params to litellm.embedding()**
```python
import os
os.environ["AWS_ACCESS_KEY_ID"] = "" # Access key
os.environ["AWS_SECRET_ACCESS_KEY"] = "" # Secret access key
os.environ["AWS_REGION_NAME"] = "" # us-east-1, us-east-2, us-west-1, us-west-2
```
### Usage
```python
from litellm import embedding
response = embedding(
model="bedrock/amazon.titan-embed-text-v1",
input=["good morning from litellm"],
)
print(response)
```
#### Titan V2 - encoding_format support
```python
from litellm import embedding
# Float format (default)
response = embedding(
model="bedrock/amazon.titan-embed-text-v2:0",
input=["good morning from litellm"],
encoding_format="float" # Returns float array
)
# Binary format
response = embedding(
model="bedrock/amazon.titan-embed-text-v2:0",
input=["good morning from litellm"],
encoding_format="base64" # Returns base64 encoded binary
)
```
## Supported AWS Bedrock Embedding Models
| Model Name | Usage | Supported Additional OpenAI params |
|----------------------|---------------------------------------------|-----|
| Titan Embeddings V2 | `embedding(model="bedrock/amazon.titan-embed-text-v2:0", input=input)` | `dimensions`, `encoding_format` |
| 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) |
| Cohere Embeddings - English | `embedding(model="bedrock/cohere.embed-english-v3", input=input)` | [here](https://github.com/BerriAI/litellm/blob/f5905e100068e7a4d61441d7453d7cf5609c2121/litellm/llms/bedrock/embed/cohere_transformation.py#L18)
| Cohere Embeddings - Multilingual | `embedding(model="bedrock/cohere.embed-multilingual-v3", input=input)` | [here](https://github.com/BerriAI/litellm/blob/f5905e100068e7a4d61441d7453d7cf5609c2121/litellm/llms/bedrock/embed/cohere_transformation.py#L18)
### Advanced - [Drop Unsupported Params](https://docs.litellm.ai/docs/completion/drop_params#openai-proxy-usage)
### Advanced - [Pass model/provider-specific Params](https://docs.litellm.ai/docs/completion/provider_specific_params#proxy-usage)
## Image Generation
Use this for stable diffusion, and amazon nova canvas on bedrock

52
litellm/llms/bedrock/embed/amazon_titan_v2_transformation.py
View file

@ -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",
)

17
litellm/types/llms/bedrock.py
View file

@ -328,15 +328,22 @@ class CohereEmbeddingResponse(TypedDict):
texts: List[str]
class AmazonTitanV2EmbeddingRequest(TypedDict):
inputText: str
class AmazonTitanV2EmbeddingRequest(TypedDict, total=False):
inputText: Required[str]
dimensions: int
normalize: bool
embeddingTypes: List[Literal["float", "binary"]]
class AmazonTitanV2EmbeddingResponse(TypedDict):
embedding: List[float]
inputTextTokenCount: int
class AmazonTitanV2EmbeddingsByType(TypedDict, total=False):
binary: List[int] # Array of integers for binary format
float: List[float] # Array of floats for float format
class AmazonTitanV2EmbeddingResponse(TypedDict, total=False):
embedding: List[float] # Legacy field - array of floats (backward compatibility)
embeddingsByType: AmazonTitanV2EmbeddingsByType # New format per AWS schema
inputTextTokenCount: Required[int] # Always present in AWS response
class AmazonTitanG1EmbeddingRequest(TypedDict):

88
tests/test_litellm/llms/bedrock/embed/test_bedrock_embedding.py
View file

@ -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 = {
@ -146,7 +147,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]}]
@ -159,4 +160,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"]