scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-10-06 02:48:13 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 103

Implemented basic support for the success path. Generates New Relic

Browse source 
custom events needed by the AI Monitorin interface.
This commit is contained in:
Josh Bonczkowski 2025-11-21 14:02:58 -05:00
parent 375545e0ca
commit 870c7662e0
2 changed files with 211 additions and 62 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

258
litellm/integrations/newrelic/newrelic.py
View file

@ -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).

15
litellm/litellm_core_utils/litellm_logging.py
View file

@ -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: