feat(bedrock-batch): route /v1/embeddings JSONL to Titan v2 modelInput - #28865
Conversation
Changes vs main:
- BedrockFilesConfig now detects OpenAI batch JSONL lines whose `url`
is /v1/embeddings (with body-shape fallback) and routes them through
a new `_map_openai_embedding_to_bedrock_params` helper instead of the
chat-completion transformer that silently produces an invalid body.
- The embedding helper currently supports Amazon Titan Text Embeddings V2
only. Other embed models (Titan G1, Titan Multimodal, Cohere Embed,
Nova Multimodal Embeddings) raise NotImplementedError with a clear
message; each will get a dedicated branch + tests in follow-up PRs to
keep schema-specific risks isolated.
- Validation refuses pre-tokenized inputs (List[int], List[List[int]])
and multi-element string lists with explicit errors so callers emit
one JSONL line per embedding instead of relying on us to fan out.
- Titan v2 model id match tolerates "bedrock/" prefix, cross-region
inference profile prefix ("us.", "eu.", etc.), and ARN forms; the
marker boundary check rejects lookalikes like "titan-embed-text-v20".
- Tests cover happy path (fixtures), dimensions/encoding_format mapping,
body-shape fallback, single-element list unwrap, error paths
(missing input, multi-element list, unsupported model, pre-tokenized),
mixed chat+embedding batch, and the model-id boundary check.
Greptile-flagged gap in `_is_embedding_record`: when an OpenAI batch JSONL line carries an explicit `url` pointing to a non-embedding endpoint (e.g. `/v1/chat/completions`) AND its body happens to have `input` without `messages`, the body-shape fallback would mis-route that record to the embedding transformer and corrupt the modelInput. Changes vs previous commit: - `_is_embedding_record` now short-circuits to NOT-embedding whenever `url` is non-empty and not equal to `/v1/embeddings`. The body-shape fallback only runs when `url` is missing or empty. Docstring updated to spell out the precedence rules. - Two new tests cover the case: direct helper assertion that an explicit chat url plus an input-bearing body returns False, plus an end-to-end check that the resulting modelInput contains no `inputText` key. A second test asserts the same short-circuit for arbitrary non-embeddings urls (`/v1/completions`, `/v1/responses`). 28/28 tests pass (was 26/26 before this commit + 2 new).
Splits the input-shape validation out of `_map_openai_embedding_to_bedrock_params` into a new static helper `_coerce_embedding_input_to_string`. Same semantics; the goal is to make the validation testable in isolation and to give future embedding-provider branches (Titan G1, Cohere) a reusable shaping function instead of duplicating type checks. - Helper accepts `str`, single-element `list[str]`, and raises `ValueError` / `NotImplementedError` with actionable messages for None, multi-element lists, pre-tokenized inputs (`list[int]` / `list[list[int]]`), and other unsupported types. - New unit test exercises the helper directly across happy paths, None / missing input, multi-element string list, multi-element int list (caught as 'one input per JSONL record' since we can't disambiguate from 'multiple strings' without more context), pre-tokenized single-element list-of-list, single-element list of bare int, and dict input. 29/29 tests in the file still pass.
Greptile SummaryThis PR adds routing support for
Confidence Score: 5/5Safe to merge — the change adds a new embedding routing branch in an isolated helper chain without touching the existing chat or passthrough paths. The routing logic is well-contained: No files require special attention.
|
| Filename | Overview |
|---|---|
| litellm/llms/bedrock/files/transformation.py | Adds embedding routing helpers and a new embedding-to-Bedrock param mapper; dispatch logic and boundary checks are correct, registry layering is sound. |
| tests/test_litellm/llms/bedrock/files/test_bedrock_files_transformation.py | Adds 17 new mocked tests in TestBedrockFilesEmbeddingTransformation covering all routing branches, error paths, registry interactions, and boundary checks; no real network calls. |
| tests/test_litellm/llms/bedrock/files/input_batch_embeddings.jsonl | Test fixture with three embedding records covering plain string, dimensions, and base64 encoding_format inputs. |
| tests/test_litellm/llms/bedrock/files/expected_bedrock_batch_embeddings.jsonl | Expected Bedrock output fixture; encoding_format=base64 correctly maps to embeddingTypes=["binary"] per AmazonTitanV2Config. |
Reviews (3): Last reviewed commit: "fix(bedrock-batch): layer registry mode ..." | Re-trigger Greptile
![greptile-apps[bot]](https://avatars.githubusercontent.com/in/867647?s=60&v=4){kind=small}
Greptile SummaryThis PR wires up
Confidence Score: 3/5The embedding transformation is logically correct and well-tested, but the model-detection logic is hardcoded directly in source rather than being driven by the model capability registry the rest of the codebase uses. The routing and parameter-mapping logic work correctly for the supported model, and the 17 unit tests give good coverage of happy and error paths. However, litellm/llms/bedrock/files/transformation.py — the
|
| Filename | Overview |
|---|---|
| litellm/llms/bedrock/files/transformation.py | Adds embedding routing via new static helpers; hardcodes Titan v2 model detection in violation of the repo rule requiring model capabilities to live in model_prices_and_context_window.json; unused provider parameter in _map_openai_embedding_to_bedrock_params. |
| tests/test_litellm/llms/bedrock/files/test_bedrock_files_transformation.py | Adds 17 new mock-only tests covering happy paths, error paths, model boundary checks, and mixed chat+embedding batches. No real network calls. |
| tests/test_litellm/llms/bedrock/files/expected_bedrock_batch_embeddings.jsonl | New fixture defining expected Bedrock modelInput for three embedding records; aligns with the encoding_format→embeddingTypes mapping in AmazonTitanV2Config. |
| tests/test_litellm/llms/bedrock/files/input_batch_embeddings.jsonl | New OpenAI-format batch JSONL fixture used as round-trip input for the fixture comparison test. |
Reviews (2): Last reviewed commit: "refactor(bedrock-batch): extract embeddi..." | Re-trigger Greptile
…param Greptile-flagged issues on Titan v2 model detection: 1. Hardcoded substring without registry consultation conflicts with the project convention of treating `model_prices_and_context_window.json` as the source of truth for model capability flags. Fix: `_is_titan_v2_embed_model` now layers a registry mode check on top of the existing marker boundary. When `get_model_info` resolves the id, we additionally require `mode == "embedding"` so a malformed id whose path-component matches the marker but whose registered mode is "chat" doesn't slip through. Registry silence (cross-region inference profile prefixes like `us.amazon.titan-embed-text-v2:0`, ARN forms) keeps the substring-only behavior because the registry genuinely can't normalize those ids today. The substring is still needed because `mode == "embedding"` alone doesn't distinguish Titan v2's InvokeModel schema from Cohere Embed, Nova Multimodal, or Titan G1 - all also embedding mode, all with incompatible bodies. 2. `provider` parameter on `_map_openai_embedding_to_bedrock_params` was accepted but never used. Fix: dropped from the signature and the call site. Also adds `_lookup_registry_mode` static helper (mirrors the one in the sibling batches transformer) so the registry try/except shape lives in one place instead of being inlined into the detector. 5 new tests pin the layered behavior: - registry mode=chat overrides the marker match (rejected) - registry mode=embedding + marker match (accepted) - registry silent + marker match for cross-region and ARN ids (accepted) - direct `_lookup_registry_mode` coverage across all return paths 33/33 tests in the file pass.
|
Addressed both Greptile flags in 5342298: Hardcoded model detection: Unused 5 new tests pin the layered behavior: registry chat overrides marker (rejected), registry embedding + marker (accepted), registry silent + marker for cross-region and ARN ids (accepted), direct helper coverage across all return paths. 33/33 tests pass. |
1 participant
Relevant issues
Reopens the use case from #15506 (closed by stale-bot, no maintainer engagement). AWS Bedrock supports batch inference for
amazon.titan-embed-text-v2:0viaCreateModelInvocationJobnatively; the gap is only on the LiteLLM side, whereBedrockFilesConfig._map_openai_to_bedrock_paramshad no embedding branch and an embedding JSONL would have silently been routed through the chat transformer.Supersedes #28862, which was opened against
main(wrong base) and accumulated force-pushes / close-reopen churn while we fixed up the base + addressed Greptile feedback. This PR is the clean, rebased version againstlitellm_oss_branch.Changes vs
litellm_oss_branchlitellm/llms/bedrock/files/transformation.pyBedrockFilesConfig._is_embedding_recordstatic helper detects whether a JSONL line is an embedding request. Strict precedence: expliciturl == \"/v1/embeddings\"-> embedding; any other non-emptyurl(e.g./v1/chat/completions) -> NOT embedding (trust the caller's signal); only fall back to body-shape (inputpresent ANDmessagesabsent) whenurlis missing.BedrockFilesConfig._is_titan_v2_embed_modelstatic helper acceptsamazon.titan-embed-text-v2:0,bedrock/amazon.titan-embed-text-v2:0,us.amazon.titan-embed-text-v2:0(cross-region inference profile), and ARN forms. Marker boundary check rejects lookalikes such astitan-embed-text-v20ortitan-embed-text-v2-experimental.BedrockFilesConfig._coerce_embedding_input_to_stringstatic helper normalizes the OpenAIinputfield into the single string Bedrock Titan v2 expects ininputText. Accepts string or single-element list; rejects None, multi-element lists, pre-tokenized inputs (List[int],List[List[int]]), and other types with actionable error messages._map_openai_embedding_to_bedrock_paramshelper builds the Bedrock InvokeModel body via the existingAmazonTitanV2Config._transform_request, mapping OpenAIdimensionsandencoding_formatthroughAmazonTitanV2Config.map_openai_paramsso this stays in sync with the synchronous/v1/embeddingspath._transform_openai_jsonl_content_to_bedrock_jsonl_contentnow dispatches per record to either the chat or the embedding transformer; the chat helper keeps its narrow contract.NotImplementedErrorwith a clear message until they get dedicated branches in follow-up PRs.tests/test_litellm/llms/bedrock/files/input_batch_embeddings.jsonlandexpected_bedrock_batch_embeddings.jsonlfixtures.TestBedrockFilesEmbeddingTransformationclass adds 17 mocked tests: happy-path round-trip; simple string input; dimensions + encoding_format mapping; body-shape fallback; single-element list unwrap; error paths (missinginput, multi-element list, unsupported model, pre-tokenized List[int] and List[List[int]]); mixed chat+embedding batch in same JSONL; model-id boundary check;_is_embedding_recordhelper coverage; ambiguous "both input and messages" routes to chat; explicit chat URL short-circuits to chat; other non-embedding URLs route to chat; isolated coverage of the input-coercion helper.Companion PR
The endpoint-propagation side (
BedrockBatchesConfig.transform_*_batch_responsereportingendpoint=\"/v1/embeddings\"correctly for embedding batches) lives in a separate PR to keep scope single-concern. That PR is also being filed under a v2 branch name.Pre-Submission checklist
tests/test_litellm/directory, Adding at least 1 test is a hard requirementmake test-unit(29/29 intests/test_litellm/llms/bedrock/files/test_bedrock_files_transformation.py)@greptileaireview after opening this PRType
New Feature