test - bedrock v2 supports dimension

This commit is contained in:
Ishaan Jaff 2024-05-03 16:23:37 -07:00
parent 4a39b95acc
commit 401bf8d67e
2 changed files with 37 additions and 2 deletions

View file

@ -4744,9 +4744,14 @@ def get_optional_params_embeddings(
message=f"Setting user/encoding format is not supported by {custom_llm_provider}. To drop it from the call, set `litellm.drop_params = True`.",
)
if custom_llm_provider == "bedrock":
if "amazon.titan-embed-text-v2" in model:
# embed-text-v2 supports the dimension param
# if dimensions is in non_default_params -> pass it for model=bedrock/amazon.titan-embed-text-v2
if (
"dimensions" in non_default_params.keys()
and "amazon.titan-embed-text-v2" in model
):
kwargs["dimensions"] = non_default_params["dimensions"]
non_default_params.pop("dimensions", None)
if len(non_default_params.keys()) > 0:
if litellm.drop_params is True: # drop the unsupported non-default values
keys = list(non_default_params.keys())
@ -4758,6 +4763,7 @@ def get_optional_params_embeddings(
status_code=500,
message=f"Setting user/encoding format is not supported by {custom_llm_provider}. To drop it from the call, set `litellm.drop_params = True`.",
)
return {**non_default_params, **kwargs}
if (
custom_llm_provider != "openai"