diff --git a/litellm/llms/bedrock/files/transformation.py b/litellm/llms/bedrock/files/transformation.py index 6669363093b2..6d1e7525d4e5 100644 --- a/litellm/llms/bedrock/files/transformation.py +++ b/litellm/llms/bedrock/files/transformation.py @@ -233,6 +233,231 @@ def map_openai_params( # example; add others here as they adopt the same schema. CONVERSE_INVOKE_PROVIDERS = ("nova",) + # OpenAI batch URL that signals an embedding request. Per OpenAI Batch API + # spec, every JSONL record carries a `url` field; we use it as the + # authoritative signal to route the line to the embedding code path + # instead of inferring from the presence of `input` vs `messages`. + OPENAI_EMBEDDINGS_URL = "/v1/embeddings" + + @staticmethod + def _is_embedding_record(openai_jsonl_record: Dict[str, Any]) -> bool: + """ + Decide whether an OpenAI batch JSONL line is an embedding request. + + Precedence (strict - any explicit `url` short-circuits): + 1. `url == "/v1/embeddings"` -> embedding. Authoritative per the + OpenAI Batch API spec. + 2. Any other non-empty `url` (e.g. `/v1/chat/completions`) -> NOT + embedding. We trust the caller's explicit signal even if the + body would otherwise suggest embedding; misrouting a chat + record into the embedding transformer would corrupt the + modelInput, while a chat-shaped body sent to the chat path + either succeeds or fails cleanly inside that transformer. + 3. `url` missing/empty -> fall back to body shape. Requires + `input` present AND `messages` absent so a malformed record + carrying both keys routes to the chat path (safer default: + Anthropic transforms ignore unknown top-level keys, whereas + the embedding transformer would silently drop the messages). + """ + url = openai_jsonl_record.get("url") + if url == BedrockFilesConfig.OPENAI_EMBEDDINGS_URL: + return True + if url: + return False + body = openai_jsonl_record.get("body", {}) + if not isinstance(body, dict): + return False + return "input" in body and "messages" not in body + + # Substring match against the model id (case-insensitive, after stripping + # any "bedrock/" routing prefix and any cross-region "." prefix + # like "us.amazon.titan-embed-text-v2:0"). Kept as a constant so future + # PRs can extend the set without touching the dispatch logic. + _TITAN_V2_EMBED_MODEL_MARKER = "titan-embed-text-v2" + + @staticmethod + def _is_titan_v2_embed_model(model: str) -> bool: + """ + True iff `model` refers to Amazon Titan Text Embeddings V2. + + Resolution order: + 1. `model_prices_and_context_window.json` via `get_model_info`. + When the registry resolves the id we trust `mode == "embedding"` + AND a matching `titan-embed-text-v2` marker in the id - the + marker is still needed because the registry's `mode` field + doesn't distinguish Titan v2's InvokeModel schema from Cohere, + Nova Multimodal, or Titan G1 (all also `mode == "embedding"` + but with incompatible bodies). + 2. Substring fallback for ids the registry can't resolve - this + catches cross-region inference profile prefixes + (`us.amazon.titan-embed-text-v2:0`) and Bedrock ARN forms. + The marker boundary check rejects lookalikes like + `titan-embed-text-v20` or `titan-embed-text-v2-experimental`. + + Tolerant of common id shapes: + - "amazon.titan-embed-text-v2:0" + - "bedrock/amazon.titan-embed-text-v2:0" + - "us.amazon.titan-embed-text-v2:0" (cross-region inference profile) + - ARN forms ending in ".../amazon.titan-embed-text-v2:0" + """ + normalized = model.lower() + if normalized.startswith("bedrock/"): + normalized = normalized[len("bedrock/") :] + marker = BedrockFilesConfig._TITAN_V2_EMBED_MODEL_MARKER + idx = normalized.find(marker) + if idx < 0: + return False + end = idx + len(marker) + if not (end == len(normalized) or normalized[end] in (":", "/")): + return False + + # Marker matches with a clean boundary. If the registry can also + # resolve this id, additionally confirm `mode == "embedding"` so a + # malformed id whose path-component is right but whose registered + # mode is, say, "chat" doesn't slip through. Registry silence + # (cross-region profiles, ARNs) is fine - the marker alone is + # authoritative there. + registry_mode = BedrockFilesConfig._lookup_registry_mode(model) + if registry_mode is not None and registry_mode != "embedding": + return False + return True + + @staticmethod + def _lookup_registry_mode(model_id: str) -> Optional[str]: + """ + Read `mode` for `model_id` from `model_prices_and_context_window.json`. + + Returns the mode string when the registry resolves the id and the + entry has a non-empty string mode, else `None`. Isolating this + means the Titan v2 detector can layer a data-driven check on top + of the marker boundary without scattering try/except shapes. + """ + try: + from litellm import get_model_info + + info = get_model_info(model_id) + except Exception: + return None + if not isinstance(info, dict): + return None + mode = info.get("mode") + return mode if isinstance(mode, str) and mode else None + + @staticmethod + def _coerce_embedding_input_to_string(raw_input: Any, model: str = "") -> str: + """ + Normalize an OpenAI /v1/embeddings `input` field into the single + string that Bedrock Titan v2 InvokeModel expects in `inputText`. + + Accepts: a string, or a single-element list containing one string. + Rejects (with actionable messages): + - None / missing -> ValueError + - Multi-element string lists -> ValueError, prompts caller to + emit one JSONL line per input + - Pre-tokenized inputs (List[int], List[List[int]]) -> NotImplementedError + - Any other type -> ValueError + + Extracted so the validation can be exercised in isolation and so + future embedding-provider branches (Titan G1, Cohere) can reuse it + without duplicating the type-shaping logic. + """ + if raw_input is None: + raise ValueError( + "Embedding batch record is missing required `input` field: " + f"model={model}" + ) + + # Bedrock InvokeModel for Titan v2 takes exactly one string `inputText` + # per call. Pre-tokenized inputs and multi-element string lists are + # explicitly unsupported so callers emit one JSONL line per embedding + # instead of relying on us to silently fan out or concatenate. + if isinstance(raw_input, list): + if len(raw_input) == 1: + candidate = raw_input[0] + else: + raise ValueError( + "Bedrock batch embedding requires one input per JSONL " + "record. Got a list with " + f"{len(raw_input)} items for model={model}; emit one " + "JSONL line per input string instead." + ) + else: + candidate = raw_input + + # Catches pre-tokenized inputs (List[int] from OpenAI spec, or a + # single int slipping past the list-unwrap above). + # NOTE: bool is a subclass of int but treating True/False as a token + # is meaningless either way, so the broad check is fine. + if isinstance(candidate, (list, int)): + raise NotImplementedError( + "Bedrock Titan v2 batch embedding does not support " + "pre-tokenized integer inputs. Pass `input` as a string " + f"(model={model})." + ) + if not isinstance(candidate, str): + raise ValueError( + "Bedrock batch embedding `input` must be a string (or a " + "single-element list of strings). Got type " + f"{type(candidate).__name__} for model={model}." + ) + return candidate + + def _map_openai_embedding_to_bedrock_params( + self, + openai_request_body: Dict[str, Any], + ) -> Dict[str, Any]: + """ + Transform an OpenAI /v1/embeddings request body into the + Bedrock InvokeModel `modelInput` for embedding models that AWS + supports via batch inference (CreateModelInvocationJob). + + Currently routes Amazon Titan Text Embeddings V2 only; other + embedding providers (Titan G1, Titan Multimodal, Cohere Embed, + Nova Multimodal Embeddings) raise NotImplementedError until they + get a dedicated branch. Splitting them keeps PR scope tight and + lets each model's request schema be exercised by its own tests. + + AWS docs (Titan v2 InvokeModel body): + https://docs.aws.amazon.com/bedrock/latest/userguide/model-parameters-titan-embed-text.html + """ + from litellm.llms.bedrock.embed.amazon_titan_v2_transformation import ( + AmazonTitanV2Config, + ) + + _model = openai_request_body.get("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 + # them here without dedicated tests would risk corrupt batches. + 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 " + "embedding models in https://github.com/BerriAI/litellm/issues." + ) + + input_text = 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 + # so this stays in sync with the synchronous /v1/embeddings path. + non_default_params = { + k: v for k, v in openai_request_body.items() if k not in ("model", "input") + } + titan_config = AmazonTitanV2Config() + inference_params = titan_config.map_openai_params( + non_default_params=non_default_params, + optional_params={}, + ) + return dict( + titan_config._transform_request( + input=input_text, inference_params=inference_params + ) + ) + def _map_openai_to_bedrock_params( self, openai_request_body: Dict[str, Any], @@ -349,10 +574,19 @@ def _transform_openai_jsonl_content_to_bedrock_jsonl_content( # Determine provider from model name provider = self.get_bedrock_invoke_provider(model) - # Transform to Bedrock modelInput format - model_input = self._map_openai_to_bedrock_params( - openai_request_body=openai_body, provider=provider - ) + # Route to the embedding transformer when the OpenAI batch line + # targets /v1/embeddings; otherwise fall back to the existing + # chat-completion path. We branch here (rather than inside + # `_map_openai_to_bedrock_params`) so the chat helper keeps its + # narrow contract and the embedding helper can evolve independently. + if self._is_embedding_record(_openai_jsonl_content): + model_input = self._map_openai_embedding_to_bedrock_params( + openai_request_body=openai_body + ) + else: + model_input = self._map_openai_to_bedrock_params( + openai_request_body=openai_body, provider=provider + ) # Create Bedrock batch record record_id = _openai_jsonl_content.get( diff --git a/tests/test_litellm/llms/bedrock/files/expected_bedrock_batch_embeddings.jsonl b/tests/test_litellm/llms/bedrock/files/expected_bedrock_batch_embeddings.jsonl new file mode 100644 index 000000000000..e798c39b7988 --- /dev/null +++ b/tests/test_litellm/llms/bedrock/files/expected_bedrock_batch_embeddings.jsonl @@ -0,0 +1,3 @@ +{"recordId": "embed-1", "modelInput": {"inputText": "Hello world"}} +{"recordId": "embed-2", "modelInput": {"inputText": "Another document to embed", "dimensions": 512}} +{"recordId": "embed-3", "modelInput": {"inputText": "Single element list", "embeddingTypes": ["binary"]}} diff --git a/tests/test_litellm/llms/bedrock/files/input_batch_embeddings.jsonl b/tests/test_litellm/llms/bedrock/files/input_batch_embeddings.jsonl new file mode 100644 index 000000000000..f87b4eba7e1a --- /dev/null +++ b/tests/test_litellm/llms/bedrock/files/input_batch_embeddings.jsonl @@ -0,0 +1,3 @@ +{"custom_id": "embed-1", "method": "POST", "url": "/v1/embeddings", "body": {"model": "bedrock/amazon.titan-embed-text-v2:0", "input": "Hello world"}} +{"custom_id": "embed-2", "method": "POST", "url": "/v1/embeddings", "body": {"model": "bedrock/amazon.titan-embed-text-v2:0", "input": "Another document to embed", "dimensions": 512}} +{"custom_id": "embed-3", "method": "POST", "url": "/v1/embeddings", "body": {"model": "bedrock/amazon.titan-embed-text-v2:0", "input": ["Single element list"], "encoding_format": "base64"}} 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 5245612e9d3a..fa5595b31740 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 @@ -426,7 +426,7 @@ def test_s3_region_name_wins_over_aws_region_name_for_signing(self): "s3_bucket_name": "litellm-batch-352026", "s3_region_name": "us-gov-west-1", } - # aws_region_name set to something different — s3_region_name must still win + # aws_region_name set to something different - s3_region_name must still win optional_params = {"aws_region_name": "us-east-1"} captured_optional_params: dict = {} @@ -482,3 +482,522 @@ def test_openai_passthrough_still_works(self): assert "messages" in model_input assert "max_tokens" in model_input assert model_input["max_tokens"] == 10 + + +class TestBedrockFilesEmbeddingTransformation: + """ + Tests for routing OpenAI /v1/embeddings batch JSONL records through the + Titan v2 transformer so AWS Bedrock's CreateModelInvocationJob receives + a valid modelInput body. + + Scope is intentionally Titan v2 only - other embedding models will get + their own follow-up PRs/tests so each schema is exercised in isolation. + """ + + def test_titan_v2_embedding_jsonl_matches_fixture(self): + """Round-trip the input fixture against the expected Bedrock output.""" + import json + import os + + from litellm.llms.bedrock.files.transformation import BedrockFilesConfig + + config = BedrockFilesConfig() + here = os.path.dirname(__file__) + with open(os.path.join(here, "input_batch_embeddings.jsonl")) as f: + openai_jsonl = [json.loads(line) for line in f if line.strip()] + with open(os.path.join(here, "expected_bedrock_batch_embeddings.jsonl")) as f: + expected = [json.loads(line) for line in f if line.strip()] + + result = config._transform_openai_jsonl_content_to_bedrock_jsonl_content( + openai_jsonl + ) + + assert result == expected + + def test_titan_v2_simple_string_input(self): + """Single string `input` maps to `{"inputText": }` with no extras.""" + from litellm.llms.bedrock.files.transformation import BedrockFilesConfig + + config = BedrockFilesConfig() + result = config._transform_openai_jsonl_content_to_bedrock_jsonl_content( + [ + { + "custom_id": "e1", + "method": "POST", + "url": "/v1/embeddings", + "body": { + "model": "bedrock/amazon.titan-embed-text-v2:0", + "input": "Hello", + }, + } + ] + ) + + assert result == [{"recordId": "e1", "modelInput": {"inputText": "Hello"}}] + + def test_titan_v2_dimensions_and_encoding_format(self): + """OpenAI `dimensions` / `encoding_format` map to Titan v2 schema.""" + from litellm.llms.bedrock.files.transformation import BedrockFilesConfig + + config = BedrockFilesConfig() + result = config._transform_openai_jsonl_content_to_bedrock_jsonl_content( + [ + { + "custom_id": "e1", + "method": "POST", + "url": "/v1/embeddings", + "body": { + "model": "bedrock/amazon.titan-embed-text-v2:0", + "input": "Hi", + "dimensions": 256, + "encoding_format": "float", + }, + } + ] + ) + + model_input = result[0]["modelInput"] + assert model_input["inputText"] == "Hi" + assert model_input["dimensions"] == 256 + assert model_input["embeddingTypes"] == ["float"] + + def test_embedding_routing_falls_back_to_body_shape(self): + """Records without `url` still route via `input` presence.""" + from litellm.llms.bedrock.files.transformation import BedrockFilesConfig + + config = BedrockFilesConfig() + result = config._transform_openai_jsonl_content_to_bedrock_jsonl_content( + [ + { + "custom_id": "e1", + "body": { + "model": "bedrock/amazon.titan-embed-text-v2:0", + "input": "Hello", + }, + } + ] + ) + + assert result[0]["modelInput"] == {"inputText": "Hello"} + + def test_embedding_single_element_list_input_is_accepted(self): + """A single-element list maps to the same shape as a bare string.""" + from litellm.llms.bedrock.files.transformation import BedrockFilesConfig + + config = BedrockFilesConfig() + result = config._transform_openai_jsonl_content_to_bedrock_jsonl_content( + [ + { + "custom_id": "e1", + "method": "POST", + "url": "/v1/embeddings", + "body": { + "model": "bedrock/amazon.titan-embed-text-v2:0", + "input": ["only one"], + }, + } + ] + ) + + assert result[0]["modelInput"]["inputText"] == "only one" + + def test_embedding_multi_input_list_raises(self): + """Multi-element `input` lists are rejected with a clear message.""" + import pytest + + from litellm.llms.bedrock.files.transformation import BedrockFilesConfig + + config = BedrockFilesConfig() + with pytest.raises(ValueError, match="one input per JSONL record"): + config._transform_openai_jsonl_content_to_bedrock_jsonl_content( + [ + { + "custom_id": "e1", + "method": "POST", + "url": "/v1/embeddings", + "body": { + "model": "bedrock/amazon.titan-embed-text-v2:0", + "input": ["a", "b"], + }, + } + ] + ) + + def test_embedding_missing_input_raises(self): + """A record routed to /v1/embeddings without `input` is an error.""" + import pytest + + from litellm.llms.bedrock.files.transformation import BedrockFilesConfig + + config = BedrockFilesConfig() + with pytest.raises(ValueError, match="missing required `input`"): + config._transform_openai_jsonl_content_to_bedrock_jsonl_content( + [ + { + "custom_id": "e1", + "method": "POST", + "url": "/v1/embeddings", + "body": {"model": "bedrock/amazon.titan-embed-text-v2:0"}, + } + ] + ) + + def test_mixed_chat_and_embedding_in_same_batch(self): + """Chat and embedding records in the same JSONL each take their path.""" + from litellm.llms.bedrock.files.transformation import BedrockFilesConfig + + config = BedrockFilesConfig() + result = config._transform_openai_jsonl_content_to_bedrock_jsonl_content( + [ + { + "custom_id": "chat-1", + "method": "POST", + "url": "/v1/chat/completions", + "body": { + "model": "bedrock/us.anthropic.claude-haiku-4-5-20251001-v1:0", + "messages": [{"role": "user", "content": "Hi"}], + "max_tokens": 5, + }, + }, + { + "custom_id": "embed-1", + "method": "POST", + "url": "/v1/embeddings", + "body": { + "model": "bedrock/amazon.titan-embed-text-v2:0", + "input": "Hi", + }, + }, + ] + ) + + assert result[0]["recordId"] == "chat-1" + assert "messages" in result[0]["modelInput"] + assert result[0]["modelInput"]["anthropic_version"] == "bedrock-2023-05-31" + + assert result[1]["recordId"] == "embed-1" + assert result[1]["modelInput"] == {"inputText": "Hi"} + + def test_unsupported_embedding_model_raises_not_implemented(self): + """Cohere/Nova/Titan-G1 embed get a clear NotImplementedError, not a corrupt body.""" + import pytest + + from litellm.llms.bedrock.files.transformation import BedrockFilesConfig + + config = BedrockFilesConfig() + for unsupported_model in ( + "bedrock/cohere.embed-english-v3", + "bedrock/amazon.titan-embed-text-v1", + "bedrock/amazon.titan-embed-image-v1", + "bedrock/amazon.nova-2-multimodal-embeddings-v1:0", + ): + with pytest.raises(NotImplementedError, match="titan-embed-text-v2"): + config._transform_openai_jsonl_content_to_bedrock_jsonl_content( + [ + { + "custom_id": "e1", + "method": "POST", + "url": "/v1/embeddings", + "body": {"model": unsupported_model, "input": "Hi"}, + } + ] + ) + + def test_titan_v2_model_name_variants_route_correctly(self): + """All common Titan v2 model id shapes route through the embedding path.""" + from litellm.llms.bedrock.files.transformation import BedrockFilesConfig + + config = BedrockFilesConfig() + for model_id in ( + "amazon.titan-embed-text-v2:0", + "bedrock/amazon.titan-embed-text-v2:0", + "us.amazon.titan-embed-text-v2:0", + "bedrock/us.amazon.titan-embed-text-v2:0", + ): + result = config._transform_openai_jsonl_content_to_bedrock_jsonl_content( + [ + { + "custom_id": "e1", + "method": "POST", + "url": "/v1/embeddings", + "body": {"model": model_id, "input": "Hi"}, + } + ] + ) + assert result[0]["modelInput"] == { + "inputText": "Hi" + }, f"model id {model_id} did not route to Titan v2 embedding path" + + def test_pretokenized_input_list_of_ints_raises(self): + """`input: List[int]` (pre-tokenized) is rejected, not silently mis-shaped.""" + import pytest + + from litellm.llms.bedrock.files.transformation import BedrockFilesConfig + + config = BedrockFilesConfig() + with pytest.raises( + (NotImplementedError, ValueError), match=r"pre-tokenized|one input per" + ): + config._transform_openai_jsonl_content_to_bedrock_jsonl_content( + [ + { + "custom_id": "e1", + "method": "POST", + "url": "/v1/embeddings", + "body": { + "model": "bedrock/amazon.titan-embed-text-v2:0", + "input": [1, 2, 3], + }, + } + ] + ) + + def test_pretokenized_single_wrapped_list_raises(self): + """`input: List[List[int]]` with one element is rejected as pre-tokenized.""" + import pytest + + from litellm.llms.bedrock.files.transformation import BedrockFilesConfig + + config = BedrockFilesConfig() + with pytest.raises(NotImplementedError, match="pre-tokenized"): + config._transform_openai_jsonl_content_to_bedrock_jsonl_content( + [ + { + "custom_id": "e1", + "method": "POST", + "url": "/v1/embeddings", + "body": { + "model": "bedrock/amazon.titan-embed-text-v2:0", + "input": [[1, 2, 3]], + }, + } + ] + ) + + def test_record_with_both_input_and_messages_routes_to_chat(self): + """If a record has both fields, chat wins (safer default - see helper docstring).""" + from litellm.llms.bedrock.files.transformation import BedrockFilesConfig + + config = BedrockFilesConfig() + result = config._transform_openai_jsonl_content_to_bedrock_jsonl_content( + [ + { + "custom_id": "ambiguous-1", + "body": { + "model": "bedrock/us.anthropic.claude-haiku-4-5-20251001-v1:0", + "messages": [{"role": "user", "content": "Hi"}], + "input": "this should be ignored by chat path", + "max_tokens": 5, + }, + } + ] + ) + + assert "messages" in result[0]["modelInput"] + assert "inputText" not in result[0]["modelInput"] + + def test_url_embeddings_with_missing_input_raises_not_chat_error(self): + """url says embed, body lacks input → embedding-path error, not chat-path crash.""" + import pytest + + from litellm.llms.bedrock.files.transformation import BedrockFilesConfig + + config = BedrockFilesConfig() + with pytest.raises(ValueError, match="missing required `input`"): + config._transform_openai_jsonl_content_to_bedrock_jsonl_content( + [ + { + "custom_id": "e1", + "method": "POST", + "url": "/v1/embeddings", + "body": {"model": "bedrock/amazon.titan-embed-text-v2:0"}, + } + ] + ) + + def test_titan_v2_marker_boundary_rejects_lookalikes(self): + """The marker must end at `:`, `/`, or end-of-string to avoid false positives.""" + from litellm.llms.bedrock.files.transformation import BedrockFilesConfig + + # Look-alikes that must NOT route through the Titan v2 path + for model in ( + "bedrock/amazon.titan-embed-text-v20:0", + "bedrock/amazon.titan-embed-text-v2-experimental:0", + "bedrock/amazon.titan-embed-text-v2foo", + ): + assert not BedrockFilesConfig._is_titan_v2_embed_model( + model + ), f"{model} unexpectedly matched the Titan v2 marker" + + # Real Titan v2 ids that MUST match + for model in ( + "amazon.titan-embed-text-v2:0", + "bedrock/amazon.titan-embed-text-v2:0", + "us.amazon.titan-embed-text-v2:0", + "arn:aws:bedrock:us-east-1:123:foundation-model/amazon.titan-embed-text-v2:0", + ): + assert BedrockFilesConfig._is_titan_v2_embed_model( + model + ), f"{model} unexpectedly missed the Titan v2 marker" + + def test_titan_v2_rejected_when_registry_mode_disagrees(self, mocker): + """If the registry resolves the id but says mode != embedding, reject.""" + from litellm.llms.bedrock.files.transformation import BedrockFilesConfig + + # Marker matches but registry claims this is chat - trust the registry + # and refuse to route through the embedding path. + mocker.patch("litellm.get_model_info", return_value={"mode": "chat"}) + assert not BedrockFilesConfig._is_titan_v2_embed_model( + "amazon.titan-embed-text-v2:0" + ), "registry mode=chat must override the marker match" + + def test_titan_v2_accepted_when_registry_confirms_embedding(self, mocker): + """Happy path: marker matches and registry says mode=embedding.""" + from litellm.llms.bedrock.files.transformation import BedrockFilesConfig + + mocker.patch("litellm.get_model_info", return_value={"mode": "embedding"}) + assert BedrockFilesConfig._is_titan_v2_embed_model( + "amazon.titan-embed-text-v2:0" + ) + + def test_titan_v2_accepted_when_registry_silent(self, mocker): + """Marker-only match is fine for ids the registry can't resolve + (cross-region profile prefixes, ARN forms).""" + from litellm.llms.bedrock.files.transformation import BedrockFilesConfig + + mocker.patch("litellm.get_model_info", side_effect=Exception("not mapped")) + assert BedrockFilesConfig._is_titan_v2_embed_model( + "us.amazon.titan-embed-text-v2:0" + ) + assert BedrockFilesConfig._is_titan_v2_embed_model( + "arn:aws:bedrock:us-east-1:123:foundation-model/amazon.titan-embed-text-v2:0" + ) + + def test_lookup_registry_mode_helper(self, mocker): + """Direct coverage of the extracted registry helper.""" + from litellm.llms.bedrock.files.transformation import BedrockFilesConfig + + # Happy path: returns the mode string + mocker.patch("litellm.get_model_info", return_value={"mode": "embedding"}) + assert BedrockFilesConfig._lookup_registry_mode("anything") == "embedding" + + # Registry raises -> None + mocker.patch("litellm.get_model_info", side_effect=Exception("not mapped")) + assert BedrockFilesConfig._lookup_registry_mode("anything") is None + + # Registry returns non-dict -> None + mocker.patch("litellm.get_model_info", return_value="not a dict") + assert BedrockFilesConfig._lookup_registry_mode("anything") is None + + # Registry returns dict without mode -> None + mocker.patch("litellm.get_model_info", return_value={}) + assert BedrockFilesConfig._lookup_registry_mode("anything") is None + + # Registry returns dict with non-string mode -> None + mocker.patch("litellm.get_model_info", return_value={"mode": 42}) + assert BedrockFilesConfig._lookup_registry_mode("anything") is None + + def test_is_embedding_record_helper(self): + """Helper detects embeddings via `url` first, then by body shape.""" + from litellm.llms.bedrock.files.transformation import BedrockFilesConfig + + assert BedrockFilesConfig._is_embedding_record( + {"url": "/v1/embeddings", "body": {"input": "x"}} + ) + # body-only fallback + assert BedrockFilesConfig._is_embedding_record({"body": {"input": "x"}}) + # chat shape + assert not BedrockFilesConfig._is_embedding_record( + {"url": "/v1/chat/completions", "body": {"messages": []}} + ) + # ambiguous body without `input` is treated as not-embedding + assert not BedrockFilesConfig._is_embedding_record({"body": {}}) + + def test_explicit_chat_url_with_input_body_short_circuits_to_chat(self): + """Explicit url=/v1/chat/completions wins even if body looks like embedding. + + Without this short-circuit, a chat record whose body happens to carry + `input` (and no `messages`) would be mis-routed to the embedding + transformer, corrupting the modelInput. + """ + from litellm.llms.bedrock.files.transformation import BedrockFilesConfig + + # Direct helper assertion + assert not BedrockFilesConfig._is_embedding_record( + { + "url": "/v1/chat/completions", + "body": { + "model": "bedrock/us.anthropic.claude-haiku-4-5-20251001-v1:0", + "input": "this would mis-route under the old precedence", + }, + } + ) + + # End-to-end: a record like this routes through the chat path. We + # just need to make sure we DON'T silently produce an inputText + # body and call it a chat completion. + config = BedrockFilesConfig() + result = config._transform_openai_jsonl_content_to_bedrock_jsonl_content( + [ + { + "custom_id": "explicit-chat-with-input", + "method": "POST", + "url": "/v1/chat/completions", + "body": { + "model": "bedrock/us.anthropic.claude-haiku-4-5-20251001-v1:0", + "messages": [{"role": "user", "content": "Hi"}], + "input": "should not become inputText", + "max_tokens": 5, + }, + } + ] + ) + + model_input = result[0]["modelInput"] + assert ( + "inputText" not in model_input + ), "explicit chat URL must not produce an embedding-shaped modelInput" + + def test_coerce_embedding_input_helper_isolated(self): + """Direct coverage of the extracted input-normalization helper.""" + import pytest + + from litellm.llms.bedrock.files.transformation import BedrockFilesConfig + + # Happy paths + assert BedrockFilesConfig._coerce_embedding_input_to_string("hello") == "hello" + assert ( + BedrockFilesConfig._coerce_embedding_input_to_string(["hello"]) == "hello" + ) + + # Error paths + with pytest.raises(ValueError, match="missing required `input`"): + BedrockFilesConfig._coerce_embedding_input_to_string(None, model="m") + with pytest.raises(ValueError, match="one input per JSONL record"): + BedrockFilesConfig._coerce_embedding_input_to_string(["a", "b"]) + # A multi-element list of ints is rejected as "one input per JSONL + # record" too - we can't tell if it's pre-tokenized or "3 strings" + # without more context, so the most-actionable error wins. + with pytest.raises(ValueError, match="one input per JSONL record"): + BedrockFilesConfig._coerce_embedding_input_to_string([1, 2, 3]) + # Single-element list wrapping a token list -> pre-tokenized error. + with pytest.raises(NotImplementedError, match="pre-tokenized"): + BedrockFilesConfig._coerce_embedding_input_to_string([[1, 2, 3]]) + # Single-element list wrapping a bare int -> pre-tokenized error. + with pytest.raises(NotImplementedError, match="pre-tokenized"): + BedrockFilesConfig._coerce_embedding_input_to_string([42]) + with pytest.raises(ValueError, match="must be a string"): + BedrockFilesConfig._coerce_embedding_input_to_string({"unsupported": True}) + + def test_other_non_embedding_urls_route_to_chat(self): + """Any non-/v1/embeddings url short-circuits to chat path.""" + from litellm.llms.bedrock.files.transformation import BedrockFilesConfig + + # /v1/completions (legacy completions endpoint) + assert not BedrockFilesConfig._is_embedding_record( + {"url": "/v1/completions", "body": {"input": "x"}} + ) + # Arbitrary unknown url - caller's explicit signal still wins + assert not BedrockFilesConfig._is_embedding_record( + {"url": "/v1/responses", "body": {"input": "x"}} + )