From fbe412a3a6a61640793a24d4c919210175e2179c Mon Sep 17 00:00:00 2001 From: Ishaan Jaff Date: Fri, 3 May 2024 15:58:05 -0700 Subject: [PATCH 1/4] feat - add amazon.titan-embed-text-v2 --- litellm/model_prices_and_context_window_backup.json | 9 +++++++++ model_prices_and_context_window.json | 9 +++++++++ 2 files changed, 18 insertions(+) diff --git a/litellm/model_prices_and_context_window_backup.json b/litellm/model_prices_and_context_window_backup.json index 7fcd425bb..c7d5aae2d 100644 --- a/litellm/model_prices_and_context_window_backup.json +++ b/litellm/model_prices_and_context_window_backup.json @@ -1832,6 +1832,15 @@ "litellm_provider": "bedrock", "mode": "embedding" }, + "amazon.titan-embed-text-v2:0": { + "max_tokens": 8192, + "max_input_tokens": 8192, + "output_vector_size": 1024, + "input_cost_per_token": 0.0000002, + "output_cost_per_token": 0.0, + "litellm_provider": "bedrock", + "mode": "embedding" + }, "mistral.mistral-7b-instruct-v0:2": { "max_tokens": 8191, "max_input_tokens": 32000, diff --git a/model_prices_and_context_window.json b/model_prices_and_context_window.json index 7fcd425bb..c7d5aae2d 100644 --- a/model_prices_and_context_window.json +++ b/model_prices_and_context_window.json @@ -1832,6 +1832,15 @@ "litellm_provider": "bedrock", "mode": "embedding" }, + "amazon.titan-embed-text-v2:0": { + "max_tokens": 8192, + "max_input_tokens": 8192, + "output_vector_size": 1024, + "input_cost_per_token": 0.0000002, + "output_cost_per_token": 0.0, + "litellm_provider": "bedrock", + "mode": "embedding" + }, "mistral.mistral-7b-instruct-v0:2": { "max_tokens": 8191, "max_input_tokens": 32000, From bf048ecda4976f2676a25de4aec69fe98c6c0b54 Mon Sep 17 00:00:00 2001 From: Ishaan Jaff Date: Fri, 3 May 2024 15:59:02 -0700 Subject: [PATCH 2/4] docs - titan embeddings v2 --- docs/my-website/docs/providers/bedrock.md | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/docs/my-website/docs/providers/bedrock.md b/docs/my-website/docs/providers/bedrock.md index 590ffc423..147c12e65 100644 --- a/docs/my-website/docs/providers/bedrock.md +++ b/docs/my-website/docs/providers/bedrock.md @@ -535,7 +535,8 @@ print(response) | Model Name | Function Call | |----------------------|---------------------------------------------| -| Titan Embeddings - G1 | `embedding(model="bedrock/amazon.titan-embed-text-v1", input=input)` | +| Titan Embeddings V2 | `embedding(model="bedrock/amazon.titan-embed-text-v2:0", input=input)` | +| Titan Embeddings - V1 | `embedding(model="bedrock/amazon.titan-embed-text-v1", input=input)` | | Cohere Embeddings - English | `embedding(model="bedrock/cohere.embed-english-v3", input=input)` | | Cohere Embeddings - Multilingual | `embedding(model="bedrock/cohere.embed-multilingual-v3", input=input)` | From 4a39b95accf25380a037ff241cb1822f412e3e22 Mon Sep 17 00:00:00 2001 From: Ishaan Jaff Date: Fri, 3 May 2024 16:13:27 -0700 Subject: [PATCH 3/4] fix - support dimension for titan embed v2 --- litellm/utils.py | 15 +++++++++++++++ 1 file changed, 15 insertions(+) diff --git a/litellm/utils.py b/litellm/utils.py index 80d26f58b..e124358b6 100644 --- a/litellm/utils.py +++ b/litellm/utils.py @@ -4743,6 +4743,21 @@ 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`.", ) + if custom_llm_provider == "bedrock": + if "amazon.titan-embed-text-v2" in model: + # embed-text-v2 supports the dimension param + 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()) + for k in keys: + non_default_params.pop(k, None) + final_params = {**non_default_params, **kwargs} + return final_params + raise UnsupportedParamsError( + 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`.", + ) if ( custom_llm_provider != "openai" From 401bf8d67ee416a177b0fe6afe1c3939d5a94ed1 Mon Sep 17 00:00:00 2001 From: Ishaan Jaff Date: Fri, 3 May 2024 16:23:37 -0700 Subject: [PATCH 4/4] test - bedrock v2 supports dimension --- .../test_get_optional_params_embeddings.py | 29 +++++++++++++++++++ litellm/utils.py | 10 +++++-- 2 files changed, 37 insertions(+), 2 deletions(-) diff --git a/litellm/tests/test_get_optional_params_embeddings.py b/litellm/tests/test_get_optional_params_embeddings.py index 41396b531..81b177030 100644 --- a/litellm/tests/test_get_optional_params_embeddings.py +++ b/litellm/tests/test_get_optional_params_embeddings.py @@ -40,3 +40,32 @@ def test_vertex_projects(): # test_vertex_projects() + + +def test_bedrock_embed_v2_regular(): + model, custom_llm_provider, _, _ = get_llm_provider( + model="bedrock/amazon.titan-embed-text-v2:0" + ) + optional_params = get_optional_params_embeddings( + model=model, + dimensions=512, + custom_llm_provider=custom_llm_provider, + ) + print(f"received optional_params: {optional_params}") + assert optional_params == {"dimensions": 512} + + +def test_bedrock_embed_v2_with_drop_params(): + litellm.drop_params = True + model, custom_llm_provider, _, _ = get_llm_provider( + model="bedrock/amazon.titan-embed-text-v2:0" + ) + optional_params = get_optional_params_embeddings( + model=model, + dimensions=512, + user="test-litellm-user-5", + encoding_format="base64", + custom_llm_provider=custom_llm_provider, + ) + print(f"received optional_params: {optional_params}") + assert optional_params == {"dimensions": 512} diff --git a/litellm/utils.py b/litellm/utils.py index e124358b6..5070e6498 100644 --- a/litellm/utils.py +++ b/litellm/utils.py @@ -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"