Spaces:
Configuration error
Configuration error
File size: 4,188 Bytes
447ebeb |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 |
import io
import os
import sys
from typing import Optional
sys.path.insert(0, os.path.abspath("../.."))
import asyncio
import gzip
import json
import logging
import time
from unittest.mock import AsyncMock, patch
import pytest
import litellm
from litellm._logging import verbose_logger
from litellm.integrations.custom_logger import CustomLogger
from litellm.types.utils import StandardLoggingPayload
class TestCustomLogger(CustomLogger):
def __init__(self, *args, **kwargs):
super().__init__(*args, **kwargs)
self.logged_standard_logging_payload: Optional[StandardLoggingPayload] = None
async def async_log_success_event(self, kwargs, response_obj, start_time, end_time):
standard_logging_payload = kwargs.get("standard_logging_object", None)
self.logged_standard_logging_payload = standard_logging_payload
@pytest.mark.asyncio
async def test_global_redaction_on():
litellm.turn_off_message_logging = True
test_custom_logger = TestCustomLogger()
litellm.callbacks = [test_custom_logger]
response = await litellm.acompletion(
model="gpt-3.5-turbo",
messages=[{"role": "user", "content": "hi"}],
mock_response="hello",
)
await asyncio.sleep(1)
standard_logging_payload = test_custom_logger.logged_standard_logging_payload
assert standard_logging_payload is not None
assert standard_logging_payload["response"] == {"text": "redacted-by-litellm"}
assert standard_logging_payload["messages"][0]["content"] == "redacted-by-litellm"
print(
"logged standard logging payload",
json.dumps(standard_logging_payload, indent=2),
)
@pytest.mark.parametrize("turn_off_message_logging", [True, False])
@pytest.mark.asyncio
async def test_global_redaction_with_dynamic_params(turn_off_message_logging):
litellm.turn_off_message_logging = True
test_custom_logger = TestCustomLogger()
litellm.callbacks = [test_custom_logger]
response = await litellm.acompletion(
model="gpt-3.5-turbo",
messages=[{"role": "user", "content": "hi"}],
turn_off_message_logging=turn_off_message_logging,
mock_response="hello",
)
await asyncio.sleep(1)
standard_logging_payload = test_custom_logger.logged_standard_logging_payload
assert standard_logging_payload is not None
print(
"logged standard logging payload",
json.dumps(standard_logging_payload, indent=2),
)
if turn_off_message_logging is True:
assert standard_logging_payload["response"] == {"text": "redacted-by-litellm"}
assert (
standard_logging_payload["messages"][0]["content"] == "redacted-by-litellm"
)
else:
assert (
standard_logging_payload["response"]["choices"][0]["message"]["content"]
== "hello"
)
assert standard_logging_payload["messages"][0]["content"] == "hi"
@pytest.mark.parametrize("turn_off_message_logging", [True, False])
@pytest.mark.asyncio
async def test_global_redaction_off_with_dynamic_params(turn_off_message_logging):
litellm.turn_off_message_logging = False
test_custom_logger = TestCustomLogger()
litellm.callbacks = [test_custom_logger]
response = await litellm.acompletion(
model="gpt-3.5-turbo",
messages=[{"role": "user", "content": "hi"}],
turn_off_message_logging=turn_off_message_logging,
mock_response="hello",
)
await asyncio.sleep(1)
standard_logging_payload = test_custom_logger.logged_standard_logging_payload
assert standard_logging_payload is not None
print(
"logged standard logging payload",
json.dumps(standard_logging_payload, indent=2),
)
if turn_off_message_logging is True:
assert standard_logging_payload["response"] == {"text": "redacted-by-litellm"}
assert (
standard_logging_payload["messages"][0]["content"] == "redacted-by-litellm"
)
else:
assert (
standard_logging_payload["response"]["choices"][0]["message"]["content"]
== "hello"
)
assert standard_logging_payload["messages"][0]["content"] == "hi"
|