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 e1171dd103..ee81822973 100644 --- a/tests/test_shadowed_module.py +++ b/tests/test_shadowed_module.py @@ -31,7 +31,6 @@ def pytest_generate_tests(metafunc): - { "clickhouse_driver", "grpc", - "litellm", "opentelemetry", "pure_eval", "ray",