feat(bedrock-batch): propagate embedding endpoint through batch responses - #28863
Conversation
|
@greptileai review |
|
|
Merging this PR will not alter performance
Comparing 
|
Greptile SummaryThis PR fixes a metadata accuracy bug in Bedrock batch processing where both
Confidence Score: 5/5Safe to merge — the change is additive metadata fidelity on non-critical batch response objects, with no effect on actual AWS wire calls or existing chat-batch behaviour. The transformation logic is narrowly scoped to two response-shaping methods. The registry-first / substring-fallback / chat-default chain preserves prior behaviour for all models that are not recognised as embeddings, so there is no regression risk for existing users. The only real gap is a test whose two inputs happen to agree rather than conflict, which does not expose a production defect. No files require special attention; the test file has one low-priority test improvement opportunity noted.
|
| Filename | Overview |
|---|---|
| litellm/llms/bedrock/batches/transformation.py | Adds _lookup_registry_mode and _infer_openai_endpoint_from_model_id helpers; replaces hardcoded /v1/chat/completions in both response transformers with a registry-first, substring-fallback inference. Existing chat batch behaviour is preserved; embedding batches now receive the correct endpoint. |
| tests/test_litellm/llms/bedrock/batches/test_endpoint_inference.py | New unit tests covering the inference helper, create-batch, and retrieve-batch paths. One test (test_create_response_uses_original_request_endpoint_when_set) uses signals that agree, so it does not verify the explicit-endpoint-over-inference priority it claims to test. |
Reviews (6): Last reviewed commit: "refactor(bedrock-batch): extract registr..." | Re-trigger Greptile
![greptile-apps[bot]](https://avatars.githubusercontent.com/in/867647?s=60&v=4){kind=small}
Greptile SummaryThis PR fixes a metadata fidelity issue where both
Confidence Score: 3/5Safe to merge for existing chat batches; embedding classification could silently regress for any future Bedrock embedding model whose id doesn't contain "embed" The core logic in litellm/llms/bedrock/batches/transformation.py — specifically the new
|
| Filename | Overview |
|---|---|
| litellm/llms/bedrock/batches/transformation.py | Adds _infer_openai_endpoint_from_model_id heuristic (string match on "embed") and wires it into both create and retrieve response transformers; hardcoded string detection violates the project rule to use model_prices_and_context_window.json + get_model_info |
| tests/test_litellm/llms/bedrock/batches/test_endpoint_inference.py | New mock-only test file covering the inference helper and both response transformers; 19 tests with good edge-case coverage, follows existing test-file patterns |
Reviews (2): Last reviewed commit: "feat(bedrock-batch): propagate embedding..." | Re-trigger Greptile
Codecov Report❌ Patch coverage is 📢 Thoughts on this report? Let us know! |
|
Addressed Greptile's registry-vs-substring concern in d462843852:
This means new Bedrock embedding models get the right endpoint as soon as their entry is added to @greptileai review |
|
Retargeted base from |
…nses Changes vs main: - BedrockBatchesConfig.transform_create_batch_response now prefers the endpoint set on `litellm_params["original_batch_request"]` and falls back to a model-id heuristic when none is present, instead of unconditionally defaulting to `/v1/chat/completions`. - BedrockBatchesConfig.transform_retrieve_batch_response replaces the hardcoded `/v1/chat/completions` with the same heuristic, fed from the caller-provided `model` argument or the `modelId` AWS returns on GetModelInvocationJob. This is necessary because the retrieve path has no original request to read from. - New static helper `_infer_openai_endpoint_from_model_id` returns `/v1/embeddings` when the lowercased model id contains "embed" (covers Titan v2 text, Titan multimodal, Cohere embed, Nova multimodal embeddings, ARN forms thereof), else `/v1/chat/completions`. Best-effort by design: chat models without "embed" in the name keep the same endpoint they reported before. - 19 new mocked unit tests across direct helper coverage, create-batch endpoint resolution (explicit override wins, model-id fallback, no-signal default), and retrieve-batch resolution (model arg wins, AWS modelId fallback, no-id default). 28 existing batch tests pass unchanged.
Greptile-flagged convention: model-capability flags should come from `model_prices_and_context_window.json` (via `litellm.get_model_info`), not from hardcoded substring matches. Every current Bedrock embedding model is registered there with `"mode": "embedding"`, so future models get classified automatically once their entry lands. Changes vs previous commit: - `_infer_openai_endpoint_from_model_id` now consults `get_model_info` first. If the registry returns `mode == "embedding"`, route to `/v1/embeddings`. If it returns any other mode, route to `/v1/chat/completions` (registry wins even if the model id name contains "embed"). - If `get_model_info` raises (cross-region inference profile ids like `us.amazon.titan-embed-text-v2:0`, Bedrock ARN forms, or unreleased models the registry can't normalize) or returns no mode, we fall back to the previous substring heuristic. This keeps the existing cross-region and ARN test cases passing without regression. - 5 new tests pin the layered behavior: registry classifies known embed ids; registry chat mode wins over substring "embed" in the name; substring fallback engages when registry raises (with and without embed marker); empty-info dict falls through to substring path. 33 batch tests pass (24 in this file + 9 in test_batch_metadata_sanitization).
…lper Splits `_lookup_registry_mode` out of `_infer_openai_endpoint_from_model_id` so the registry-vs-fallback split is unit-testable in isolation and so future call sites (e.g. cost calculation, billing routing) can reuse the same defensive lookup without duplicating the try/except shape. - New `_lookup_registry_mode` returns the registry's `mode` string when `get_model_info` resolves and the entry has a non-empty string mode, else `None`. Catches the three real-world failure modes: registry raises, registry returns a non-dict, registry returns a dict without (or with empty / non-string) `mode`. - `_infer_openai_endpoint_from_model_id` now reads from the helper and the dispatch is linear: registry-says-embedding -> embeddings, registry-says-anything-else -> chat, registry-silent -> substring fallback -> chat default. - 6 new isolated tests pin the helper's behavior: returns mode for known ids, returns None on each of the four failure modes (raise, non-dict, missing mode, empty string mode, non-string mode). 30/30 tests in the file still pass.
57c4645 to
dcc81d6
Compare
|
Superseded by #28867 (rebased clean against |
2 participants
Relevant issues
Companion to #28862 (Titan v2 embedding JSONL input transform). That PR makes embedding JSONL records produce a valid Bedrock modelInput; this PR makes the resulting
LiteLLMBatchobject correctly reportendpoint=\"/v1/embeddings\"to OpenAI clients that introspect it.Reopens the use case from #15506.
Problem
BedrockBatchesConfig.transform_create_batch_responseandtransform_retrieve_batch_responseboth hardcodedendpoint=\"/v1/chat/completions\". For an embedding batch the create path could readoriginal_batch_request.endpointbut defaulted to chat-completions when the caller didn't set it; the retrieve path had no fallback at all because it has no original request to read from.Changes vs
mainlitellm/llms/bedrock/batches/transformation.py_infer_openai_endpoint_from_model_idreturns/v1/embeddingsiff the lowercased Bedrock model id contains "embed" (Titan v2 text, Titan multimodal, Cohere Embed, Nova Multimodal Embeddings, ARN forms thereof). Otherwise defaults to/v1/chat/completionsso the existing chat batches keep the endpoint they always reported.transform_create_batch_responsenow prefersoriginal_batch_request.endpointand falls back to the heuristic instead of the chat-completions default.transform_retrieve_batch_responsereplaces the hardcoded value with the same heuristic, fed from the caller-suppliedmodelargument or themodelIdAWS returns on GetModelInvocationJob.tests/test_litellm/llms/bedrock/batches/test_endpoint_inference.py(new)modelarg wins over AWS modelId, default-chat when neither present.What this does NOT change
CreateModelInvocationJobis endpoint-agnostic at the wire level: only themodelIdand the JSONLmodelInputschema vary. So this PR is metadata fidelity, not feature unlock; feat(bedrock-batch): route /v1/embeddings JSONL to Titan v2 modelInput #28862 is the schema fix.modelOutput-> OpenAI batch embedding response shape) is intentionally not in this PR. It will be a follow-up so the schema-specific risk stays isolated.Pre-Submission checklist
tests/test_litellm/directory, Adding at least 1 test is a hard requirementmake test-unit(19 new + 28 existing batch tests pass; the 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