From 45f9e26760ff5582fc563bc9b1a22741851b8380 Mon Sep 17 00:00:00 2001 From: Josh Bonczkowski Date: Thu, 11 Jun 2026 07:52:14 -0400 Subject: [PATCH] feat(newrelic): Add New Relic extension (#26989) * initial New Relic integration. * Minor fixes for basic observability. * Implemented basic support for the success path. Generates New Relic custom events needed by the AI Monitorin interface. * Supportability metric is sent on first request. * Emit supportability metric every hour instead of once a day. * Add the start/end times to the messages before sending them so that the start time and end time reflect the correct time and both are not set to 'now'. * Make use of `turn_off_message_logging` configuration that is available by default from CustomLogger. * Enabling New Relic agent to be wired when docker container starts if an environment variable is set. * If we cannot find trace information, send the AI events without the trace ID attached. * Use a fake trace_id if we cannot find one. * Implementing a configuration so that users can use litellm configuration to disable sending LLM messages to New Relic. There is a second method to do this via New Relic env var. * Mised file. * Cleaning up logic to turn off recording content via either the LiteLLM configuration or an env var. * Removing debugging. Fixed logic / comments around how often to send supportability metric. * Initial version of public doc for New Relic. * Use a proper name for the doc file. * Updating newrelic.md document. * Updating LiteLLM documentation for New Relic extension. * Moving New Relic imports into the methods to support unit tests. * Adding unit tests for the New Relic extension. * Updating linting and the unit tests that are not running in the CI environment. * Address reviewer feedback on New Relic integration. - Fix _record_error_metric to use app.record_custom_metric() instead of module-level newrelic.agent.record_custom_metric() so the call works outside of an active transaction context - Remove unreachable except ImportError block in _get_trace_context - Update stale "23 hours" comment to "27 hours" (matches 97200s threshold) - Remove commented-out debug code from _process_success - Fix docs typo: NEW_RELIC_CUSTOM_INSIGHTS_EVENTS_MAX_SAMPLES_STOREDA -> NEW_RELIC_CUSTOM_INSIGHTS_EVENTS_MAX_SAMPLES_STORED - Update TestRecordErrorMetric to verify app.record_custom_metric call Co-Authored-By: Claude Sonnet 4.6 * Reformating for the linter. * Addressing additional automated feedback. - Removed a legacy comment about the New Relic header - Reordered imports in one file - Switched another file to use the import at the top of the file instead of inline when used - Added unit tests for untested methods that were identified * Addressing new feedback. - Proper handling of time to floats. Created a util method and updated code to use it. - added the missing guard to ensure the app is enabled * Addressing feedback. - When an error occurs, still check if the periodic supportability metric should be emitted - Added a check to ensure the extension is ready in the error handler to match _process_success * Updating the NR event timestamps to more accurately reflect when the messages were generated. * Addressing feedback for potential better practice. * Addressing feedback on accessing default values. Added tests for most of these cases. * Adding a new catch exception block based on feedback. * Addressing feedback about a potential issue around a timestamp for the supportability metric. * Addressing minor feedback on length of generated, fallback traceId. * Addressing feedback. - A few more cases were found where the dictionary access might not return the correct value. - Handling cases where `traceparent` is not lower cased * Addressed feedback where the newrelic options might not apply correctly. * Addressing some feedback. * Addressing feedback. * Validating testing / formatting for our changes. * Updating linting, adding tests, defining data type for UI. * Configuration for the logging callback definition. * Adding a newrelic image for the UI to use. * Putting the New Relic callback in proper alphabetic order. * Copying the logo to a committed output directory so it shows up in a locally built container. * Adding missing definition of new env vars that were causing a build failure. * Addressing automated feedback from greptile. * Adding a few more unit tests to increase the code coverage just a bit more. * Additional unit tests to push coverage to almost 90%. * Adding a custom newrelic docker image build process. This removes the need to add the newrelic agent to the core litellm container or dependencies. * Clarifying message when the New Relic agent is not installed and someone is trying to use the newrelic extension. Either use the proper image when using docker, or install the agent manually when running from source. * Ensuring pip is available to install the New Relic agent. * Updating the definition and handling of traceId (no spanId). Clarifying behavior of env vars vs UI configuration for the newrelic extension. * Removing entries from the New Relic logger configuraiton UI as these values must be set as part of running the image. * Removing a stale doc file that has moved to the litellm-docs repo. Cleanup of Dockerfile to remove a LABEL that was incorrect. * Updating container image name to be the best guess for the new name. * Addressing feedback from greptile. - Added a comment around token_count=0 - Updated the boolean parser to allow a wider set of options which matches existing patterns in other parts of LiteLLM. * Removing option for a separate New Relic container image. The agreement is to handle this in the New Relic integration docs. * Updating error message when New Relic agent is not available. * Wiring in the test message from the LiteLLM callback UX. * Missed saving one of the file conflicts. * Fixed a lint error I introduced. Somehow, I dropped another string and now added it back. * Adding newrelic to the schema definition. * Added an admin check on the call before sending test message as mentioned by the AI code review. * Updating to use should_redact_message_logging(kwargs) as part of the logic to determine if message content should be sent to New Relic or not. This still uses the `record_content` property as well, but both have to be true in order for content to be included. --------- Co-authored-by: Claude Sonnet 4.6 --- litellm/__init__.py | 3 + litellm/integrations/callback_configs.json | 15 + litellm/integrations/newrelic/__init__.py | 10 + litellm/integrations/newrelic/newrelic.py | 926 +++++++++++ .../custom_logger_registry.py | 2 + litellm/litellm_core_utils/litellm_logging.py | 12 + .../out/assets/logos/newrelic.png | Bin 0 -> 862 bytes litellm/proxy/_types.py | 8 + .../health_endpoints/_health_endpoints.py | 22 + litellm/types/integrations/newrelic.py | 9 + .../integrations/newrelic/test_newrelic.py | 1351 +++++++++++++++++ .../health_endpoints/test_health_endpoints.py | 79 + .../public/assets/logos/newrelic.png | Bin 0 -> 862 bytes ui/litellm-dashboard/src/lib/http/schema.d.ts | 2 +- 14 files changed, 2438 insertions(+), 1 deletion(-) create mode 100644 litellm/integrations/newrelic/__init__.py create mode 100644 litellm/integrations/newrelic/newrelic.py create mode 100644 litellm/proxy/_experimental/out/assets/logos/newrelic.png create mode 100644 litellm/types/integrations/newrelic.py create mode 100644 tests/test_litellm/integrations/newrelic/test_newrelic.py create mode 100644 ui/litellm-dashboard/public/assets/logos/newrelic.png diff --git a/litellm/__init__.py b/litellm/__init__.py index 514e2519504..e8422952ab4 100644 --- a/litellm/__init__.py +++ b/litellm/__init__.py @@ -43,6 +43,7 @@ from typing import ( Type, ) from litellm.types.integrations.datadog import DatadogInitParams +from litellm.types.integrations.newrelic import NewRelicInitParams from litellm._logging import ( set_verbose, _turn_on_debug, @@ -159,6 +160,7 @@ _custom_logger_compatible_callbacks_literal = Literal[ "posthog", "levo", "compression_interception", + "newrelic", ] cold_storage_custom_logger: Optional[_custom_logger_compatible_callbacks_literal] = None logged_real_time_event_types: Optional[Union[List[str], Literal["*"]]] = None @@ -413,6 +415,7 @@ s3_callback_params: Optional[Dict] = None s3_audit_callback_params: Optional[Dict] = None datadog_llm_observability_params: Optional[Union[DatadogLLMObsInitParams, Dict]] = None datadog_params: Optional[Union[DatadogInitParams, Dict]] = None +newrelic_params: Optional[Union[NewRelicInitParams, Dict]] = None aws_sqs_callback_params: Optional[Dict] = None generic_logger_headers: Optional[Dict] = None default_key_generate_params: Optional[Dict] = None diff --git a/litellm/integrations/callback_configs.json b/litellm/integrations/callback_configs.json index 3a69c9a7936..590c848767a 100644 --- a/litellm/integrations/callback_configs.json +++ b/litellm/integrations/callback_configs.json @@ -290,6 +290,21 @@ }, "description": "Langsmith Logging Integration" }, + { + "id": "newrelic", + "displayName": "New Relic", + "logo": "newrelic.png", + "supports_key_team_logging": false, + "dynamic_params": { + "NEW_RELIC_AI_MONITORING_RECORD_CONTENT_ENABLED": { + "type": "text", + "ui_name": "Record AI Content (default: true)", + "description": "Whether to record AI message content. Set to false to disable.", + "required": false + } + }, + "description": "New Relic AI Monitoring Integration" + }, { "id": "openmeter", "displayName": "OpenMeter", diff --git a/litellm/integrations/newrelic/__init__.py b/litellm/integrations/newrelic/__init__.py new file mode 100644 index 00000000000..5b0f5b9cb24 --- /dev/null +++ b/litellm/integrations/newrelic/__init__.py @@ -0,0 +1,10 @@ +""" +New Relic AI Monitoring Integration for LiteLLM + +This module provides integration with New Relic's AI Monitoring feature to track +LLM requests, responses, and usage metrics. +""" + +from litellm.integrations.newrelic.newrelic import NewRelicLogger + +__all__ = ["NewRelicLogger"] diff --git a/litellm/integrations/newrelic/newrelic.py b/litellm/integrations/newrelic/newrelic.py new file mode 100644 index 00000000000..753b8520337 --- /dev/null +++ b/litellm/integrations/newrelic/newrelic.py @@ -0,0 +1,926 @@ +""" +New Relic AI Monitoring Integration for LiteLLM + +This module provides integration with New Relic's AI Monitoring feature to track +LLM requests, responses, and usage metrics. + +Environment Variables (consumed by the New Relic agent at process bootstrap - +set via container env, or before invoking `newrelic-admin run-program`): + NEW_RELIC_LICENSE_KEY: Your New Relic license key (required) + NEW_RELIC_APP_NAME: Your application name (required) + +UI- and runtime-toggleable: + NEW_RELIC_AI_MONITORING_RECORD_CONTENT_ENABLED: Whether to record message + content (optional, default: true) + +Configuration: + Message logging can be controlled via (both must agree to record): + 1. turn_off_message_logging parameter - pass via callback initialization or config YAML + 2. NEW_RELIC_AI_MONITORING_RECORD_CONTENT_ENABLED env var + + Default behavior: Messages ARE recorded unless explicitly disabled by either method + Either method can disable recording - both must enable for recording to occur + +Usage - Python SDK: + import litellm + litellm.callbacks = ["newrelic"] + + # Or with explicit configuration: + from litellm.integrations.newrelic import NewRelicLogger + litellm.callbacks = [NewRelicLogger(turn_off_message_logging=True)] + +Usage - Proxy Server (config.yaml): + litellm_settings: + callbacks: ["newrelic"] + newrelic_params: + turn_off_message_logging: true # Disable message content recording + + # Or disable via environment variable: + # export NEW_RELIC_AI_MONITORING_RECORD_CONTENT_ENABLED=false + + # Ensure New Relic agent is initialized (use newrelic-admin or initialize manually) + # newrelic-admin run-program python your_app.py +""" + +import json +import os +import threading +import time +import uuid +from typing import Any, Dict, List, Optional, Tuple, Union + +import litellm +from litellm._logging import verbose_logger +from litellm.integrations.custom_logger import CustomLogger +from litellm.litellm_core_utils.redact_messages import should_redact_message_logging +from litellm.types.integrations.newrelic import NewRelicInitParams +from litellm.types.integrations.base_health_check import IntegrationHealthCheckStatus +from litellm.types.utils import ModelResponse, Message, StandardLoggingPayload + +try: + import newrelic.agent as _newrelic_agent +except ImportError: + _newrelic_agent = None # type: ignore + + +class NewRelicLogger(CustomLogger): + """ + New Relic logger for LiteLLM to send AI monitoring events. + + This logger creates two types of New Relic custom events: + 1. LlmChatCompletionSummary - One per completion request + 2. LlmChatCompletionMessage - One per message (request and response) + """ + + # Class-level state for supportability metric emission, shared across all instances. + # Protected by _metric_lock to ensure thread-safe access. + _last_metric_emission_time: float = 0.0 + _metric_lock = threading.Lock() + + def __init__(self, **kwargs): + ######################################################### + # Handle newrelic_params set as litellm.newrelic_params + ######################################################### + dict_newrelic_params = self._get_newrelic_params() + + # Use setdefault so constructor kwargs take priority over global params. + # model_dump() always returns all fields (including defaults), so update() + # would silently overwrite explicit constructor args like turn_off_message_logging=True. + for k, v in dict_newrelic_params.items(): + kwargs.setdefault(k, v) + + # CustomLogger.__init__ will set self.turn_off_message_logging from kwargs + super().__init__(**kwargs) + + # Check for required environment variables + self.license_key = os.getenv("NEW_RELIC_LICENSE_KEY") + self.app_name = os.getenv("NEW_RELIC_APP_NAME") + + # Validate configuration + if not self.license_key or not self.app_name: + verbose_logger.warning( + "New Relic integration requires NEW_RELIC_LICENSE_KEY and " + "NEW_RELIC_APP_NAME environment variables. Integration will be disabled." + ) + self.enabled = False + elif _newrelic_agent is None: + verbose_logger.error( + "New Relic Python agent not installed. Review the New Relic integration documentation at https://docs.litellm.ai/docs/observability/newrelic." + ) + self.enabled = False + else: + try: + # timeout=0 forces non-blocking startup: the agent connects in a + # background thread regardless of newrelic.ini / NEW_RELIC_STARTUP_TIMEOUT. + _newrelic_agent.register_application(timeout=0) + + self.enabled = True + verbose_logger.info( + f"New Relic AI Monitoring initialized for app: {self.app_name}, " + f"content recording: {self.record_content}" + ) + except Exception as e: + verbose_logger.error( + f"Failed to initialize New Relic agent: {e}. " + "Integration will be disabled." + ) + self.enabled = False + + def _get_newrelic_params(self) -> Dict: + """ + Get the newrelic_params from litellm.newrelic_params + + These are params specific to initializing the NewRelicLogger e.g. turn_off_message_logging + """ + dict_newrelic_params: Dict = {} + if litellm.newrelic_params is not None: + if isinstance(litellm.newrelic_params, NewRelicInitParams): + dict_newrelic_params = litellm.newrelic_params.model_dump() + elif isinstance(litellm.newrelic_params, Dict): + # only allow params that are of NewRelicInitParams + dict_newrelic_params = NewRelicInitParams( + **litellm.newrelic_params + ).model_dump() + return dict_newrelic_params + + @property + def record_content(self) -> bool: + """Whether to record message content in New Relic. + + Both turn_off_message_logging param AND NEW_RELIC_AI_MONITORING_RECORD_CONTENT_ENABLED + env var must agree to record content. If either disables recording, content will not + be recorded. Read at call time so UI config changes take effect without a restart. + Default: True (record content) unless explicitly disabled by either method. + """ + return (not self.turn_off_message_logging) and self._parse_bool_env( + "NEW_RELIC_AI_MONITORING_RECORD_CONTENT_ENABLED", True + ) + + def _parse_bool_env(self, var_name: str, default: bool = False) -> bool: + """Parse a boolean environment variable. + + Accepts true/false, 1/0, yes/no, on/off (case-insensitive, + whitespace-tolerant) — matching the convention used in + ``litellm/__init__.py`` and the standard library's + ``configparser.BOOLEAN_STATES``. Unrecognised values log a + warning and fall back to ``default`` rather than silently + flipping user intent. + """ + raw = os.getenv(var_name) + if not raw: + return default + value = raw.strip().lower() + if value in ("1", "true", "yes", "on"): + return True + if value in ("0", "false", "no", "off"): + return False + verbose_logger.warning( + f"{var_name}={raw!r} is not a recognised boolean " + f"(accepts true/false, 1/0, yes/no, on/off). " + f"Falling back to default ({default})." + ) + return default + + def _get_litellm_version(self) -> str: + """ + Get litellm version for supportability metrics. + + Returns: + Version string (e.g., "1.80.0") or "unknown" if unable to determine + """ + try: + from importlib.metadata import version + + return version("litellm") + except Exception as e: + verbose_logger.warning(f"Unable to determine litellm version: {e}") + return "unknown" + + def _emit_supportability_metric(self): + """ + Emit New Relic supportability metric for LiteLLM usage. + + Per spec, this metric should be emitted at least once every 27 hours + to indicate the library is in use. Format: + Supportability/Python/ML/LiteLLM/{version} + + This method updates _last_metric_emission_time and should + be called within a lock when checking periodic emission. + """ + try: + litellm_version = self._get_litellm_version() + metric_name = f"Supportability/Python/ML/LiteLLM/{litellm_version}" + + # Record metric with value of 1 (will be aggregated by New Relic) + app = _newrelic_agent.application() + + # Always update the timestamp so the 27-hour back-off applies + # regardless of whether the app is ready, preventing lock contention + # on every request when the agent is slow to register or never starts. + NewRelicLogger._last_metric_emission_time = time.time() + + if app and app.enabled: + app.record_custom_metric(metric_name, 1) + verbose_logger.info( + f"Emitted New Relic supportability metric: {metric_name}" + ) + else: + verbose_logger.info( + "New Relic application is not enabled; skipping metric recording." + ) + + except Exception as e: + verbose_logger.warning(f"Failed to emit supportability metric: {e}") + + def _check_and_emit_periodic_metric(self): + """ + Check if 27 hours have passed since last metric emission and re-emit if needed. + + Uses a mutex to ensure only one thread emits the metric even if multiple + requests are being processed concurrently. + """ + # Quick check without lock to avoid unnecessary locking + current_time = time.time() + time_since_last_emission = ( + current_time - NewRelicLogger._last_metric_emission_time + ) + + if time_since_last_emission >= 97200: # 27 hours = 97200 seconds + # Acquire lock to ensure only one thread emits + with NewRelicLogger._metric_lock: + # Double-check inside lock in case another thread just emitted + current_time = time.time() + time_since_last_emission = ( + current_time - NewRelicLogger._last_metric_emission_time + ) + + if time_since_last_emission >= 97200: + self._emit_supportability_metric() + + def _get_trace_context( + self, + kwargs: Dict, + standard_logging_object: Optional[StandardLoggingPayload] = None, + ) -> str: + """ + Get the New Relic trace ID for AI monitoring events. + + This integration runs in LiteLLM's async logging worker, outside the + New Relic agent's current transaction. Because we can't call + `newrelic.agent.current_trace_id()` to let the agent populate the + trace_id on AIM custom events, we manually simulate what the agent + would do. An AIM event without a trace_id is malformed per the NR + schema, so this method always returns a valid string. + + Resolution order: + 1. W3C traceparent header (litellm_params.metadata.headers.traceparent) - + what the agent would link to if we were in-transaction. + 2. StandardLoggingPayload.trace_id - LiteLLM's internal trace for + retry/fallback grouping. + 3. Generated UUID - synthetic grouping key when upstream context is + absent or parsing it fails. + + Span IDs are intentionally not emitted: any span ID recoverable from + the inbound traceparent is the caller's parent span, not ours. + + Returns: + trace_id: always a non-empty string. + """ + trace_id: Optional[str] = None + try: + litellm_params = kwargs.get("litellm_params") or {} + metadata = litellm_params.get("metadata") or {} + headers = metadata.get("headers") or {} + # Normalize header key lookup to be case-insensitive per W3C spec + traceparent = next( + (v for k, v in headers.items() if k.lower() == "traceparent"), None + ) + + if traceparent: + # Extract trace_id from traceparent header if available + # traceparent format: "00-4bf92f3577b34da6a3ce929d0e0e4736-00f067aa0ba902b7-00" + parts = traceparent.split("-") + if len(parts) == 4: + trace_id = parts[1] + + if not trace_id and standard_logging_object: + slo_trace_id = standard_logging_object.get("trace_id") + if slo_trace_id: + trace_id = slo_trace_id + + except Exception as e: + verbose_logger.warning( + f"Unable to parse New Relic trace context from upstream sources: {e}" + ) + + if not trace_id: + trace_id = uuid.uuid4().hex + verbose_logger.debug( + f"New Relic trace_id not available from distributed tracing headers or " + f"StandardLoggingPayload. Generated trace_id={trace_id} for AI monitoring " + f"event grouping." + ) + + return trace_id + + def _extract_completion_id(self, kwargs: Dict, response_obj: ModelResponse) -> str: + """ + Extract completion ID from kwargs or response_obj, or generate one. + """ + completion_id = None + + if response_obj: + completion_id = response_obj.get("id") + + if not completion_id: + completion_id = kwargs.get("litellm_call_id") + + # If still not found, generate UUID and log warning per spec + if not completion_id: + completion_id = str(uuid.uuid4()) + + return completion_id + + def _get_vendor( + self, + kwargs: Dict, + standard_logging_object: Optional[StandardLoggingPayload] = None, + ) -> str: + """Extract vendor/provider, preferring StandardLoggingPayload.""" + if standard_logging_object: + vendor = standard_logging_object.get("custom_llm_provider") + if vendor: + return vendor + litellm_params = kwargs.get("litellm_params", {}) or {} + return litellm_params.get("custom_llm_provider") or "litellm" + + def _get_model_names( + self, + kwargs: Dict, + response_obj: ModelResponse, + standard_logging_object: Optional[StandardLoggingPayload] = None, + ) -> Tuple[str, str]: + """ + Extract request and response model names, preferring StandardLoggingPayload + for the request model. + + Returns: + Tuple of (request_model, response_model) + """ + request_model = None + if standard_logging_object: + slo_model = standard_logging_object.get("model") + if slo_model: + request_model = str(slo_model) + if not request_model: + request_model = str(kwargs.get("model") or "unknown") + response_model: str = str(response_obj.get("model") or request_model) + return request_model, response_model + + def _extract_usage( + self, + response_obj: ModelResponse, + standard_logging_object: Optional[StandardLoggingPayload] = None, + ) -> Dict[str, int]: + """Extract usage statistics, preferring StandardLoggingPayload.""" + if standard_logging_object: + prompt = standard_logging_object.get("prompt_tokens") + completion = standard_logging_object.get("completion_tokens") + total = standard_logging_object.get("total_tokens") + if any(x is not None for x in [prompt, completion, total]): + return { + "prompt_tokens": prompt or 0, + "completion_tokens": completion or 0, + "total_tokens": total or 0, + } + + usage = response_obj.get("usage", None) + if not usage: + return {"prompt_tokens": 0, "completion_tokens": 0, "total_tokens": 0} + + return { + "prompt_tokens": usage.get("prompt_tokens") or 0, + "completion_tokens": usage.get("completion_tokens") or 0, + "total_tokens": usage.get("total_tokens") or 0, + } + + def _get_finish_reason(self, response_obj: ModelResponse) -> str: + """ + Extract finish reason from first choice in the response. + + Returns "unknown" if choices are not present or finish_reason is not found. + """ + choices = response_obj.get("choices") or [] + if choices and len(choices) > 0: + return choices[0].get("finish_reason") or "unknown" + return "unknown" + + def _to_epoch_ms(self, t: Any) -> float: + """Convert a datetime or float timestamp to epoch milliseconds.""" + if hasattr(t, "timestamp"): + return t.timestamp() * 1000.0 + return float(t) * 1000.0 + + def _get_duration( + self, + kwargs: Dict, + start_time: Any, + end_time: Any, + standard_logging_object: Optional[StandardLoggingPayload] = None, + ) -> Optional[float]: + """ + Extract duration in milliseconds. + + Resolution order: + 1. StandardLoggingPayload.response_time (already computed by LiteLLM) + 2. llm_api_duration_ms from kwargs + 3. Calculated from start_time and end_time + """ + if standard_logging_object: + response_time = standard_logging_object.get("response_time") + if response_time is not None: + return ( + float(response_time) * 1000.0 + ) # SLO stores seconds; convert to ms + + duration_ms = kwargs.get("llm_api_duration_ms") + if duration_ms is not None: + return float(duration_ms) + + if start_time is not None and end_time is not None: + return self._to_epoch_ms(end_time) - self._to_epoch_ms(start_time) + + return None + + def _get_request_params( + self, + kwargs: Dict, + standard_logging_object: Optional[StandardLoggingPayload] = None, + ) -> Dict[str, Any]: + """ + Extract request parameters like temperature and max_tokens, preferring + StandardLoggingPayload.model_parameters. + + Returns dict with available parameters, omitting those not present. + """ + if standard_logging_object: + source_params = standard_logging_object.get("model_parameters") or {} + else: + source_params = kwargs.get("optional_params") or {} + + params = {} + + temperature = source_params.get("temperature") + if temperature is not None: + params["temperature"] = temperature + + max_tokens = source_params.get("max_tokens") + if max_tokens is not None: + params["max_tokens"] = max_tokens + + return params + + def _extract_message_content(self, message: Union[Message, Dict]) -> str: + """ + Extract content from a message, handling various formats. + + Handles tool calls, multimodal content (as JSON), and standard text content. + Returns empty string if content is None or missing. + """ + content = message.get("content") + + # Handle tool calls + if message.get("tool_calls"): + try: + return json.dumps(message["tool_calls"]) + except Exception: + return str(message["tool_calls"]) + + # Handle None or missing content + if content is None: + return "" + + # Handle list content (multimodal) + if isinstance(content, list): + try: + return json.dumps(content) + except Exception: + return str(content) + + # Handle non-string content + if not isinstance(content, str): + return str(content) + + return content + + def _extract_all_messages( + self, + kwargs: Dict, + response_obj: ModelResponse, + response_model: str, + vendor: str, + standard_logging_object: Optional[StandardLoggingPayload] = None, + ) -> List[Dict[str, Any]]: + """ + Extract all messages (request + response) with sequence numbers and timestamps. + + Processes request messages from StandardLoggingPayload.messages (preferred) or + kwargs["messages"] (fallback), and response messages from response_obj["choices"]. + Assigns sequential numbers starting at 0. + Adds timestamps from StandardLoggingPayload (preferred) or kwargs if available + (converted to epoch milliseconds). + """ + messages = [] + sequence = 0 + + # Extract timestamps, preferring StandardLoggingPayload + start_time = None + if standard_logging_object: + start_time = standard_logging_object.get("startTime") + if not start_time: + start_time = kwargs.get("start_time") + + end_time = None + if standard_logging_object: + end_time = standard_logging_object.get("endTime") + if not end_time: + end_time = kwargs.get("end_time") + + # Content is recorded only when the NR-specific switches allow it AND + # LiteLLM's wider redaction decision (turn_off_message_logging, dynamic + # params, headers) does not require redaction. Async streaming hands the + # callback an unredacted async_complete_streaming_response, so without + # this gate generated content would still reach NR even when the user + # has globally disabled message logging. + record_content = self.record_content and not should_redact_message_logging( + kwargs + ) + + # Extract request messages, preferring StandardLoggingPayload. + # SLO messages can be a string (serialized/redacted), so only use it when it's a list. + slo_messages = ( + standard_logging_object.get("messages") if standard_logging_object else None + ) + if isinstance(slo_messages, list): + request_messages = slo_messages + else: + request_messages = kwargs.get("messages") or [] + for msg in request_messages: + message_data = { + "role": msg.get("role") or "user", + "sequence": sequence, + "response.model": response_model, + "vendor": vendor, + } + + # Add timestamp for request message if available (convert to milliseconds) + if start_time is not None: + message_data["timestamp"] = int(self._to_epoch_ms(start_time)) + + if record_content: + message_data["content"] = self._extract_message_content(msg) + + messages.append(message_data) + sequence += 1 + + # Extract response messages from choices + choices = response_obj.get("choices") or [] + if choices and len(choices) > 0: + for choice in choices: + # Prefer "message" (non-streaming); fall back to "delta" (streaming-assembled) + message = choice.get("message", None) or choice.get("delta", None) + if message: + message_data = { + "role": message.get("role") or "assistant", + "sequence": sequence, + "response.model": response_model, + "vendor": vendor, + "is_response": True, + } + + # Add timestamp for response message if available (convert to milliseconds) + if end_time is not None: + message_data["timestamp"] = int(self._to_epoch_ms(end_time)) + + if record_content: + message_data["content"] = self._extract_message_content(message) + + messages.append(message_data) + sequence += 1 + + return messages + + def _record_summary_event( + self, + request_id: str, + trace_id: Optional[str], + request_model: str, + response_model: str, + vendor: str, + finish_reason: str, + num_messages: int, + usage: Dict[str, int], + duration: Optional[float] = None, + request_params: Optional[Dict[str, Any]] = None, + ): + """Record LlmChatCompletionSummary event to New Relic.""" + try: + event_data = { + "id": request_id, + "request_id": request_id, + "request.model": request_model, + "response.model": response_model, + "response.choices.finish_reason": finish_reason, + "response.number_of_messages": num_messages, + "vendor": vendor, + "ingest_source": "litellm", + "response.usage.prompt_tokens": usage["prompt_tokens"], + "response.usage.completion_tokens": usage["completion_tokens"], + "response.usage.total_tokens": usage["total_tokens"], + } + + # Add optional attributes if present + if trace_id: + event_data["trace_id"] = trace_id + + if duration is not None: + event_data["duration"] = duration + + # Add request parameters if present + if request_params: + if "temperature" in request_params: + event_data["request.temperature"] = request_params["temperature"] + if "max_tokens" in request_params: + event_data["request.max_tokens"] = request_params["max_tokens"] + + app = _newrelic_agent.application() + + if app and app.enabled: + app.record_custom_event("LlmChatCompletionSummary", event_data) + else: + verbose_logger.warning( + "New Relic application is not enabled; skipping summary event recording." + ) + + except Exception as e: + verbose_logger.warning(f"Failed to record New Relic summary event: {e}") + self.handle_callback_failure("newrelic") + + def _record_message_events( + self, + request_id: str, + llm_response_id: str, + trace_id: Optional[str], + messages: List[Dict[str, Any]], + ): + """Record LlmChatCompletionMessage events to New Relic. + + Args: + request_id: Agent-generated UUID that links to Summary event's id + llm_response_id: LLM's response ID (e.g., "chatcmpl-...") for message id format + trace_id: Trace ID for distributed tracing (None if not available) + messages: List of message dicts to record + """ + try: + app = _newrelic_agent.application() + + if not (app and app.enabled): + verbose_logger.warning( + "New Relic application is not enabled; skipping message event recording." + ) + return + + for message in messages: + sequence = message["sequence"] + event_data = { + "id": f"{llm_response_id}-{sequence}", + "request_id": request_id, + "completion_id": request_id, + "role": message["role"], + "sequence": sequence, + "response.model": message["response.model"], + "vendor": message["vendor"], + "ingest_source": "litellm", + "token_count": 0, # Per-message token counts are not available from LiteLLM + } + + # Add trace context if available + if trace_id: + event_data["trace_id"] = trace_id + + # Add content only if it was included in the message data + if "content" in message: + event_data["content"] = message["content"] + + # Add is_response only if True (per spec, omit for request messages) + if message.get("is_response"): + event_data["is_response"] = True + + # Forward actual request/response timestamp (ms) so NR uses the + # real LLM call window rather than the async-logger fire time. + # Requires newrelic>=11.2.0 which reads params["timestamp"] as + # the intrinsic event timestamp. + if "timestamp" in message: + event_data["timestamp"] = message["timestamp"] + + app.record_custom_event("LlmChatCompletionMessage", event_data) + + except Exception as e: + verbose_logger.warning(f"Failed to record New Relic message events: {e}") + self.handle_callback_failure("newrelic") + + def _record_error_metric(self): + """Record error metric to New Relic.""" + try: + if not self.enabled: + return + + self._check_and_emit_periodic_metric() + + app = _newrelic_agent.application() + if app and app.enabled: + app.record_custom_metric("LLM/LiteLLM/Error", 1) + except Exception as e: + verbose_logger.warning(f"Failed to record New Relic error metric: {e}") + self.handle_callback_failure("newrelic") + + def _process_success( + self, + kwargs: Dict, + response_obj: ModelResponse, + start_time: Optional[float] = None, + end_time: Optional[float] = None, + ): + """ + Core logic for processing successful LLM calls. + Used by both sync and async success event handlers. + """ + # Early exit if not enabled + if not self.enabled: + return + + # Check and emit periodic supportability metric if 27 hours have passed + self._check_and_emit_periodic_metric() + + # Use StandardLoggingPayload where available for normalized, pre-computed values + standard_logging_object: Optional[StandardLoggingPayload] = kwargs.get( + "standard_logging_object" + ) + + # Get trace context + trace_id = self._get_trace_context(kwargs, standard_logging_object) + + # Generate unique request ID for this request (used as Summary event id) + request_id = str(uuid.uuid4()) + + # Extract data from response + llm_response_id = self._extract_completion_id(kwargs, response_obj) + vendor = self._get_vendor(kwargs, standard_logging_object) + request_model, response_model = self._get_model_names( + kwargs, response_obj, standard_logging_object + ) + usage = self._extract_usage(response_obj, standard_logging_object) + finish_reason = self._get_finish_reason(response_obj) + + # Extract additional summary event fields + duration = self._get_duration( + kwargs, start_time, end_time, standard_logging_object + ) + request_params = self._get_request_params(kwargs, standard_logging_object) + + # Extract all messages + messages = self._extract_all_messages( + kwargs, response_obj, response_model, vendor, standard_logging_object + ) + + # Record summary event + self._record_summary_event( + request_id=request_id, + trace_id=trace_id, + request_model=request_model, + response_model=response_model, + vendor=vendor, + finish_reason=finish_reason, + num_messages=len(messages), + usage=usage, + duration=duration, + request_params=request_params, + ) + + # Record message events + self._record_message_events( + request_id=request_id, + llm_response_id=llm_response_id, + trace_id=trace_id, + messages=messages, + ) + + async def async_health_check(self) -> IntegrationHealthCheckStatus: + """ + Check if the New Relic integration is healthy. + + Verifies that the integration is enabled and the New Relic agent + has an active, connected application, then records a small + `LiteLLMConnectionTest` custom event so the user can confirm the + end-to-end pipeline in the New Relic UI via NRQL: + `SELECT * FROM LiteLLMConnectionTest SINCE 1 hour ago`. + + The `LiteLLMConnectionTest` event type is intentionally outside the + `Llm*` family that AI Monitoring queries, so test events do not + appear in AI Monitoring dashboards. + """ + if not self.enabled: + return IntegrationHealthCheckStatus( + status="unhealthy", + error_message="New Relic integration is disabled. Check that " + "NEW_RELIC_LICENSE_KEY and NEW_RELIC_APP_NAME are set and the " + "newrelic package is installed.", + ) + + try: + app = _newrelic_agent.application() + if not (app and app.enabled): + return IntegrationHealthCheckStatus( + status="unhealthy", + error_message=( + "New Relic Python agent not installed. Review the New Relic integration documentation at https://docs.litellm.ai/docs/observability/newrelic." + ), + ) + + app.record_custom_event( + "LiteLLMConnectionTest", + { + "is_test_event": True, + "app_name": self.app_name, + "source": "litellm-proxy", + "timestamp": time.time(), + }, + ) + return IntegrationHealthCheckStatus(status="healthy", error_message=None) + except Exception as e: + return IntegrationHealthCheckStatus( + status="unhealthy", + error_message=str(e), + ) + + # CustomLogger interface implementation + + def log_pre_api_call(self, model, messages, kwargs): + """Unused per spec.""" + pass + + def log_post_api_call(self, kwargs, response_obj, start_time, end_time): + """Unused per spec.""" + pass + + def log_success_event(self, kwargs, response_obj, start_time, end_time): + """ + Main success path for non-streaming requests. + + Note: New Relic's record_custom_event is synchronous but non-blocking + (in-memory operation), so it's safe to call from sync context. + """ + try: + self._process_success(kwargs, response_obj, start_time, end_time) + except Exception as e: + verbose_logger.warning(f"Error in New Relic log_success_event: {e}") + self.handle_callback_failure("newrelic") + + async def async_log_success_event(self, kwargs, response_obj, start_time, end_time): + """ + Main success path for async/streaming requests. + + Note: New Relic's SDK is thread-safe and record_custom_event is fast, + so we can call it directly without asyncio.to_thread(). + """ + try: + self._process_success(kwargs, response_obj, start_time, end_time) + except Exception as e: + verbose_logger.warning(f"Error in New Relic async_log_success_event: {e}") + self.handle_callback_failure("newrelic") + + def log_failure_event(self, kwargs, response_obj, start_time, end_time): + """ + Log error metric for failed LLM calls (sync). + + Per spec: Do not send AI events on failure, only record error metric. + """ + try: + self._record_error_metric() + + except Exception as e: + verbose_logger.warning(f"Error in New Relic log_failure_event: {e}") + self.handle_callback_failure("newrelic") + + async def async_log_failure_event(self, kwargs, response_obj, start_time, end_time): + """ + Log error metric for failed LLM calls (async). + + Per spec: Do not send AI events on failure, only record error metric. + """ + try: + self._record_error_metric() + + except Exception as e: + verbose_logger.warning(f"Error in New Relic async_log_failure_event: {e}") + self.handle_callback_failure("newrelic") diff --git a/litellm/litellm_core_utils/custom_logger_registry.py b/litellm/litellm_core_utils/custom_logger_registry.py index ff7f77cf130..a7fae104c92 100644 --- a/litellm/litellm_core_utils/custom_logger_registry.py +++ b/litellm/litellm_core_utils/custom_logger_registry.py @@ -40,6 +40,7 @@ from litellm.integrations.langsmith import LangsmithLogger from litellm.integrations.litellm_agent import LiteLLMAgentModelResolver from litellm.integrations.literal_ai import LiteralAILogger from litellm.integrations.mlflow import MlflowLogger +from litellm.integrations.newrelic import NewRelicLogger from litellm.integrations.openmeter import OpenMeterLogger from litellm.integrations.opentelemetry import OpenTelemetry from litellm.integrations.opik.opik import OpikLogger @@ -106,6 +107,7 @@ class CustomLoggerRegistry: "mavvrik": MavvrikFocusLogger, "vantage": VantageLogger, "posthog": PostHogLogger, + "newrelic": NewRelicLogger, } try: diff --git a/litellm/litellm_core_utils/litellm_logging.py b/litellm/litellm_core_utils/litellm_logging.py index a5bb1c2d56d..087d1682553 100644 --- a/litellm/litellm_core_utils/litellm_logging.py +++ b/litellm/litellm_core_utils/litellm_logging.py @@ -158,6 +158,7 @@ from ..integrations.litellm_agent import LiteLLMAgentModelResolver from ..integrations.literal_ai import LiteralAILogger from ..integrations.logfire_logger import LogfireLevel, LogfireLogger from ..integrations.lunary import LunaryLogger +from ..integrations.newrelic import NewRelicLogger from ..integrations.openmeter import OpenMeterLogger from ..integrations.opik.opik import OpikLogger from ..integrations.posthog import PostHogLogger @@ -4430,6 +4431,13 @@ def _init_custom_logger_compatible_class( # noqa: PLR0915 gitlab_logger = GitLabPromptManager(gitlab_config=gitlab_config) _in_memory_loggers.append(gitlab_logger) return gitlab_logger # type: ignore + elif logging_integration == "newrelic": + for callback in _in_memory_loggers: + if isinstance(callback, NewRelicLogger): + return callback # type: ignore + newrelic_logger = NewRelicLogger() + _in_memory_loggers.append(newrelic_logger) + return newrelic_logger # type: ignore return None except Exception as e: verbose_logger.exception( @@ -4731,6 +4739,10 @@ def get_custom_logger_compatible_class( # noqa: PLR0915 for callback in _in_memory_loggers: if isinstance(callback, SMTPEmailLogger): return callback + elif logging_integration == "newrelic": + for callback in _in_memory_loggers: + if isinstance(callback, NewRelicLogger): + return callback return None except Exception as e: diff --git a/litellm/proxy/_experimental/out/assets/logos/newrelic.png b/litellm/proxy/_experimental/out/assets/logos/newrelic.png new file mode 100644 index 0000000000000000000000000000000000000000..c841e3e71365fb00db681b80ef149c812998dede GIT binary patch literal 862 zcmV-k1EKthP)Px&7)eAyR9Hv7mp@DsaTLeD@2(9tIA9FY+F*0g`AXufX z?VUDF5Ceh9s2dB)?5ZvrCyk2+5=glAN*WUe4Z(7K=J;d(UGJ{6cl5V*CvOqcV*u@TP8i66IFkPG?Io_hHobs&0#FRK2Sk%#obh55gfCzo)TP*aH+e^Z zpb-!GAdUmn3$Q2HqC;W^&zmquvQm7x7Kp9Uj_i1=59?uKVHn_m+bLZBuQTz)CIYF| z;K=Vv*>V5|DTdk)PznUu(J*6;0Tlv3Kn#pF7x^^6VO@&Ng2217WdL0ZcUsmD9vLOz zK7iI522CcsFjW#w*w#Y4kH6s}IE8$j`fEHPu$>&b1?l1Ypa30Cnc`Md%YRUC8g6d5_IJs?K1 zQ2h^m{iD!V$ta?o-&{MDPKP|1yfS%{wAUp$DhKd!Ch5jiF(SAZBDd#MFF=Un2l9q+&vPkmCcvJWQ8NIx+c8Z7 zxL9{PZUjKqlZW9mpENX(Gy>3RXg3Amm7@)Y`$*pF07*qoM6N<$g8i0&3jhEB literal 0 HcmV?d00001 diff --git a/litellm/proxy/_types.py b/litellm/proxy/_types.py index 1b594e20d32..99ea28381c5 100644 --- a/litellm/proxy/_types.py +++ b/litellm/proxy/_types.py @@ -3090,6 +3090,14 @@ class AllCallbacks(LiteLLMPydanticObjectBase): ui_callback_name="Galileo", ) + newrelic: CallbackOnUI = CallbackOnUI( + litellm_callback_name="newrelic", + ui_callback_name="New Relic", + litellm_callback_params=[ + "NEW_RELIC_AI_MONITORING_RECORD_CONTENT_ENABLED", + ], + ) + class SpendLogsMetadata(TypedDict): """ diff --git a/litellm/proxy/health_endpoints/_health_endpoints.py b/litellm/proxy/health_endpoints/_health_endpoints.py index 6ef8bbc4006..507e8e4d4da 100644 --- a/litellm/proxy/health_endpoints/_health_endpoints.py +++ b/litellm/proxy/health_endpoints/_health_endpoints.py @@ -130,6 +130,7 @@ services = Union[ "generic_api", "arize", "galileo", + "newrelic", "sqs", ], str, @@ -208,6 +209,7 @@ async def health_services_endpoint( # noqa: PLR0915 "generic_api", "arize", "galileo", + "newrelic", "sqs", ]: raise HTTPException( @@ -325,6 +327,26 @@ async def health_services_endpoint( # noqa: PLR0915 "status": "success", "message": "Mock LLM request made - check langfuse.", } + elif service == "newrelic": + if not _is_proxy_admin(user_api_key_dict): + raise HTTPException( + status_code=status.HTTP_403_FORBIDDEN, + detail={ + "error": "Only proxy admins can trigger the New Relic test event." + }, + ) + from litellm.integrations.newrelic.newrelic import NewRelicLogger + + newrelic_logger = NewRelicLogger() + response = await newrelic_logger.async_health_check() + return { + "status": response["status"], + "message": ( + response["error_message"] + if response["status"] == "unhealthy" + else "New Relic is healthy — test event sent" + ), + } if service == "webhook": user_info = CallInfo( diff --git a/litellm/types/integrations/newrelic.py b/litellm/types/integrations/newrelic.py new file mode 100644 index 00000000000..2de9769b181 --- /dev/null +++ b/litellm/types/integrations/newrelic.py @@ -0,0 +1,9 @@ +from litellm.types.integrations.custom_logger import StandardCustomLoggerInitParams + + +class NewRelicInitParams(StandardCustomLoggerInitParams): + """ + Params for initializing a New Relic logger on litellm + """ + + pass diff --git a/tests/test_litellm/integrations/newrelic/test_newrelic.py b/tests/test_litellm/integrations/newrelic/test_newrelic.py new file mode 100644 index 00000000000..541b271cb77 --- /dev/null +++ b/tests/test_litellm/integrations/newrelic/test_newrelic.py @@ -0,0 +1,1351 @@ +import os +import sys +from datetime import datetime, timezone +from unittest.mock import MagicMock, patch + +import pytest + +# newrelic is a proxy-runtime dependency (pyproject.toml) 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 + +import litellm +import litellm.integrations.newrelic.newrelic as nr_module +from litellm.integrations.newrelic.newrelic import NewRelicLogger + +# The module may have been imported before sys.modules was patched (e.g. via +# litellm's own startup imports), leaving _newrelic_agent=None. Point it at +# the mock agent so all tests see a non-None agent. +nr_module._newrelic_agent = _mock_newrelic_agent + + +# --------------------------------------------------------------------------- +# Shared fixtures +# --------------------------------------------------------------------------- + +NR_ENV = { + "NEW_RELIC_LICENSE_KEY": "test-license-key", + "NEW_RELIC_APP_NAME": "test-app", +} + + +def make_logger(**kwargs) -> NewRelicLogger: + """Instantiate NewRelicLogger with NR agent calls mocked out.""" + with patch.dict(os.environ, NR_ENV): + return NewRelicLogger(**kwargs) + + +def make_kwargs( + model="gpt-4", + provider="openai", + messages=None, + optional_params=None, + traceparent=None, +) -> dict: + """Build a minimal kwargs dict representative of a litellm callback invocation.""" + headers = {} + if traceparent: + headers["traceparent"] = traceparent + + return { + "model": model, + "messages": messages or [{"role": "user", "content": "Hello"}], + "optional_params": optional_params or {}, + "litellm_params": { + "custom_llm_provider": provider, + "metadata": {"headers": headers}, + }, + "start_time": 1_000_000.0, + "end_time": 1_000_001.5, + "llm_api_duration_ms": 1500.0, + } + + +def make_response( + model="gpt-4", + response_id="chatcmpl-abc123", + content="Hello there!", + finish_reason="stop", + prompt_tokens=10, + completion_tokens=20, +): + """Build a minimal ModelResponse-like dict.""" + return { + "id": response_id, + "model": model, + "choices": [ + { + "message": {"role": "assistant", "content": content}, + "finish_reason": finish_reason, + } + ], + "usage": { + "prompt_tokens": prompt_tokens, + "completion_tokens": completion_tokens, + "total_tokens": prompt_tokens + completion_tokens, + }, + } + + +def make_slo(**overrides): + """Build a StandardLoggingPayload-like dict with sentinel values distinct from + make_kwargs/make_response defaults, so tests can prove the SLO branch won.""" + base = { + "trace_id": "slo-trace-abc", + "custom_llm_provider": "slo-provider", + "model": "slo-model", + "prompt_tokens": 100, + "completion_tokens": 200, + "total_tokens": 300, + "response_time": 1.5, # seconds; converted to ms by _get_duration + "model_parameters": {"temperature": 0.7, "max_tokens": 500}, + "startTime": 2_000_000.0, + "endTime": 2_000_001.5, + "messages": [{"role": "user", "content": "from-slo"}], + } + base.update(overrides) + return base + + +# --------------------------------------------------------------------------- +# Init / configuration +# --------------------------------------------------------------------------- + + +class TestNewRelicLoggerInit: + def test_disabled_when_license_key_missing(self): + with patch("newrelic.agent.register_application"): + with patch.dict(os.environ, {"NEW_RELIC_APP_NAME": "app"}, clear=True): + logger = NewRelicLogger() + assert logger.enabled is False + + def test_disabled_when_app_name_missing(self): + with patch("newrelic.agent.register_application"): + with patch.dict(os.environ, {"NEW_RELIC_LICENSE_KEY": "key"}, clear=True): + logger = NewRelicLogger() + assert logger.enabled is False + + def test_enabled_with_valid_env_vars(self): + with patch("newrelic.agent.register_application"): + with patch.dict(os.environ, NR_ENV): + logger = NewRelicLogger() + assert logger.enabled is True + + def test_disabled_on_import_error(self): + with patch.object( + _mock_newrelic_agent, "register_application", side_effect=ImportError + ): + with patch.dict(os.environ, NR_ENV): + logger = NewRelicLogger() + assert logger.enabled is False + + def test_disabled_on_agent_startup_error(self): + with patch.object( + _mock_newrelic_agent, + "register_application", + side_effect=RuntimeError("agent startup failed"), + ): + with patch.dict(os.environ, NR_ENV): + logger = NewRelicLogger() + assert logger.enabled is False + + def test_disabled_when_agent_package_missing(self): + with patch.object(nr_module, "_newrelic_agent", None): + with patch.dict(os.environ, NR_ENV): + logger = NewRelicLogger() + assert logger.enabled is False + + def test_record_content_default_true(self): + logger = make_logger() + assert logger.record_content is True + + def test_record_content_disabled_by_param(self): + logger = make_logger(turn_off_message_logging=True) + assert logger.record_content is False + + def test_record_content_disabled_by_env_var(self): + with patch("newrelic.agent.register_application"): + with patch.dict( + os.environ, + {**NR_ENV, "NEW_RELIC_AI_MONITORING_RECORD_CONTENT_ENABLED": "false"}, + ): + logger = NewRelicLogger() + assert logger.record_content is False + + def test_record_content_requires_both_enabled(self): + """param says record, but env var says no — result is False.""" + with patch("newrelic.agent.register_application"): + with patch.dict( + os.environ, + {**NR_ENV, "NEW_RELIC_AI_MONITORING_RECORD_CONTENT_ENABLED": "false"}, + ): + logger = NewRelicLogger(turn_off_message_logging=False) + assert logger.record_content is False + + def test_constructor_kwargs_take_priority_over_global_params(self): + """Constructor turn_off_message_logging=True must not be overwritten by + litellm.newrelic_params which defaults turn_off_message_logging to False.""" + from litellm.types.integrations.newrelic import NewRelicInitParams + + with patch("newrelic.agent.register_application"): + with patch.dict(os.environ, NR_ENV): + with patch( + "litellm.newrelic_params", + NewRelicInitParams(turn_off_message_logging=False), + ): + logger = NewRelicLogger(turn_off_message_logging=True) + assert logger.record_content is False + + def test_newrelic_params_plain_dict_branch(self): + """litellm.newrelic_params can be a plain dict; it should be validated + through NewRelicInitParams and its values applied to the logger.""" + with patch("newrelic.agent.register_application"): + with patch.dict(os.environ, NR_ENV): + with patch( + "litellm.newrelic_params", + {"turn_off_message_logging": True}, + ): + logger = NewRelicLogger() + assert logger.turn_off_message_logging is True + + +# --------------------------------------------------------------------------- +# _parse_bool_env +# --------------------------------------------------------------------------- + + +class TestParseBoolEnv: + def setup_method(self): + self.logger = make_logger() + + @pytest.mark.parametrize("raw", ["true", "TRUE", "True", "1", "yes", "on", "ON"]) + def test_truthy_values(self, raw): + with patch.dict(os.environ, {"MY_VAR": raw}): + assert self.logger._parse_bool_env("MY_VAR") is True + + @pytest.mark.parametrize("raw", ["false", "FALSE", "0", "no", "off", "Off"]) + def test_falsy_values(self, raw): + with patch.dict(os.environ, {"MY_VAR": raw}): + assert self.logger._parse_bool_env("MY_VAR") is False + + @pytest.mark.parametrize("raw", [" true ", " 1\t", "\nyes"]) + def test_whitespace_tolerance_truthy(self, raw): + with patch.dict(os.environ, {"MY_VAR": raw}): + assert self.logger._parse_bool_env("MY_VAR") is True + + @pytest.mark.parametrize("raw", [" false ", " 0\t", "\nno"]) + def test_whitespace_tolerance_falsy(self, raw): + with patch.dict(os.environ, {"MY_VAR": raw}): + assert self.logger._parse_bool_env("MY_VAR") is False + + def test_missing_uses_default(self): + with patch.dict(os.environ, {}, clear=True): + assert self.logger._parse_bool_env("MY_VAR", default=True) is True + assert self.logger._parse_bool_env("MY_VAR", default=False) is False + + def test_empty_string_uses_default(self): + with patch.dict(os.environ, {"MY_VAR": ""}): + assert self.logger._parse_bool_env("MY_VAR", default=True) is True + assert self.logger._parse_bool_env("MY_VAR", default=False) is False + + @pytest.mark.parametrize("raw", ["maybe", "2", "enabled", "tru"]) + def test_unrecognised_value_falls_back_to_default_with_warning(self, raw): + with ( + patch.dict(os.environ, {"MY_VAR": raw}), + patch.object(nr_module.verbose_logger, "warning") as mock_warn, + ): + assert self.logger._parse_bool_env("MY_VAR", default=True) is True + assert self.logger._parse_bool_env("MY_VAR", default=False) is False + assert mock_warn.call_count == 2 + # Warning should mention the variable name and the raw value + for call in mock_warn.call_args_list: + assert "MY_VAR" in call.args[0] + assert repr(raw) in call.args[0] + + +# --------------------------------------------------------------------------- +# _get_trace_context +# --------------------------------------------------------------------------- + + +class TestGetTraceContext: + def setup_method(self): + self.logger = make_logger() + + def test_extracts_trace_id_from_traceparent(self): + kwargs = make_kwargs( + traceparent="00-4bf92f3577b34da6a3ce929d0e0e4736-00f067aa0ba902b7-00" + ) + trace_id = self.logger._get_trace_context(kwargs) + assert trace_id == "4bf92f3577b34da6a3ce929d0e0e4736" + + def test_generates_uuid_when_no_headers(self): + kwargs = make_kwargs() + trace_id = self.logger._get_trace_context(kwargs) + assert trace_id is not None + assert ( + len(trace_id) == 32 + ) # 32-char lowercase hex, matches W3C traceparent format + + def test_generates_uuid_when_traceparent_malformed(self): + kwargs = make_kwargs(traceparent="not-valid") + trace_id = self.logger._get_trace_context(kwargs) + # Falls back to a 32-char lowercase hex, matching W3C traceparent format + assert trace_id is not None + assert len(trace_id) == 32 + + def test_extracts_trace_id_from_mixed_case_traceparent_header(self): + # Callers passing headers directly may not normalise case; per W3C spec + # header names are case-insensitive, so "Traceparent" must work too. + kwargs = make_kwargs() + kwargs["litellm_params"]["metadata"]["headers"] = { + "Traceparent": "00-4bf92f3577b34da6a3ce929d0e0e4736-00f067aa0ba902b7-00" + } + trace_id = self.logger._get_trace_context(kwargs) + assert trace_id == "4bf92f3577b34da6a3ce929d0e0e4736" + + def test_parse_failure_falls_through_to_synthetic_uuid(self): + """When parsing upstream sources raises, emit a synthetic UUID rather + than dropping the event. NR schema requires every AIM event carry a + trace_id; this method's contract is to always return a valid string. + """ + # Non-dict headers value forces .items() to raise inside the try + kwargs = {"litellm_params": {"metadata": {"headers": "not-a-dict"}}} + trace_id = self.logger._get_trace_context(kwargs) + assert trace_id is not None + assert len(trace_id) == 32 # 32-char lowercase hex fallback + + +# --------------------------------------------------------------------------- +# _extract_message_content edge cases +# --------------------------------------------------------------------------- + + +class TestExtractMessageContent: + def setup_method(self): + self.logger = make_logger() + + def test_plain_text(self): + assert self.logger._extract_message_content({"content": "hello"}) == "hello" + + def test_none_content_returns_empty_string(self): + assert self.logger._extract_message_content({"content": None}) == "" + + def test_missing_content_returns_empty_string(self): + assert self.logger._extract_message_content({}) == "" + + def test_tool_calls_serialized_as_json(self): + msg = { + "content": None, + "tool_calls": [{"id": "call_1", "function": {"name": "get_weather"}}], + } + result = self.logger._extract_message_content(msg) + assert "get_weather" in result + assert "call_1" in result + + def test_multimodal_list_serialized_as_json(self): + msg = { + "content": [ + {"type": "text", "text": "describe this"}, + {"type": "image_url"}, + ] + } + result = self.logger._extract_message_content(msg) + assert "describe this" in result + assert "image_url" in result + + def test_non_string_content_coerced_to_str(self): + """Numeric/bool content passes the None and list guards; final branch coerces to str.""" + assert self.logger._extract_message_content({"content": 123}) == "123" + assert self.logger._extract_message_content({"content": True}) == "True" + + +# --------------------------------------------------------------------------- +# _extract_all_messages — record_content=False path +# --------------------------------------------------------------------------- + + +class TestExtractAllMessagesContentDisabled: + def test_no_content_key_when_recording_disabled(self): + logger = make_logger(turn_off_message_logging=True) + kwargs = make_kwargs(messages=[{"role": "user", "content": "secret"}]) + response = make_response(content="also secret") + + messages = logger._extract_all_messages( + kwargs, response, response_model="gpt-4", vendor="openai" + ) + + for msg in messages: + assert "content" not in msg + + +class TestExtractAllMessagesRespectsLitellmRedaction: + """Regression tests for the async-streaming redaction bypass. + + NR-specific switches alone are insufficient: when + ``litellm.turn_off_message_logging=True`` (or the per-request equivalents), + async streaming callbacks receive an unredacted + ``async_complete_streaming_response``. Without consulting LiteLLM's + redaction decision the integration would still write generated content + into NR events. + """ + + def _assert_no_content(self, logger, kwargs): + response = make_response(content="streamed assistant text") + + messages = logger._extract_all_messages( + kwargs, response, response_model="gpt-4", vendor="openai" + ) + + # All extracted messages must carry no content payload + for msg in messages: + assert ( + "content" not in msg + ), f"content leaked despite redaction signal: {msg}" + # And there must actually be at least one user + one assistant entry, + # otherwise the test would pass vacuously. + assert any(not m.get("is_response") for m in messages) + assert any(m.get("is_response") for m in messages) + + def test_global_turn_off_message_logging_blocks_content(self, monkeypatch): + monkeypatch.setattr(litellm, "turn_off_message_logging", True) + logger = make_logger() + assert logger.record_content is True + + kwargs = make_kwargs(messages=[{"role": "user", "content": "user prompt"}]) + self._assert_no_content(logger, kwargs) + + def test_dynamic_param_turn_off_message_logging_blocks_content(self): + logger = make_logger() + assert logger.record_content is True + + kwargs = make_kwargs(messages=[{"role": "user", "content": "user prompt"}]) + kwargs["standard_callback_dynamic_params"] = { + "turn_off_message_logging": True, + } + self._assert_no_content(logger, kwargs) + + def test_enable_redaction_header_blocks_content(self): + logger = make_logger() + assert logger.record_content is True + + kwargs = make_kwargs(messages=[{"role": "user", "content": "user prompt"}]) + kwargs["litellm_params"]["metadata"]["headers"] = { + "x-litellm-enable-message-redaction": True, + } + self._assert_no_content(logger, kwargs) + + def test_dynamic_param_explicit_false_overrides_global_redaction(self, monkeypatch): + """The dynamic param has higher priority than the global flag (see + should_redact_message_logging). When a caller explicitly opts back into + message logging per-request, NR must record content again.""" + monkeypatch.setattr(litellm, "turn_off_message_logging", True) + logger = make_logger() + + kwargs = make_kwargs(messages=[{"role": "user", "content": "ok to log"}]) + kwargs["standard_callback_dynamic_params"] = { + "turn_off_message_logging": False, + } + response = make_response(content="response text") + + messages = logger._extract_all_messages( + kwargs, response, response_model="gpt-4", vendor="openai" + ) + + request_msg = next(m for m in messages if not m.get("is_response")) + response_msg = next(m for m in messages if m.get("is_response")) + assert request_msg["content"] == "ok to log" + assert response_msg["content"] == "response text" + + +class TestExtractAllMessagesTimestamps: + def setup_method(self): + self.logger = make_logger() + + def test_input_messages_get_start_time_timestamp(self): + kwargs = make_kwargs(messages=[{"role": "user", "content": "Hi"}]) + # make_kwargs sets start_time=1_000_000.0 and end_time=1_000_001.5 + response = make_response() + + messages = self.logger._extract_all_messages( + kwargs, response, response_model="gpt-4", vendor="openai" + ) + + input_msg = next(m for m in messages if not m.get("is_response")) + assert input_msg["timestamp"] == int(1_000_000.0 * 1000.0) + + def test_output_messages_get_end_time_timestamp(self): + kwargs = make_kwargs(messages=[{"role": "user", "content": "Hi"}]) + response = make_response() + + messages = self.logger._extract_all_messages( + kwargs, response, response_model="gpt-4", vendor="openai" + ) + + output_msg = next(m for m in messages if m.get("is_response")) + assert output_msg["timestamp"] == int(1_000_001.5 * 1000.0) + + def test_timestamp_forwarded_to_event_data(self): + logger = make_logger() + mock_app = MagicMock() + mock_app.enabled = True + + kwargs = make_kwargs( + traceparent="00-aabbccddeeff00112233445566778899-0011223344556677-01", + messages=[{"role": "user", "content": "Hi"}], + ) + response = make_response() + + with patch("newrelic.agent.application", return_value=mock_app): + logger._process_success(kwargs, response, start_time=1.0, end_time=2.5) + + calls = mock_app.record_custom_event.call_args_list + message_events = [ + c[0][1] for c in calls if c[0][0] == "LlmChatCompletionMessage" + ] + for event in message_events: + assert "timestamp" in event + + +# --------------------------------------------------------------------------- +# Streaming response handling +# --------------------------------------------------------------------------- + + +def make_streaming_response( + model="gpt-4", + response_id="chatcmpl-stream123", + content="Hello from streaming!", + finish_reason="stop", + prompt_tokens=8, + completion_tokens=15, +): + """Build a streaming-assembled response dict using 'delta' instead of 'message'.""" + return { + "id": response_id, + "model": model, + "choices": [ + { + "delta": {"role": "assistant", "content": content}, + "finish_reason": finish_reason, + } + ], + "usage": { + "prompt_tokens": prompt_tokens, + "completion_tokens": completion_tokens, + "total_tokens": prompt_tokens + completion_tokens, + }, + } + + +class TestStreamingResponse: + """Verify graceful handling of streaming-assembled responses. + + When LiteLLM assembles a streaming response, some providers produce a + final choice dict with a 'delta' key instead of 'message'. The integration + must extract content from either key without raising. + """ + + def setup_method(self): + self.logger = make_logger() + + def test_extracts_content_from_delta_key(self): + kwargs = make_kwargs(messages=[{"role": "user", "content": "Hi"}]) + response = make_streaming_response(content="Streamed reply") + + messages = self.logger._extract_all_messages( + kwargs, response, response_model="gpt-4", vendor="openai" + ) + + response_msgs = [m for m in messages if m.get("is_response")] + assert len(response_msgs) == 1 + assert response_msgs[0]["content"] == "Streamed reply" + assert response_msgs[0]["role"] == "assistant" + + def test_streaming_response_records_summary_and_message_events(self): + mock_app = MagicMock() + mock_app.enabled = True + + kwargs = make_kwargs( + traceparent="00-aabbccddeeff00112233445566778899-0011223344556677-01", + messages=[{"role": "user", "content": "Hi"}], + ) + response = make_streaming_response( + response_id="chatcmpl-stream123", + content="Streamed reply", + finish_reason="stop", + prompt_tokens=8, + completion_tokens=15, + ) + + with patch("newrelic.agent.application", return_value=mock_app): + self.logger._process_success(kwargs, response, start_time=1.0, end_time=2.0) + + calls = mock_app.record_custom_event.call_args_list + event_types = [c[0][0] for c in calls] + assert "LlmChatCompletionSummary" in event_types + assert "LlmChatCompletionMessage" in event_types + + message_events = [ + c[0][1] for c in calls if c[0][0] == "LlmChatCompletionMessage" + ] + response_msg = next((e for e in message_events if e.get("is_response")), None) + assert response_msg is not None + assert response_msg["content"] == "Streamed reply" + + @pytest.mark.asyncio + async def test_async_log_success_event_streaming(self): + """async_log_success_event is the primary entry point for streaming calls.""" + mock_app = MagicMock() + mock_app.enabled = True + + kwargs = make_kwargs(messages=[{"role": "user", "content": "Hi"}]) + response = make_streaming_response() + + with patch("newrelic.agent.application", return_value=mock_app): + await self.logger.async_log_success_event( + kwargs, response, start_time=1.0, end_time=2.0 + ) + + calls = mock_app.record_custom_event.call_args_list + event_types = [c[0][0] for c in calls] + assert "LlmChatCompletionSummary" in event_types + assert "LlmChatCompletionMessage" in event_types + + def test_no_content_when_recording_disabled_streaming(self): + logger = make_logger(turn_off_message_logging=True) + kwargs = make_kwargs(messages=[{"role": "user", "content": "secret"}]) + response = make_streaming_response(content="also secret") + + messages = logger._extract_all_messages( + kwargs, response, response_model="gpt-4", vendor="openai" + ) + + for msg in messages: + assert "content" not in msg + + +# --------------------------------------------------------------------------- +# Explicit-None defensive tests +# --------------------------------------------------------------------------- + + +class TestExplicitNoneValues: + """Verify that explicitly None values in kwargs/response don't raise or silently drop events.""" + + def setup_method(self): + self.logger = make_logger() + + # _get_trace_context — chained dict lookups + def test_trace_context_litellm_params_none(self): + kwargs = make_kwargs() + kwargs["litellm_params"] = None + trace_id = self.logger._get_trace_context(kwargs) + assert trace_id is not None # falls back to UUID + + def test_trace_context_metadata_none(self): + kwargs = make_kwargs() + kwargs["litellm_params"] = {"metadata": None} + trace_id = self.logger._get_trace_context(kwargs) + assert trace_id is not None + + def test_trace_context_headers_none(self): + kwargs = make_kwargs() + kwargs["litellm_params"] = {"metadata": {"headers": None}} + trace_id = self.logger._get_trace_context(kwargs) + assert trace_id is not None + + # _get_request_params + def test_request_params_optional_params_none(self): + assert self.logger._get_request_params({"optional_params": None}) == {} + + # _get_model_names + def test_model_names_model_none_in_kwargs(self): + request_model, _ = self.logger._get_model_names( + {"model": None}, make_response() + ) + assert request_model == "unknown" + + def test_model_names_model_none_in_response(self): + response = make_response() + response["model"] = None + _, response_model = self.logger._get_model_names(make_kwargs(), response) + assert response_model == "gpt-4" # falls back to request_model from kwargs + + # _extract_all_messages + def test_extract_messages_messages_none(self): + kwargs = make_kwargs() + kwargs["messages"] = None + response = make_response() + messages = self.logger._extract_all_messages( + kwargs, response, response_model="gpt-4", vendor="openai" + ) + # No request messages, but response message should still be extracted + assert any(m.get("is_response") for m in messages) + + def test_extract_messages_choices_none(self): + kwargs = make_kwargs(messages=[{"role": "user", "content": "Hi"}]) + response = make_response() + response["choices"] = None + messages = self.logger._extract_all_messages( + kwargs, response, response_model="gpt-4", vendor="openai" + ) + # No response messages, but request message should still be extracted + assert any(not m.get("is_response") for m in messages) + + +# --------------------------------------------------------------------------- +# Helper edge cases +# --------------------------------------------------------------------------- + + +class TestExtractUsage: + def setup_method(self): + self.logger = make_logger() + + def test_missing_usage_returns_zeros(self): + response = {"id": "r1", "model": "gpt-4", "choices": []} + usage = self.logger._extract_usage(response) + assert usage == {"prompt_tokens": 0, "completion_tokens": 0, "total_tokens": 0} + + def test_explicit_none_token_fields_return_zeros(self): + response = { + "usage": { + "prompt_tokens": None, + "completion_tokens": None, + "total_tokens": None, + } + } + usage = self.logger._extract_usage(response) + assert usage == {"prompt_tokens": 0, "completion_tokens": 0, "total_tokens": 0} + + +class TestGetFinishReason: + def setup_method(self): + self.logger = make_logger() + + def test_returns_unknown_when_no_choices(self): + response = {"choices": []} + assert self.logger._get_finish_reason(response) == "unknown" + + def test_returns_unknown_when_choices_missing(self): + assert self.logger._get_finish_reason({}) == "unknown" + + def test_returns_unknown_when_finish_reason_explicitly_none(self): + response = {"choices": [{"finish_reason": None}]} + assert self.logger._get_finish_reason(response) == "unknown" + + +class TestToEpochMs: + def setup_method(self): + self.logger = make_logger() + + def test_float_passthrough(self): + assert self.logger._to_epoch_ms(1.0) == pytest.approx(1000.0) + + def test_datetime_converted(self): + dt = datetime(2024, 1, 1, 0, 0, 0, tzinfo=timezone.utc) + assert self.logger._to_epoch_ms(dt) == pytest.approx(dt.timestamp() * 1000.0) + + +class TestGetDuration: + def setup_method(self): + self.logger = make_logger() + + def test_uses_kwargs_value_when_present(self): + kwargs = {"llm_api_duration_ms": 750.0} + assert self.logger._get_duration(kwargs, 0.0, 1.0) == 750.0 + + def test_calculates_from_float_timestamps(self): + kwargs = {} + result = self.logger._get_duration(kwargs, 1.0, 2.5) + assert result == pytest.approx(1500.0) + + def test_calculates_from_datetime_timestamps(self): + kwargs = {} + start = datetime(2024, 1, 1, 0, 0, 0, tzinfo=timezone.utc) + end = datetime(2024, 1, 1, 0, 0, 1, 500000, tzinfo=timezone.utc) # +1.5s + result = self.logger._get_duration(kwargs, start, end) + assert result == pytest.approx(1500.0) + + def test_returns_none_when_nothing_available(self): + assert self.logger._get_duration({}, None, None) is None + + +class TestGetRequestParams: + def setup_method(self): + self.logger = make_logger() + + def test_includes_only_present_params(self): + kwargs = {"optional_params": {"temperature": 0.7}} + params = self.logger._get_request_params(kwargs) + assert params == {"temperature": 0.7} + assert "max_tokens" not in params + + def test_empty_when_no_optional_params(self): + assert self.logger._get_request_params({}) == {} + + +# --------------------------------------------------------------------------- +# _process_success — comprehensive happy-path +# --------------------------------------------------------------------------- + + +class TestProcessSuccess: + def test_records_summary_and_message_events(self): + logger = make_logger() + mock_app = MagicMock() + mock_app.enabled = True + + kwargs = make_kwargs( + traceparent="00-aabbccddeeff00112233445566778899-0011223344556677-01", + messages=[{"role": "user", "content": "Hello"}], + optional_params={"temperature": 0.5, "max_tokens": 100}, + ) + response = make_response( + response_id="chatcmpl-xyz", + content="Hi there!", + finish_reason="stop", + prompt_tokens=5, + completion_tokens=10, + ) + + with patch("newrelic.agent.application", return_value=mock_app): + logger._process_success(kwargs, response, start_time=1.0, end_time=2.5) + + calls = mock_app.record_custom_event.call_args_list + event_types = [c[0][0] for c in calls] + + assert "LlmChatCompletionSummary" in event_types + assert "LlmChatCompletionMessage" in event_types + + # Verify summary event fields + summary_data = next( + c[0][1] for c in calls if c[0][0] == "LlmChatCompletionSummary" + ) + assert summary_data["vendor"] == "openai" + assert summary_data["request.model"] == "gpt-4" + assert summary_data["response.model"] == "gpt-4" + assert summary_data["response.choices.finish_reason"] == "stop" + assert summary_data["response.usage.prompt_tokens"] == 5 + assert summary_data["response.usage.completion_tokens"] == 10 + assert summary_data["response.usage.total_tokens"] == 15 + assert summary_data["request.temperature"] == 0.5 + assert summary_data["request.max_tokens"] == 100 + assert summary_data["ingest_source"] == "litellm" + assert summary_data["trace_id"] == "aabbccddeeff00112233445566778899" + + # Verify message event id format: "{llm_response_id}-{sequence}" + message_events = [ + c[0][1] for c in calls if c[0][0] == "LlmChatCompletionMessage" + ] + assert any(e["id"].startswith("chatcmpl-xyz-") for e in message_events) + response_msg = next(e for e in message_events if e.get("is_response")) + assert response_msg["content"] == "Hi there!" + assert response_msg["role"] == "assistant" + + def test_skips_when_disabled(self): + logger = make_logger() + logger.enabled = False + + with patch("newrelic.agent.application") as mock_app: + logger._process_success(make_kwargs(), make_response()) + + mock_app.assert_not_called() + + +# --------------------------------------------------------------------------- +# _record_error_metric +# --------------------------------------------------------------------------- + + +class TestRecordErrorMetric: + def setup_method(self): + self.logger = make_logger() + + def test_calls_record_custom_metric(self): + mock_app = MagicMock() + mock_app.enabled = True + + with patch.object(self.logger, "_check_and_emit_periodic_metric"): + with patch("newrelic.agent.application", return_value=mock_app): + self.logger._record_error_metric() + + mock_app.record_custom_metric.assert_called_once_with("LLM/LiteLLM/Error", 1) + + def test_skips_when_app_disabled(self): + mock_app = MagicMock() + mock_app.enabled = False + + with patch.object(self.logger, "_check_and_emit_periodic_metric"): + with patch("newrelic.agent.application", return_value=mock_app): + self.logger._record_error_metric() + + mock_app.record_custom_metric.assert_not_called() + + def test_calls_check_and_emit_periodic_metric(self): + with patch.object( + self.logger, "_check_and_emit_periodic_metric" + ) as mock_periodic: + with patch("newrelic.agent.application", return_value=MagicMock()): + self.logger._record_error_metric() + + mock_periodic.assert_called_once() + + def test_skips_when_logger_disabled(self): + self.logger.enabled = False + with patch("newrelic.agent.application") as mock_app: + self.logger._record_error_metric() + mock_app.assert_not_called() + + def test_handles_exception(self): + with patch( + "newrelic.agent.application", side_effect=RuntimeError("agent down") + ): + self.logger._record_error_metric() # must not raise + + +# --------------------------------------------------------------------------- +# _emit_supportability_metric +# --------------------------------------------------------------------------- + + +class TestEmitSupportabilityMetric: + def setup_method(self): + self.logger = make_logger() + NewRelicLogger._last_metric_emission_time = 0.0 + + def test_records_metric_with_correct_name_and_value(self): + mock_app = MagicMock() + mock_app.enabled = True + with patch("newrelic.agent.application", return_value=mock_app): + with patch.object( + self.logger, "_get_litellm_version", return_value="1.80.0" + ): + self.logger._emit_supportability_metric() + mock_app.record_custom_metric.assert_called_once_with( + "Supportability/Python/ML/LiteLLM/1.80.0", 1 + ) + + def test_updates_last_emission_time(self): + mock_app = MagicMock() + mock_app.enabled = True + fake_now = 9_999_999.0 + with patch("newrelic.agent.application", return_value=mock_app): + with patch( + "litellm.integrations.newrelic.newrelic.time.time", + return_value=fake_now, + ): + self.logger._emit_supportability_metric() + assert NewRelicLogger._last_metric_emission_time == fake_now + + def test_skips_when_app_disabled(self): + mock_app = MagicMock() + mock_app.enabled = False + with patch("newrelic.agent.application", return_value=mock_app): + self.logger._emit_supportability_metric() + mock_app.record_custom_metric.assert_not_called() + # Timestamp is still updated to back off lock contention during registration. + assert NewRelicLogger._last_metric_emission_time != 0.0 + + def test_skips_when_no_app(self): + with patch("newrelic.agent.application", return_value=None): + self.logger._emit_supportability_metric() + # Timestamp is updated even when app is None to back off lock contention + # if the agent never starts or is slow to initialise. + assert NewRelicLogger._last_metric_emission_time != 0.0 + + def test_handles_exception(self): + with patch( + "newrelic.agent.application", side_effect=RuntimeError("agent down") + ): + self.logger._emit_supportability_metric() # must not raise + + +# --------------------------------------------------------------------------- +# _check_and_emit_periodic_metric +# --------------------------------------------------------------------------- + + +class TestCheckAndEmitPeriodicMetric: + def setup_method(self): + self.logger = make_logger() + NewRelicLogger._last_metric_emission_time = 0.0 + + def test_emits_on_first_call(self): + """_last_metric_emission_time starts at 0.0; any real time satisfies 27-hour window.""" + with patch.object(self.logger, "_emit_supportability_metric") as mock_emit: + with patch( + "litellm.integrations.newrelic.newrelic.time.time", + return_value=100_000.0, + ): + self.logger._check_and_emit_periodic_metric() + mock_emit.assert_called_once() + + def test_does_not_re_emit_within_27_hours(self): + recent = 1_000_000.0 + NewRelicLogger._last_metric_emission_time = recent + with patch.object(self.logger, "_emit_supportability_metric") as mock_emit: + with patch( + "litellm.integrations.newrelic.newrelic.time.time", + return_value=recent + 3600, # 1 hour later + ): + self.logger._check_and_emit_periodic_metric() + mock_emit.assert_not_called() + + def test_re_emits_after_27_hours(self): + old = 1_000_000.0 + NewRelicLogger._last_metric_emission_time = old + with patch.object(self.logger, "_emit_supportability_metric") as mock_emit: + with patch( + "litellm.integrations.newrelic.newrelic.time.time", + return_value=old + 97201, # 27 hours + 1 second + ): + self.logger._check_and_emit_periodic_metric() + mock_emit.assert_called_once() + + def test_boundary_exactly_27_hours_triggers_emission(self): + old = 1_000_000.0 + NewRelicLogger._last_metric_emission_time = old + with patch.object(self.logger, "_emit_supportability_metric") as mock_emit: + with patch( + "litellm.integrations.newrelic.newrelic.time.time", + return_value=old + 97200, + ): + self.logger._check_and_emit_periodic_metric() + mock_emit.assert_called_once() + + +# --------------------------------------------------------------------------- +# _get_litellm_version +# --------------------------------------------------------------------------- + + +class TestGetLitellmVersion: + def setup_method(self): + self.logger = make_logger() + + def test_returns_unknown_on_exception(self): + with patch("importlib.metadata.version", side_effect=Exception("no package")): + result = self.logger._get_litellm_version() + assert result == "unknown" + + +# --------------------------------------------------------------------------- +# _record_summary_event — disabled-app and exception paths +# --------------------------------------------------------------------------- + +_USAGE = {"prompt_tokens": 5, "completion_tokens": 10, "total_tokens": 15} + + +class TestRecordSummaryEvent: + def setup_method(self): + self.logger = make_logger() + + def _call(self, **kwargs): + self.logger._record_summary_event( + request_id="req-1", + trace_id="trace-abc", + request_model="gpt-4", + response_model="gpt-4", + vendor="openai", + finish_reason="stop", + num_messages=2, + usage=_USAGE, + **kwargs, + ) + + def test_skips_when_app_disabled(self): + mock_app = MagicMock() + mock_app.enabled = False + with patch("newrelic.agent.application", return_value=mock_app): + self._call() + mock_app.record_custom_event.assert_not_called() + + def test_handles_exception(self): + with patch( + "newrelic.agent.application", side_effect=RuntimeError("agent down") + ): + self._call() # must not raise + + +# --------------------------------------------------------------------------- +# _record_message_events — disabled-app and exception paths +# --------------------------------------------------------------------------- + +_MESSAGES = [ + {"role": "user", "sequence": 0, "response.model": "gpt-4", "vendor": "openai"} +] + + +class TestRecordMessageEvents: + def setup_method(self): + self.logger = make_logger() + + def _call(self): + self.logger._record_message_events( + request_id="req-1", + llm_response_id="resp-1", + trace_id="trace-abc", + messages=_MESSAGES, + ) + + def test_skips_when_app_disabled(self): + mock_app = MagicMock() + mock_app.enabled = False + with patch("newrelic.agent.application", return_value=mock_app): + self._call() + mock_app.record_custom_event.assert_not_called() + + def test_handles_exception(self): + with patch( + "newrelic.agent.application", side_effect=RuntimeError("agent down") + ): + self._call() # must not raise + + +# --------------------------------------------------------------------------- +# CustomLogger interface entry points +# --------------------------------------------------------------------------- + + +class TestLogSuccessEvent: + def test_delegates_to_process_success(self): + logger = make_logger() + with patch.object(logger, "_process_success") as mock_process: + logger.log_success_event(make_kwargs(), make_response(), 1.0, 2.0) + mock_process.assert_called_once() + + def test_exception_is_handled(self): + logger = make_logger() + with patch.object(logger, "_process_success", side_effect=RuntimeError("boom")): + logger.log_success_event(make_kwargs(), make_response(), 1.0, 2.0) + + @pytest.mark.asyncio + async def test_async_delegates_to_process_success(self): + logger = make_logger() + with patch.object(logger, "_process_success") as mock_process: + await logger.async_log_success_event( + make_kwargs(), make_response(), 1.0, 2.0 + ) + mock_process.assert_called_once() + + @pytest.mark.asyncio + async def test_async_exception_is_handled(self): + logger = make_logger() + with patch.object(logger, "_process_success", side_effect=RuntimeError("boom")): + await logger.async_log_success_event( + make_kwargs(), make_response(), 1.0, 2.0 + ) + + +class TestLogFailureEvent: + def test_sync_records_error_metric(self): + logger = make_logger() + with patch.object(logger, "_record_error_metric") as mock_metric: + logger.log_failure_event(make_kwargs(), None, 1.0, 2.0) + mock_metric.assert_called_once() + + def test_sync_exception_is_handled(self): + logger = make_logger() + with patch.object( + logger, "_record_error_metric", side_effect=RuntimeError("boom") + ): + logger.log_failure_event(make_kwargs(), None, 1.0, 2.0) + + @pytest.mark.asyncio + async def test_async_records_error_metric(self): + logger = make_logger() + with patch.object(logger, "_record_error_metric") as mock_metric: + await logger.async_log_failure_event(make_kwargs(), None, 1.0, 2.0) + mock_metric.assert_called_once() + + @pytest.mark.asyncio + async def test_async_exception_is_handled(self): + logger = make_logger() + with patch.object( + logger, "_record_error_metric", side_effect=RuntimeError("boom") + ): + await logger.async_log_failure_event(make_kwargs(), None, 1.0, 2.0) + + +# --------------------------------------------------------------------------- +# async_health_check +# --------------------------------------------------------------------------- + + +class TestAsyncHealthCheck: + @pytest.mark.asyncio + async def test_unhealthy_when_disabled(self): + logger = make_logger() + logger.enabled = False + result = await logger.async_health_check() + assert result["status"] == "unhealthy" + assert result["error_message"] is not None + + @pytest.mark.asyncio + async def test_healthy_when_app_enabled_records_test_event(self): + logger = make_logger() + mock_app = MagicMock() + mock_app.enabled = True + with patch("newrelic.agent.application", return_value=mock_app): + result = await logger.async_health_check() + assert result["status"] == "healthy" + assert result["error_message"] is None + + mock_app.record_custom_event.assert_called_once() + event_type, event_data = mock_app.record_custom_event.call_args[0] + assert event_type == "LiteLLMConnectionTest" + assert event_data["is_test_event"] is True + assert event_data["app_name"] == logger.app_name + assert event_data["source"] == "litellm-proxy" + assert isinstance(event_data["timestamp"], float) + + @pytest.mark.asyncio + async def test_unhealthy_when_app_disabled(self): + logger = make_logger() + mock_app = MagicMock() + mock_app.enabled = False + with patch("newrelic.agent.application", return_value=mock_app): + result = await logger.async_health_check() + assert result["status"] == "unhealthy" + assert result["error_message"] is not None + mock_app.record_custom_event.assert_not_called() + + @pytest.mark.asyncio + async def test_exception_returns_unhealthy(self): + logger = make_logger() + with patch( + "newrelic.agent.application", side_effect=RuntimeError("agent down") + ): + result = await logger.async_health_check() + assert result["status"] == "unhealthy" + assert "agent down" in result["error_message"] + + @pytest.mark.asyncio + async def test_record_custom_event_failure_returns_unhealthy(self): + logger = make_logger() + mock_app = MagicMock() + mock_app.enabled = True + mock_app.record_custom_event.side_effect = RuntimeError("intake unreachable") + with patch("newrelic.agent.application", return_value=mock_app): + result = await logger.async_health_check() + assert result["status"] == "unhealthy" + assert "intake unreachable" in result["error_message"] + + +# --------------------------------------------------------------------------- +# _extract_completion_id fallback chain +# --------------------------------------------------------------------------- + + +class TestExtractCompletionId: + def setup_method(self): + self.logger = make_logger() + + def test_uses_litellm_call_id_when_response_has_no_id(self): + result = self.logger._extract_completion_id( + kwargs={"litellm_call_id": "call-abc-123"}, + response_obj={}, + ) + assert result == "call-abc-123" + + def test_generates_uuid_when_neither_id_present(self): + result = self.logger._extract_completion_id(kwargs={}, response_obj={}) + # UUID4 hex-with-dashes is 36 chars; just confirm shape and uniqueness + assert isinstance(result, str) + assert len(result) == 36 + second = self.logger._extract_completion_id(kwargs={}, response_obj={}) + assert result != second + + +# --------------------------------------------------------------------------- +# StandardLoggingPayload preference across extractors +# --------------------------------------------------------------------------- + + +class TestStandardLoggingPayloadPreference: + """Each extractor that accepts a StandardLoggingPayload must prefer its + values over the raw kwargs/response fallbacks.""" + + def setup_method(self): + self.logger = make_logger() + + def test_trace_context_uses_slo_trace_id_when_no_traceparent(self): + kwargs = {"litellm_params": {"metadata": {"headers": {}}}} + trace_id = self.logger._get_trace_context( + kwargs, standard_logging_object=make_slo() + ) + assert trace_id == "slo-trace-abc" + + def test_vendor_from_slo(self): + # kwargs carries a different provider; SLO must win. + kwargs = {"litellm_params": {"custom_llm_provider": "kwargs-provider"}} + assert ( + self.logger._get_vendor(kwargs, standard_logging_object=make_slo()) + == "slo-provider" + ) + + def test_model_names_uses_slo_model(self): + request_model, _ = self.logger._get_model_names( + {"model": "kwargs-model"}, + make_response(model="response-model"), + standard_logging_object=make_slo(), + ) + assert request_model == "slo-model" + + def test_usage_from_slo_when_any_token_field_present(self): + # make_response defaults to 10/20/30 tokens; SLO sentinels are 100/200/300. + usage = self.logger._extract_usage( + make_response(), standard_logging_object=make_slo() + ) + assert usage == { + "prompt_tokens": 100, + "completion_tokens": 200, + "total_tokens": 300, + } + + def test_duration_from_slo_response_time_converted_to_ms(self): + # SLO response_time is 1.5 seconds; expected 1500.0 ms. + # Pass start/end that would compute a different value to prove SLO won. + duration = self.logger._get_duration( + kwargs={"llm_api_duration_ms": 9999.0}, + start_time=1.0, + end_time=2.0, + standard_logging_object=make_slo(), + ) + assert duration == 1500.0 + + def test_request_params_from_slo_model_parameters(self): + params = self.logger._get_request_params( + {"optional_params": {"temperature": 0.1}}, + standard_logging_object=make_slo(), + ) + assert params == {"temperature": 0.7, "max_tokens": 500} + + def test_extract_all_messages_sources_timestamps_and_messages_from_slo(self): + """Covers three SLO branches at once: startTime, endTime, and messages list.""" + kwargs = make_kwargs(messages=[{"role": "user", "content": "from-kwargs"}]) + messages = self.logger._extract_all_messages( + kwargs, + make_response(), + response_model="gpt-4", + vendor="openai", + standard_logging_object=make_slo(), + ) + + request = next(m for m in messages if not m.get("is_response")) + assert request["content"] == "from-slo" # SLO messages list wins + assert request["timestamp"] == int(2_000_000.0 * 1000.0) # SLO startTime + + response = next(m for m in messages if m.get("is_response")) + assert response["timestamp"] == int(2_000_001.5 * 1000.0) # SLO endTime diff --git a/tests/test_litellm/proxy/health_endpoints/test_health_endpoints.py b/tests/test_litellm/proxy/health_endpoints/test_health_endpoints.py index 64c57ab90e3..d31cfdc39bd 100644 --- a/tests/test_litellm/proxy/health_endpoints/test_health_endpoints.py +++ b/tests/test_litellm/proxy/health_endpoints/test_health_endpoints.py @@ -756,6 +756,85 @@ async def test_health_services_endpoint_rejects_unknown_service(): await health_services_endpoint(service="totally_unknown_service_xyz") +@pytest.mark.asyncio +@pytest.mark.parametrize( + "role", + [ + None, + LitellmUserRoles.INTERNAL_USER, + LitellmUserRoles.INTERNAL_USER_VIEW_ONLY, + LitellmUserRoles.TEAM, + LitellmUserRoles.CUSTOMER, + ], +) +async def test_health_services_endpoint_newrelic_blocks_non_admin(role): + """ + /health/services?service=newrelic emits a real LiteLLMConnectionTest event + to the configured New Relic account. Only proxy admins (full or view-only) + should be able to trigger it; every other caller must be rejected before + the external event is recorded. + """ + from litellm.proxy._types import ProxyException + + user_api_key_dict = UserAPIKeyAuth( + token="non-admin-token", + user_id="non-admin-user", + user_role=role, + ) + + with patch( + "litellm.integrations.newrelic.newrelic.NewRelicLogger" + ) as MockNewRelicLogger: + mock_instance = MagicMock() + mock_instance.async_health_check = AsyncMock( + return_value={"status": "healthy", "error_message": ""} + ) + MockNewRelicLogger.return_value = mock_instance + + with pytest.raises(ProxyException) as exc_info: + await health_services_endpoint( + user_api_key_dict=user_api_key_dict, + service="newrelic", + ) + + assert str(exc_info.value.code) == "403" + mock_instance.async_health_check.assert_not_awaited() + MockNewRelicLogger.assert_not_called() + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "admin_role", + [LitellmUserRoles.PROXY_ADMIN, LitellmUserRoles.PROXY_ADMIN_VIEW_ONLY], +) +async def test_health_services_endpoint_newrelic_allows_proxy_admin(admin_role): + """ + Proxy admins (full and view-only) can trigger the New Relic test event. + """ + user_api_key_dict = UserAPIKeyAuth( + token="admin-token", + user_id="admin-user", + user_role=admin_role, + ) + + with patch( + "litellm.integrations.newrelic.newrelic.NewRelicLogger" + ) as MockNewRelicLogger: + mock_instance = MagicMock() + mock_instance.async_health_check = AsyncMock( + return_value={"status": "healthy", "error_message": ""} + ) + MockNewRelicLogger.return_value = mock_instance + + result = await health_services_endpoint( + user_api_key_dict=user_api_key_dict, + service="newrelic", + ) + + assert result["status"] == "healthy" + mock_instance.async_health_check.assert_awaited_once() + + @pytest.fixture(scope="function") def proxy_client(monkeypatch): """ diff --git a/ui/litellm-dashboard/public/assets/logos/newrelic.png b/ui/litellm-dashboard/public/assets/logos/newrelic.png new file mode 100644 index 0000000000000000000000000000000000000000..c841e3e71365fb00db681b80ef149c812998dede GIT binary patch literal 862 zcmV-k1EKthP)Px&7)eAyR9Hv7mp@DsaTLeD@2(9tIA9FY+F*0g`AXufX z?VUDF5Ceh9s2dB)?5ZvrCyk2+5=glAN*WUe4Z(7K=J;d(UGJ{6cl5V*CvOqcV*u@TP8i66IFkPG?Io_hHobs&0#FRK2Sk%#obh55gfCzo)TP*aH+e^Z zpb-!GAdUmn3$Q2HqC;W^&zmquvQm7x7Kp9Uj_i1=59?uKVHn_m+bLZBuQTz)CIYF| z;K=Vv*>V5|DTdk)PznUu(J*6;0Tlv3Kn#pF7x^^6VO@&Ng2217WdL0ZcUsmD9vLOz zK7iI522CcsFjW#w*w#Y4kH6s}IE8$j`fEHPu$>&b1?l1Ypa30Cnc`Md%YRUC8g6d5_IJs?K1 zQ2h^m{iD!V$ta?o-&{MDPKP|1yfS%{wAUp$DhKd!Ch5jiF(SAZBDd#MFF=Un2l9q+&vPkmCcvJWQ8NIx+c8Z7 zxL9{PZUjKqlZW9mpENX(Gy>3RXg3Amm7@)Y`$*pF07*qoM6N<$g8i0&3jhEB literal 0 HcmV?d00001 diff --git a/ui/litellm-dashboard/src/lib/http/schema.d.ts b/ui/litellm-dashboard/src/lib/http/schema.d.ts index 15123bcdbf8..21e93d05a66 100644 --- a/ui/litellm-dashboard/src/lib/http/schema.d.ts +++ b/ui/litellm-dashboard/src/lib/http/schema.d.ts @@ -40443,7 +40443,7 @@ export interface operations { parameters: { query: { /** @description Specify the service being hit. */ - service: ("slack_budget_alerts" | "langfuse" | "langfuse_otel" | "slack" | "openmeter" | "webhook" | "email" | "braintrust" | "datadog" | "datadog_llm_observability" | "generic_api" | "arize" | "galileo" | "sqs") | string; + service: ("slack_budget_alerts" | "langfuse" | "langfuse_otel" | "slack" | "openmeter" | "webhook" | "email" | "braintrust" | "datadog" | "datadog_llm_observability" | "generic_api" | "arize" | "galileo" | "newrelic" | "sqs") | string; }; header?: never; path?: never;