diff --git a/litellm/llms/bedrock/batches/transformation.py b/litellm/llms/bedrock/batches/transformation.py index 17bb7012900..5bedc3ab7c5 100644 --- a/litellm/llms/bedrock/batches/transformation.py +++ b/litellm/llms/bedrock/batches/transformation.py @@ -91,10 +91,13 @@ def openai_endpoint_for_bedrock_job_model( if not model_id: return "/v1/chat/completions" try: - mode: Final = get_cached_model_info()(model=model_id, custom_llm_provider="bedrock").get("mode") + model_info: Final = get_cached_model_info()(model=model_id, custom_llm_provider="bedrock") except ModelNotMappedError: return "/v1/chat/completions" - return "/v1/embeddings" if mode == "embedding" else "/v1/chat/completions" + is_embedding_model: Final = ( + model_info.get("mode") == "embedding" or model_info.get("output_vector_size") is not None + ) + return "/v1/embeddings" if is_embedding_model else "/v1/chat/completions" def titan_embedding_usage_from_batch_output(model_output: Mapping[str, object]) -> Usage | None: diff --git a/tests/unit/llms/bedrock/batches/test_transformation.py b/tests/unit/llms/bedrock/batches/test_transformation.py index 21157ff72f0..e774082668b 100644 --- a/tests/unit/llms/bedrock/batches/test_transformation.py +++ b/tests/unit/llms/bedrock/batches/test_transformation.py @@ -24,6 +24,7 @@ from botocore.auth import SigV4Auth from botocore.awsrequest import AWSRequest from botocore.credentials import Credentials +import litellm from litellm.llms.bedrock.batches.transformation import BedrockBatchesConfig from litellm.types.utils import LlmProviders @@ -940,6 +941,20 @@ def test_retrieve_response_endpoint_follows_the_job_models_mode( assert out.endpoint == expected_endpoint +def test_retrieve_response_endpoint_survives_a_batch_deployment_overriding_the_models_mode( + config: BedrockBatchesConfig, monkeypatch: pytest.MonkeyPatch +) -> 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( + model=None, + raw_response=_raw({"jobArn": ARN, "status": "InProgress", "modelId": job_model_id}), + logging_obj=MagicMock(), + litellm_params={}, + ) + assert out.endpoint == "/v1/embeddings" + + def test_retrieve_response_enriches_metadata(config): body = { "jobArn": ARN,