diff --git a/headroom/providers/litellm.py b/headroom/providers/litellm.py index cc374d5e6..95d13311a 100644 --- a/headroom/providers/litellm.py +++ b/headroom/providers/litellm.py @@ -22,6 +22,7 @@ Requires: pip install litellm from __future__ import annotations import logging +import os from typing import Any from headroom.tokenizers import EstimatingTokenCounter @@ -32,6 +33,10 @@ logger = logging.getLogger(__name__) # Check if litellm is available try: + # LiteLLM can print its provider-list banner during import, before the + # module-level suppression flags below can be set. + os.environ.setdefault("LITELLM_SUPPRESS_DEBUG_INFO", "True") + import litellm # Suppress litellm's startup banner ("Provider List: https://...") and diff --git a/tests/test_startup_log_noise.py b/tests/test_startup_log_noise.py index 0da69ddeb..b04c72b93 100644 --- a/tests/test_startup_log_noise.py +++ b/tests/test_startup_log_noise.py @@ -9,8 +9,12 @@ Covers the fixes in: from __future__ import annotations +import builtins +import importlib import logging +import sys import warnings +from types import ModuleType class TestAnthropicWarnParameter: @@ -110,6 +114,40 @@ class TestEmbedderLogLevels: class TestLiteLLMLogSuppression: """litellm startup banner suppression must be applied at import time.""" + def test_litellm_suppress_env_is_set_before_import(self, monkeypatch): + """The env flag must exist before litellm itself is imported.""" + import os + + monkeypatch.delenv("LITELLM_SUPPRESS_DEBUG_INFO", raising=False) + sys.modules.pop("headroom.providers.litellm", None) + sys.modules.pop("litellm", None) + + original_import = builtins.__import__ + fake_litellm = ModuleType("litellm") + fake_litellm.suppress_debug_info = False + fake_litellm.set_verbose = True + fake_litellm.get_model_info = lambda _model: {} + fake_litellm.model_cost = {} + fake_litellm.token_counter = lambda **_kwargs: 0 + observed_env: list[str | None] = [] + + def import_spy(name, globals=None, locals=None, fromlist=(), level=0): + if name == "litellm": + observed_env.append(os.environ.get("LITELLM_SUPPRESS_DEBUG_INFO")) + sys.modules["litellm"] = fake_litellm + return fake_litellm + return original_import(name, globals, locals, fromlist, level) + + monkeypatch.setattr(builtins, "__import__", import_spy) + try: + importlib.import_module("headroom.providers.litellm") + finally: + sys.modules.pop("headroom.providers.litellm", None) + sys.modules.pop("litellm", None) + + assert observed_env + assert all(value == "True" for value in observed_env) + def test_litellm_suppress_debug_info_is_set(self): """litellm.suppress_debug_info must be True after importing the litellm provider.""" litellm = pytest_importorskip_litellm()