mirror of
https://github.com/BerriAI/litellm.git
synced 2026-09-14 23:21:35 +00:00
Fix Rubrik batch drain and SSE recovery mutation
This commit is contained in:
parent
ab7d7d016d
commit
87c9e3bdd3
4 changed files with 67 additions and 13 deletions
|
|
@ -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",
|
||||
|
|
|
|||
|
|
@ -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 -------------------------------------------------
|
||||
|
||||
|
|
|
|||
|
|
@ -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,
|
||||
|
|
|
|||
|
|
@ -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()
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue