From 3588a1e87760ebec0b2f7cd4894e02d361a1c898 Mon Sep 17 00:00:00 2001 From: mateo Date: Sun, 1 Nov 2026 00:17:01 +0000 Subject: [PATCH] fix(responses): carry streamed logprobs into the completed snapshot and drop them on guardrail rewrites Greptile review follow-ups on the bridge logprobs change. The response.completed snapshot is rebuilt through stream_chunk_builder, which drops chat logprobs, so the text part now takes the logprobs the stream carried. The per-chunk accumulator appends and flattens once instead of recopying every earlier token. CarriesOutputTextLogprobs carries a frozen-ok reason since guardrails rewrite text parts and stream events in place. A guardrail rewrite of output_text (part, delta, done) now also clears the logprobs that would spell out the original text. Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .../guardrail_translation/handler.py | 13 ++- .../streaming_iterator.py | 37 ++++++-- litellm/types/llms/base.py | 2 +- ...test_openai_responses_guardrail_handler.py | 89 +++++++++++++++++++ .../test_streaming_iterator_transformation.py | 2 + 5 files changed, 132 insertions(+), 11 deletions(-) diff --git a/litellm/llms/openai/responses/guardrail_translation/handler.py b/litellm/llms/openai/responses/guardrail_translation/handler.py index a905bff3521..7db75645e88 100644 --- a/litellm/llms/openai/responses/guardrail_translation/handler.py +++ b/litellm/llms/openai/responses/guardrail_translation/handler.py @@ -1425,11 +1425,18 @@ class OpenAIResponsesHandler(BaseTranslation): else: setattr(event, field, value) + @staticmethod + def _write_rewritten_text(event: object, field: str, text: str) -> None: + """A rewrite replaces the text, so the logprobs that spell out the original text go with it.""" + OpenAIResponsesHandler._write_event_field(event, field, text) + if stream_item_field(event, "logprobs") is not None: + OpenAIResponsesHandler._write_event_field(event, "logprobs", None) + @staticmethod def _write_content_part_text(part: object, text: str) -> None: field: Final = _content_part_text_field(part) if _has_writable_field(part, field): - OpenAIResponsesHandler._write_event_field(part, field, text) + OpenAIResponsesHandler._write_rewritten_text(part, field, text) @staticmethod def _write_compaction_summary(item: object, summary: str) -> None: @@ -1468,9 +1475,9 @@ class OpenAIResponsesHandler(BaseTranslation): continue position = (output_index, content_index) if event_type in _TEXT_DELTA_EVENT_TYPES and position in delta_replacements: - self._write_event_field(event, "delta", next(delta_replacements[position])) + self._write_rewritten_text(event, "delta", next(delta_replacements[position])) elif event_type == "response.output_text.done" and position in rewrites_by_position: - self._write_event_field(event, "text", rewrites_by_position[position]) + self._write_rewritten_text(event, "text", rewrites_by_position[position]) elif event_type == "response.refusal.done" and position in rewrites_by_position: self._write_event_field(event, "refusal", rewrites_by_position[position]) elif event_type == "response.content_part.done" and position in rewrites_by_position: diff --git a/litellm/responses/litellm_completion_transformation/streaming_iterator.py b/litellm/responses/litellm_completion_transformation/streaming_iterator.py index 7364eabb602..bf137629553 100644 --- a/litellm/responses/litellm_completion_transformation/streaming_iterator.py +++ b/litellm/responses/litellm_completion_transformation/streaming_iterator.py @@ -1,7 +1,7 @@ import time import uuid from collections.abc import Mapping, Sequence -from itertools import filterfalse +from itertools import chain, filterfalse from types import MappingProxyType from typing import Any, Final, cast @@ -63,7 +63,11 @@ from litellm.types.llms.openai import ( WebSearchCallInProgressEvent, WebSearchCallSearchingEvent, ) -from litellm.types.responses.main import GenericResponseOutputItemContentAnnotation +from litellm.types.responses.main import ( + GenericResponseOutputItem, + GenericResponseOutputItemContentAnnotation, + OutputText, +) from litellm.types.utils import Delta as ChatCompletionDelta from litellm.types.utils import ( ModelResponse, @@ -213,7 +217,9 @@ class LiteLLMCompletionStreamingIterator(ResponsesAPIStreamingIterator): self._sent_reasoning_summary_text_done_event: bool = False self._sent_reasoning_summary_part_done_event: bool = False self._reasoning_summary_text: str = "" - self._text_logprobs: tuple[OutputTextLogprob, ...] = () + self._text_logprob_chunks: list[ + tuple[OutputTextLogprob, ...] + ] = [] # mutable-ok: appended per chunk, flattened once at done # -- GENERIC RESPONSE-EVENTS PENDING QUEUE as required by fix -- self._pending_response_events: list[BaseLiteLLMOpenAIResponseObject] = [] self._reasoning_active = False @@ -1021,9 +1027,15 @@ class LiteLLMCompletionStreamingIterator(ResponsesAPIStreamingIterator): logprobs=self._final_text_logprobs(litellm_complete_object), ) + def _streamed_text_logprobs(self) -> tuple[OutputTextLogprob, ...] | None: + if not self._text_logprob_chunks: + return None + return tuple(chain.from_iterable(self._text_logprob_chunks)) + def _final_text_logprobs(self, litellm_complete_object: ModelResponse) -> tuple[OutputTextLogprob, ...] | None: - if self._text_logprobs: - return self._text_logprobs + streamed: Final = self._streamed_text_logprobs() + if streamed is not None: + return streamed return output_text_logprobs(getattr(litellm_complete_object.choices[0], "logprobs", None)) def create_output_content_part_done_event(self, litellm_complete_object: ModelResponse) -> ContentPartDoneEvent: @@ -1525,7 +1537,7 @@ class LiteLLMCompletionStreamingIterator(ResponsesAPIStreamingIterator): self._sequence_number += 1 delta_logprobs: Final = output_text_logprobs(getattr(chunk.choices[0], "logprobs", None)) if delta_logprobs is not None: - self._text_logprobs = (*self._text_logprobs, *delta_logprobs) + self._text_logprob_chunks.append(delta_logprobs) text_delta_event: Final = OutputTextDeltaEvent( type=ResponsesAPIStreamEvents.OUTPUT_TEXT_DELTA, item_id=item_id, @@ -1582,7 +1594,18 @@ class LiteLLMCompletionStreamingIterator(ResponsesAPIStreamingIterator): ) compaction_aligned: Final = _output_items_with_id(reasoning_aligned, "compaction", self._compaction_item_id) streamed_items: Final = filterfalse(self._is_unstreamed_empty_message, compaction_aligned) - return tuple(sorted(streamed_items, key=self._streamed_output_position)) + return tuple(map(self._with_streamed_logprobs, sorted(streamed_items, key=self._streamed_output_position))) + + def _with_streamed_logprobs(self, item: object) -> object: + """stream_chunk_builder drops chat logprobs, so the snapshot's text part takes the ones the stream carried.""" + logprobs: Final = self._streamed_text_logprobs() + if logprobs is None or not isinstance(item, GenericResponseOutputItem): + return item + content: Final = tuple( + part.model_copy(update={"logprobs": logprobs}) if isinstance(part, OutputText) else part + for part in item.content + ) + return item.model_copy(update={"content": list(content)}) def _is_unstreamed_empty_message(self, item: object) -> bool: if getattr(item, "type", None) != "message" or self._message_output_index is not None: diff --git a/litellm/types/llms/base.py b/litellm/types/llms/base.py index 6af1d4d7046..231bfff2036 100644 --- a/litellm/types/llms/base.py +++ b/litellm/types/llms/base.py @@ -100,7 +100,7 @@ class OutputTextLogprob(TypedDict): top_logprobs: ReadOnly[Sequence[OutputTextTopLogprob]] -class CarriesOutputTextLogprobs(BaseLiteLLMOpenAIResponseObject): +class CarriesOutputTextLogprobs(BaseLiteLLMOpenAIResponseObject): # frozen-ok: guardrails rewrite these in place """A Responses text payload carries ``logprobs`` only when the chat choice behind it sent them.""" logprobs: Sequence[OutputTextLogprob] | None = None diff --git a/tests/unit/llms/openai/responses/test_openai_responses_guardrail_handler.py b/tests/unit/llms/openai/responses/test_openai_responses_guardrail_handler.py index d80a5e835c9..5c8b147f842 100644 --- a/tests/unit/llms/openai/responses/test_openai_responses_guardrail_handler.py +++ b/tests/unit/llms/openai/responses/test_openai_responses_guardrail_handler.py @@ -36,6 +36,8 @@ from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLogging from litellm.llms import get_guardrail_translation_mapping from litellm.llms.openai.responses.guardrail_translation.handler import ( OpenAIResponsesHandler, + stream_item_field, + stream_item_items, ) from litellm.llms.openai.responses.guardrail_translation.tool_merge import merge_guardrailed_tools from litellm.proxy.guardrails.guardrail_hooks.generic_guardrail_api import GenericGuardrailAPI @@ -46,8 +48,11 @@ from litellm.types.guardrails import ContentFilterAction, ContentFilterPattern from litellm.types.llms.openai import ( ChatCompletionToolCallChunk, ContentPartDoneEvent, + ContentPartDonePartOutputText, ContentPartDonePartRefusal, OutputItemDoneEvent, + OutputTextDeltaEvent, + OutputTextDoneEvent, RefusalDeltaEvent, RefusalDoneEvent, ResponseIncompleteEvent, @@ -5229,3 +5234,87 @@ class TestRefusalPartsReachTextGuardrails: ) assert (events[0].delta, events[1].refusal) == (_MASKED_REFUSAL, _MASKED_REFUSAL) + + +class _PersimmonTextMaskingGuardrail(CustomGuardrail): + async def apply_guardrail( + self, + inputs: GenericGuardrailAPIInputs, + request_data: dict, + input_type: Literal["request", "response"], + logging_obj: Optional[LiteLLMLoggingObj] = None, + ) -> GenericGuardrailAPIInputs: + return {**inputs, "texts": [text.replace("persimmon", "[MASKED]") for text in inputs.get("texts", [])]} + + +_PERSIMMON_LOGPROBS: Final = ( + {"token": "ripe", "logprob": -0.1, "bytes": None, "top_logprobs": []}, + {"token": " persimmon", "logprob": -0.2, "bytes": None, "top_logprobs": []}, +) + + +@pytest.mark.asyncio +async def test_masked_output_text_drops_the_logprobs_that_spell_the_original_text() -> None: + response = ResponsesAPIResponse.model_validate( + { + "id": "resp_1", + "created_at": 1, + "model": "gpt-5.6", + "object": "response", + "status": "completed", + "output": [ + { + "type": "message", + "id": "msg_1", + "status": "completed", + "role": "assistant", + "content": [ + {"type": "output_text", "text": "ripe persimmon", "annotations": [], "logprobs": list(_PERSIMMON_LOGPROBS)} + ], + } + ], + } + ) + + result = await OpenAIResponsesHandler().process_output_response( + response, _PersimmonTextMaskingGuardrail(guardrail_name="mask") + ) + + part = stream_item_items(result.output[0], "content")[0] + assert (stream_item_field(part, "text"), stream_item_field(part, "logprobs")) == ("ripe [MASKED]", None) + + +def test_stream_rewrite_drops_logprobs_from_the_delta_done_and_part_events() -> None: + delta = OutputTextDeltaEvent( + type=ResponsesAPIStreamEvents.OUTPUT_TEXT_DELTA, + item_id="msg_1", + output_index=0, + content_index=0, + delta="ripe persimmon", + logprobs=_PERSIMMON_LOGPROBS, + ) + done = OutputTextDoneEvent( + type=ResponsesAPIStreamEvents.OUTPUT_TEXT_DONE, + item_id="msg_1", + output_index=0, + content_index=0, + text="ripe persimmon", + logprobs=_PERSIMMON_LOGPROBS, + ) + part_done = ContentPartDoneEvent( + type=ResponsesAPIStreamEvents.CONTENT_PART_DONE, + item_id="msg_1", + output_index=0, + content_index=0, + part=ContentPartDonePartOutputText( + type="output_text", text="ripe persimmon", annotations=[], logprobs=_PERSIMMON_LOGPROBS + ), + ) + + OpenAIResponsesHandler()._sync_stream_events_with_rewrites( + stream_events=(delta, done, part_done), rewrites_by_position={(0, 0): "ripe [MASKED]"} + ) + + assert (delta.delta, delta.logprobs) == ("ripe [MASKED]", None) + assert (done.text, done.logprobs) == ("ripe [MASKED]", None) + assert (part_done.part.text, part_done.part.logprobs) == ("ripe [MASKED]", None) diff --git a/tests/unit/responses/litellm_completion_transformation/test_streaming_iterator_transformation.py b/tests/unit/responses/litellm_completion_transformation/test_streaming_iterator_transformation.py index fd7a91d12f6..c0fcf509456 100644 --- a/tests/unit/responses/litellm_completion_transformation/test_streaming_iterator_transformation.py +++ b/tests/unit/responses/litellm_completion_transformation/test_streaming_iterator_transformation.py @@ -2023,6 +2023,8 @@ def test_streamed_chat_logprobs_reach_the_text_delta_and_done_events() -> None: assert done["logprobs"] == [hello, period] assert part_done["part"]["logprobs"] == [hello, period] assert item_done["item"]["content"][0]["logprobs"] == [hello, period] + completed: Final = next(event for event in events if event["type"] == "response.completed") + assert completed["response"]["output"][0]["content"][0]["logprobs"] == [hello, period] def test_streamed_chat_without_logprobs_emits_text_events_without_a_logprobs_key() -> None: