Repository navigation
fix(streaming): keep n>1 choices separate in stream_chunk_builder - #36208
vineethsaivs wants to merge 9 commits into
Conversation
 |
![greptile-apps[bot]](https://avatars.githubusercontent.com/in/867647?s=60&v=4){kind=small}
| return TextCompletionResponse(**response) | ||
|
|
||
|
|
||
| def _streaming_choice_index(choice: Any) -> int | None: |
There was a problem hiding this comment.
The new helpers accept Any and list[Any], preventing static validation of supported chunk shapes and increasing the repository's explicit-Any budget; use strongly typed chunk and choice unions throughout these helpers.
Context Used: CLAUDE.md (source)
Note: If this suggestion doesn't match your team's coding style, reply to this and let me know. I'll remember it for next time!
![veria-ai[bot]](https://avatars.githubusercontent.com/in/1723776?s=60&v=4){kind=small}
PR overviewAll previously flagged issues have been addressed. No open security concerns remain on this pull request. Security reviewNo open security issues remain on this pull request. Fixed/addressed: 1 · PR risk: 0/10 |
Codecov Report❌ Patch coverage is
📢 Thoughts on this report? Let us know! |
|
The bot is right about the mechanism, and I would rather flag it clearly than tick a box past it. It is also wider than Noma, so here is what I found
content = None
for choice in response.choices:
if isinstance(choice, litellm.Choices) and choice.message.content:
content = choice.message.content
breakThat is not new, and this PR does not touch it. What this PR changes is how much that loop happens to cover on the streaming path. Before, a streamed Two things worth knowing before deciding what to do about it The same Noma's write side has the matching problem in the other direction. # Replace content in all choices
for choice in response.choices:
if isinstance(choice, litellm.Choices) and choice.message.content:
choice.message.content = anonymized_contentso with On what to do, I do not think I should quietly redesign an enterprise vendor integration inside a PR about the stream builder, so I will follow your call. Three options as I see them Restore the previous coverage cheaply, here: change the read side to scan every non-empty choice joined together rather than the first, which is exactly what Noma saw before this PR for streamed Fix it properly and separately: scan and anonymize per choice across all six guardrails, which changes the payload shape those vendors receive and is a real design decision. Happy to open that as its own PR with an issue first Or land this as is and track the guardrail behaviour separately, on the grounds that the single-choice scan is the actual defect and My preference is the first, because it means this PR does not reduce anyone's coverage relative to today, and it leaves the real fix free to land later. Say the word and I will push it Separately, on the Codecov note: the uncovered lines are the |
|
Two review threads outstanding here, taking both. The Noma finding (@veria-ai) is real about Noma and not caused by this PRVerified the mechanism you describe. content = None
for choice in response.choices:
if isinstance(choice, litellm.Choices) and choice.message.content:
content = choice.message.content
breakso with Where I disagree is the attribution. That same method is what the non-streaming path uses: What this PR changes is that the streaming path now produces the same shape as the non-streaming path. Before it, So the fix belongs in The
|
|
Both responses are sound. Noma attribution: Confirmed. content = None
for choice in response.choices:
if isinstance(choice, litellm.Choices) and choice.message.content:
content = choice.message.content
breakAnd
|
build_base_response emits a single hardcoded choice and every assembly step reads choices[0], while get_combined_content joins the delta content of every choice of every chunk into one string. A streamed n=2 call therefore came back as one choice holding both completions concatenated, which is text no model produced, plus whichever finish_reason arrived last. The reassembled response is what feeds cost and token accounting, what the caller gets with complete_response=True, and what goes into the response cache, so a cached streamed n=2 call later replays the concatenation. Assemble each choice index on its own by recursing over the chunks filtered to that index, then swap the results in after usage has been calculated. Usage is still computed over the full chunk list, so it is unchanged. A single-choice stream sees no behaviour change: the recursion only runs when the chunks carry more than one distinct index.
The strict gate flagged the new helpers for ANN401 (bare `Any` on `choice`, `chunk` and the `_replace_chunk_choices` return) and ANN001 (`start_time` and `end_time` unannotated). Each one has a concrete type: the chunk and choice helpers already branch on `isinstance(..., dict)`, so the union of the pydantic model and a dict is what they actually take, and the two times are the `datetime.datetime | None` pair threaded into stream_chunk_builder. That also puts the already-imported StreamingChoices to use, clearing an F401. Net ruff difference on the file is 7 findings removed and none added.
…pers Giving the helpers concrete types traded ANN401 for LIT001, since `dict[...]` and `list[...]` are mutable collections in an annotation. Eleven of the fourteen are genuinely read-only and become `Mapping` / `Sequence`. `_replace_chunk_choices` narrowed with `isinstance(chunk, dict)`, which does not narrow a `Mapping`, so it now narrows on `ModelResponseStream` and the mapping branch stays a mapping. Same two expected inputs, same result for each. The remaining three have to be real lists and say why: two are forwarded to `stream_chunk_builder`, whose `chunks` and `messages` parameters are lists, and the third is assigned to `ModelResponse.choices`. `tests/test_litellm/test_main.py` is 98 passed before and 98 passed after.
LIT002 covers construction as well as annotation, so the lists, the local set and the two mapping literals in the per-index helpers each carry a `# mutable-ok:` reason now. Nothing moves at runtime: tests/test_litellm/test_main.py is 98 passed before and after.
The per-index builder asserted its result with two nested cast() calls, which the repo's LIT006 budget counts as unchecked assertions. Move the narrowing into a small helper that isinstance-checks the response and its first choice and returns None when either is not what the merge needs, so the caller's fallback to the single-response path is reached by a real check rather than by faith. The two one-shot locals in the index/choices readers pick up the Final annotation LIT010 asks for. No behaviour change: the same inputs produce the same choices, and the same inputs fall back. Signed-off-by: Vineeth Sai <vineethsai4444@gmail.com>
… Any The helper isinstance-narrows its argument, so it does not need a dynamically typed annotation; object works and keeps the repo's ANN401 budget unchanged. Signed-off-by: Vineeth Sai <vineethsai4444@gmail.com>
ModelResponse.choices is declared list[Choices], so the second isinstance was an unnecessary check basedpyright counts against the budget; only the response itself can be the streaming variant. Signed-off-by: Vineeth Sai <vineethsai4444@gmail.com>
4dd061f to
4771bba
Compare
| import json | ||
| from collections.abc import AsyncIterator | ||
| from typing import Final | ||
|
|
||
| import httpx | ||
| import pytest | ||
| import pytest_asyncio | ||
| import respx | ||
|
|
||
| import litellm | ||
| from litellm import ModelResponse | ||
| from litellm import Choices, ModelResponse | ||
| from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler | ||
| from litellm.proxy._types import UserAPIKeyAuth | ||
| from litellm.proxy.guardrails.guardrail_hooks.noma import ( | ||
| NomaGuardrail, | ||
| NomaV2Guardrail, | ||
| initialize_guardrail, | ||
| ) | ||
| import litellm.proxy.guardrails.guardrail_hooks.noma.noma as noma_legacy_module | ||
| from litellm.proxy.guardrails.guardrail_hooks.noma.noma import NomaBlockedMessage | ||
| from litellm.proxy.guardrails.init_guardrails import init_guardrails_v2 | ||
| from litellm.types.llms.openai import AllMessageValues | ||
| from litellm.types.utils import Choices, Message | ||
| from litellm.proxy.guardrails.guardrail_hooks.noma.noma import NomaBlockedMessage, NomaGuardrail | ||
| from litellm.types.utils import Delta, ModelResponseStream, StreamingChoices |
There was a problem hiding this comment.
This deletes existing configuration, failure-mode, and lifecycle tests. Repository rules prohibit weakening modified tests, so this requirement must be satisfied before merging
Rule Used: What: Flag any modifications to existing tests and verify they don't weaken test coverage or mask regressions. Why: Developers may alter tests to make failing code pass rather than fix the actual bug, hiding regressions. Good: ``` // Test updated t... (source)
Note: If this suggestion doesn't match your team's coding style, reply to this and let me know. I'll remember it for next time!
2 participants
TLDR
Problem this solves:
How it solves it:
User Flow
Before: multiple streamed completions could escape complete guardrail checking
stream=trueandn=2After: all choices are checked before the stream is returned
stream=trueandn=2Relevant issues
Addresses the guardrail review thread in this PR
Linear ticket
Pre-Submission checklist
make checkpassesScreenshots / Proof of Fix
Live proxy/vendor QA is pending because Noma credentials are unavailable
Type
Bug Fix
Caveats
Medium
Final Attestation