diff --git a/litellm/integrations/custom_batch_logger.py b/litellm/integrations/custom_batch_logger.py index f9d4496c21f..0549c3aeaaf 100644 --- a/litellm/integrations/custom_batch_logger.py +++ b/litellm/integrations/custom_batch_logger.py @@ -14,6 +14,8 @@ from litellm.integrations.custom_logger import CustomLogger class CustomBatchLogger(CustomLogger): + preserve_events_added_during_flush = False + def __init__( self, flush_lock: Optional[asyncio.Lock] = None, @@ -47,11 +49,15 @@ class CustomBatchLogger(CustomLogger): async with self.flush_lock: if self.log_queue: + log_queue_length = len(self.log_queue) verbose_logger.debug( "CustomLogger: Flushing batch of %s events", len(self.log_queue) ) await self.async_send_batch() - self.log_queue.clear() + if self.preserve_events_added_during_flush: + del self.log_queue[:log_queue_length] + else: + self.log_queue.clear() self.last_flush_time = time.time() async def async_send_batch(self, *args, **kwargs): diff --git a/litellm/integrations/rubrik.py b/litellm/integrations/rubrik.py index a342b95dc0f..ff2d811e74a 100644 --- a/litellm/integrations/rubrik.py +++ b/litellm/integrations/rubrik.py @@ -48,6 +48,8 @@ class BlockedToolsResult: class RubrikLogger(CustomGuardrail, CustomBatchLogger): + preserve_events_added_during_flush = True + def __init__( self, api_key: str | None = None, @@ -391,15 +393,6 @@ class RubrikLogger(CustomGuardrail, CustomBatchLogger): 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: - await self.async_send_batch() # -- Tool blocking service ------------------------------------------------- diff --git a/litellm/router.py b/litellm/router.py index 12032789155..b42b957141a 100644 --- a/litellm/router.py +++ b/litellm/router.py @@ -6809,7 +6809,7 @@ class Router: # Multiple aliases can point at the same provider/model backend, # but their deployment-level overrides should not downgrade the # backend from responses -> chat via last-write-wins registration. - _shared_model_info.pop("mode", None) + _shared_model_info["mode"] = _existing_shared_mode litellm.register_model( model_cost={ _model_name: _shared_model_info, diff --git a/tests/test_litellm/integrations/test_rubrik.py b/tests/test_litellm/integrations/test_rubrik.py index 06942b5eb8b..414ae05d721 100644 --- a/tests/test_litellm/integrations/test_rubrik.py +++ b/tests/test_litellm/integrations/test_rubrik.py @@ -12,6 +12,7 @@ from unittest.mock import AsyncMock, Mock, patch import httpx import pytest +from litellm.integrations.custom_batch_logger import CustomBatchLogger from litellm.integrations.custom_guardrail import ModifyResponseException from litellm.integrations.rubrik import RubrikLogger from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler @@ -236,7 +237,7 @@ class TestBatchLogging: assert handler.log_queue == [{"msg": "c"}] - async def test_async_send_batch_drains_sent_events(self, handler): + async def test_async_send_batch_does_not_drain_events(self, handler): handler.log_queue = [{"msg": "a"}, {"msg": "b"}] async def mock_post(*_args, **_kwargs): @@ -250,6 +251,22 @@ class TestBatchLogging: await handler.async_send_batch() + assert handler.log_queue == [{"msg": "a"}, {"msg": "b"}, {"msg": "c"}] + + async def test_parent_flush_queue_preserves_events_added_during_send(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 CustomBatchLogger.flush_queue(handler) + assert handler.log_queue == [{"msg": "c"}] async def test_log_batch_error_does_not_crash(self, handler): diff --git a/tests/test_litellm/test_router_model_cost_isolation.py b/tests/test_litellm/test_router_model_cost_isolation.py index 5de4027382f..f9ca60a2d94 100644 --- a/tests/test_litellm/test_router_model_cost_isolation.py +++ b/tests/test_litellm/test_router_model_cost_isolation.py @@ -335,6 +335,7 @@ def test_should_not_downgrade_chatgpt_shared_key_mode_with_alias_override(): ) assert litellm.model_cost[backend_model]["mode"] == "responses" + assert "mode" in litellm.model_cost[backend_model] bridge_model_info, bridge_model = responses_api_bridge_check( model="gpt-5.4",