fix: restore chat completion multi-choice support

[nabinchha]

📋 Summary

Restores chat completion multi-choice compatibility by reintroducing the n request field and preserving all returned choices in the canonical response. Existing single-choice callers can continue using response.message, while callers that request multiple completions can use response.choices.

🔗 Related Issue

Fixes #620

🔄 Changes

• Add ChatCompletionRequest.n and export a new ChatCompletionChoice response type.
• Preserve every returned chat completion choice in ChatCompletionResponse.choices, while keeping response.message as the first choice for existing callers.
• Parse all OpenAI-compatible chat completion choices instead of discarding everything after index 0.
• Forward n through ModelFacade.completion() and OpenAI-compatible transport bodies, while excluding it from Anthropic request forwarding.
• Strip n from generate() and agenerate() before delegating to completion because those APIs return one parsed result.
• Add tests for request forwarding, multi-choice parsing, async parsing, compatibility access, Anthropic n exclusion, and generate(..., n=...) stripping.

🔍 Attention Areas

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

• types.py — updates the canonical chat completion response contract while preserving response.message.
• facade.py — completion(..., n=...) exposes multiple choices, while generate(..., n=...) strips n because it returns a single parsed result.

🧪 Testing

• make test passes (not run; focused model suite was run instead)
• PYTHONPATH=packages/data-designer-config/src:packages/data-designer-engine/src:packages/data-designer/src uv run --group dev pytest packages/data-designer-engine/tests/engine/models -q passes (532 passed)
• uv run --group dev ruff check ... passes for touched files
• Unit tests added/updated
• E2E tests added/updated (N/A — no E2E surface)

✅ Checklist

• Follows commit message conventions
• Commits are signed off (DCO)
• Architecture docs updated (N/A — no architecture docs needed)

[greptile-apps]

Greptile Summary

This PR restores multi-choice chat completion support by adding ChatCompletionRequest.n, preserving all returned choices in a new ChatCompletionResponse.choices list, and stripping n in generate/agenerate paths that only consume a single parsed result. Backward compatibility is maintained via response.message (first choice) and the new __post_init__ fallback for callers that construct ChatCompletionResponse without an explicit choices argument.

• types.py — introduces ChatCompletionChoice, adds n to ChatCompletionRequest, adds choices field and messages property to ChatCompletionResponse, with __post_init__ ensuring single-choice back-compat.
• parsing.py — refactors the response-parsing path into per-choice helpers (parse_chat_completion_choice / aparse_chat_completion_choice), aggregating images across all choices for usage tracking.
• facade.py — forwards n through completion(); generate()/agenerate() strip n (both top-level and nested in extra_body) via the new _drop_multi_choice_request_fields helper before delegating to completion(allow_multiple_choices=False).

Confidence Score: 5/5

The change is safe to merge; all code paths produce consistent ChatCompletionResponse objects and existing single-choice callers are unaffected.

The post_init guard ensures backward compatibility for any code that constructs ChatCompletionResponse without an explicit choices argument. The _drop_multi_choice_request_fields helper correctly strips n from both top-level kwargs and nested extra_body after consolidate_kwargs runs, and the allow_multiple_choices pop prevents duplicate-keyword errors. The Anthropic exclusion list update is correctly scoped. Test coverage is thorough, spanning sync/async parsing, request forwarding, and config-sourced n stripping.

No files require special attention.

Important Files Changed

Filename	Overview
packages/data-designer-engine/src/data_designer/engine/models/clients/types.py	Adds ChatCompletionChoice dataclass, n field to ChatCompletionRequest, and choices/messages to ChatCompletionResponse with post_init backward-compat guard; logic is correct in all construction paths.
packages/data-designer-engine/src/data_designer/engine/models/clients/parsing.py	Refactors response parsing into per-choice helpers; normalize_choice_list handles None/non-list gracefully; image count aggregation across all choices is consistent with the multi-choice intent.
packages/data-designer-engine/src/data_designer/engine/models/facade.py	Correctly threads n through completion() and strips it for generate()/agenerate() via _drop_multi_choice_request_fields; allow_multiple_choices is popped from kwargs before being passed explicitly, preventing duplicate-keyword errors.
packages/data-designer-engine/src/data_designer/engine/models/clients/adapters/anthropic.py	Adds n to the Anthropic exclusion list so it is never forwarded to an API that does not support it; change is minimal and correct.
packages/data-designer-engine/tests/engine/models/test_facade.py	New tests cover n forwarding in completion(), n stripping in generate()/agenerate() for top-level, extra_body, and config-sourced n values; async variants are included.
packages/data-designer-engine/tests/engine/models/clients/test_parsing.py	Tests cover single-choice compat, multi-choice preservation (sync and async), n forwarding into transport body, and the ChatCompletionResponse.messages property; good coverage of the changed parsing paths.
packages/data-designer-engine/tests/engine/models/clients/test_openai_compatible.py	Verifies n is forwarded into the OpenAI-compatible request body; straightforward and correct.
packages/data-designer-engine/tests/engine/models/clients/test_anthropic.py	Verifies n is excluded from Anthropic payloads; minimal and correct addition to existing exclusion test.
packages/data-designer-engine/src/data_designer/engine/models/clients/init.py	Exports the new ChatCompletionChoice type alongside the existing public API; no issues.

