phoenix/litellm-mirror
0
0
Fork 1
mirror of https://github.com/BerriAI/litellm.git synced 2025-04-26 11:14:04 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions

fix(utils.py): move adding custom logger callback to success event in… (#7905)

Browse source 
* fix(utils.py): move adding custom logger callback to success event into separate function + don't add success callback to failure event

if user is explicitly choosing 'success' callback, don't log failure as well

* test(test_utils.py): add unit test to ensure custom logger callback only adds callback to specific event

* fix(utils.py): remove string from list of callbacks once corresponding callback class is added

prevents floating values - simplifies testing

* fix(utils.py): fix linting error

* test: cleanup args before test

* test: fix test

* test: update test

* test: fix test
This commit is contained in:
Krish Dholakia 2025-01-22 21:49:09 -08:00 • committed by GitHub
parent cefbada875
commit 4911cd80a1
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
5 changed files with 98 additions and 49 deletions
Download patch file Download diff file Expand all files Collapse all files

48
tests/local_testing/test_utils.py
View file

@ -1496,24 +1496,34 @@ def test_get_num_retries(num_retries):
)
@pytest.mark.parametrize("filter_invalid_headers", [True, False])
@pytest.mark.parametrize(
"custom_llm_provider, expected_result",
[("anthropic", {"anthropic-beta": "123"}), ("bedrock", {}), ("vertex_ai", {})],
)
def test_get_clean_extra_headers(
filter_invalid_headers, custom_llm_provider, expected_result, monkeypatch
):
from litellm.utils import get_clean_extra_headers
def test_add_custom_logger_callback_to_specific_event(monkeypatch):
from litellm.utils import _add_custom_logger_callback_to_specific_event
monkeypatch.setattr(litellm, "filter_invalid_headers", filter_invalid_headers)
monkeypatch.setattr(litellm, "success_callback", [])
monkeypatch.setattr(litellm, "failure_callback", [])
if filter_invalid_headers:
assert (
get_clean_extra_headers({"anthropic-beta": "123"}, custom_llm_provider)
== expected_result
)
else:
assert get_clean_extra_headers(
{"anthropic-beta": "123"}, custom_llm_provider
) == {"anthropic-beta": "123"}
_add_custom_logger_callback_to_specific_event("langfuse", "success")
assert len(litellm.success_callback) == 1
assert len(litellm.failure_callback) == 0
def test_add_custom_logger_callback_to_specific_event_e2e(monkeypatch):
monkeypatch.setattr(litellm, "success_callback", [])
monkeypatch.setattr(litellm, "failure_callback", [])
monkeypatch.setattr(litellm, "callbacks", [])
litellm.success_callback = ["humanloop"]
curr_len_success_callback = len(litellm.success_callback)
curr_len_failure_callback = len(litellm.failure_callback)
litellm.completion(
model="gpt-4o-mini",
messages=[{"role": "user", "content": "Hello, world!"}],
mock_response="Testing langfuse",
)
assert len(litellm.success_callback) == curr_len_success_callback
assert len(litellm.failure_callback) == curr_len_failure_callback