fix: preserve extra_body for LiteLLM to avoid UnsupportedParamsError (#409)

[nabinchha]

📋 Summary

TransportKwargs.from_request() flattened extra_body keys into top-level kwargs, causing LiteLLM to reject provider-specific params like reasoning_effort via its per-provider allowlist validation. This adds a flatten_extra_body flag so the LiteLLM bridge can preserve extra_body as a distinct kwarg that LiteLLM forwards without validation.

🔄 Changes

🐛 Fixed

• LiteLLMBridgeClient now passes flatten_extra_body=False to TransportKwargs.from_request() across all 6 methods (completion, embedding, image generation — both sync and async), so extra_body keys like reasoning_effort are preserved as a nested dict rather than flattened into top-level kwargs where LiteLLM rejects them

🔧 Changed

• TransportKwargs.from_request() gained a flatten_extra_body parameter — True (default) merges extra_body into the top-level body dict for future native adapters, False preserves it as a distinct extra_body kwarg for LiteLLM

🧪 Tests

• Updated existing bridge test assertions to expect extra_body={"foo": "bar"} instead of flattened foo="bar"
• Added test_completion_passes_extra_body_as_distinct_kwarg to explicitly verify reasoning_effort flows through as extra_body
• Added test_extra_body_preserved_when_flatten_disabled for the new TransportKwargs code path

🔍 Attention Areas

⚠️ Reviewers: Please pay special attention to the following:

• types.py — TransportKwargs.from_request() — New branching logic for flatten_extra_body. Default True preserves existing behavior for future native adapters; False is the LiteLLM-specific path.

---

🤖 Generated with AI

Closes #409

[greptile-apps]

Greptile Summary

This PR fixes a bug where extra_body fields like reasoning_effort were flattened into top-level LiteLLM kwargs, causing UnsupportedParamsError due to LiteLLM's per-provider allowlist validation. The fix adds a flatten_extra_body flag to TransportKwargs.from_request() and sets it to False across all 6 LiteLLMBridgeClient methods so that extra_body is preserved as a nested dict that LiteLLM can forward without validation.

Key changes:

• TransportKwargs.from_request() gains flatten_extra_body: bool = True — default preserves backward compatibility for future native adapters; False is the LiteLLM-specific path
• All 6 LiteLLMBridgeClient methods pass flatten_extra_body=False, meaning any field in extra_body (including n for image generation) is now forwarded as extra_body={...} rather than as individual top-level kwargs
• Two new unit tests cover the flatten_extra_body=False path in test_parsing.py, and one new integration-level test in test_litellm_bridge.py verifies reasoning_effort flows through correctly
• The _META_FIELDS exclusion in _collect_optional_fields already prevents key collisions between optional_fields and the extra_body dict, so the new branching logic is safe

Confidence Score: 4/5

• This PR is safe to merge; the core logic is correct with one noteworthy behavioral side-effect worth confirming before ship.
• The TransportKwargs change is logically sound — _META_FIELDS already excludes extra_body from optional_fields, so there are no key collisions. The flatten_extra_body=False path is fully tested and the True default preserves backward compatibility. Score is 4 rather than 5 because the uniform application of flatten_extra_body=False to image-generation calls changes how n (and any other extra_body field) is forwarded to LiteLLM's image_generation / aimage_generation endpoints — previously as a top-level kwarg, now nested inside extra_body — and while LiteLLM supports this pattern for completions very explicitly, the same guarantee for image-generation endpoints is worth a quick smoke-test against a real provider before merging.
• packages/data-designer-engine/src/data_designer/engine/models/clients/adapters/litellm_bridge.py — generate_image / agenerate_image paths where extra_body params (e.g., n) are now forwarded as extra_body={...} to the LiteLLM image generation router instead of as top-level kwargs.

Important Files Changed

Filename	Overview
packages/data-designer-engine/src/data_designer/engine/models/clients/types.py	Adds flatten_extra_body boolean parameter to TransportKwargs.from_request(). Default True preserves existing flatten behavior; False preserves extra_body as a nested dict. Logic is correct: _META_FIELDS already excludes extra_body/extra_headers from optional_fields, preventing any key collisions. Edge cases (None, empty dict) handled consistently.
packages/data-designer-engine/src/data_designer/engine/models/clients/adapters/litellm_bridge.py	All 6 bridge methods (completion, acompletion, embeddings, aembeddings, generate_image, agenerate_image) updated to pass flatten_extra_body=False. This is a behavioral change: any field previously in extra_body (e.g., n for image generation) is now passed as extra_body={...} to LiteLLM instead of being flattened to a top-level kwarg — which is the intended fix.
packages/data-designer-engine/src/data_designer/engine/models/facade.py	Comment-only change updating the inline documentation for TransportKwargs behavior. No functional changes.
packages/data-designer-engine/tests/engine/models/clients/test_litellm_bridge.py	Existing assertions updated to reflect extra_body={"foo": "bar"} instead of flattened foo="bar". Two new tests added: one verifying reasoning_effort flows through as a nested extra_body kwarg, and corrected parameter ordering in other assertions. All test logic is sound and accurately validates the new behavior.
packages/data-designer-engine/tests/engine/models/clients/test_parsing.py	Two new unit tests added for the flatten_extra_body=False path: one verifying preservation of multi-key extra_body, and one verifying an empty extra_body is not injected. Existing default-path tests are unmodified and remain correct.

Sequence Diagram

sequenceDiagram
    participant Caller
    participant LiteLLMBridgeClient
    participant TransportKwargs
    participant LiteLLMRouter

    Caller->>LiteLLMBridgeClient: completion(request)<br/>extra_body={"reasoning_effort": "high"}

    LiteLLMBridgeClient->>TransportKwargs: from_request(request, flatten_extra_body=False)
    note over TransportKwargs: optional_fields = {temperature, ...}<br/>(extra_body excluded via _META_FIELDS)
    note over TransportKwargs: body = {temperature: ...}<br/>body["extra_body"] = {"reasoning_effort": "high"}
    TransportKwargs-->>LiteLLMBridgeClient: TransportKwargs(body={..., extra_body={...}}, headers={})

    LiteLLMBridgeClient->>LiteLLMRouter: completion(model=..., messages=...,<br/>extra_headers=None, **transport.body)<br/>→ extra_body={"reasoning_effort":"high"} as kwarg

    note over LiteLLMRouter: LiteLLM forwards extra_body<br/>to provider WITHOUT allowlist validation
    LiteLLMRouter-->>LiteLLMBridgeClient: response
    LiteLLMBridgeClient-->>Caller: ChatCompletionResponse

Loading

_{Last reviewed commit: 8cfdce3}

[andreatgretel]

nice fix. maybe worth adding one regression test a bit closer to the original config / health check path too, just to lock this in end to end.

[andreatgretel]

packages/data-designer-engine/src/data_designer/engine/models/clients/types.py:121

- ``body``: API-level keyword arguments with ``extra_body`` keys merged
  into the top level (mirroring how LiteLLM flattens them).

small doc nit: could we update this comment too? it still describes the old flattened extra_body behavior.

[andreatgretel]

packages/data-designer-engine/src/data_designer/engine/models/facade.py:53

# `TransportKwargs` (extra_body is flattened into the request body, extra_headers
# are forwarded as HTTP headers).  They are NOT regular model parameters.

tiny follow-up: this note still says TransportKwargs flattens extra_body. might be nice to keep it in sync with the new behavior here.

[nabinchha]

nice fix. maybe worth adding one regression test a bit closer to the original config / health check path too, just to lock this in end to end.

good catch! updated stale docs in 8cfdce3