From 3884855a6e789794d5c0be6693081350648b1298 Mon Sep 17 00:00:00 2001 From: Henry Su Date: Sun, 6 Sep 2026 12:10:57 -0500 Subject: [PATCH] fix(streaming): keep chat moderation results in the accumulated snapshot `ChatCompletionStreamState._accumulate_chunk` carried `usage` and `system_fingerprint` forward from each chunk but ignored `moderation`. Moderation results arrive on their own chunk after the content chunks, so `get_final_completion()` and `current_completion_snapshot` reported `moderation=None` for streams that requested moderated completions, while the equivalent non-streaming response kept the field. Carry `moderation` into the snapshot when a chunk provides it. It is copied through `construct_type` because `ChatCompletionChunk.Moderation` and `ChatCompletion.Moderation` are distinct models, and it is only applied when present so a later chunk without moderation cannot clear it. --- src/openai/lib/streaming/chat/_completions.py | 9 ++- tests/lib/chat/test_completions_streaming.py | 78 +++++++++++++++++++ 2 files changed, 86 insertions(+), 1 deletion(-) diff --git a/src/openai/lib/streaming/chat/_completions.py b/src/openai/lib/streaming/chat/_completions.py index f9dec645b6..7050fad965 100644 --- a/src/openai/lib/streaming/chat/_completions.py +++ b/src/openai/lib/streaming/chat/_completions.py @@ -38,7 +38,7 @@ from ...._streaming import Stream, AsyncStream from ....types.chat import ChatCompletionChunk, ParsedChatCompletion, ChatCompletionToolUnionParam from ...._exceptions import LengthFinishReasonError, ContentFilterFinishReasonError -from ....types.chat.chat_completion import ChoiceLogprobs +from ....types.chat.chat_completion import Moderation, ChoiceLogprobs from ....types.chat.chat_completion_chunk import Choice as ChoiceChunk from ....types.chat.completion_create_params import ResponseFormat as ResponseFormatParam @@ -487,6 +487,13 @@ def _accumulate_chunk(self, chunk: ChatCompletionChunk) -> ParsedChatCompletionS completion_snapshot.usage = chunk.usage completion_snapshot.system_fingerprint = chunk.system_fingerprint + if chunk.moderation is not None: + # moderation results arrive on their own chunk, so they have to be carried + # into the snapshot instead of being dropped along with the rest of that chunk + completion_snapshot.moderation = cast( + Moderation, + construct_type(type_=Moderation, value=chunk.moderation.to_dict()), + ) return completion_snapshot diff --git a/tests/lib/chat/test_completions_streaming.py b/tests/lib/chat/test_completions_streaming.py index 25d2942639..25be04a755 100644 --- a/tests/lib/chat/test_completions_streaming.py +++ b/tests/lib/chat/test_completions_streaming.py @@ -1060,6 +1060,84 @@ def test_stream_obfuscation_stays_on_raw_chunks(padding: tuple[str | None, str | assert "obfuscation" not in completion.to_json() +_MODERATION_RESULT = { + "categories": {"violence": False}, + "category_applied_input_types": {"violence": ["text"]}, + "category_scores": {"violence": 0.01}, + "flagged": False, + "model": "omni-moderation-latest", + "type": "moderation_result", +} +_MODERATION = { + "input": {"type": "moderation_results", "model": "omni-moderation-latest", "results": [_MODERATION_RESULT]}, + "output": {"type": "moderation_results", "model": "omni-moderation-latest", "results": [_MODERATION_RESULT]}, +} +_MODERATION_CONTENT_CHOICE = { + "index": 0, + "delta": {"role": "assistant", "content": "Hello"}, + "finish_reason": "stop", + "logprobs": None, +} + + +def _moderation_test_chunk(**extra: Any) -> ChatCompletionChunk: + return model_parse( + ChatCompletionChunk, + { + "id": "chatcmpl-test", + "object": "chat.completion.chunk", + "created": 0, + "model": "gpt-test", + "choices": [], + **extra, + }, + ) + + +def test_stream_snapshot_keeps_moderation_from_a_later_chunk() -> None: + state = ChatCompletionStreamState() + state.handle_chunk(_moderation_test_chunk(choices=[_MODERATION_CONTENT_CHOICE])) + state.handle_chunk( + _moderation_test_chunk( + moderation=_MODERATION, + system_fingerprint="fp_test", + usage={"prompt_tokens": 1, "completion_tokens": 1, "total_tokens": 2}, + ) + ) + + completion = state.get_final_completion() + assert completion.choices[0].message.content == "Hello" + assert completion.system_fingerprint == "fp_test" + assert completion.usage is not None + assert completion.moderation is not None + assert completion.moderation.input.type == "moderation_results" + assert completion.moderation.output.type == "moderation_results" + assert completion.to_dict()["moderation"] == _MODERATION + + +def test_stream_snapshot_keeps_moderation_from_the_first_chunk() -> None: + state = ChatCompletionStreamState() + state.handle_chunk(_moderation_test_chunk(moderation=_MODERATION, choices=[_MODERATION_CONTENT_CHOICE])) + + assert state.current_completion_snapshot.moderation is not None + assert state.get_final_completion().moderation is not None + + +def test_stream_snapshot_moderation_survives_a_later_chunk_without_moderation() -> None: + state = ChatCompletionStreamState() + state.handle_chunk(_moderation_test_chunk(moderation=_MODERATION, choices=[_MODERATION_CONTENT_CHOICE])) + state.handle_chunk(_moderation_test_chunk(usage={"prompt_tokens": 1, "completion_tokens": 1, "total_tokens": 2})) + + assert state.get_final_completion().moderation is not None + + +def test_stream_snapshot_has_no_moderation_when_the_stream_has_none() -> None: + state = ChatCompletionStreamState() + state.handle_chunk(_moderation_test_chunk(choices=[_MODERATION_CONTENT_CHOICE])) + + assert state.get_final_completion().moderation is None + + @pytest.mark.respx2(base_url=base_url) def test_chat_completion_state_helper(client: OpenAI, respx2_mock: MockRouter, monkeypatch: pytest.MonkeyPatch) -> None: state = ChatCompletionStreamState()