Merge pull request #14974 from deepanshululla/feature/sqs_pushes_errors

Error logging in SQS
This commit is contained in:
Krish Dholakia 2025-09-27 07:31:23 -07:00 committed by GitHub
commit 0f0b34e936
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
2 changed files with 83 additions and 0 deletions

View file

@ -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)}"

View file

@ -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=<url_encoded_json>"
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"