11import io
22import os
3- from re import M
43import sys
54
6- from litellm .integrations .custom_logger import CustomLogger
7-
85
96sys .path .insert (0 , os .path .abspath ("../.." ))
107
@@ -395,28 +392,22 @@ async def test_datadog_payload_environment_variables():
395392 pytest .fail (f"Test failed with exception: { str (e )} " )
396393
397394
398- class TestDDLogger (CustomLogger ):
399- def __init__ (self ):
400- self .standard_logging_object : Optional [StandardLoggingPayload ] = None
401-
402- async def async_log_success_event (self , kwargs , response_obj , start_time , end_time ):
403- self .standard_logging_object = kwargs ["standard_logging_object" ]
404-
405-
406395@pytest .mark .asyncio
407396async def test_datadog_payload_content_truncation ():
408397 """
409398 Test that DataDog payload correctly truncates long content
410399
411400 DataDog has a limit of 1MB for the logged payload size.
412401 """
413- dd_logger = TestDDLogger ()
414- litellm .callbacks = [dd_logger ]
402+ dd_logger = DataDogLogger ()
415403
404+ # Create a standard payload with very long content
405+ standard_payload = create_standard_logging_payload ()
416406 long_content = "x" * 80_000 # Create string longer than MAX_STR_LENGTH (10_000)
417407
418- # messages with long content
419- messages = [
408+ # Modify payload with long content
409+ standard_payload ["error_str" ] = long_content
410+ standard_payload ["messages" ] = [
420411 {
421412 "role" : "user" ,
422413 "content" : [
@@ -430,26 +421,28 @@ async def test_datadog_payload_content_truncation():
430421 ],
431422 }
432423 ]
433- await litellm .acompletion (
434- model = "gpt-3.5-turbo" ,
435- messages = messages ,
436- temperature = 0.2 ,
437- mock_response = long_content ,
438- )
439-
440- await asyncio .sleep (2 )
424+ standard_payload ["response" ] = {"choices" : [{"message" : {"content" : long_content }}]}
441425
442426 # Create the payload
443- standard_logging_payload = dd_logger .standard_logging_object
427+ dd_payload = dd_logger .create_datadog_logging_payload (
428+ kwargs = {"standard_logging_object" : standard_payload },
429+ response_obj = None ,
430+ start_time = datetime .now (),
431+ end_time = datetime .now (),
432+ )
433+
434+ print ("dd_payload" , json .dumps (dd_payload , indent = 2 ))
444435
445- print ("standard_logging_payload" , json .dumps (standard_logging_payload , indent = 2 ))
436+ # Parse the message back to dict to verify truncation
437+ message_dict = json .loads (dd_payload ["message" ])
446438
447439 # Verify truncation of fields
440+ assert len (message_dict ["error_str" ]) < 10_100 , "error_str not truncated correctly"
448441 assert (
449- len (str (standard_logging_payload ["messages" ])) < 10_100
442+ len (str (message_dict ["messages" ])) < 10_100
450443 ), "messages not truncated correctly"
451444 assert (
452- len (str (standard_logging_payload ["response" ])) < 10_100
445+ len (str (message_dict ["response" ])) < 10_100
453446 ), "response not truncated correctly"
454447
455448
0 commit comments