Merge pull request #3432 from BerriAI/litellm_add_bedrock_titan_embed

Feat - add bedrock titan embed-v2

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)` |
 

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,

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}

litellm/utils.py

@@ -4754,6 +4754,27 @@ 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 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())
+                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`.",
+            )
+        return {**non_default_params, **kwargs}
 
     if (
         custom_llm_provider != "openai"

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,