diff --git a/litellm/completion_extras/litellm_responses_transformation/transformation.py b/litellm/completion_extras/litellm_responses_transformation/transformation.py index f58512ec930..de16680f1bb 100644 --- a/litellm/completion_extras/litellm_responses_transformation/transformation.py +++ b/litellm/completion_extras/litellm_responses_transformation/transformation.py @@ -652,9 +652,10 @@ class LiteLLMResponsesTransformationHandler(CompletionTransformationBridge): except (TypeError, ValueError): output_index = len(recovered_text_only_items) - item = recovered_output_items.get( - output_index - ) or recovered_text_only_items.get(output_index) + if output_index in recovered_output_items: + return + + item = recovered_text_only_items.get(output_index) if item is None: item = { "type": "message", diff --git a/litellm/integrations/rubrik.py b/litellm/integrations/rubrik.py index 80979e48ae8..a342b95dc0f 100644 --- a/litellm/integrations/rubrik.py +++ b/litellm/integrations/rubrik.py @@ -384,23 +384,22 @@ class RubrikLogger(CustomGuardrail, CustomBatchLogger): if not self.log_queue: return - await self._log_batch_to_rubrik( - data=list(self.log_queue), + log_queue_snapshot = list(self.log_queue) + verbose_logger.debug( + "Rubrik: Flushing batch of %s events", len(log_queue_snapshot) ) + await self._log_batch_to_rubrik( + data=log_queue_snapshot, + ) + del self.log_queue[: len(log_queue_snapshot)] + self.last_flush_time = time.time() async def flush_queue(self): if self.flush_lock is None: return async with self.flush_lock: - if self.log_queue: - log_queue_snapshot = list(self.log_queue) - verbose_logger.debug( - "Rubrik: Flushing batch of %s events", len(log_queue_snapshot) - ) - await self._log_batch_to_rubrik(data=log_queue_snapshot) - del self.log_queue[: len(log_queue_snapshot)] - self.last_flush_time = time.time() + await self.async_send_batch() # -- Tool blocking service ------------------------------------------------- diff --git a/tests/test_litellm/completion_extras/litellm_responses_transformation/test_completion_extras_litellm_responses_transformation_transformation.py b/tests/test_litellm/completion_extras/litellm_responses_transformation/test_completion_extras_litellm_responses_transformation_transformation.py index 77126a868cf..903d27056de 100644 --- a/tests/test_litellm/completion_extras/litellm_responses_transformation/test_completion_extras_litellm_responses_transformation_transformation.py +++ b/tests/test_litellm/completion_extras/litellm_responses_transformation/test_completion_extras_litellm_responses_transformation_transformation.py @@ -705,6 +705,44 @@ def test_transform_response_recovers_output_item_done_from_whitespace_padded_raw assert result.choices[0].message.content == "Recovered from padded output item" +def test_transform_response_preserves_output_item_when_text_done_arrives_later(): + from litellm.completion_extras.litellm_responses_transformation.transformation import ( + LiteLLMResponsesTransformationHandler, + ) + + handler = LiteLLMResponsesTransformationHandler() + + raw_sse = "\n".join( + [ + 'data: {"type":"response.output_item.done","output_index":0,"item":{"type":"message","id":"msg_from_item","role":"assistant","status":"completed","content":[{"type":"output_text","text":"Complete output item text","annotations":[]}]}}', + 'data: {"type":"response.output_text.done","output_index":0,"content_index":0,"item_id":"msg_from_stream","text":"Late text event"}', + 'data: {"type":"response.completed","response":{"id":"resp_from_stream","object":"response","created_at":1760144904,"status":"completed","model":"gpt-5.4","output":[]}}', + "data: [DONE]", + "", + ] + ) + + raw_response = _make_empty_responses_api_response() + model_response = _make_empty_model_response() + logging_obj = Mock() + logging_obj.model_call_details = {"original_response": raw_sse} + + result = handler.transform_response( + model="gpt-5.4", + raw_response=raw_response, + model_response=model_response, + logging_obj=logging_obj, + request_data={"model": "gpt-5.4"}, + messages=[{"role": "user", "content": "Reply with exactly: ok"}], + optional_params={}, + litellm_params={}, + encoding=Mock(), + ) + + assert len(result.choices) == 1 + assert result.choices[0].message.content == "Complete output item text" + + def test_transform_response_prefers_completed_output_from_raw_sse(): from litellm.completion_extras.litellm_responses_transformation.transformation import ( LiteLLMResponsesTransformationHandler, diff --git a/tests/test_litellm/integrations/test_rubrik.py b/tests/test_litellm/integrations/test_rubrik.py index 6e51106f632..06942b5eb8b 100644 --- a/tests/test_litellm/integrations/test_rubrik.py +++ b/tests/test_litellm/integrations/test_rubrik.py @@ -236,6 +236,22 @@ class TestBatchLogging: assert handler.log_queue == [{"msg": "c"}] + async def test_async_send_batch_drains_sent_events(self, handler): + handler.log_queue = [{"msg": "a"}, {"msg": "b"}] + + async def mock_post(*_args, **_kwargs): + handler.log_queue.append({"msg": "c"}) + mock_response = Mock() + mock_response.raise_for_status = Mock() + return mock_response + + handler.async_httpx_client = AsyncMock() + handler.async_httpx_client.post = mock_post + + await handler.async_send_batch() + + assert handler.log_queue == [{"msg": "c"}] + async def test_log_batch_error_does_not_crash(self, handler): handler.log_queue = [{"msg": "a"}] mock_response = Mock()