mirror of
https://github.com/BerriAI/litellm.git
synced 2026-09-10 22:41:41 +00:00
Merge pull request #14974 from deepanshululla/feature/sqs_pushes_errors
Error logging in SQS
This commit is contained in:
commit
0f0b34e936
2 changed files with 83 additions and 0 deletions
|
|
@ -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)}"
|
||||
|
|
|
|||
|
|
@ -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"
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue