diff --git a/litellm/integrations/sqs.py b/litellm/integrations/sqs.py index 2a0c73dfdbf..8a2ebf8d344 100644 --- a/litellm/integrations/sqs.py +++ b/litellm/integrations/sqs.py @@ -7,6 +7,7 @@ This logger sends ``StandardLoggingPayload`` entries to an AWS SQS queue. from __future__ import annotations import asyncio +import traceback from typing import List, Optional import litellm @@ -200,6 +201,25 @@ class SQSLogger(CustomBatchLogger, BaseAWSLLM): except Exception as e: verbose_logger.exception(f"sqs Layer Error - {str(e)}") + async def async_log_failure_event(self, kwargs, response_obj, start_time, end_time): + try: + standard_logging_payload = kwargs.get("standard_logging_object") + if standard_logging_payload is None: + raise ValueError("standard_logging_payload is None") + + self.log_queue.append(standard_logging_payload) + verbose_logger.debug( + "sqs logging: queue length %s, batch size %s", + len(self.log_queue), + self.batch_size, + ) + + except Exception as e: + verbose_logger.exception( + f"Datadog Layer Error - {str(e)}\n{traceback.format_exc()}" + ) + pass + async def async_send_batch(self) -> None: verbose_logger.debug( f"sqs logger - sending batch of {len(self.log_queue)}" diff --git a/tests/logging_callback_tests/test_sqs_logger.py b/tests/logging_callback_tests/test_sqs_logger.py index a3fe09b2a51..38d730712f3 100644 --- a/tests/logging_callback_tests/test_sqs_logger.py +++ b/tests/logging_callback_tests/test_sqs_logger.py @@ -71,3 +71,66 @@ async def test_async_sqs_logger_flush(): assert len(payload_data["messages"]) == 1 assert payload_data["messages"][0]["role"] == "user" assert payload_data["messages"][0]["content"] == "hello" + + +@pytest.mark.asyncio +async def test_async_sqs_logger_error_flush(): + expected_queue_url = "https://sqs.us-east-1.amazonaws.com/123456789012/test-queue" + expected_region = "us-east-1" + + sqs_logger = SQSLogger( + sqs_queue_url=expected_queue_url, + sqs_region_name=expected_region, + sqs_flush_interval=1, + ) + + # Mock the httpx client + mock_response = MagicMock() + mock_response.raise_for_status = Exception("Something went wrong") + sqs_logger.async_httpx_client.post = AsyncMock(return_value=mock_response) + + litellm.callbacks = [sqs_logger] + + await litellm.acompletion( + model="gpt-4o", + messages=[{"role": "user", "content": "hello"}], + mock_response="Error occurred" + ) + + await asyncio.sleep(2) + + # Verify that httpx post was called + sqs_logger.async_httpx_client.post.assert_called() + + # Get the call arguments + call_args = sqs_logger.async_httpx_client.post.call_args + + # Verify the URL is correct + called_url = call_args[0][0] # First positional argument + assert called_url == expected_queue_url, f"Expected URL {expected_queue_url}, got {called_url}" + + # Verify the payload contains StandardLoggingPayload data + called_data = call_args.kwargs['data'] + + # Extract the MessageBody from the URL-encoded data + # Format: "Action=SendMessage&Version=2012-11-05&MessageBody=" + assert "Action=SendMessage" in called_data + assert "Version=2012-11-05" in called_data + assert "MessageBody=" in called_data + + # Extract and decode the message body + message_body_start = called_data.find("MessageBody=") + len("MessageBody=") + message_body_encoded = called_data[message_body_start:] + message_body_json = unquote(message_body_encoded) + + # Parse the JSON to verify it's a StandardLoggingPayload + payload_data = json.loads(message_body_json) + + # Verify it has the expected StandardLoggingPayload structure + assert "model" in payload_data + assert "messages" in payload_data + assert "response" in payload_data + assert payload_data["model"] == "gpt-4o" + assert len(payload_data["messages"]) == 1 + assert payload_data["messages"][0]["role"] == "user" + assert payload_data["messages"][0]["content"] == "hello"