diff --git a/litellm/llms/bedrock/files/transformation.py b/litellm/llms/bedrock/files/transformation.py index b50a9ae04d1..7d13ae82a6c 100644 --- a/litellm/llms/bedrock/files/transformation.py +++ b/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") diff --git a/tests/test_litellm/llms/bedrock/files/test_bedrock_files_transformation.py b/tests/test_litellm/llms/bedrock/files/test_bedrock_files_transformation.py index 841736acd73..2445bae97cd 100644 --- a/tests/test_litellm/llms/bedrock/files/test_bedrock_files_transformation.py +++ b/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: """