diff --git a/litellm-rust/crates/python-bridge/src/routes/messages.rs b/litellm-rust/crates/python-bridge/src/routes/messages.rs index b85057e9d75..07f15f1cebd 100644 --- a/litellm-rust/crates/python-bridge/src/routes/messages.rs +++ b/litellm-rust/crates/python-bridge/src/routes/messages.rs @@ -340,6 +340,7 @@ class Host: self.current = arguments self.asynchronous = asynchronous self.logger = arguments.get('litellm_logging_obj') + self.deployment_hooks_owned = self.logger is None self.state = None self.response = None self.error = None @@ -353,7 +354,9 @@ class Host: self.streaming = self.logger.stream is True async def deployment_pre(self): - modified = await utils.async_pre_call_deployment_hook(self.current, 'amessages') + if not self.deployment_hooks_owned: + return + modified = await utils.async_pre_call_deployment_hook(self.current, 'anthropic_messages') if modified is not None: self.current = modified self.current['litellm_logging_obj'] = self.logger @@ -370,7 +373,8 @@ class Host: self.end = datetime.now() async def deployment_success(self): - self.response = await utils.async_post_call_success_deployment_hook(self.current, self.response, CallTypes.aanthropic_messages) + if self.deployment_hooks_owned: + self.response = await utils.async_post_call_success_deployment_hook(self.current, self.response, CallTypes.aanthropic_messages) if self.streaming: self.response = retain_stream_response( self.response, @@ -380,7 +384,8 @@ class Host: ) async def deployment_failure(self): - await utils.async_post_call_failure_deployment_hook(self.current, self.error, 'amessages') + if self.deployment_hooks_owned: + await utils.async_post_call_failure_deployment_hook(self.current, self.error, 'anthropic_messages') def terminal(self, action, value): if self.streaming: diff --git a/tests/test_litellm_rust/test_messages_callbacks.py b/tests/test_litellm_rust/test_messages_callbacks.py index fdf90d2fbf4..e21ad4661db 100644 --- a/tests/test_litellm_rust/test_messages_callbacks.py +++ b/tests/test_litellm_rust/test_messages_callbacks.py @@ -8,6 +8,8 @@ import pytest import litellm from litellm.integrations.custom_logger import CustomLogger +from litellm.litellm_core_utils.logging_worker import GLOBAL_LOGGING_WORKER +from litellm.types.utils import CallTypes from tests.test_litellm_rust.callback_recorder import RecordingLogger from tests.test_litellm_rust.contracts import ( MESSAGES, @@ -219,3 +221,132 @@ async def test_messages_stream_logs_success_after_exhaustion(messages_server: Re assert len(events) == 1 assert "async_log_stream_event" not in recorder.names assert events[0].kwargs["complete_streaming_response"] is not None + + +@pytest.mark.asyncio +async def test_messages_compression_hook_replaces_messages_sent_to_provider(messages_server: RecordingServer) -> None: + compressed_messages: Final = [{"role": "user", "content": "Compressed context"}] + call_types: Final = [] + + class CompressMessages(CustomLogger): + async def async_pre_call_deployment_hook(self, kwargs, call_type): + call_types.append(call_type) + kwargs["messages"] = compressed_messages + return kwargs + + litellm.callbacks.append(CompressMessages()) + + await call_messages(messages_server, []) + + assert call_types == [CallTypes.anthropic_messages] + assert messages_server.requests[0].body["messages"] == compressed_messages + + +@pytest.mark.asyncio +async def test_messages_post_call_guardrail_replacement_reaches_caller_and_logging( + messages_server: RecordingServer, +) -> None: + recorder: Final = RecordingLogger() + call_types: Final = [] + + class ReviewResponse(CustomLogger): + async def async_post_call_success_deployment_hook(self, request_data, response, call_type): + call_types.append(call_type) + response["content"][0]["text"] = "Reviewed response" + return response + + litellm.callbacks.append(ReviewResponse()) + + response: Final = await call_messages(messages_server, [recorder]) + events: Final = await recorder.wait_for_async("async_log_success_event") + + assert call_types == [CallTypes.anthropic_messages] + assert response["content"][0]["text"] == "Reviewed response" + assert events[0].response.choices[0].message.content == "Reviewed response" + + +@pytest.mark.asyncio +async def test_messages_logging_hook_replacement_reaches_later_loggers_only(messages_server: RecordingServer) -> None: + observations: Final = [] + + class RecordGuardrailVerdict(CustomLogger): + async def async_logging_hook(self, kwargs, result, call_type): + observations.append("guardrail") + return {**kwargs, "guardrail-verdict": "allowed"}, result + + class ExportLog(CustomLogger): + async def async_logging_hook(self, kwargs, result, call_type): + return kwargs, result + + async def async_log_success_event(self, kwargs, response_obj, start_time, end_time): + observations.append(("export", kwargs["guardrail-verdict"])) + + response: Final = await call_messages(messages_server, [RecordGuardrailVerdict(), ExportLog()]) + await asyncio.wait_for(GLOBAL_LOGGING_WORKER.flush(), timeout=10) + + assert observations == ["guardrail", ("export", "allowed")] + assert response["content"][0]["text"] == "Hello from native Messages" + + +@pytest.mark.asyncio +async def test_messages_success_callback_failure_does_not_skip_later_loggers( + messages_server: RecordingServer, +) -> None: + recorder: Final = RecordingLogger() + + class UnavailableExporter(CustomLogger): + async def async_log_success_event(self, kwargs, response_obj, start_time, end_time): + raise RuntimeError("exporter unavailable") + + response: Final = await call_messages(messages_server, [UnavailableExporter(), recorder]) + events: Final = await recorder.wait_for_async("async_log_success_event") + + assert response["content"][0]["text"] == "Hello from native Messages" + assert len(events) == 1 + assert "async_log_failure_event" not in recorder.names + + +@pytest.mark.asyncio +async def test_messages_concurrent_calls_keep_callback_state_isolated(messages_server: RecordingServer) -> None: + messages_server.expected_requests = 4 + tokens: Final = {f"messages-{index}": object() for index in range(4)} + terminal_state: Final = [] + + class CorrelateCallState(RecordingLogger): + def log_pre_api_call(self, model, messages, kwargs): + kwargs["correlation-token"] = tokens[kwargs["litellm_call_id"]] + super().log_pre_api_call(model, messages, kwargs) + + async def async_log_success_event(self, kwargs, response_obj, start_time, end_time): + terminal_state.append((kwargs["litellm_call_id"], kwargs["correlation-token"])) + await super().async_log_success_event(kwargs, response_obj, start_time, end_time) + + correlate: Final = CorrelateCallState() + + await asyncio.gather(*(call_messages(messages_server, [correlate], litellm_call_id=call_id) for call_id in tokens)) + await correlate.wait_for_async("async_log_success_event", count=4) + + assert len(terminal_state) == 4 + assert all(token is tokens[call_id] for call_id, token in terminal_state) + + +@pytest.mark.asyncio +async def test_messages_cancelled_call_runs_no_terminal_callbacks(messages_server: RecordingServer) -> None: + messages_server.default_response = ResponseSpec(body=MESSAGES_RESPONSE, delay=0.5) + recorder: Final = RecordingLogger() + task: Final = asyncio.create_task(call_messages(messages_server, [recorder])) + + async with asyncio.timeout(10): + while not messages_server.requests: + await asyncio.sleep(0.01) + task.cancel() + + with pytest.raises(asyncio.CancelledError): + await task + await asyncio.wait_for(GLOBAL_LOGGING_WORKER.flush(), timeout=10) + + assert recorder.names.count("log_pre_api_call") == 1 + assert "log_success_event" not in recorder.names + assert "async_log_success_event" not in recorder.names + assert "log_failure_event" not in recorder.names + assert "async_log_failure_event" not in recorder.names diff --git a/tests/test_litellm_rust/test_ocr_callbacks.py b/tests/test_litellm_rust/test_ocr_callbacks.py index 6cbf0ce8eb0..e0cbb446c59 100644 --- a/tests/test_litellm_rust/test_ocr_callbacks.py +++ b/tests/test_litellm_rust/test_ocr_callbacks.py @@ -11,7 +11,6 @@ from litellm.integrations.custom_logger import CustomLogger from tests.test_litellm_rust.callback_recorder import RecordingLogger from tests.test_litellm_rust.contracts import ( OCR_DOCUMENT, - OCR_MODEL, OCR_RESPONSE, call_native_aocr, call_native_ocr, @@ -247,3 +246,69 @@ async def test_pre_call_runs_in_callers_execution_context(ocr_server: RecordingS assert len(events) == 1 assert events[0].loop is caller_loop assert events[0].thread is caller_thread + + +@pytest.mark.asyncio +async def test_ocr_failure_callbacks_receive_pre_call_state(ocr_server: RecordingServer) -> None: + ocr_server.enqueue(ResponseSpec(body={"message": "provider unavailable"}, status=500)) + token: Final = object() + observed: Final = [] + + class TrackInFlightRequest(CustomLogger): + def log_pre_api_call(self, model, messages, kwargs): + kwargs["request-token"] = token + + def log_failure_event(self, kwargs, response_obj, start_time, end_time): + observed.append(("sync", kwargs["request-token"])) + + async def async_log_failure_event(self, kwargs, response_obj, start_time, end_time): + observed.append(("async", kwargs["request-token"])) + + with pytest.raises(litellm.InternalServerError): + await call_aocr(ocr_server, [TrackInFlightRequest()]) + + assert [event for event, _ in observed] == ["sync", "async"] + assert all(observed_token is token for _, observed_token in observed) + + +@pytest.mark.asyncio +async def test_ocr_failure_callback_error_does_not_mask_provider_error_or_later_callbacks( + ocr_server: RecordingServer, +) -> None: + ocr_server.enqueue(ResponseSpec(body={"message": "provider unavailable"}, status=500)) + recorder: Final = RecordingLogger() + + class UnavailableExporter(CustomLogger): + def log_failure_event(self, kwargs, response_obj, start_time, end_time): + raise RuntimeError("exporter unavailable") + + async def async_log_failure_event(self, kwargs, response_obj, start_time, end_time): + raise RuntimeError("exporter unavailable") + + with pytest.raises(litellm.InternalServerError) as caught: + await call_aocr(ocr_server, [UnavailableExporter(), recorder]) + + sync_events: Final = tuple(event for event in recorder.events if event.name == "log_failure_event") + async_events: Final = tuple(event for event in recorder.events if event.name == "async_log_failure_event") + assert len(sync_events) == 1 + assert len(async_events) == 1 + assert sync_events[0].kwargs["exception"] is caught.value + assert async_events[0].kwargs["exception"] is caught.value + assert "async_log_success_event" not in recorder.names + + +def test_ocr_duplicate_callback_registration_dispatches_once(ocr_server: RecordingServer) -> None: + recorder: Final = RecordingLogger() + + call_ocr( + ocr_server, + [recorder, recorder], + success_callback=[recorder], + failure_callback=[recorder], + ) + recorder.wait_for("log_success_event") + + assert recorder.names.count("log_pre_api_call") == 1 + assert recorder.names.count("logging_hook") == 1 + assert recorder.names.count("log_success_event") == 1 + assert "log_failure_event" not in recorder.names