feat(bedrock-batch): route /v1/embeddings JSONL to Titan v2 modelInput - #28862
Conversation
|
@greptileai review |
Merging this PR will not alter performance
Comparing 
|
Codecov Report❌ Patch coverage is 📢 Thoughts on this report? Let us know! |
Greptile SummaryThis PR adds JSONL batch embedding support for
Confidence Score: 4/5Safe to merge for the input-transform path; the hardcoded Titan v2 model-ID detection noted in earlier review threads remains unaddressed and will require a code change when AWS adds new Titan v2 variants. The transformation logic is correct and well-tested: URL-precedence routing works, input coercion covers all OpenAI input shapes, and the AmazonTitanV2Config delegation keeps the batch path in sync with the synchronous embedding path. The outstanding concern (hardcoded model identity check rather than a data-driven lookup) was raised in earlier threads and has not yet been resolved. litellm/llms/bedrock/files/transformation.py — specifically _is_titan_v2_embed_model and _TITAN_V2_EMBED_MODEL_MARKER
|
| Filename | Overview |
|---|---|
| litellm/llms/bedrock/files/transformation.py | Core change: adds embedding routing and Titan v2 transformation helpers. URL-precedence logic in _is_embedding_record is correct; _coerce_embedding_input_to_string covers all input shapes including pre-tokenized rejection. _is_titan_v2_embed_model hard-codes model-ID detection rather than querying get_model_info/model_prices_and_context_window.json per project convention (flagged in previous review threads). |
| tests/test_litellm/llms/bedrock/files/test_bedrock_files_transformation.py | 14 new mock-only tests covering happy-path, error paths, mixed batch, boundary checks, and helper isolation. No existing tests modified. |
| tests/test_litellm/llms/bedrock/files/input_batch_embeddings.jsonl | Input fixture with 3 records covering plain string, dimensions override, and base64 encoding_format with single-element list. |
| tests/test_litellm/llms/bedrock/files/expected_bedrock_batch_embeddings.jsonl | Expected Bedrock output fixture; correctly maps base64 to embeddingTypes: ["binary"] and dimensions to dimensions, matching AmazonTitanV2Config.map_openai_params behavior. |
Reviews (6): Last reviewed commit: "refactor(bedrock-batch): extract embeddi..." | Re-trigger Greptile
![greptile-apps[bot]](https://avatars.githubusercontent.com/in/867647?s=60&v=4){kind=small}
Greptile SummaryThis PR adds input-side batch-embedding support for Amazon Titan Text Embeddings V2 on AWS Bedrock, routing
Confidence Score: 3/5The core Titan v2 transformation is correct, but the record-routing helper has a gap where a non-embedding URL combined with certain body shapes can misroute a record to the embedding transformer. The routing guard in litellm/llms/bedrock/files/transformation.py — specifically the
|
| Filename | Overview |
|---|---|
| litellm/llms/bedrock/files/transformation.py | Adds embedding routing helpers and Titan v2 transformer; the _is_embedding_record fallback misroutes records that carry an explicit non-embedding url with an input field, and model detection is hardcoded rather than data-driven. |
| tests/test_litellm/llms/bedrock/files/test_bedrock_files_transformation.py | Adds 14 new mocked unit tests for the embedding transformation path; test coverage is thorough but the url-present-non-embedding misrouting case is not covered. |
| tests/test_litellm/llms/bedrock/files/input_batch_embeddings.jsonl | New fixture file providing representative Titan v2 embedding batch inputs; content is correct and matches the expected output fixture. |
| tests/test_litellm/llms/bedrock/files/expected_bedrock_batch_embeddings.jsonl | New fixture with expected Bedrock-format output records; values are consistent with the Titan v2 InvokeModel schema. |
Reviews (2): Last reviewed commit: "feat(bedrock-batch): route /v1/embedding..." | Re-trigger Greptile
|
Addressed Greptile's routing concern in defb2e5d21:
On the second comment (hardcoded @greptileai review |
|
|
|
Retargeted base from |
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.
23f7bdf to
8105b7e
Compare
|
Superseded by #28865 (rebased clean against |
2 participants
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.Changes vs
mainlitellm/llms/bedrock/files/transformation.pyBedrockFilesConfig._is_embedding_recordstatic helper detects whether a JSONL line is an embedding request. Precedence: expliciturl == \"/v1/embeddings\"first, then body-shape fallback (inputpresent ANDmessagesabsent).BedrockFilesConfig._is_titan_v2_embed_modelstatic helper accepts "amazon.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 as "titan-embed-text-v20" or "titan-embed-text-v2-experimental"._map_openai_embedding_to_bedrock_paramshelper builds the Bedrock InvokeModel body via the existingAmazonTitanV2Config._transform_request, mapping OpenAI `dimensions` and `encoding_format` throughAmazonTitanV2Config.map_openai_paramsso this stays in sync with the synchronous `/v1/embeddings` path._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.List[int],List[List[int]]) and multi-element string lists are explicitly rejected so callers emit one JSONL line per embedding instead of relying on us to fan out or concatenate.tests/test_litellm/llms/bedrock/files/input_batch_embeddings.jsonlandexpected_bedrock_batch_embeddings.jsonlfixtures.TestBedrockFilesEmbeddingTransformationclass adds 14 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" record routing to chat.Why split into two PRs
The companion PR (incoming) replaces the hardcoded
endpoint=\"/v1/chat/completions\"inBedrockBatchesConfig.transform_create_batch_response/transform_retrieve_batch_responseand adds the output-file transform that maps Titan v2modelOutputback to the OpenAI embeddings batch shape. Splitting keeps each PR single-concern per the template, makes review easier, and lets the input-transform land independently if the output-transform needs iteration.Out of scope (will follow up)
transform_create_batch_response/transform_retrieve_batch_responseendpoint propagation (companion PR).Pre-Submission checklist
tests/test_litellm/directory, Adding at least 1 test is a hard requirementmake test-unit(26/26 intests/test_litellm/llms/bedrock/files/test_bedrock_files_transformation.py; baseline 71 fails + 19 errors in the broader bedrock suite are pre-existing onmaindue to missing botocore in the test env, unchanged on this branch)@greptileaireview after opening this PRType
New Feature