Sequence Diagram

sequenceDiagram
    participant Caller
    participant ModelFacade
    participant ModelClient
    participant Parser

    Note over Caller,Parser: completion() path — n is preserved
    Caller->>ModelFacade: "completion(messages, n=4)"
    ModelFacade->>ModelFacade: "consolidate_kwargs(n=4)"
    ModelFacade->>ModelFacade: "_build_chat_completion_request(n=4)"
    ModelFacade->>ModelClient: client.completion(request)
    ModelClient->>Parser: parse_chat_completion_response(raw)
    Parser->>Parser: normalize_choice_list(raw["choices"])
    loop for each choice [0..3]
        Parser->>Parser: parse_chat_completion_choice(choice)
    end
    Parser-->>ModelClient: "ChatCompletionResponse(message=choices[0].message, choices=[...])"
    ModelClient-->>ModelFacade: response
    ModelFacade-->>Caller: response (response.message + response.choices)

    Note over Caller,Parser: generate() path — n is stripped
    Caller->>ModelFacade: "generate(prompt, n=4)"
    ModelFacade->>ModelFacade: "consolidate_kwargs(n=4)"
    ModelFacade->>ModelFacade: _drop_multi_choice_request_fields (removes n)
    ModelFacade->>ModelFacade: "completion(..., allow_multiple_choices=False)"
    ModelFacade->>ModelClient: "client.completion(request, n=None)"
    ModelClient-->>ModelFacade: ChatCompletionResponse(single choice)
    ModelFacade->>ModelFacade: parser(response.message.content)
    ModelFacade-->>Caller: (parsed_output, messages)

Loading

_{Reviews (6): Last reviewed commit: "Merge branch 'main' into nmulepati/fix-6..." | Re-trigger Greptile}

[github-actions]

Review: PR #672 — fix: restore chat completion multi-choice support

Summary

