fix(streaming): finish each streamed text completion choice on its own index and null token_ids in redacted logs

litellm/litellm_core_utils/redact_messages.py

@@ -196,6 +196,11 @@ def _redact_function_call(function_call) -> None:
 
 def _redact_choice_content(choice):
     """Helper to redact content in a choice (message or delta)."""
+    if (
+        isinstance(choice, (litellm.Choices, litellm.utils.StreamingChoices))
+        and getattr(choice, "token_ids", None) is not None
+    ):
+        setattr(choice, "token_ids", None)  # noqa: B010  # pydantic extra, not a declared field
     if isinstance(choice, litellm.Choices):
         if choice.message.content is not None:
             choice.message.content = REDACTED_BY_LITELLM
@@ -250,6 +255,7 @@ def _redacted_provider_specific_fields(fields: Mapping[str, object]) -> Mapping[
                 "citations",
                 "compaction_blocks",
                 "compaction_delta",
+                "token_ids",
             )
             if key in fields
         },
@@ -421,6 +427,8 @@ def _redact_model_response_dict_choices(choices, redacted_str: str):
     for choice in choices:
         if isinstance(choice, dict):
             _redact_provider_specific_fields_dict(choice)
+            if choice.get("token_ids") is not None:
+                choice["token_ids"] = None
             if "message" in choice and isinstance(choice["message"], dict):
                 if choice["message"].get("content") is not None:
                     choice["message"]["content"] = redacted_str

litellm/litellm_core_utils/streaming_handler.py

@@ -891,13 +891,15 @@ class CustomStreamWrapper:
                     is_finished = True
                     finish_reason = choices[0].finish_reason
             usage = getattr(chunk, "usage", None)
+            choice_index: Final = _upstream_choice_index(chunk)
             return {
                 "text": text,
                 "is_finished": is_finished,
                 "finish_reason": finish_reason,
                 "usage": usage,
                 "chunk_choice_fields": _non_null_upstream_choice_fields(chunk, STREAM_CHUNK_CHOICE_FIELDS),
-                "choice_index": _upstream_choice_index(chunk),
+                "choice_index": choice_index,
+                "sibling_finish_reason": finish_reason if choice_index != 0 else None,
             }
 
         except Exception as e:
@@ -1212,6 +1214,7 @@ class CustomStreamWrapper:
             "chunk_choice_fields", _NO_UPSTREAM_CHOICE_FIELDS
         )
         chunk_choice_fields: Final = self._take_carried_chunk_choice_fields(response_obj, upstream_chunk_choice_fields)
+        sibling_finish_reason: Final[str | None] = response_obj.get("sibling_finish_reason")
 
         if is_chunk_non_empty:  # cannot set content of an OpenAI Object to be an empty string
             self.raise_on_model_repetition()
@@ -1305,11 +1308,19 @@ class CustomStreamWrapper:
                         model_response.choices[0].index = _index
 
                 _apply_choice_fields(model_response.choices[0], chunk_choice_fields)
+                if sibling_finish_reason is not None:
+                    self._finish_sibling_choice(model_response.choices[0], sibling_finish_reason)
                 self._optional_combine_thinking_block_in_choices(model_response=model_response)
 
                 return model_response
             else:
                 return
