Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

[Fix] Graceful rejection of token input for AWS Embeddings API #1685

Merged
merged 3 commits into from
Jan 30, 2024
Merged

[Fix] Graceful rejection of token input for AWS Embeddings API #1685

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
3 commits
Select commit Hold shift + click to select a range
6488de3
(test) bedrock input validation - exceptions
ishaan-jaff Jan 30, 2024
b9a117f
(feat) Bedrock embedding - raise correct exception BadRequest
ishaan-jaff Jan 30, 2024
f941c57
(fix) use isinstance to check types
ishaan-jaff Jan 30, 2024
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
18 changes: 15 additions & 3 deletions litellm/llms/bedrock.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -702,6 +702,11 @@ def _embedding_func_single(
encoding=None,
logging_obj=None,
):
if isinstance(input, str) is False:
raise BedrockError(
message="Bedrock Embedding API input must be type str | List[str]",
status_code=400,
)
# logic for parsing in - calling - parsing out model embedding calls
## FORMAT EMBEDDING INPUT ##
provider = model.split(".")[0]
Expand Down Expand Up @@ -795,7 +800,8 @@ def embedding(
aws_role_name=aws_role_name,
aws_session_name=aws_session_name,
)
if type(input) == str:
if isinstance(input, str):
## Embedding Call
embeddings = [
_embedding_func_single(
model,
Expand All @@ -805,8 +811,8 @@ def embedding(
logging_obj=logging_obj,
)
]
else:
## Embedding Call
elif isinstance(input, list):
## Embedding Call - assuming this is a List[str]
embeddings = [
_embedding_func_single(
model,
Expand All @@ -817,6 +823,12 @@ def embedding(
)
for i in input
] # [TODO]: make these parallel calls
else:
# enters this branch if input = int, ex. input=2
raise BedrockError(
message="Bedrock Embedding API input must be type str | List[str]",
status_code=400,
)

## Populate OpenAI compliant dictionary
embedding_response = []
Expand Down
19 changes: 19 additions & 0 deletions litellm/tests/test_embedding.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -302,6 +302,25 @@ def test_bedrock_embedding_cohere():
# test_bedrock_embedding_cohere()


def test_demo_tokens_as_input_to_embeddings_fails_for_titan():
litellm.set_verbose = True

with pytest.raises(
litellm.BadRequestError,
match="BedrockException - Bedrock Embedding API input must be type str | List[str]",
):
litellm.embedding(model="amazon.titan-embed-text-v1", input=[[1]])

with pytest.raises(
litellm.BadRequestError,
match="BedrockException - Bedrock Embedding API input must be type str | List[str]",
):
litellm.embedding(
model="amazon.titan-embed-text-v1",
input=[1],
)

Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

added your unit test to our ci/cd @rmann-nflx

Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Awesome!


# comment out hf tests - since hf endpoints are unstable
def test_hf_embedding():
try:
Expand Down