The PR restores chat-completion multi-choice support (issue #620) by:

• Adding n to ChatCompletionRequest and forwarding it through ModelFacade.completion() and OpenAI-compatible transport bodies; explicitly excluded from Anthropic forwarding.
• Adding a new ChatCompletionChoice dataclass and a choices: list[ChatCompletionChoice] field on ChatCompletionResponse. response.message is preserved as the first choice for backward compatibility via __post_init__.
• Refactoring parse_chat_completion_response / aparse_chat_completion_response to iterate every returned choice (instead of discarding everything after index 0), with shared helpers parse_chat_completion_choice, parse_assistant_message, parse_choice_index, parse_choice_finish_reason, and normalize_choice_list.
• Adding tests for request n forwarding, multi-choice parsing, single-choice compatibility, and generate(..., n=...) forwarding.

The change is well-scoped and backward-compatible. Existing single-choice callers continue to work via response.message; new callers can iterate response.choices.

Findings

Correctness

• generate(..., n=...) is a footgun. The n keyword is now on the _COMPLETION_REQUEST_FIELDS allowlist (facade.py:83), so callers can pass it to generate() — and the test test_generate_forwards_n_to_completion exercises that. But generate() only ever consumes completion_response.message.content (facade.py:372, facade.py:475); the remaining n-1 choices are silently dropped. Callers pay for n× output tokens for no benefit. Options worth considering:
  • Reject n>1 in generate() with a clear error, or
  • Document explicitly in the docstring that n is forwarded but only the first choice is parsed (the PR description notes this, but the public API doesn't), or
  • Pop n out of kwargs in generate() before delegating, since generate semantics don't currently accommodate it.
• Anthropic silently drops n. n is added to AnthropicClient._TRANSPORT_EXCLUDE (correct — Anthropic Messages API doesn't accept it), so a caller who passes n=4 to an Anthropic-backed model gets a single response with no warning. Consider logging at debug or warning level when n>1 is supplied to a provider that doesn't honor it; otherwise this is an easy-to-miss surprise for users switching providers.
• generated_images semantics changed for multi-choice. Previously extract_usage received len(images) from the first choice only; now it receives sum(len(c.message.images) for c in choices) (parsing.py:41, parsing.py:55). Identical for n=1 (the only currently exercised path) but worth confirming this aggregation is what billing/usage downstream expects when multi-choice image generation is added.
• No validation of n. A caller can pass n=0 or a negative integer; the dataclass accepts it and forwards. Fine if the upstream provider is the source of truth, but a small __post_init__ check on ChatCompletionRequest would fail fast.

API design

• ChatCompletionResponse.message and choices[0].message can drift. If both are passed at construction (as some external callers might), __post_init__ only backfills choices when empty — it does not assert consistency. The internal parsers always keep them in sync, but a stricter invariant (assert/normalize) would prevent subtle bugs later. Lightweight: in __post_init__, after backfill, assert self.message is self.choices[0].message.
• messages property name collides conceptually with the request messages field. ChatCompletionResponse.messages returns a list[AssistantMessage], while ChatCompletionRequest.messages is the conversation history. assistant_messages or just steering callers toward choices would reduce ambiguity. Minor.

Tests

• Coverage of the sync path is solid: request forwarding (test_chat_completion_request_n_is_forwarded_into_body), multi-choice parsing (test_parse_chat_completion_response_preserves_all_choices), backward-compat single-message (test_chat_completion_response_exposes_choices_for_single_message), and facade forwarding (test_completion_forwards_n_to_request, test_generate_forwards_n_to_completion).
• No async multi-choice parsing test. aparse_chat_completion_response and aparse_chat_completion_choice mirror the sync logic but are uncovered. A 1-line async equivalent of test_parse_chat_completion_response_preserves_all_choices would close that gap.
• No test asserting n is excluded from Anthropic payloads. Adding n to _TRANSPORT_EXCLUDE is a correctness-critical change (Anthropic would 400 on it); a test that builds an AnthropicClient payload from a request with n=4 and asserts "n" not in payload would lock that in.
• No test for normalize_choice_list edge cases. Specifically, the "single non-list choice" branch (return [raw_choices]) is untested. Worth a unit test if any provider actually returns that shape — otherwise consider deleting the branch and just handling None + list.

Style / conventions

• Follows project conventions: from __future__ import annotations, modern type syntax (int | None), absolute imports, no comments on obvious code. Good.
• parse_chat_completion_choice and aparse_chat_completion_choice differ only in sync vs async image extraction — acceptable duplication given the project's existing sync/async split pattern (extract_images_from_chat_message vs aextract_images_from_chat_message).

Performance / security

• No performance impact — multi-choice parsing iterates choices once. No new I/O.
• No security implications.

Structural Impact

No pre-computed structural impact analysis was available (/tmp/structural-impact-672.md not present). Manual assessment: the change is contained to data_designer.engine.models.clients and data_designer.engine.models.facade. Import direction (interface → engine → config) is preserved; no new cross-package edges. The ChatCompletionResponse shape change is additive (new field with default factory + backfilling __post_init__), so external consumers reading response.message are unaffected. Risk: low.

Verdict

Approve with minor follow-ups. The core change is clean, backward-compatible, and well-tested for the sync OpenAI-compatible path. The two items I'd most like to see addressed before merge:

1. Decide on the generate(..., n=...) story — either reject, document, or strip n in generate(). Tests currently encode the silent-drop behavior, which makes it harder to change later.
2. Add a test asserting n is excluded from the Anthropic payload, since that's the only thing preventing a 400 when an Anthropic-backed model is used.

Nice-to-have: async multi-choice parsing test, soft validation of n>=1, and a clarifying note on ChatCompletionResponse.messages vs choices.

[nabinchha]

Addressed the feedback in b2e22f7f:

• generate() and agenerate() now strip n before delegating to completion, since those APIs only expose one parsed result.
• Replaced the generate(..., n=...) forwarding test with sync/async tests that assert n is dropped.
• Added async multi-choice parsing coverage.
• Added Anthropic coverage confirming n is excluded from the payload.

Validation:

• PYTHONPATH=packages/data-designer-config/src:packages/data-designer-engine/src:packages/data-designer/src uv run --group dev pytest packages/data-designer-engine/tests/engine/models -q → 532 passed
• uv run --group dev ruff check ... on touched files → passed

[johnnygreco]

I found one remaining issue after the direct n stripping update: generate()/agenerate() can still forward n when it comes from model/provider configuration or nested extra_body, so the single-result APIs may still request multiple choices and discard all but the first.

Validation I ran locally on the reviewed diff:

• uv run pytest packages/data-designer-engine/tests/engine/models/clients/test_parsing.py packages/data-designer-engine/tests/engine/models/clients/test_openai_compatible.py packages/data-designer-engine/tests/engine/models/clients/test_anthropic.py packages/data-designer-engine/tests/engine/models/test_facade.py -> 188 passed
• git diff --check 71997624b31045c8b0268055f4a14eb02acce905 -> passed

[johnnygreco]

Approved. I verified the follow-up fix addresses the remaining multi-choice sanitization issue: generate()/agenerate() force allow_multiple_choices=False, sanitization runs after consolidate_kwargs(), configured and nested extra_body.n are stripped, and completion(..., n=...) still preserves multi-choice behavior. Focused model tests passed locally (191 passed).