Spaces:
Configuration error
Configuration error
File size: 6,739 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 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 196 197 198 199 200 201 202 203 204 205 206 207 208 209 210 211 212 213 |
import json
import os
import sys
import time
from datetime import datetime
from unittest.mock import AsyncMock, patch, MagicMock
import pytest
sys.path.insert(
0, os.path.abspath("../..")
) # Adds the parent directory to the system path
import litellm
from litellm.integrations.custom_logger import CustomLogger
from litellm.litellm_core_utils.logging_callback_manager import LoggingCallbackManager
from litellm.integrations.langfuse.langfuse_prompt_management import (
LangfusePromptManagement,
)
from litellm.integrations.opentelemetry import OpenTelemetry
# Test fixtures
@pytest.fixture
def callback_manager():
manager = LoggingCallbackManager()
# Reset callbacks before each test
manager._reset_all_callbacks()
return manager
@pytest.fixture
def mock_custom_logger():
class TestLogger(CustomLogger):
def log_success_event(self, kwargs, response_obj, start_time, end_time):
pass
return TestLogger()
# Test cases
def test_add_string_callback():
"""
Test adding a string callback to litellm.callbacks - only 1 instance of the string callback should be added
"""
manager = LoggingCallbackManager()
test_callback = "test_callback"
# Add string callback
manager.add_litellm_callback(test_callback)
assert test_callback in litellm.callbacks
# Test duplicate prevention
manager.add_litellm_callback(test_callback)
assert litellm.callbacks.count(test_callback) == 1
def test_duplicate_langfuse_logger_test():
manager = LoggingCallbackManager()
for _ in range(10):
langfuse_logger = LangfusePromptManagement()
manager.add_litellm_success_callback(langfuse_logger)
print("litellm.success_callback: ", litellm.success_callback)
assert len(litellm.success_callback) == 1
def test_duplicate_multiple_loggers_test():
manager = LoggingCallbackManager()
for _ in range(10):
langfuse_logger = LangfusePromptManagement()
otel_logger = OpenTelemetry()
manager.add_litellm_success_callback(langfuse_logger)
manager.add_litellm_success_callback(otel_logger)
print("litellm.success_callback: ", litellm.success_callback)
assert len(litellm.success_callback) == 2
# Check exactly one instance of each logger type
langfuse_count = sum(
1
for callback in litellm.success_callback
if isinstance(callback, LangfusePromptManagement)
)
otel_count = sum(
1
for callback in litellm.success_callback
if isinstance(callback, OpenTelemetry)
)
assert (
langfuse_count == 1
), "Should have exactly one LangfusePromptManagement instance"
assert otel_count == 1, "Should have exactly one OpenTelemetry instance"
def test_add_function_callback():
manager = LoggingCallbackManager()
def test_func(kwargs):
pass
# Add function callback
manager.add_litellm_callback(test_func)
assert test_func in litellm.callbacks
# Test duplicate prevention
manager.add_litellm_callback(test_func)
assert litellm.callbacks.count(test_func) == 1
def test_add_custom_logger(mock_custom_logger):
manager = LoggingCallbackManager()
# Add custom logger
manager.add_litellm_callback(mock_custom_logger)
assert mock_custom_logger in litellm.callbacks
def test_add_multiple_callback_types(mock_custom_logger):
manager = LoggingCallbackManager()
def test_func(kwargs):
pass
string_callback = "test_callback"
# Add different types of callbacks
manager.add_litellm_callback(string_callback)
manager.add_litellm_callback(test_func)
manager.add_litellm_callback(mock_custom_logger)
assert string_callback in litellm.callbacks
assert test_func in litellm.callbacks
assert mock_custom_logger in litellm.callbacks
assert len(litellm.callbacks) == 3
def test_success_failure_callbacks():
manager = LoggingCallbackManager()
success_callback = "success_callback"
failure_callback = "failure_callback"
# Add callbacks
manager.add_litellm_success_callback(success_callback)
manager.add_litellm_failure_callback(failure_callback)
assert success_callback in litellm.success_callback
assert failure_callback in litellm.failure_callback
def test_async_callbacks():
manager = LoggingCallbackManager()
async_success = "async_success"
async_failure = "async_failure"
# Add async callbacks
manager.add_litellm_async_success_callback(async_success)
manager.add_litellm_async_failure_callback(async_failure)
assert async_success in litellm._async_success_callback
assert async_failure in litellm._async_failure_callback
def test_remove_callback_from_list_by_object():
manager = LoggingCallbackManager()
# Reset all callbacks
manager._reset_all_callbacks()
def TestObject():
def __init__(self):
manager.add_litellm_callback(self.callback)
manager.add_litellm_success_callback(self.callback)
manager.add_litellm_failure_callback(self.callback)
manager.add_litellm_async_success_callback(self.callback)
manager.add_litellm_async_failure_callback(self.callback)
def callback(self):
pass
obj = TestObject()
manager.remove_callback_from_list_by_object(litellm.callbacks, obj)
manager.remove_callback_from_list_by_object(litellm.success_callback, obj)
manager.remove_callback_from_list_by_object(litellm.failure_callback, obj)
manager.remove_callback_from_list_by_object(litellm._async_success_callback, obj)
manager.remove_callback_from_list_by_object(litellm._async_failure_callback, obj)
# Verify all callback lists are empty
assert len(litellm.callbacks) == 0
assert len(litellm.success_callback) == 0
assert len(litellm.failure_callback) == 0
assert len(litellm._async_success_callback) == 0
assert len(litellm._async_failure_callback) == 0
def test_reset_callbacks(callback_manager):
# Add various callbacks
callback_manager.add_litellm_callback("test")
callback_manager.add_litellm_success_callback("success")
callback_manager.add_litellm_failure_callback("failure")
callback_manager.add_litellm_async_success_callback("async_success")
callback_manager.add_litellm_async_failure_callback("async_failure")
# Reset all callbacks
callback_manager._reset_all_callbacks()
# Verify all callback lists are empty
assert len(litellm.callbacks) == 0
assert len(litellm.success_callback) == 0
assert len(litellm.failure_callback) == 0
assert len(litellm._async_success_callback) == 0
assert len(litellm._async_failure_callback) == 0
|