feat(bedrock-batch): route /v1/embeddings JSONL to Titan v2 modelInput

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 "<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(

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"]}}

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"}}