Skip to content

fix(specs): stream every choice in full when OpenAISpec n > 1 - #793

Open
RaphaelFakhri wants to merge 1 commit into
Lightning-AI:mainfrom
RaphaelFakhri:fix/openai-stream-n-choices
Open

RaphaelFakhri wants to merge 1 commit into
Lightning-AI:mainfrom
RaphaelFakhri:fix/openai-stream-n-choices

Conversation

@RaphaelFakhri

Copy link
Copy Markdown

What does this PR do?

Fixes #792

With stream: true and n > 1, OpenAISpec merged the per-choice streams with azip, which stops at the shortest stream. Longer choices were cut off without any error, and the response still ended with finish_reason: "stop".

This PR:

  • Adds azip_longest to litserve.utils. It keeps yielding until every iterable is exhausted and fills finished ones with None. azip keeps its existing zip semantics.
  • Uses azip_longest in OpenAISpec.streaming_completion. Each chunk carries only the choices that still produce output, with their original index.
  • Computes the final usage from the last item of each choice, so a choice that finishes early still counts. For equal-length choices, the result is the same as before.
  • Adds test_openai_spec_streaming_choices_with_different_lengths (fails on main with ['a b ', 'a b '] == ['a b ', 'a b c d ']) and a unit test for azip_longest.
Before submitting
  • Was this discussed/agreed via a Github issue? (no need for typos and docs improvements)
  • Did you read the contributor guideline, Pull Request section?
  • Did you make sure to update the docs? (no docs change needed)
  • Did you write any new necessary tests?

How does this PR impact the user?

As a user serving an LLM through the OpenAI-compatible endpoint, you request several completions at once (n > 1) and stream them. Sampled completions have different lengths, so today every completion except the shortest is silently cut short. With this PR, each completion streams in full.

PR review

Anyone in the community is free to review the PR once the tests have passed.
If we didn't discuss your PR in GitHub issues there's a high chance it will not be merged.

Did you have fun?

Make sure you had fun coding 🙃

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

OpenAISpec streaming truncates choices when n > 1 and the choices have different lengths

1 participant