diff --git a/litellm/integrations/newrelic/newrelic.py b/litellm/integrations/newrelic/newrelic.py index e5c40a4b064..6f2ac539b86 100644 --- a/litellm/integrations/newrelic/newrelic.py +++ b/litellm/integrations/newrelic/newrelic.py @@ -243,7 +243,6 @@ class NewRelicLogger(CustomLogger): litellm_params = kwargs.get("litellm_params", {}) metadata = litellm_params.get("metadata", {}) headers = metadata.get("headers", {}) - newrelic = headers.get("newrelic", None) traceparent = headers.get("traceparent", None) trace_id = None @@ -256,10 +255,6 @@ class NewRelicLogger(CustomLogger): if len(parts) == 4: trace_id = parts[1] - if not trace_id: - # attempt to pull from newrelic header - pass - if not trace_id: # Generate a random trace_id for grouping AI monitoring events trace_id = str(uuid.uuid4()) diff --git a/litellm/types/integrations/newrelic.py b/litellm/types/integrations/newrelic.py index d3589867cc6..2de9769b181 100644 --- a/litellm/types/integrations/newrelic.py +++ b/litellm/types/integrations/newrelic.py @@ -1,5 +1,3 @@ -from typing import Optional - from litellm.types.integrations.custom_logger import StandardCustomLoggerInitParams diff --git a/tests/test_litellm/integrations/newrelic/test_newrelic.py b/tests/test_litellm/integrations/newrelic/test_newrelic.py index d6bfbea6908..fce9db561e5 100644 --- a/tests/test_litellm/integrations/newrelic/test_newrelic.py +++ b/tests/test_litellm/integrations/newrelic/test_newrelic.py @@ -2,9 +2,21 @@ import os import sys from unittest.mock import MagicMock, patch -import newrelic.agent # ensure submodule is loaded so patch() can resolve it import pytest +# newrelic is a container-only dependency (requirements.txt) and is not installed +# in the CI Python environment. Mock it in sys.modules before importing the +# integration so that deferred `import newrelic.agent` calls inside NewRelicLogger +# methods resolve to these mocks rather than failing with ModuleNotFoundError. +_mock_newrelic = MagicMock() +_mock_newrelic_agent = MagicMock() +# Explicitly link so _mock_newrelic.agent IS _mock_newrelic_agent. Without this, +# the first getattr(_mock_newrelic, 'agent') auto-creates a different child mock, +# causing patch("newrelic.agent.xxx") to patch the wrong object. +_mock_newrelic.agent = _mock_newrelic_agent +sys.modules["newrelic"] = _mock_newrelic +sys.modules["newrelic.agent"] = _mock_newrelic_agent + sys.path.insert(0, os.path.abspath("../..")) from litellm.integrations.newrelic.newrelic import NewRelicLogger @@ -22,9 +34,8 @@ NR_ENV = { def make_logger(**kwargs) -> NewRelicLogger: """Instantiate NewRelicLogger with NR agent calls mocked out.""" - with patch("newrelic.agent.register_application"): - with patch.dict(os.environ, NR_ENV): - return NewRelicLogger(**kwargs) + with patch.dict(os.environ, NR_ENV): + return NewRelicLogger(**kwargs) def make_kwargs( @@ -106,8 +117,8 @@ class TestNewRelicLoggerInit: assert logger.enabled is True def test_disabled_on_import_error(self): - with patch( - "newrelic.agent.register_application", side_effect=ImportError + with patch.object( + _mock_newrelic_agent, "register_application", side_effect=ImportError ): with patch.dict(os.environ, NR_ENV): logger = NewRelicLogger()