scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-10-27 06:13:07 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 72

fix(proxy): share the billing logger across callback registries (#47510)

Browse source 
* fix(logging): deduplicate equivalent completed success callbacks

* fix(proxy): share the billing logger across callback registries

* test(proxy): use fixed timestamps for billing regressions

---------

Co-authored-by: Joshua Valluru <326636767+joshua-berri@users.noreply.github.com>
This commit is contained in:
joshua-berri 2026-10-20 08:21:49 -07:00 • committed by GitHub
parent 91724bb0a7
commit 02199b670b
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
3 changed files with 112 additions and 37 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

5
litellm/proxy/proxy_server.py
View file

@ -3118,8 +3118,9 @@ def cost_tracking():
from litellm.integrations.shadow_eval_logger import ShadowEvalLogger
spend_event_producer = build_spend_event_producer(CollectorSettings(), fallback=run_spend_event)
litellm.logging_callback_manager.add_litellm_callback(ProxyDBLogger(spend_event_producer))
litellm.logging_callback_manager.add_litellm_async_success_callback(ProxyDBLogger(spend_event_producer))
db_logger: Final = ProxyDBLogger(spend_event_producer)
litellm.logging_callback_manager.add_litellm_callback(db_logger)
litellm.logging_callback_manager.add_litellm_async_success_callback(db_logger)
litellm.logging_callback_manager.add_litellm_callback(ShadowEvalLogger())

72
tests/unit/litellm_core_utils/test_litellm_logging.py
View file

@ -2462,7 +2462,41 @@ def test_success_handler_runs_sync_callbacks_for_sync_requests(
@pytest.mark.asyncio
@pytest.mark.parametrize("registration", ["sync", "async", "overlapping"])
@pytest.mark.parametrize("different_headers", [False, True])
async def test_completed_success_preserves_generic_logger_configuration(
monkeypatch: pytest.MonkeyPatch, different_headers: bool
) -> None:
from litellm.integrations.generic_api.generic_api_callback import GenericAPILogger
first: Final = GenericAPILogger(
endpoint="https://logs.example.test/events",
headers={"Authorization": "Bearer first"},
event_types=["llm_api_success"],
)
second: Final = GenericAPILogger(
endpoint="https://logs.example.test/events",
headers={"Authorization": "Bearer second" if different_headers else "Bearer first"},
event_types=["llm_api_success"] if different_headers else ["llm_api_failure"],
)
monkeypatch.setattr(litellm, "success_callback", [first])
monkeypatch.setattr(litellm, "_async_success_callback", [second])
monkeypatch.setattr(litellm, "generic_api_use_v1", False)
logging_obj: Final = _async_request_logging_obj("distinct-generic-loggers", [])
response: Final = ModelResponse(id="chatcmpl-distinct-generic-loggers")
now: Final = datetime.datetime(2026, 1, 1)
await logging_obj.async_success_handler(response, now, now)
assert len(first.log_queue) == 1
assert first.log_queue[0]["id"] == response.id
if different_headers:
assert second.log_queue == first.log_queue
else:
assert second.log_queue == []
@pytest.mark.asyncio
@pytest.mark.parametrize("registration", ["sync", "async", "overlapping", "equivalent", "distinct"])
@pytest.mark.parametrize("kind", ["sync", "async", "dual"])
@pytest.mark.parametrize("compatibility_first", [False, True])
async def test_completed_success_selects_one_implemented_hook(
@ -2517,9 +2551,15 @@ async def test_completed_success_selects_one_implemented_hook(
) -> None:
record("async", kwargs, response_obj)
callback: Final = {"sync": SyncRecorder, "async": AsyncRecorder, "dual": DualRecorder}[kind]()
recorder: Final = {"sync": SyncRecorder, "async": AsyncRecorder, "dual": DualRecorder}[kind]
callback: Final = recorder()
async_callback: Final = (
recorder(turn_off_message_logging=registration == "distinct")
if registration in ("equivalent", "distinct")
else callback
)
monkeypatch.setattr(litellm, "success_callback", [callback] if registration != "async" else [])
monkeypatch.setattr(litellm, "_async_success_callback", [callback] if registration != "sync" else [])
monkeypatch.setattr(litellm, "_async_success_callback", [async_callback] if registration != "sync" else [])
logging_obj: Final = _async_request_logging_obj("completed-hook", [])
response: Final = ModelResponse(id="chatcmpl-completed-hook")
now: Final = datetime.datetime.now()
@ -2531,20 +2571,26 @@ async def test_completed_success_selects_one_implemented_hook(
await logging_obj.async_success_handler(response, now, now)
logging_obj.handle_sync_success_callbacks_for_async_calls(response, now, now)
assert events.qsize() == 1
hook, response_id, cost, thread_id = events.get_nowait()
assert hook == ("sync" if kind == "sync" else "async")
assert response_id == response.id
assert cost == logging_obj.model_call_details["response_cost"]
assert (thread_id == loop_thread) is (kind != "sync")
expected_count: Final = 2 if registration in ("equivalent", "distinct") else 1
assert events.qsize() == expected_count
for _ in range(expected_count):
hook, response_id, cost, thread_id = events.get_nowait()
assert hook == ("sync" if kind == "sync" else "async")
assert response_id == response.id
assert cost == logging_obj.model_call_details["response_cost"]
assert (thread_id == loop_thread) is (kind != "sync")
@pytest.mark.parametrize("for_async_request", [False, True])
@pytest.mark.parametrize("for_async_request,split_registries", [(False, False), (True, False), (True, True)])
@pytest.mark.parametrize("private_first", [False, True])
@pytest.mark.parametrize("response_kind", ["responses", "completion"])
def test_sync_success_redacts_only_opted_out_loggers(
monkeypatch: pytest.MonkeyPatch, for_async_request: bool, private_first: bool, response_kind: str
monkeypatch: pytest.MonkeyPatch,
for_async_request: bool,
split_registries: bool,
private_first: bool,
response_kind: str,
) -> None:
from queue import SimpleQueue
@ -2563,8 +2609,8 @@ def test_sync_success_redacts_only_opted_out_loggers(
private_logger: Final = PayloadRecorder(turn_off_message_logging=True)
public_logger: Final = PayloadRecorder()
callbacks: Final = [private_logger, public_logger] if private_first else [public_logger, private_logger]
monkeypatch.setattr(litellm, "success_callback", callbacks)
monkeypatch.setattr(litellm, "_async_success_callback", [])
monkeypatch.setattr(litellm, "success_callback", callbacks[:1] if split_registries else callbacks)
monkeypatch.setattr(litellm, "_async_success_callback", callbacks[1:] if split_registries else [])
monkeypatch.setattr(litellm, "turn_off_message_logging", False)
logging_obj: Final = LitellmLogging(
model="openai/gpt-6.1-sol",

72
tests/unit/proxy/proxy_server/test_lifecycle.py
View file

@ -697,33 +697,61 @@ def test_load_from_azure_key_vault_missing_uri_failure_is_swallowed(monkeypatch)
# ---------------------------------------------------------------------------
def test_cost_tracking_adds_db_and_shadow_eval_callbacks_when_prisma_set(monkeypatch):
@pytest.mark.asyncio
@pytest.mark.parametrize("setup_count", [1, 2])
async def test_cost_tracking_adds_db_and_shadow_eval_callbacks_when_prisma_set(
monkeypatch: pytest.MonkeyPatch, setup_count: int
) -> None:
from datetime import datetime
import litellm
from litellm.integrations.shadow_eval_logger import ShadowEvalLogger
from litellm.proxy.spend_tracking.spend_event import SpendEventDecodeError, decode_spend_event
from litellm.proxy.spend_tracking.spend_event_producer import SpendEventProducer
from litellm.types.utils import ModelResponse, Usage
from litellm.utils import Rules, function_setup
fake_prisma = MagicMock()
monkeypatch.setattr(ps, "prisma_client", fake_prisma, raising=False)
monkeypatch.setattr(litellm, "callbacks", [], raising=False)
monkeypatch.setattr(litellm, "_async_success_callback", [], raising=False)
producer: Final = AsyncMock(spec=SpendEventProducer)
monkeypatch.setattr(ps, "prisma_client", MagicMock())
monkeypatch.setattr(ps, "spend_event_producer", None)
monkeypatch.setattr(ps, "build_spend_event_producer", lambda settings, fallback: producer)
for name in (
"callbacks",
"input_callback",
"success_callback",
"failure_callback",
"_async_success_callback",
"_async_failure_callback",
):
monkeypatch.setattr(litellm, name, [])
before_callbacks = len(litellm.callbacks)
before_async = len(litellm._async_success_callback)
for _ in range(setup_count):
cost_tracking()
cost_tracking()
cost_tracking()
observed = {
"added_to_callbacks": len(litellm.callbacks) - before_callbacks,
"added_to_async_success": len(litellm._async_success_callback) - before_async,
"shadow_eval_loggers": sum(isinstance(cb, ShadowEvalLogger) for cb in litellm.callbacks),
"prisma_was_set": True,
}
assert normalize(observed) == {
"added_to_callbacks": 2,
"added_to_async_success": 1,
"shadow_eval_loggers": 1,
"prisma_was_set": True,
}
assert len(litellm.callbacks) == 2
assert sum(isinstance(cb, ShadowEvalLogger) for cb in litellm.callbacks) == 1
assert len(litellm._async_success_callback) == 1
now: Final = datetime(2026, 1, 1)
logging_obj, _ = function_setup(
"acompletion",
Rules(),
now,
model="openai/gpt-4o-mini",
messages=[{"role": "user", "content": "hello"}],
litellm_call_id="single-accounting-event",
)
logging_obj.update_environment_variables(litellm_params={"acompletion": True, "metadata": {}}, optional_params={})
response: Final = ModelResponse(
id="chatcmpl-single-accounting-event",
model="gpt-4o-mini",
usage=Usage(prompt_tokens=10, completion_tokens=5, total_tokens=15),
)
await logging_obj.async_success_handler(response, now, now)
producer.publish.assert_awaited_once()
event: Final = decode_spend_event(producer.publish.call_args.args[0])
assert not isinstance(event, SpendEventDecodeError)
assert event.response_cost == logging_obj.model_call_details["response_cost"]
assert event.response_cost > 0
def test_cost_tracking_no_op_when_prisma_missing(monkeypatch):