From 870c7662e069ff46ffe70a62218116d9d5361a76 Mon Sep 17 00:00:00 2001 From: Josh Bonczkowski Date: Fri, 21 Nov 2025 14:02:58 -0500 Subject: [PATCH] Implemented basic support for the success path. Generates New Relic custom events needed by the AI Monitorin interface. --- litellm/integrations/newrelic/newrelic.py | 258 +++++++++++++----- litellm/litellm_core_utils/litellm_logging.py | 15 + 2 files changed, 211 insertions(+), 62 deletions(-) diff --git a/litellm/integrations/newrelic/newrelic.py b/litellm/integrations/newrelic/newrelic.py index cf9b453b859..3ca22dba046 100644 --- a/litellm/integrations/newrelic/newrelic.py +++ b/litellm/integrations/newrelic/newrelic.py @@ -20,10 +20,13 @@ Usage: import json import os import uuid -from typing import Any, Dict, List, Optional, Tuple +from typing import Any, Dict, List, Optional, Tuple, Union from litellm._logging import verbose_logger from litellm.integrations.custom_logger import CustomLogger +from litellm.types.utils import ModelResponse, Message + +import newrelic.agent class NewRelicLogger(CustomLogger): @@ -81,20 +84,45 @@ class NewRelicLogger(CustomLogger): """Check if message content should be recorded.""" return self.record_content - def _get_trace_context(self) -> Tuple[Optional[str], Optional[str]]: + def _get_trace_context(self, kwargs: Dict) -> Tuple[Optional[str], Optional[str]]: """ Get current New Relic trace ID and span ID. + This integration runs asynchronously from the actual request, so we + cannot just use the New Relic agent to pull the current traceId and spanId. + + For the traceID, we can look in the kwargs for litellm_params.metadata.headers.newrelic + or litellm_params.metadata.headers.traceparent. If either are set, we can pull the traceId + from those values. + + We do not have access to the SpanID for the method in the litellm server where the + request to the LLM was being invoked. We may be able to pass that through later + by including litellm instrumentation with the New Relic agent. Returns: - Tuple of (trace_id, span_id) or (None, None) if not available + Tuple of (trace_id, None) or (None, None) if not available """ try: - import newrelic.agent + litellm_params = kwargs.get("litellm_params", {}) + metadata = litellm_params.get("metadata", {}) + headers = metadata.get("headers", {}) + newrelic = headers.get("newrelic", None) + traceparent = headers.get("traceparent", None) - trace_id = newrelic.agent.current_trace_id() - span_id = newrelic.agent.current_span_id() + trace_id = None + span_id = None - if not trace_id or not span_id: + 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: + # attempt to pull from newrelic header + pass + + if not trace_id: verbose_logger.warning( "New Relic trace_id or span_id not available. " "Skipping New Relic event recording." @@ -112,19 +140,18 @@ class NewRelicLogger(CustomLogger): verbose_logger.warning("Unable to get New Relic trace context.") return None, None - def _extract_completion_id(self, kwargs: Dict, response_obj: Dict) -> str: + def _extract_completion_id(self, kwargs: Dict, response_obj: ModelResponse) -> str: """ Extract completion ID from kwargs or response_obj, or generate one. - - Per spec: Check kwargs first, then response_obj, then generate UUID. """ - # Check kwargs first per spec - completion_id = kwargs.get("id") + completion_id = None - # If not in kwargs, check response_obj - if not completion_id: + 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()) @@ -137,22 +164,22 @@ class NewRelicLogger(CustomLogger): def _get_vendor(self, kwargs: Dict) -> str: """Extract vendor/provider from kwargs.""" litellm_params = kwargs.get("litellm_params", {}) or {} - return litellm_params.get("custom_llm_provider", "unknown") + return litellm_params.get("custom_llm_provider", "litellm") - def _get_model_names(self, kwargs: Dict, response_obj: Dict) -> Tuple[str, str]: + def _get_model_names(self, kwargs: Dict, response_obj: ModelResponse) -> Tuple[str, str]: """ Extract request and response model names. Returns: Tuple of (request_model, response_model) """ - request_model = kwargs.get("model", "unknown") - response_model = response_obj.get("model", request_model) + request_model: str = str(kwargs.get("model", "unknown")) + response_model: str = str(response_obj.get("model", request_model)) return request_model, response_model - def _extract_usage(self, response_obj: Dict) -> Dict[str, int]: + def _extract_usage(self, response_obj: ModelResponse) -> Dict[str, int]: """Extract usage statistics from response.""" - usage = response_obj.get("usage", {}) + usage = response_obj.get("usage", None) if not usage: return { "prompt_tokens": 0, @@ -166,7 +193,7 @@ class NewRelicLogger(CustomLogger): "total_tokens": usage.get("total_tokens", 0) } - def _get_finish_reason(self, response_obj: Dict) -> str: + def _get_finish_reason(self, response_obj: ModelResponse) -> str: """ Extract finish reason from first choice in the response. @@ -177,7 +204,44 @@ class NewRelicLogger(CustomLogger): return choices[0].get("finish_reason", "unknown") return "unknown" - def _extract_message_content(self, message: Dict) -> str: + def _get_duration(self, kwargs: Dict, start_time: Optional[float], end_time: Optional[float]) -> Optional[float]: + """ + Extract duration in milliseconds. + + First tries to get llm_api_duration_ms from kwargs, then falls back to + calculating from start_time and end_time if available. + """ + # Try to get pre-calculated duration from kwargs + duration_ms = kwargs.get("llm_api_duration_ms") + if duration_ms is not None: + return float(duration_ms) + + # Fall back to calculating from timestamps + if start_time is not None and end_time is not None: + return (end_time - start_time) * 1000.0 # Convert to milliseconds + + return None + + def _get_request_params(self, kwargs: Dict) -> Dict[str, Any]: + """ + Extract request parameters like temperature and max_tokens. + + Returns dict with available parameters, omitting those not present. + """ + optional_params = kwargs.get("optional_params", {}) + params = {} + + temperature = optional_params.get("temperature") + if temperature is not None: + params["temperature"] = temperature + + max_tokens = optional_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. @@ -213,7 +277,7 @@ class NewRelicLogger(CustomLogger): def _extract_all_messages( self, kwargs: Dict, - response_obj: Dict, + response_obj: ModelResponse, response_model: str, vendor: str ) -> List[Dict[str, Any]]: @@ -245,57 +309,77 @@ class NewRelicLogger(CustomLogger): # Extract response messages from choices choices = response_obj.get("choices", []) - for choice in choices: - message = choice.get("message", {}) - if message: - message_data = { - "role": message.get("role", "assistant"), - "sequence": sequence, - "response.model": response_model, - "vendor": vendor - } + if choices and len(choices) > 0: + for choice in choices: + message = choice.get("message", None) + if message: + message_data = { + "role": message.get("role", "assistant"), + "sequence": sequence, + "response.model": response_model, + "vendor": vendor, + "is_response": True + } - # Only add content if recording is enabled - if self._should_record_content(): - message_data["content"] = self._extract_message_content(message) + # Only add content if recording is enabled + if self._should_record_content(): + message_data["content"] = self._extract_message_content(message) - messages.append(message_data) - sequence += 1 + messages.append(message_data) + sequence += 1 return messages def _record_summary_event( self, - completion_id: str, + request_id: str, trace_id: str, - span_id: str, + span_id: Optional[str], request_model: str, response_model: str, vendor: str, finish_reason: str, num_messages: int, - usage: Dict[str, int] + usage: Dict[str, int], + duration: Optional[float] = None, + request_params: Optional[Dict[str, Any]] = None ): """Record LlmChatCompletionSummary event to New Relic.""" try: import newrelic.agent event_data = { - "id": completion_id, + "id": request_id, + "request_id": request_id, "trace_id": trace_id, - "span_id": span_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 span_id: + event_data["span_id"] = span_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"] + newrelic.agent.record_custom_event("LlmChatCompletionSummary", event_data) - verbose_logger.debug("Recorded LlmChatCompletionSummary event") + import pprint + verbose_logger.info(f"Recorded LlmChatCompletionSummary event: {pprint.pformat(event_data)}") except Exception as e: verbose_logger.warning(f"Failed to record New Relic summary event: {e}") @@ -303,34 +387,55 @@ class NewRelicLogger(CustomLogger): def _record_message_events( self, - completion_id: str, + request_id: str, + llm_response_id: str, trace_id: str, - span_id: str, + span_id: Optional[str], messages: List[Dict[str, Any]] ): - """Record LlmChatCompletionMessage events to New Relic.""" + """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 + span_id: Span ID for distributed tracing + messages: List of message dicts to record + """ try: import newrelic.agent for message in messages: + sequence = message["sequence"] event_data = { - "completion_id": completion_id, + "id": f"{llm_response_id}-{sequence}", + "request_id": request_id, + "completion_id": request_id, "trace_id": trace_id, - "span_id": span_id, "role": message["role"], - "sequence": message["sequence"], + "sequence": sequence, "response.model": message["response.model"], - "vendor": message["vendor"] + "vendor": message["vendor"], + "ingest_source": "litellm", + "token_count": 0 } # 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 + + if span_id: + event_data["span_id"] = span_id + newrelic.agent.record_custom_event("LlmChatCompletionMessage", event_data) - verbose_logger.debug( - f"Recorded {len(messages)} LlmChatCompletionMessage events" + import pprint + verbose_logger.info( + f"Recorded {len(messages)} LlmChatCompletionMessage events: {pprint.pformat(messages)}" ) except Exception as e: @@ -343,13 +448,19 @@ class NewRelicLogger(CustomLogger): import newrelic.agent newrelic.agent.record_custom_metric("LLM/LiteLLM/Error", 1) - verbose_logger.debug("Recorded LLM/LiteLLM/Error metric") + verbose_logger.info("Recorded LLM/LiteLLM/Error metric") 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: Dict): + 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. @@ -358,18 +469,30 @@ class NewRelicLogger(CustomLogger): if not self.enabled: return + import pprint + verbose_logger.info(f"newrelic._process_success called, kwargs=\n{pprint.pformat(kwargs)}, \nresponse_obj=\n{pprint.pformat(response_obj)}") + # Get trace context - trace_id, span_id = self._get_trace_context() - if not trace_id or not span_id: + trace_id, span_id = self._get_trace_context(kwargs) + if not trace_id: return + verbose_logger.info(f"Trace ID: {trace_id}, Span ID: {span_id}") + + # Generate unique request ID for this request (used as Summary event id) + request_id = str(uuid.uuid4()) + # Extract data from response - completion_id = self._extract_completion_id(kwargs, response_obj) + llm_response_id = self._extract_completion_id(kwargs, response_obj) vendor = self._get_vendor(kwargs) request_model, response_model = self._get_model_names(kwargs, response_obj) usage = self._extract_usage(response_obj) finish_reason = self._get_finish_reason(response_obj) + # Extract additional summary event fields + duration = self._get_duration(kwargs, start_time, end_time) + request_params = self._get_request_params(kwargs) + # Extract all messages messages = self._extract_all_messages( kwargs, response_obj, response_model, vendor @@ -377,7 +500,7 @@ class NewRelicLogger(CustomLogger): # Record summary event self._record_summary_event( - completion_id=completion_id, + request_id=request_id, trace_id=trace_id, span_id=span_id, request_model=request_model, @@ -385,12 +508,15 @@ class NewRelicLogger(CustomLogger): vendor=vendor, finish_reason=finish_reason, num_messages=len(messages), - usage=usage + usage=usage, + duration=duration, + request_params=request_params ) # Record message events self._record_message_events( - completion_id=completion_id, + request_id=request_id, + llm_response_id=llm_response_id, trace_id=trace_id, span_id=span_id, messages=messages @@ -406,6 +532,8 @@ class NewRelicLogger(CustomLogger): """Unused per spec.""" pass + @newrelic.agent.background_task() + @newrelic.agent.function_trace(name="log_success_event") def log_success_event(self, kwargs, response_obj, start_time, end_time): """ Main success path for non-streaming requests. @@ -414,11 +542,13 @@ class NewRelicLogger(CustomLogger): (in-memory operation), so it's safe to call from sync context. """ try: - self._process_success(kwargs, response_obj) + 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") + @newrelic.agent.background_task() + @newrelic.agent.function_trace(name="async_log_success_event") async def async_log_success_event(self, kwargs, response_obj, start_time, end_time): """ Main success path for async/streaming requests. @@ -427,11 +557,13 @@ class NewRelicLogger(CustomLogger): so we can call it directly without asyncio.to_thread(). """ try: - self._process_success(kwargs, response_obj) + 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") + @newrelic.agent.background_task() + @newrelic.agent.function_trace(name="log_failure_event") def log_failure_event(self, kwargs, response_obj, start_time, end_time): """ Log error metric for failed LLM calls (sync). @@ -448,6 +580,8 @@ class NewRelicLogger(CustomLogger): verbose_logger.warning(f"Error in New Relic log_failure_event: {e}") self.handle_callback_failure("newrelic") + @newrelic.agent.background_task() + @newrelic.agent.function_trace(name="async_log_failure_event") async def async_log_failure_event(self, kwargs, response_obj, start_time, end_time): """ Log error metric for failed LLM calls (async). diff --git a/litellm/litellm_core_utils/litellm_logging.py b/litellm/litellm_core_utils/litellm_logging.py index 6bad7ee29e2..fe97d5c95bf 100644 --- a/litellm/litellm_core_utils/litellm_logging.py +++ b/litellm/litellm_core_utils/litellm_logging.py @@ -3788,6 +3788,15 @@ 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": + from litellm.integrations.newrelic import NewRelicLogger + + 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( @@ -3972,6 +3981,12 @@ def get_custom_logger_compatible_class( # noqa: PLR0915 for callback in _in_memory_loggers: if isinstance(callback, SMTPEmailLogger): return callback + elif logging_integration == "newrelic": + from litellm.integrations.newrelic import NewRelicLogger + + for callback in _in_memory_loggers: + if isinstance(callback, NewRelicLogger): + return callback return None except Exception as e: