Skip to content

fix: preserve logprobs in OpenAIChatCompletionClient.create_stream#7967

Open
sohumt123 wants to merge 1 commit into
microsoft:mainfrom
sohumt123:fix-stream-logprobs
Open

fix: preserve logprobs in OpenAIChatCompletionClient.create_stream#7967
sohumt123 wants to merge 1 commit into
microsoft:mainfrom
sohumt123:fix-stream-logprobs

Conversation

@sohumt123

@sohumt123 sohumt123 commented Jul 17, 2026

Copy link
Copy Markdown

Why are these changes needed?

When streaming with logprobs enabled (extra_create_args={"logprobs": True, "top_logprobs": N}), OpenAIChatCompletionClient.create_stream returns a CreateResult whose logprobs is always None for text responses, while the non-streaming create() 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):

  1. The content fast-path skips logprobs collection. Any chunk carrying delta.content hits continue before the if 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 and CreateResult.logprobs is always None.
  2. The collection overwrites instead of accumulating. Even when the block is reached (tool-call streams), logprobs = [...] reassigns the list on every chunk, so only the last chunk's tokens survive.

Fix

  • Move the logprobs collection above the content fast-path so it runs for every non-empty choice.
  • Accumulate with extend across chunks instead of reassigning.
  • Semantics preserved: logprobs stays None when the stream carried none.

The construction of ChatCompletionTokenLogprob/TopLogprob entries is unchanged and matches the non-streaming path.

How it was tested

Added test_openai_chat_completion_client_create_stream_logprobs in python/packages/autogen-ext/tests/models/test_openai_model_client.py, mirroring the existing mocked-chunk streaming tests: three content chunks each carrying ChoiceLogprobs for their delta token plus a finish chunk. The test fails on current main (result.logprobs is None) and passes with this fix (all three tokens present, in order, with their top_logprobs).

pytest tests/models/test_openai_model_client.py shows no regressions (failure list identical before/after this change), and ruff 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

  • I've included any doc changes needed for https://microsoft.github.io/autogen/. (No doc changes needed: bug fix restoring documented CreateResult.logprobs behavior; no API change.)
  • I've added tests (if relevant) corresponding to the changes introduced in this PR.
  • I've made sure all auto checks have passed.

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.
@sohumt123

Copy link
Copy Markdown
Author

@microsoft-github-policy-service agree

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.

1 participant