Merge pull request #36159 from daleselaji-dev/codex/bedrock-alias-36156

fix(bedrock): resolve aliases in batch file records

litellm/llms/bedrock/files/transformation.py

@@ -17,6 +17,7 @@ from typing_extensions import ReadOnly
 
 from litellm._logging import verbose_logger
 from litellm._uuid import uuid
+from litellm.constants import BEDROCK_INVOKE_PROVIDERS_LITERAL
 from litellm.files.utils import FilesAPIUtils
 from litellm.litellm_core_utils.cloud_storage_security import (
     BEDROCK_MANAGED_S3_BATCH_PREFIX,
@@ -68,6 +69,18 @@ def _frozen_mapping(items: Iterable[tuple[str, object]]) -> Mapping[str, object]
     return MappingProxyType(dict(items))
 
 
+def _strip_llm_routing_prefix(model: str) -> str:
+    try:
+        stripped_model, _, _, _ = get_llm_provider(model=model, custom_llm_provider=None)
+    except Exception as e:
+        verbose_logger.exception(
+            "litellm.llms.bedrock.files.transformation.py::_strip_llm_routing_prefix() - Error inferring custom_llm_provider - %s",
+            e,
+        )
+        return model
+    return stripped_model
+
+
 _EmbeddingBatchInput: TypeAlias = (
     str | int | float | Sequence[str] | Sequence[int] | Sequence[Sequence[int]] | Mapping[str, object]
 )
@@ -572,6 +585,7 @@ class BedrockFilesConfig(BaseAWSLLM, BaseFilesConfig):
     def _map_openai_embedding_to_bedrock_params(
         self,
         openai_request_body: _OpenAIBatchRecordBody,
+        model: str,
     ) -> dict[str, object]:
         """
         Transform an OpenAI /v1/embeddings request body into the
@@ -591,8 +605,7 @@ class BedrockFilesConfig(BaseAWSLLM, BaseFilesConfig):
             AmazonTitanV2Config,
         )
 
-        _model: Final = openai_request_body.get("model", "")
-        if not self._is_titan_v2_embed_model(_model):
+        if not self._is_titan_v2_embed_model(model):
             # Refuse early instead of silently shaping the body for the wrong
             # provider. The synchronous /v1/embeddings path supports more
             # models, but each has a different InvokeModel schema; mapping
@@ -600,11 +613,11 @@ class BedrockFilesConfig(BaseAWSLLM, BaseFilesConfig):
             raise NotImplementedError(
                 "Bedrock batch embedding currently supports only Amazon "
                 "Titan Text Embeddings V2 (model id contains "
-                f"'titan-embed-text-v2'). Got model={_model!r}. Track other "
+                f"'titan-embed-text-v2'). Got model={model!r}. Track other "
                 "embedding models in https://github.com/BerriAI/litellm/issues."
             )
 
-        input_text: Final = self._coerce_embedding_input_to_string(openai_request_body.get("input"), model=_model)
+        input_text: Final = self._coerce_embedding_input_to_string(openai_request_body.get("input"), model=model)
 
         # Map OpenAI-style params (dimensions, encoding_format) onto the
         # Titan v2 schema (dimensions, embeddingTypes) via the embed config
@@ -699,6 +712,7 @@ class BedrockFilesConfig(BaseAWSLLM, BaseFilesConfig):
     def _map_openai_to_bedrock_params(
         self,
         openai_request_body: Mapping[str, Any],
+        model: str,
         provider: str | None = None,
     ) -> dict[str, object]:
         """
@@ -711,7 +725,6 @@ class BedrockFilesConfig(BaseAWSLLM, BaseFilesConfig):
         """
         from litellm.types.utils import LlmProviders
 
-        _model: Final[str] = openai_request_body.get("model", "")
         messages: Final = openai_request_body.get("messages", [])
         optional_params: Final = {k: v for k, v in openai_request_body.items() if k not in ["model", "messages"]}
 
@@ -725,11 +738,11 @@ class BedrockFilesConfig(BaseAWSLLM, BaseFilesConfig):
             mapped_params = config.map_openai_params(
                 non_default_params={},
                 optional_params=optional_params,
-                model=_model,
+                model=model,
                 drop_params=False,
             )
             return config.transform_request(
-                model=_model,
+                model=model,
                 messages=messages,
                 optional_params=mapped_params,
                 litellm_params={},
@@ -748,11 +761,11 @@ class BedrockFilesConfig(BaseAWSLLM, BaseFilesConfig):
             mapped_params = converse_config.map_openai_params(
                 non_default_params=optional_params,
                 optional_params={},
-                model=_model,
+                model=model,
                 drop_params=False,
             )
             return converse_config.transform_request(
-                model=_model,
+                model=model,
                 messages=messages,
                 optional_params=mapped_params,
                 litellm_params={},
@@ -766,8 +779,21 @@ class BedrockFilesConfig(BaseAWSLLM, BaseFilesConfig):
             **optional_params,
         }
 
+    def _resolve_batch_record_model_and_provider(
+        self,
+        record_model: str,
+        target_model: str,
+    ) -> tuple[str, BEDROCK_INVOKE_PROVIDERS_LITERAL | None]:
+        record_provider: Final = self.get_bedrock_invoke_provider(_strip_llm_routing_prefix(record_model))
+        if record_provider is not None or not target_model:
+            return record_model, record_provider
+        target_provider: Final = self.get_bedrock_invoke_provider(_strip_llm_routing_prefix(target_model))
+        if target_provider is None:
+            return record_model, record_provider
+        return target_model, target_provider
+
     def _transform_openai_jsonl_content_to_bedrock_jsonl_content(
-        self, openai_jsonl_content: Sequence[_OpenAIBatchRecord]
+        self, openai_jsonl_content: Sequence[_OpenAIBatchRecord], target_model: str = ""
     ) -> list[_BedrockBatchRecord]:
         """
         Transforms OpenAI JSONL content to Bedrock batch format
@@ -789,25 +815,17 @@ class BedrockFilesConfig(BaseAWSLLM, BaseFilesConfig):
         }
         """
 
+        import litellm
+
         bedrock_jsonl_content: Final = []
         for idx, _openai_jsonl_content in enumerate(openai_jsonl_content):
             # Extract the request body from OpenAI format
             openai_body = _openai_jsonl_content.get("body", {})
-            model = openai_body.get("model", "")
-
-            try:
-                model, _, _, _ = get_llm_provider(
-                    model=model,
-                    custom_llm_provider=None,
-                )
-            except Exception as e:
-                verbose_logger.exception(
-                    "litellm.llms.bedrock.files.transformation.py::_transform_openai_jsonl_content_to_bedrock_jsonl_content() - Error inferring custom_llm_provider - %s",
-                    e,
-                )
-
-            # Determine provider from model name
-            provider = self.get_bedrock_invoke_provider(model)
+            record_model = openai_body.get("model", "")
+            resolved_model = litellm.model_alias_map.get(record_model, record_model)
+            model_for_transform, provider = self._resolve_batch_record_model_and_provider(
+                record_model=resolved_model, target_model=target_model
+            )
 
             # Route to the embedding transformer when the OpenAI batch line
             # targets /v1/embeddings; every other endpoint shape is normalized
@@ -816,10 +834,13 @@ class BedrockFilesConfig(BaseAWSLLM, BaseFilesConfig):
             # narrow contract and the embedding helper can evolve independently.
             record_kind = self._classify_batch_record(_openai_jsonl_content)
             if record_kind is BedrockBatchRecordKind.EMBEDDING:
-                model_input = self._map_openai_embedding_to_bedrock_params(openai_request_body=openai_body)
+                model_input = self._map_openai_embedding_to_bedrock_params(
+                    openai_request_body=openai_body, model=model_for_transform
+                )
             else:
                 model_input = self._map_openai_to_bedrock_params(
                     openai_request_body=self._transform_batch_body_to_chat_body(openai_body, record_kind),
+                    model=model_for_transform,
                     provider=provider,
                 )
 
@@ -858,7 +879,11 @@ class BedrockFilesConfig(BaseAWSLLM, BaseFilesConfig):
             ## Transform JSONL content to Bedrock format
             original_file_content: Final = self._get_content_from_openai_file(extracted_file_data_content)
             openai_jsonl_content = [json.loads(line) for line in original_file_content.splitlines() if line.strip()]
-            bedrock_jsonl_content = self._transform_openai_jsonl_content_to_bedrock_jsonl_content(openai_jsonl_content)
+            litellm_params_model: Final = litellm_params.get("model")
+            target_model: Final = model or (litellm_params_model if isinstance(litellm_params_model, str) else "")
+            bedrock_jsonl_content = self._transform_openai_jsonl_content_to_bedrock_jsonl_content(
+                openai_jsonl_content, target_model=target_model
+            )
             file_content = "\n".join(json.dumps(item) for item in bedrock_jsonl_content)
         elif isinstance(extracted_file_data_content, bytes):
             file_content = extracted_file_data_content.decode("utf-8")

tests/test_litellm/llms/bedrock/files/test_bedrock_files_transformation.py

@@ -622,6 +622,200 @@ class TestBedrockFilesTransformation:
         assert "max_tokens" in model_input
         assert model_input["max_tokens"] == 10
 
+    def test_resolves_model_alias_before_provider_mapping(self, monkeypatch):
+        import litellm
+        from litellm.llms.bedrock.files.transformation import BedrockFilesConfig
+
+        monkeypatch.setitem(
+            litellm.model_alias_map,
+            "bedrock-batch",
+            "bedrock/anthropic.claude-haiku-4-5-20251001-v1:0",
+        )
+
+        result = BedrockFilesConfig()._transform_openai_jsonl_content_to_bedrock_jsonl_content(
+            [
+                {
+                    "custom_id": "req-1",
+                    "body": {
+                        "model": "bedrock-batch",
+                        "messages": [{"role": "user", "content": "hi"}],
+                        "max_tokens": 16,
+                    },
+                }
+            ]
+        )
+
+        assert result == [
+            {
+                "recordId": "req-1",
+                "modelInput": {
+                    "messages": [{"role": "user", "content": [{"type": "text", "text": "hi"}]}],
+                    "max_tokens": 16,
+                    "anthropic_version": "bedrock-2023-05-31",
+                },
+            }
+        ]
+
+    def test_resolves_model_alias_before_embedding_mapping(self, monkeypatch):
+        import litellm
+        from litellm.llms.bedrock.files.transformation import BedrockFilesConfig
+
+        monkeypatch.setitem(
+            litellm.model_alias_map,
+            "bedrock-embedding-batch",
+            "bedrock/amazon.titan-embed-text-v2:0",
+        )
+
+        result = BedrockFilesConfig()._transform_openai_jsonl_content_to_bedrock_jsonl_content(
+            [
+                {
+                    "custom_id": "embedding-1",
+                    "url": "/v1/embeddings",
+                    "body": {
+                        "model": "bedrock-embedding-batch",
+                        "input": "hello",
+                    },
+                }
+            ]
+        )
+
+        assert result == [
+            {
+                "recordId": "embedding-1",
+                "modelInput": {"inputText": "hello"},
+            }
+        ]
+
+    def test_unmapped_alias_falls_back_to_target_model(self):
+        from litellm.llms.bedrock.files.transformation import BedrockFilesConfig
+
+        result = BedrockFilesConfig()._transform_openai_jsonl_content_to_bedrock_jsonl_content(
+            [
+                {
+                    "custom_id": "req-1",
+                    "body": {
+                        "model": "bedrock-batch",
+                        "messages": [{"role": "user", "content": "hi"}],
+                        "max_tokens": 16,
+                    },
+                },
+                {
+                    "custom_id": "req-2",
+                    "body": {
+                        "messages": [{"role": "user", "content": "hi"}],
+                        "max_tokens": 16,
+                    },
+                },
+            ],
+            target_model="bedrock/us.anthropic.claude-haiku-4-5-20251001-v1:0",
+        )
+
+        expected_model_input = {
+            "messages": [{"role": "user", "content": [{"type": "text", "text": "hi"}]}],
+            "max_tokens": 16,
+            "anthropic_version": "bedrock-2023-05-31",
+        }
+        assert result == [
+            {"recordId": "req-1", "modelInput": expected_model_input},
+            {"recordId": "req-2", "modelInput": expected_model_input},
+        ]
+
+    def test_record_provider_wins_over_target_model(self):
+        from litellm.llms.bedrock.files.transformation import BedrockFilesConfig
+
+        result = BedrockFilesConfig()._transform_openai_jsonl_content_to_bedrock_jsonl_content(
+            [
+                {
+                    "custom_id": "openai-1",
+                    "url": "/v1/chat/completions",
+                    "body": {
+                        "model": "openai.gpt-oss-120b-1:0",
+                        "messages": [{"role": "user", "content": "Hello!"}],
+                        "max_tokens": 10,
+                    },
+                }
+            ],
+            target_model="bedrock/us.anthropic.claude-haiku-4-5-20251001-v1:0",
+        )
+
+        assert result == [
+            {
+                "recordId": "openai-1",
+                "modelInput": {
+                    "messages": [{"role": "user", "content": "Hello!"}],
+                    "max_tokens": 10,
+                },
+            }
+        ]
+
+    def test_embedding_alias_falls_back_to_target_model(self):
+        from litellm.llms.bedrock.files.transformation import BedrockFilesConfig
+
+        result = BedrockFilesConfig()._transform_openai_jsonl_content_to_bedrock_jsonl_content(
+            [
+                {
+                    "custom_id": "embedding-1",
+                    "url": "/v1/embeddings",
+                    "body": {
+                        "model": "bedrock-embedding-batch",
+                        "input": "hello",
+                    },
+                }
+            ],
+            target_model="bedrock/amazon.titan-embed-text-v2:0",
+        )
+
+        assert result == [
+            {
+                "recordId": "embedding-1",
+                "modelInput": {"inputText": "hello"},
+            }
+        ]
+
+    def test_create_file_request_threads_deployment_model_to_alias_records(self):
+        from litellm.llms.bedrock.files.transformation import BedrockFilesConfig
+
+        class CapturingSignConfig(BedrockFilesConfig):
+            def __init__(self):
+                super().__init__()
+                self.signed_content: str | None = None
+
+            def _sign_s3_request(self, content, api_base, optional_params, s3_encryption_key_id=None):
+                self.signed_content = content
+                return {"Authorization": "fake"}, content
+
+        config = CapturingSignConfig()
+        jsonl_content = json.dumps(
+            {
+                "custom_id": "req-1",
+                "method": "POST",
+                "url": "/v1/chat/completions",
+                "body": {
+                    "model": "bedrock-batch",
+                    "messages": [{"role": "user", "content": "Hello"}],
+                    "max_tokens": 10,
+                },
+            }
+        ).encode()
+
+        config.transform_create_file_request(
+            model="",
+            create_file_data={
+                "file": ("batch.jsonl", jsonl_content, "application/jsonl"),
+                "purpose": "batch",
+            },
+            optional_params={},
+            litellm_params={
+                "s3_bucket_name": "litellm-batch-352026",
+                "model": "bedrock/us.anthropic.claude-haiku-4-5-20251001-v1:0",
+            },
+        )
+
+        assert config.signed_content is not None
+        record = json.loads(config.signed_content)
+        assert record["modelInput"]["anthropic_version"] == "bedrock-2023-05-31"
+        assert "model" not in record["modelInput"]
+
 
 class TestBedrockFilesEmbeddingTransformation:
     """