diff --git a/sentry_sdk/integrations/litellm.py b/sentry_sdk/integrations/litellm.py index 08cb217962..06ae8a0782 100644 --- a/sentry_sdk/integrations/litellm.py +++ b/sentry_sdk/integrations/litellm.py @@ -19,6 +19,7 @@ try: import litellm # type: ignore[import-not-found] + from litellm import input_callback, success_callback, failure_callback except ImportError: raise DidNotEnable("LiteLLM not installed") @@ -278,14 +279,14 @@ def __init__(self: "LiteLLMIntegration", include_prompts: bool = True) -> None: @staticmethod def setup_once() -> None: """Set up LiteLLM callbacks for monitoring.""" - litellm.input_callback = litellm.input_callback or [] + litellm.input_callback = input_callback or [] if _input_callback not in litellm.input_callback: litellm.input_callback.append(_input_callback) - litellm.success_callback = litellm.success_callback or [] + litellm.success_callback = success_callback or [] if _success_callback not in litellm.success_callback: litellm.success_callback.append(_success_callback) - litellm.failure_callback = litellm.failure_callback or [] + litellm.failure_callback = failure_callback or [] if _failure_callback not in litellm.failure_callback: litellm.failure_callback.append(_failure_callback) diff --git a/tests/test_shadowed_module.py b/tests/test_shadowed_module.py index fed3abe24a..10d86f285b 100644 --- a/tests/test_shadowed_module.py +++ b/tests/test_shadowed_module.py @@ -26,11 +26,7 @@ def pytest_generate_tests(metafunc): metafunc.parametrize( "integration_submodule_name", - # Temporarily skip some integrations - submodule_names - - { - "litellm", - }, + submodule_names, )