fix(bedrock): correct a saved chat endpoint on embeddings batches and cover Cohere Embed v3

A Bedrock answer of /v1/embeddings now replaces a saved /v1/chat/completions, so a batch read before the upgrade or by an old pod during a rolling upgrade reports the right endpoint on its next provider sync. The two Cohere Embed v3 cost-map rows gain output_vector_size so they stay embeddings when a batch deployment overrides the model's mode

litellm/model_prices_and_context_window_backup.json

@@ -17536,6 +17536,7 @@
         "max_input_tokens": 512,
         "mode": "embedding",
         "output_cost_per_token": 0.0,
+        "output_vector_size": 1024,
         "source": "https://aws.amazon.com/bedrock/pricing/",
         "supports_embedding_image_input": true
     },
@@ -17545,6 +17546,7 @@
         "max_input_tokens": 512,
         "mode": "embedding",
         "output_cost_per_token": 0.0,
+        "output_vector_size": 1024,
         "source": "https://aws.amazon.com/bedrock/pricing/",
         "supports_embedding_image_input": true
     },

litellm/proxy/openai_files_endpoints/common_utils.py

@@ -1608,7 +1608,10 @@ def batch_cost_poller_is_active() -> bool:
 
 
 def created_batch_endpoint(provider_batch: "LiteLLMBatch", stored_batch: "LiteLLMBatch | None") -> str:
-    return stored_batch.endpoint if stored_batch is not None and stored_batch.endpoint else provider_batch.endpoint
+    if stored_batch is None or not stored_batch.endpoint:
+        return provider_batch.endpoint
+    provider_named_the_endpoint: Final = provider_batch.endpoint not in ("", "/v1/chat/completions")
+    return provider_batch.endpoint if provider_named_the_endpoint else stored_batch.endpoint
 
 
 def _completed_batch_safe_to_retire(response: "LiteLLMBatch") -> bool:

model_prices_and_context_window.json

@@ -17536,6 +17536,7 @@
         "max_input_tokens": 512,
         "mode": "embedding",
         "output_cost_per_token": 0.0,
+        "output_vector_size": 1024,
         "source": "https://aws.amazon.com/bedrock/pricing/",
         "supports_embedding_image_input": true
     },
@@ -17545,6 +17546,7 @@
         "max_input_tokens": 512,
         "mode": "embedding",
         "output_cost_per_token": 0.0,
+        "output_vector_size": 1024,
         "source": "https://aws.amazon.com/bedrock/pricing/",
         "supports_embedding_image_input": true
     },

tests/unit/llms/bedrock/batches/test_handler.py

@@ -178,7 +178,7 @@ def test_job_status_endpoint_follows_the_job_models_mode(
     fake_client, _ = patched_boto3
     fake_client.get_model_invocation_job.return_value = {**_fake_boto3_response(), "modelId": job_model_id}
 
-    batch = BedrockBatchesHandler.handle_model_invocation_job_status(batch_id=JOB_ARN)
+    batch: Final = BedrockBatchesHandler.handle_model_invocation_job_status(batch_id=JOB_ARN)
 
     assert batch.endpoint == expected_endpoint
 

tests/unit/llms/bedrock/batches/test_transformation.py

@@ -14,6 +14,7 @@ URL/ARN handling, and the error class. AWS auth/sigv4 is the only external seam
 we mock; everything else runs for real.
 """
 
+from collections.abc import Mapping
 from datetime import datetime, timezone
 from typing import Final
 from unittest.mock import MagicMock, patch
@@ -932,7 +933,7 @@ def test_retrieve_response_no_errors_when_no_message(config):
 def test_retrieve_response_endpoint_follows_the_job_models_mode(
     config: BedrockBatchesConfig, job_model_id: str | None, expected_endpoint: str
 ) -> None:
-    out = config.transform_retrieve_batch_response(
+    out: Final = config.transform_retrieve_batch_response(
         model=None,
         raw_response=_raw({"jobArn": ARN, "status": "InProgress", "modelId": job_model_id}),
         logging_obj=MagicMock(),
@@ -941,12 +942,19 @@ def test_retrieve_response_endpoint_follows_the_job_models_mode(
     assert out.endpoint == expected_endpoint
 
 
+def _is_bedrock_embedding_model(cost_map_entry: Mapping[str, object]) -> bool:
+    return cost_map_entry.get("litellm_provider") == "bedrock" and cost_map_entry.get("mode") == "embedding"
+
+
+@pytest.mark.parametrize(
+    "job_model_id",
+    sorted(model_id for model_id, entry in litellm.model_cost.items() if _is_bedrock_embedding_model(entry)),
+)
 def test_retrieve_response_endpoint_survives_a_batch_deployment_overriding_the_models_mode(
-    config: BedrockBatchesConfig, monkeypatch: pytest.MonkeyPatch
+    config: BedrockBatchesConfig, monkeypatch: pytest.MonkeyPatch, job_model_id: str
 ) -> None:
-    job_model_id = "amazon.titan-embed-text-v2:0"
     monkeypatch.setitem(litellm.model_cost, job_model_id, {**litellm.model_cost[job_model_id], "mode": "batch"})
-    out = config.transform_retrieve_batch_response(
+    out: Final = config.transform_retrieve_batch_response(
         model=None,
         raw_response=_raw({"jobArn": ARN, "status": "InProgress", "modelId": job_model_id}),
         logging_obj=MagicMock(),

tests/unit/proxy/batches_endpoints/test_endpoints.py

@@ -1952,6 +1952,24 @@ async def test_retrieve__a_saved_batch_without_an_endpoint_keeps_the_providers(r
     assert retrieve_harness.update_batch_in_db.call_args.kwargs["response"].endpoint == "/v1/embeddings"
 
 
+@pytest.mark.asyncio
+async def test_retrieve__an_embeddings_endpoint_from_the_provider_corrects_a_saved_chat_endpoint(retrieve_harness):
+    stored = make_batch(id="batch-unified-blob", status="validating").model_copy(
+        update={"endpoint": "/v1/chat/completions"}
+    )
+    provider_batch = make_batch(id="batch-raw-xyz", status="in_progress").model_copy(
+        update={"endpoint": "/v1/embeddings"}
+    )
+    retrieve_harness.get_batch_from_db.return_value = (MagicMock(status="validating"), stored)
+    retrieve_harness.router.aretrieve_batch.return_value = provider_batch
+
+    with patch.object(endpoints, "is_base64_encoded_unified_file_id", return_value=UNIFIED_BATCH_ID):
+        resp = await call_retrieve(retrieve_harness, "batch-unified-blob")
+
+    assert resp.endpoint == "/v1/embeddings"
+    assert retrieve_harness.update_batch_in_db.call_args.kwargs["response"].endpoint == "/v1/embeddings"
+
+
 @pytest.mark.asyncio
 @pytest.mark.parametrize("status", ["validating", "in_progress", "finalizing", "cancelling"])
 async def test_retrieve__executed_batch_served_from_db_in_every_status(retrieve_harness, status):