From 7f51cb1ce79f2cce39a684713069ce5600428ade Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Mon, 9 Nov 2026 21:05:03 -0800 Subject: [PATCH] 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 | 7 ++++-- .../llms/bedrock/batches/transformation.py | 16 +++++++++++++- .../unit/llms/bedrock/batches/test_handler.py | 18 +++++++++++++++ .../bedrock/batches/test_transformation.py | 22 +++++++++++++++++++ .../proxy/batches_endpoints/test_endpoints.py | 9 ++++---- 5 files changed, 65 insertions(+), 7 deletions(-) diff --git a/litellm/llms/bedrock/batches/handler.py b/litellm/llms/bedrock/batches/handler.py index f37abacc883..df9dca50d3d 100644 --- a/litellm/llms/bedrock/batches/handler.py +++ b/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, ) diff --git a/litellm/llms/bedrock/batches/transformation.py b/litellm/llms/bedrock/batches/transformation.py index 84f40bec9ca..17bb7012900 100644 --- a/litellm/llms/bedrock/batches/transformation.py +++ b/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", diff --git a/tests/unit/llms/bedrock/batches/test_handler.py b/tests/unit/llms/bedrock/batches/test_handler.py index 602504f0c49..ce507e99f4b 100644 --- a/tests/unit/llms/bedrock/batches/test_handler.py +++ b/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 diff --git a/tests/unit/llms/bedrock/batches/test_transformation.py b/tests/unit/llms/bedrock/batches/test_transformation.py index 5360ac00018..21157ff72f0 100644 --- a/tests/unit/llms/bedrock/batches/test_transformation.py +++ b/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, diff --git a/tests/unit/proxy/batches_endpoints/test_endpoints.py b/tests/unit/proxy/batches_endpoints/test_endpoints.py index 00ff9161062..01c5ef8c3bd 100644 --- a/tests/unit/proxy/batches_endpoints/test_endpoints.py +++ b/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