fix(bedrock): send max_tokens and priority in Cohere embedding requests (#51435)

* fix(bedrock): send max_tokens and priority in Cohere embedding requests

Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>

* test(sagemaker): mark Cohere transform result Final

Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>

* ci: rerun after unrelated checkout timeout and langsmith flake

Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>

---------

Co-authored-by: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>

litellm/types/llms/bedrock.py

@@ -505,6 +505,8 @@ class CohereEmbeddingRequest(TypedDict, total=False):
     truncate: Literal["NONE", "START", "END"]
     embedding_types: Literal["float", "int8", "uint8", "binary", "ubinary"]
     output_dimension: int
+    max_tokens: ReadOnly[int]
+    priority: ReadOnly[int]
 
 
 class CohereEmbeddingRequestWithModel(CohereEmbeddingRequest):

tests/unit/llms/bedrock/embed/test_bedrock_embedding.py

@@ -190,6 +190,34 @@ def test_bedrock_cohere_embedding_returns_the_scripted_vector(respx_mock: respx.
     assert response.data == [{"object": "embedding", "index": 0, "embedding": [0.3, 0.2, 0.1]}]
 
 
+def test_bedrock_cohere_embedding_sends_max_tokens_and_priority(respx_mock: respx.MockRouter) -> None:
+    route = respx_mock.post("https://bedrock-runtime.us-east-1.amazonaws.com/model/cohere.embed-v4:0/invoke").mock(
+        return_value=httpx.Response(
+            200,
+            json={"embeddings": {"float": [[0.3, 0.2, 0.1]]}, "response_type": "embeddings_by_type"},
+        )
+    )
+
+    litellm.embedding(
+        model="bedrock/cohere.embed-v4:0",
+        input=["hello"],
+        api_key="test-bearer-token",
+        aws_region_name="us-east-1",
+        aws_bedrock_runtime_endpoint="https://bedrock-runtime.us-east-1.amazonaws.com",
+        max_tokens=16,
+        priority=5,
+        truncate="END",
+    )
+
+    assert json.loads(route.calls.last.request.read()) == {
+        "input_type": "search_document",
+        "texts": ["hello"],
+        "max_tokens": 16,
+        "priority": 5,
+        "truncate": "END",
+    }
+
+
 GOVCLOUD_WEST_RUNTIME: Final = "https://bedrock-runtime.us-gov-west-1.amazonaws.com"
 
 

tests/unit/llms/sagemaker/test_sagemaker_embedding_voyage.py

@@ -7,6 +7,7 @@ transformation, and model type detection.
 """
 
 import json
+from typing import Final
 from unittest.mock import MagicMock, patch
 
 import httpx
@@ -105,6 +106,21 @@ class TestSagemakerCohereEmbeddingConfig:
         )
         assert result["texts"] == ["hello"]
 
+    def test_transform_request_keeps_max_tokens_and_priority(self):
+        result: Final = self.config.transform_embedding_request(
+            model="cohere-embed-v4",
+            input=["hello"],
+            optional_params={"max_tokens": 16, "priority": 5, "truncate": "END"},
+            headers={},
+        )
+        assert result == {
+            "texts": ["hello"],
+            "input_type": "search_document",
+            "max_tokens": 16,
+            "priority": 5,
+            "truncate": "END",
+        }
+
     def test_map_openai_params_dimensions_to_output_dimension(self):
         params = self.config.map_openai_params(
             non_default_params={"dimensions": 512, "encoding_format": "float"},