+        elif sibling_finish_reason is not None:
+            sibling_chunk: Final = _choice_fields_chunk(
+                model_response, _response_choice_index(response_obj), chunk_choice_fields
+            )
+            self._finish_sibling_choice(sibling_chunk.choices[0], sibling_finish_reason)
+            return sibling_chunk
         elif self.received_finish_reason is not None:
             chunk_upstream_index: Final = _response_choice_index(response_obj)
             if chunk_choice_fields and self._streams_past_the_pending_finish(
@@ -1597,8 +1608,9 @@ class CustomStreamWrapper:
             completion_obj["content"] = response_obj["text"]
             print_verbose(f"completion obj content: {completion_obj['content']}")
             if response_obj["is_finished"]:
-                self.received_finish_reason = response_obj["finish_reason"]
                 self._record_finish_choice_fields(chunk)
+            if response_obj["is_finished"] and response_obj["sibling_finish_reason"] is None:
+                self.received_finish_reason = response_obj["finish_reason"]
             if response_obj["usage"] is not None:
                 _text_completion_usage: Final[Usage] = response_obj["usage"]
                 setattr(
@@ -2039,6 +2051,10 @@ class CustomStreamWrapper:
     def _apply_finish_choice_fields(self, choice: StreamingChoices) -> None:
         _apply_choice_fields(choice, self.received_finish_choice_fields.get(choice.index, _NO_UPSTREAM_CHOICE_FIELDS))
 
+    def _finish_sibling_choice(self, choice: StreamingChoices, finish_reason: str) -> None:
+        choice.finish_reason = map_finish_reason(finish_reason)  # rebind-ok: chunks are built in place here
+        self._apply_finish_choice_fields(choice)
+
     def _streams_past_the_pending_finish(
         self, model_response: ModelResponseStream, response_obj: Mapping[str, object], upstream_index: int
     ) -> bool:

tests/unit/litellm_core_utils/test_redact_messages.py

@@ -560,6 +560,34 @@ class TestPerformRedaction:
             {"text": "redacted-by-litellm", "finish_reason": "length", "index": 1},
         ]
 
+    def test_redacts_token_ids_that_decode_to_the_completion(self):
+        details = {
+            "stream": True,
+            "complete_streaming_response": litellm.ModelResponse(
+                choices=[
+                    litellm.Choices(message=litellm.Message(content="secret", role="assistant"), token_ids=[20, 11])
+                ]
+            ),
+            "standard_logging_object": {
+                "response": {
+                    "choices": [
+                        {
+                            "token_ids": [20, 11],
+                            "message": {"content": "secret", "provider_specific_fields": {"token_ids": [20, 11]}},
+                        },
+                        {"text": " Paris.", "token_ids": [20, 11], "index": 1},
+                    ]
+                }
+            },
+        }
+
+        perform_redaction(details, None)
+
+        logged_choices = details["standard_logging_object"]["response"]["choices"]
+        assert [choice["token_ids"] for choice in logged_choices] == [None, None]
+        assert logged_choices[0]["message"]["provider_specific_fields"]["token_ids"] is None
+        assert details["complete_streaming_response"].choices[0].get("token_ids") is None
+
     def test_redacts_object_choices_inside_model_response_dict(self):
         result = {
             "choices": [

tests/unit/text_completions/test_streaming_iterator.py

@@ -187,6 +187,36 @@ def test_text_completion_stream_keeps_each_choices_text_and_token_ids_on_its_own
     } == {0: (" 5, 6", [20, 11, 21, 11, 220]), 1: (" 7, 8", [22, 11, 220, 23])}
 
 
+@pytest.mark.parametrize(
+    "upstream",
+    [
+        (
+            _completion_chunk(" a", None, index=0),
+            _completion_chunk(" b", None, index=1),
+            _completion_chunk("", "length", index=2),
+            _completion_chunk(" red", "stop", index=1, stop_reason=","),
+            _completion_chunk(" s", "stop", index=0, stop_reason=","),
+        ),
+        (
+            _completion_chunk(" a", None, index=0),
+            _completion_chunk(" b", None, index=1),
+            _completion_chunk(" s", "stop", index=0, stop_reason=","),
+            _completion_chunk("", "length", index=2),
+            _completion_chunk(" red", "stop", index=1, stop_reason=","),
+        ),
+    ],
+    ids=["other-choices-finish-first", "choice-0-finishes-first"],
+)
+def test_text_completion_stream_finishes_each_choice_on_its_own_index(upstream: Sequence[Completion]) -> None:
+    finished: Final = [
+        (choice.index, choice.finish_reason, choice.get("stop_reason"))
+        for choice in chain.from_iterable(response.choices for response in _text_completion_stream(upstream))
+        if choice.finish_reason is not None
+    ]
+
+    assert sorted(finished) == [(0, "stop", ","), (1, "stop", ","), (2, "length", None)]
+
+
 def test_text_completion_conversion_rejects_non_streaming_choices() -> None:
     chunk: Final = Choices()
     wrapper: Final = TextCompletionStreamWrapper(completion_stream=None, model="test-model")