scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-11-02 07:13:02 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 48

fix(responses): carry streamed logprobs into the completed snapshot and drop them on guardrail rewrites

Browse source 
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>
This commit is contained in:
mateo 2026-11-01 00:17:01 +00:00
parent 45805384e5
commit 3588a1e877
5 changed files with 132 additions and 11 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

13
litellm/llms/openai/responses/guardrail_translation/handler.py
View file

@ -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:

37
litellm/responses/litellm_completion_transformation/streaming_iterator.py
View file

@ -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:

2
litellm/types/llms/base.py
View file

@ -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

89
tests/unit/llms/openai/responses/test_openai_responses_guardrail_handler.py
View file

@ -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)

2
tests/unit/responses/litellm_completion_transformation/test_streaming_iterator_transformation.py
View file

@ -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: