diff --git a/headroom/integrations/litellm_callback.py b/headroom/integrations/litellm_callback.py index 3d5e137a1..4c2e4e271 100644 --- a/headroom/integrations/litellm_callback.py +++ b/headroom/integrations/litellm_callback.py @@ -24,10 +24,18 @@ logger = logging.getLogger(__name__) _DEFAULT_CLOUD_URL = "https://api.headroomlabs.ai" -class HeadroomCallback: +try: + from litellm.integrations.custom_logger import CustomLogger as _CustomLogger +except ImportError: # litellm not installed — fall back to plain object + _CustomLogger = object # type: ignore[assignment,misc] + + +class HeadroomCallback(_CustomLogger): """LiteLLM callback that compresses messages before each API call. - Implements LiteLLM's CustomLogger interface (async_pre_call_hook). + Inherits from litellm.integrations.custom_logger.CustomLogger so that + any hook LiteLLM adds in future versions (e.g. async_post_call_success_hook + added in 1.89.x) has a no-op default and won't raise AttributeError (#1114). Two modes: - Local (default): Compresses in-process using headroom.compress(). @@ -56,6 +64,7 @@ class HeadroomCallback: api_key: str | None = None, api_url: str | None = None, ) -> None: + super().__init__() self._min_tokens = min_tokens self._model_limit = model_limit self._hooks = hooks diff --git a/tests/test_litellm_callback.py b/tests/test_litellm_callback.py new file mode 100644 index 000000000..5ad40d265 --- /dev/null +++ b/tests/test_litellm_callback.py @@ -0,0 +1,72 @@ +"""Tests for HeadroomCallback LiteLLM integration. + +Regression for #1114: HeadroomCallback did not inherit CustomLogger, so any +hook LiteLLM added post-1.89.x (e.g. async_post_call_success_hook) raised +AttributeError and crashed the LiteLLM proxy. +""" + +from __future__ import annotations + +import asyncio + +import pytest + +from tests._dotenv import importorskip_no_env_leak + +importorskip_no_env_leak("litellm") + +from headroom.integrations.litellm_callback import HeadroomCallback # noqa: E402 + + +class TestHeadroomCallbackCustomLoggerInheritance: + def test_instantiates_without_error(self) -> None: + cb = HeadroomCallback() + assert cb is not None + + def test_has_async_post_call_success_hook(self) -> None: + """Regression: AttributeError: 'HeadroomCallback' has no attr 'async_post_call_success_hook'.""" + cb = HeadroomCallback() + assert hasattr(cb, "async_post_call_success_hook"), ( + "async_post_call_success_hook must exist (added in litellm 1.89.x)" + ) + + def test_async_post_call_success_hook_is_callable(self) -> None: + """LiteLLM must be able to await the hook without exception.""" + cb = HeadroomCallback() + hook = cb.async_post_call_success_hook + assert callable(hook) + + def test_async_post_call_success_hook_does_not_raise(self) -> None: + """Calling the hook (no-op from CustomLogger) must not raise.""" + cb = HeadroomCallback() + + async def _run() -> None: + await cb.async_post_call_success_hook( + data={"model": "gpt-4o", "messages": []}, + user_api_key_dict={}, + response=None, + ) + + asyncio.run(_run()) + + def test_all_current_litellm_async_hooks_present(self) -> None: + """HeadroomCallback must expose every async hook CustomLogger defines.""" + from litellm.integrations.custom_logger import CustomLogger + + cb = HeadroomCallback() + missing = [ + name + for name in dir(CustomLogger) + if name.startswith("async_") and not hasattr(cb, name) + ] + assert not missing, f"Missing CustomLogger hooks: {missing}" + + def test_async_pre_call_hook_still_works(self) -> None: + """Inheritance must not break the existing compression hook.""" + cb = HeadroomCallback() + assert hasattr(cb, "async_pre_call_hook") + assert callable(cb.async_pre_call_hook) + + def test_total_tokens_saved_property(self) -> None: + cb = HeadroomCallback() + assert cb.total_tokens_saved == 0