diff --git a/litellm/litellm_core_utils/redact_messages.py b/litellm/litellm_core_utils/redact_messages.py index 548d088d833..91f4d144b47 100644 --- a/litellm/litellm_core_utils/redact_messages.py +++ b/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 diff --git a/litellm/litellm_core_utils/streaming_handler.py b/litellm/litellm_core_utils/streaming_handler.py index 5c1f354fe4b..b248aebd121 100644 --- a/litellm/litellm_core_utils/streaming_handler.py +++ b/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: diff --git a/tests/unit/litellm_core_utils/test_redact_messages.py b/tests/unit/litellm_core_utils/test_redact_messages.py index 5b64e42d8a8..387289eeea6 100644 --- a/tests/unit/litellm_core_utils/test_redact_messages.py +++ b/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": [ diff --git a/tests/unit/text_completions/test_streaming_iterator.py b/tests/unit/text_completions/test_streaming_iterator.py index 402e3302eeb..8f7fd6fa327 100644 --- a/tests/unit/text_completions/test_streaming_iterator.py +++ b/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")