Fix Rubrik batch drain and SSE recovery mutation

This commit is contained in:
Cursor Agent 2026-05-05 00:35:08 +00:00
parent ab7d7d016d
commit 87c9e3bdd3
No known key found for this signature in database
4 changed files with 67 additions and 13 deletions

View file

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

View file

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

View file

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

View file

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