fix: preserve logprobs in OpenAIChatCompletionClient.create_stream#7967
Open
sohumt123 wants to merge 1 commit into
Open
fix: preserve logprobs in OpenAIChatCompletionClient.create_stream#7967sohumt123 wants to merge 1 commit into
sohumt123 wants to merge 1 commit into
Conversation
When streaming with logprobs enabled, CreateResult.logprobs was always None for text responses: the content fast-path's `continue` in the chunk loop ran before the logprobs collection block, and content-bearing chunks are exactly the ones that carry the logprobs for their delta tokens. Even when the block was reached (tool-call streams), it reassigned the list on every chunk instead of extending it, keeping only the last chunk's tokens. Collect logprobs before the content fast-path and accumulate them across chunks, so create_stream returns the same logprobs as the non-streaming create() path. Returns None as before when no logprobs were received. Adds a regression test with mocked chunks carrying ChoiceLogprobs.
Author
|
@microsoft-github-policy-service agree |
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Why are these changes needed?
When streaming with logprobs enabled (
extra_create_args={"logprobs": True, "top_logprobs": N}),OpenAIChatCompletionClient.create_streamreturns aCreateResultwhoselogprobsis alwaysNonefor text responses, while the non-streamingcreate()path returns full logprobs for the same response.Two issues in the chunk loop of
create_stream(python/packages/autogen-ext/src/autogen_ext/models/openai/_openai_client.py):delta.contenthitscontinuebefore theif choice.logprobs and choice.logprobs.content:block — but content-bearing chunks are exactly the ones that carry the logprobs for those tokens. So for text streams the block never runs andCreateResult.logprobsis alwaysNone.logprobs = [...]reassigns the list on every chunk, so only the last chunk's tokens survive.Fix
extendacross chunks instead of reassigning.logprobsstaysNonewhen the stream carried none.The construction of
ChatCompletionTokenLogprob/TopLogprobentries is unchanged and matches the non-streaming path.How it was tested
Added
test_openai_chat_completion_client_create_stream_logprobsinpython/packages/autogen-ext/tests/models/test_openai_model_client.py, mirroring the existing mocked-chunk streaming tests: three content chunks each carryingChoiceLogprobsfor their delta token plus a finish chunk. The test fails on currentmain(result.logprobs is None) and passes with this fix (all three tokens present, in order, with theirtop_logprobs).pytest tests/models/test_openai_model_client.pyshows no regressions (failure list identical before/after this change), andruff check/ruff format --check(pinned 0.4.8) pass on both touched files.Related issue number
None found — searched open/closed issues and PRs for "logprobs", "create_stream logprobs", "logprobs stream"; closest hits (#6192, #1044, #4213) are unrelated.
Checks
CreateResult.logprobsbehavior; no API change.)