fix(bedrock): report the embeddings endpoint from the batch job's model

Bedrock's GetModelInvocationJob has no endpoint field, so a retrieve that
has no saved batch to read still answered /v1/chat/completions for an
embeddings job. The answer now follows the mode of the job's modelId.

litellm/llms/bedrock/batches/handler.py

@@ -327,7 +327,10 @@ class BedrockBatchesHandler:
         # provider uses, so model_list / env / role-assumption configs are
         # honored. We instantiate BedrockBatchesConfig (which extends
         # BaseAWSLLM) lazily to avoid a circular import at module load.
-        from litellm.llms.bedrock.batches.transformation import BedrockBatchesConfig
+        from litellm.llms.bedrock.batches.transformation import (
+            BedrockBatchesConfig,
+            openai_endpoint_for_bedrock_job_model,
+        )
 
         creds: Final = BedrockBatchesConfig().resolve_credentials(AwsAuthParams.model_validate(kwargs), region)
         control_plane_endpoint: Final = configured_bedrock_control_plane_endpoint(
@@ -419,7 +422,7 @@ class BedrockBatchesHandler:
             request_counts=_record_counts_from_response(response),
             metadata=openai_batch_metadata,
             completion_window="24h",
-            endpoint="/v1/chat/completions",
+            endpoint=openai_endpoint_for_bedrock_job_model(response.get("modelId")),
             input_file_id=input_uri,
             output_file_id=output_file_uri if openai_status == "completed" else None,
         )

litellm/llms/bedrock/batches/transformation.py

@@ -7,6 +7,7 @@ from typing import TYPE_CHECKING, Final, Literal, cast
 from httpx import Headers, Response
 from pydantic import TypeAdapter, ValidationError
 
+from litellm.exceptions import ModelNotMappedError
 from litellm.litellm_core_utils.aws_partition import is_bedrock_arn
 from litellm.litellm_core_utils.cloud_storage_security import (
     BEDROCK_MANAGED_S3_BATCH_PREFIX,
@@ -33,6 +34,7 @@ from ..base_aws_llm import BaseAWSLLM
 from ..common_utils import (
     CommonBatchFilesUtils,
     bedrock_control_plane_base_url,
+    get_cached_model_info,
     merge_bedrock_aws_request_params,
     resolve_s3_bucket_owner,
     resolve_s3_encryption_key_id,
@@ -83,6 +85,18 @@ def _validate_bedrock_tags(raw_tags: object) -> list[BedrockTag]:
         ) from e
 
 
+def openai_endpoint_for_bedrock_job_model(
+    model_id: str | None,
+) -> Literal["/v1/chat/completions", "/v1/embeddings"]:
+    if not model_id:
+        return "/v1/chat/completions"
+    try:
+        mode: Final = get_cached_model_info()(model=model_id, custom_llm_provider="bedrock").get("mode")
+    except ModelNotMappedError:
+        return "/v1/chat/completions"
+    return "/v1/embeddings" if mode == "embedding" else "/v1/chat/completions"
+
+
 def titan_embedding_usage_from_batch_output(model_output: Mapping[str, object]) -> Usage | None:
     """Titan embedding batch lines report usage as a top-level inputTextTokenCount, not a usage block."""
     if "embedding" not in model_output and "embeddingsByType" not in model_output:
@@ -641,7 +655,7 @@ class BedrockBatchesConfig(BaseAWSLLM, BaseBatchesConfig):
         return LiteLLMBatch(
             id=job_arn,
             object="batch",
-            endpoint="/v1/chat/completions",
+            endpoint=openai_endpoint_for_bedrock_job_model(response_data.get("modelId")),
             errors=errors,
             input_file_id=input_file_id,
             completion_window="24h",

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

@@ -165,6 +165,24 @@ def test_handle_model_invocation_job_status_completed(patched_boto3):
     assert batch.metadata["output_s3_uri"] == OUTPUT_PREFIX
 
 
+@pytest.mark.parametrize(
+    "job_model_id,expected_endpoint",
+    [
+        ("amazon.titan-embed-text-v2:0", "/v1/embeddings"),
+        ("anthropic.claude-3-5-sonnet-20240620-v1:0", "/v1/chat/completions"),
+    ],
+)
+def test_job_status_endpoint_follows_the_job_models_mode(
+    patched_boto3: tuple[MagicMock, MagicMock], job_model_id: str, expected_endpoint: str
+) -> None:
+    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)
+
+    assert batch.endpoint == expected_endpoint
+
+
 @pytest.mark.parametrize("success_count,error_count", [(100, 0), (86, 14)])
 def test_completed_job_maps_provider_record_counts(patched_boto3, success_count, error_count):
     fake_client, _ = patched_boto3

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

@@ -918,6 +918,28 @@ def test_retrieve_response_no_errors_when_no_message(config):
     assert out.errors is None
 
 
+@pytest.mark.parametrize(
+    "job_model_id,expected_endpoint",
+    [
+        ("amazon.titan-embed-text-v2:0", "/v1/embeddings"),
+        ("arn:aws:bedrock:us-west-2::foundation-model/amazon.titan-embed-text-v2:0", "/v1/embeddings"),
+        ("anthropic.claude-3-5-sonnet-20240620-v1:0", "/v1/chat/completions"),
+        ("a-model-the-cost-map-does-not-list", "/v1/chat/completions"),
+        (None, "/v1/chat/completions"),
+    ],
+)
+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(
+        model=None,
+        raw_response=_raw({"jobArn": ARN, "status": "InProgress", "modelId": job_model_id}),
+        logging_obj=MagicMock(),
+        litellm_params={},
+    )
+    assert out.endpoint == expected_endpoint
+
+
 def test_retrieve_response_enriches_metadata(config):
     body = {
         "jobArn": ARN,

tests/unit/proxy/batches_endpoints/test_endpoints.py

@@ -36,6 +36,7 @@ from contextlib import ExitStack
 from dataclasses import dataclass
 from datetime import datetime, timedelta, timezone
 from pathlib import Path
+from collections.abc import Callable
 from typing import Any, Dict, Final, Optional
 from unittest.mock import DEFAULT, AsyncMock, MagicMock, patch
 
@@ -1909,10 +1910,10 @@ async def test_retrieve__db_non_terminal_state_syncs_with_provider(retrieve_harn
     assert retrieve_harness.update_batch_in_db.call_count == 1
 
 
-def _hook_rewriting_the_batch_id_to(batch_id):
-    def hook(**kwargs):
-        kwargs["response"].id = batch_id
-        return kwargs["response"]
+def _hook_rewriting_the_batch_id_to(batch_id: str) -> Callable[..., LiteLLMBatch]:
+    def hook(*, response: LiteLLMBatch, **_: object) -> LiteLLMBatch:
+        response.id = batch_id
+        return response
 
     return hook