scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-08-28 05:25:59 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 36

feat(newrelic): per-team cost and usage metrics via team callbacks (#37610)

Browse source 
* feat(newrelic): per-team cost and usage metrics via team callbacks

Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>

* fix(newrelic): retry transient 429/408 metric posts instead of dropping

* fix(newrelic): drop only records queued when the drain began, not mid-drain arrivals

---------

Co-authored-by: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>
This commit is contained in:
yucheng-berri 2026-08-26 23:42:02 -07:00 • committed by GitHub
parent 807ee7f232
commit 8ebcb3e181
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
8 changed files with 1764 additions and 20 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

2
litellm/integrations/custom_batch_logger.py
View file

@ -45,7 +45,7 @@ class CustomBatchLogger(CustomLogger):
super().__init__(**kwargs)
async def periodic_flush(self):
async def periodic_flush(self) -> None:
while True:
await asyncio.sleep(self.flush_interval)
verbose_logger.debug("CustomLogger periodic flush after %s seconds", self.flush_interval)

395
litellm/integrations/newrelic/newrelic_metrics.py Normal file
View file

@ -0,0 +1,395 @@
"""
New Relic Metric API Integration - sends per-team cost/usage metrics to /metric/v1
NR Reference API: https://docs.newrelic.com/docs/data-apis/ingest-apis/metric-api/introduction-metric-api/
`async_log_success_event` / `async_log_failure_event` queue one record per request;
at flush the queue is aggregated by (team, model group, model, provider, status)
into count/summary metrics. `interval.ms` is the real window between flushes,
computed at flush time.
Team-scoped by construction: the ingest key is injected explicitly and there is
deliberately no environment-variable fallback, so a team's metrics are never sent
with the proxy operator's credentials (mirrors ``allow_env_credentials=False`` on
the Datadog team logger).
Error policy on flush: 4xx drops the batch (a retry would fail identically; 403
is a permanent credential failure), 5xx/network re-queues capped at
``max_queue_size`` records with the oldest dropped.
For batching specific details see CustomBatchLogger class
"""
import asyncio
import gzip
import time
import traceback
from collections.abc import Mapping
from math import ceil
from types import MappingProxyType
from typing import Final
from httpx import HTTPStatusError, Response
from litellm._logging import verbose_logger
from litellm.integrations.custom_batch_logger import CustomBatchLogger
from litellm.litellm_core_utils.safe_json_dumps import safe_dumps
from litellm.llms.custom_httpx.http_handler import (
get_async_httpx_client,
httpxSpecialProvider,
)
from litellm.types.integrations.newrelic import (
NEWRELIC_DEFAULT_REGION,
NEWRELIC_METRIC_ATTRIBUTE_MAX_LEN,
NEWRELIC_METRIC_COMPLETION_TOKENS,
NEWRELIC_METRIC_COST_USD,
NEWRELIC_METRIC_ENDPOINT_BY_REGION,
NEWRELIC_METRIC_PROMPT_TOKENS,
NEWRELIC_METRIC_REQUEST_DURATION_MS,
NEWRELIC_METRIC_REQUESTS,
NEWRELIC_METRIC_TOTAL_TOKENS,
NEWRELIC_METRICS_MAX_BATCH_SIZE,
NEWRELIC_METRICS_MAX_DRAIN_PASSES,
NEWRELIC_METRICS_MAX_RETRY_QUEUE_SIZE,
NewRelicCountMetric,
NewRelicMetric,
NewRelicMetricCommon,
NewRelicMetricEnvelope,
NewRelicMetricRecord,
NewRelicSummaryMetric,
NewRelicSummaryValue,
)
from litellm.types.utils import StandardLoggingPayload
# 408 (request timeout) and 429 (rate limit) are transient client errors the
# Metric API expects a retry on, unlike 400/403 which a retry would only repeat.
_RETRYABLE_CLIENT_STATUSES: Final = frozenset({408, 429})
def resolve_newrelic_metric_endpoint(newrelic_region: str | None) -> str:
if not newrelic_region:
return NEWRELIC_METRIC_ENDPOINT_BY_REGION[NEWRELIC_DEFAULT_REGION]
endpoint: Final = NEWRELIC_METRIC_ENDPOINT_BY_REGION.get(newrelic_region.lower())
if endpoint is None:
verbose_logger.warning(
"New Relic: unknown newrelic_region %r; supported regions: %s. Using the default (US) endpoint.",
newrelic_region,
", ".join(sorted(NEWRELIC_METRIC_ENDPOINT_BY_REGION)),
)
return NEWRELIC_METRIC_ENDPOINT_BY_REGION[NEWRELIC_DEFAULT_REGION]
return endpoint
def _metric_record_from_payload(standard_logging_object: StandardLoggingPayload) -> NewRelicMetricRecord:
metadata: Final = standard_logging_object.get("metadata")
team_id: Final = ((metadata.get("user_api_key_team_id") or metadata.get("team_id")) if metadata else None) or ""
team_alias: Final = (
(metadata.get("user_api_key_team_alias") or metadata.get("team_alias")) if metadata else None
) or ""
return NewRelicMetricRecord(
team_id=team_id,
team_alias=team_alias,
model_group=standard_logging_object.get("model_group") or "",
model=standard_logging_object.get("model") or "",
custom_llm_provider=standard_logging_object.get("custom_llm_provider") or "",
status=str(standard_logging_object.get("status") or "success"),
response_cost=float(standard_logging_object.get("response_cost") or 0.0),
prompt_tokens=int(standard_logging_object.get("prompt_tokens") or 0),
completion_tokens=int(standard_logging_object.get("completion_tokens") or 0),
total_tokens=int(standard_logging_object.get("total_tokens") or 0),
duration_ms=float(standard_logging_object.get("response_time") or 0.0) * 1000.0,
)
def _bucket_metrics(bucket_records: tuple[NewRelicMetricRecord, ...]) -> tuple[NewRelicMetric, ...]:
first: Final = bucket_records[0]
attributes: Final[Mapping[str, str]] = { # mutable-ok: JSON leaf; safe_dumps stringifies MappingProxyType
key: value[:NEWRELIC_METRIC_ATTRIBUTE_MAX_LEN]
for key, value in (
("team_id", first.team_id),
("team_alias", first.team_alias),
("model_group", first.model_group),
("model", first.model),
("custom_llm_provider", first.custom_llm_provider),
("status", first.status),
)
if value
}
durations: Final = tuple(record.duration_ms for record in bucket_records)
counts: Final[tuple[tuple[str, float], ...]] = (
(NEWRELIC_METRIC_REQUESTS, float(len(bucket_records))),
(NEWRELIC_METRIC_COST_USD, sum(record.response_cost for record in bucket_records)),
(NEWRELIC_METRIC_PROMPT_TOKENS, float(sum(record.prompt_tokens for record in bucket_records))),
(NEWRELIC_METRIC_COMPLETION_TOKENS, float(sum(record.completion_tokens for record in bucket_records))),
(NEWRELIC_METRIC_TOTAL_TOKENS, float(sum(record.total_tokens for record in bucket_records))),
)
count_metrics: Final[tuple[NewRelicMetric, ...]] = tuple(
NewRelicCountMetric(name=name, type="count", value=value, attributes=attributes) for name, value in counts
)
summary_metric: Final = NewRelicSummaryMetric(
name=NEWRELIC_METRIC_REQUEST_DURATION_MS,
type="summary",
value=NewRelicSummaryValue(
count=len(durations),
sum=sum(durations),
min=min(durations),
max=max(durations),
),
attributes=attributes,
)
return (*count_metrics, summary_metric)
def build_metric_payload(
records: tuple[NewRelicMetricRecord, ...],
*,
window_start: float,
now: float,
) -> tuple[NewRelicMetricEnvelope, ...]:
"""Aggregates records into one Metric API envelope for the flush window."""
interval_ms: Final = max(1, int((now - window_start) * 1000))
bucket_keys: Final = tuple(dict.fromkeys(record.bucket_key for record in records))
metrics: Final = tuple(
metric
for key in bucket_keys
for metric in _bucket_metrics(tuple(record for record in records if record.bucket_key == key))
)
common: Final[NewRelicMetricCommon] = {
"timestamp": int(window_start * 1000),
"interval.ms": interval_ms,
}
return (NewRelicMetricEnvelope(common=common, metrics=metrics),)
class NewRelicMetricsLogger(CustomBatchLogger):
def __init__(
self,
newrelic_api_key: str,
newrelic_region: str | None = None,
) -> None:
if not newrelic_api_key:
raise ValueError(
"newrelic_api_key is required for NewRelicMetricsLogger; "
"team-scoped metrics never fall back to environment credentials"
)
self.newrelic_api_key: Final = newrelic_api_key
self.metric_api_url: Final = resolve_newrelic_metric_endpoint(newrelic_region)
self.async_client = get_async_httpx_client(llm_provider=httpxSpecialProvider.LoggingCallback)
self._stopped: bool = False
self._drain_lock = asyncio.Lock()
asyncio.create_task(self.periodic_flush())
self.flush_lock = asyncio.Lock()
super().__init__(
flush_lock=self.flush_lock,
batch_size=NEWRELIC_METRICS_MAX_BATCH_SIZE,
max_queue_size=NEWRELIC_METRICS_MAX_RETRY_QUEUE_SIZE,
)
def stop(self) -> None:
"""Ends the periodic flush loop; called on DynamicLoggingCache eviction.
Schedules one final drain of anything still queued, so eviction never
silently discards records. Guarded so it can never raise into the
cache's eviction path.
"""
self._stopped = True
try:
asyncio.get_running_loop().create_task(self._final_drain())
except Exception: # noqa: BLE001 # no running loop / shutdown; the periodic loop's final drain still runs
verbose_logger.debug("New Relic Metrics: could not schedule final drain on stop()", exc_info=True)
async def _drain_with_retry(self) -> None:
"""Deliver everything queued on a stopped logger, or drop it with a log.
A stopped logger has no periodic loop left, so every post-stop path
funnels through here. ``_drain_lock`` serializes drains: a callback that
appends and starts its own drain queues behind the running one instead
of racing it. Each pass attempts the whole current queue in
``batch_size`` chunks, unlike the periodic path it does not stop at the
first failing chunk, so a persistently failing head never starves the
tail. Only after ``_MAX_DRAIN_PASSES`` against a permanently failing
destination is the remainder dropped, and then only the records that were
queued when this drain began, so every dropped record got the full retry
budget: a record a callback appended mid-drain is not in that snapshot,
so it is left for its own serialized drain rather than dropped after
fewer attempts, and is never stranded.
"""
async with self._drain_lock:
attempted: Final = tuple(self.log_queue)
for _pass in range(NEWRELIC_METRICS_MAX_DRAIN_PASSES):
await self._drain_flush_once()
if not self.log_queue:
return
if _pass < NEWRELIC_METRICS_MAX_DRAIN_PASSES - 1:
await asyncio.sleep(2**_pass)
async with self.flush_lock:
tried_ids: Final = frozenset(id(record) for record in attempted)
survivors: Final = tuple(record for record in self.log_queue if id(record) not in tried_ids)
dropped: Final = len(self.log_queue) - len(survivors)
if dropped:
verbose_logger.warning(
"New Relic Metrics: dropping %s records after %s drain passes",
dropped,
NEWRELIC_METRICS_MAX_DRAIN_PASSES,
)
self.log_queue[:] = list(survivors) # mutable-ok: leave late arrivals for the next serialized drain
async def _drain_flush_once(self) -> None:
"""Attempt every queued record once, in ``batch_size`` chunks, without
stopping at the first failing chunk so a persistently failing head does
not starve the tail (the periodic ``flush_queue`` deliberately stops
instead). Takes the queue under ``flush_lock`` and re-queues only the
chunks a 5xx/network error left undelivered, so records a concurrent
request appends during the sends survive for the next pass."""
async with self.flush_lock:
pending: Final = tuple(self.log_queue)
window_start: Final = self.last_flush_time
self.last_flush_time = time.time()
del self.log_queue[:]
if not pending:
return
chunks: Final = tuple(
pending[start : start + self.batch_size] for start in range(0, len(pending), self.batch_size)
)
delivered: Final = tuple([await self._classify_and_send(chunk, window_start) for chunk in chunks])
failed: Final = tuple(record for chunk, ok in zip(chunks, delivered) for record in (() if ok else chunk))
if failed:
self._requeue(failed)
async def _final_drain(self) -> None:
await self._drain_with_retry()
async def periodic_flush(self) -> None:
while not self._stopped:
await asyncio.sleep(self.flush_interval)
if self._stopped:
break
await self.flush_queue()
await self._final_drain()
async def async_log_success_event(self, kwargs, response_obj, start_time, end_time) -> None:
try:
await self._log_async_event(standard_logging_object=kwargs.get("standard_logging_object", None))
except Exception as e: # noqa: BLE001 # logging must never break the request path
verbose_logger.exception("New Relic Metrics Layer Error - %s\n%s", e, traceback.format_exc())
async def async_log_failure_event(self, kwargs, response_obj, start_time, end_time) -> None:
try:
await self._log_async_event(standard_logging_object=kwargs.get("standard_logging_object", None))
except Exception as e: # noqa: BLE001 # logging must never break the request path
verbose_logger.exception("New Relic Metrics Layer Error - %s\n%s", e, traceback.format_exc())
async def _log_async_event(self, standard_logging_object: StandardLoggingPayload | None) -> None:
if standard_logging_object is None:
raise ValueError("standard_logging_object not found in kwargs")
self.log_queue.append(_metric_record_from_payload(standard_logging_object))
if self._stopped:
# A stopped logger has no periodic loop left; an in-flight callback
# that appends after the eviction drain delivers its own record.
await self._drain_with_retry()
return
if len(self.log_queue) >= self.batch_size:
await self.flush_queue()
async def flush_queue(self) -> None:
async with self.flush_lock:
window_start: Final = self.last_flush_time
self.last_flush_time = time.time()
queued: Final = len(self.log_queue)
if not queued:
return
verbose_logger.debug("New Relic Metrics: Flushing %s queued records", queued)
# Bounded by what is queued now: records appended mid-flush belong to
# the next window, and looping until empty would never end under load.
for _chunk in range(ceil(queued / self.batch_size)):
if not await self.async_send_batch(window_start=window_start):
return
async def async_send_batch(self, window_start: float | None = None) -> bool:
"""Sends the oldest ``batch_size`` records only, so a queue grown past that
by re-queues cannot breach the Metric API data point cap in one request.
Returns False once a chunk fails and is re-queued, so the caller stops."""
if not self.log_queue:
return False
batch_to_send: Final[tuple[NewRelicMetricRecord, ...]] = tuple(self.log_queue[: self.batch_size])
del self.log_queue[: len(batch_to_send)]
delivered: Final = await self._classify_and_send(
batch_to_send, window_start if window_start is not None else self.last_flush_time
)
if not delivered:
self._requeue(batch_to_send)
return delivered
async def _classify_and_send(self, batch: tuple[NewRelicMetricRecord, ...], window_start: float) -> bool:
"""Send one chunk and classify the outcome, never touching the queue.
Returns True when the batch is done with (delivered on any 2xx, or a 4xx
a retry would only repeat, 403 being a permanent bad-key rejection), and
False when a 5xx or network error means the caller should re-queue it.
``AsyncHTTPHandler.post`` raises ``HTTPStatusError`` on any non-2xx, so a
4xx never returns a response here; the status is read off the raised
error to keep the client-error path (drop) distinct from 5xx (retry)."""
payload: Final = build_metric_payload(records=batch, window_start=window_start, now=time.time())
try:
status = (
await self.async_send_compressed_data(payload)
).status_code # rebind-ok: reassigned from the raised HTTPStatusError below
except HTTPStatusError as e:
status = e.response.status_code
except Exception as e: # noqa: BLE001 # transport/network failure re-queues the batch
verbose_logger.warning(
"New Relic Metrics: network error sending %s records, will retry - %s",
len(batch),
e,
)
return False
if 200 <= status < 300:
return True
if 400 <= status < 500 and status not in _RETRYABLE_CLIENT_STATUSES:
verbose_logger.warning(
"New Relic Metrics: %s from Metric API%s, dropping %s records.",
status,
" (permanent credential failure: invalid or revoked team ingest key)" if status == 403 else "",
len(batch),
)
return True
verbose_logger.warning(
"New Relic Metrics: %s from Metric API, will retry %s records",
status,
len(batch),
)
return False
def _requeue(self, batch: tuple[NewRelicMetricRecord, ...]) -> None:
"""Prepends ``batch`` in place (never by assignment: records appended by
concurrent requests during the flush await must survive), keeping
chronological order so the cap drops the oldest records first."""
self.log_queue[:0] = batch
overflow: Final = len(self.log_queue) - self.max_queue_size
if overflow > 0:
del self.log_queue[:overflow]
verbose_logger.warning(
"New Relic Metrics: retry queue exceeded max_queue_size=%s; dropped %s oldest records.",
self.max_queue_size,
overflow,
)
async def async_send_compressed_data(self, payload: tuple[NewRelicMetricEnvelope, ...]) -> Response:
compressed_data: Final = gzip.compress(safe_dumps(payload).encode("utf-8"))
headers: Final[Mapping[str, str]] = MappingProxyType(
{
"Content-Type": "application/json",
"Content-Encoding": "gzip",
"Api-Key": self.newrelic_api_key,
}
)
return await self.async_client.post(
url=self.metric_api_url,
data=compressed_data,
headers=headers,
)

90
litellm/integrations/newrelic/newrelic_team_handler.py Normal file
View file

@ -0,0 +1,90 @@
"""
New Relic Team Handler
Used to get the NewRelicMetricsLogger for a given request.
Handles Key/Team Based New Relic metrics, following the same pattern as DataDogHandler.
"""
from typing import TYPE_CHECKING, Final
from typing_extensions import ReadOnly, TypedDict
from litellm._logging import verbose_logger
from litellm.litellm_core_utils.litellm_logging import StandardCallbackDynamicParams
from .newrelic_metrics import NewRelicMetricsLogger
if TYPE_CHECKING:
from litellm.litellm_core_utils.litellm_logging import DynamicLoggingCache
class NewRelicLoggingConfig(TypedDict):
newrelic_api_key: ReadOnly[str | None]
newrelic_region: ReadOnly[str | None]
class NewRelicHandler:
@staticmethod
def get_newrelic_logger_for_request(
standard_callback_dynamic_params: StandardCallbackDynamicParams,
in_memory_dynamic_logger_cache: "DynamicLoggingCache",
) -> NewRelicMetricsLogger:
"""
Get a team-scoped NewRelicMetricsLogger for a given request.
Resolves and caches per-team NewRelicMetricsLogger instances using
DynamicLoggingCache, keyed by the team's New Relic credentials. Each unique
set of credentials gets its own logger instance with its own batch/flush loop.
Note: This handler is only called when a team-scoped newrelic_api_key is
present. The trace logger for the ``newrelic`` callback (OTel v2 / legacy
agent) is managed separately by _init_custom_logger_compatible_class via
_in_memory_loggers.
"""
_credentials: Final = NewRelicHandler.get_dynamic_newrelic_logging_config(
standard_callback_dynamic_params=standard_callback_dynamic_params,
)
temp_newrelic_logger = in_memory_dynamic_logger_cache.get_cache(
credentials=_credentials, service_name="newrelic"
)
if temp_newrelic_logger is None:
temp_newrelic_logger = NewRelicHandler._create_newrelic_logger_from_credentials(
credentials=_credentials,
in_memory_dynamic_logger_cache=in_memory_dynamic_logger_cache,
)
return temp_newrelic_logger
@staticmethod
def _create_newrelic_logger_from_credentials(
credentials: NewRelicLoggingConfig,
in_memory_dynamic_logger_cache: "DynamicLoggingCache",
) -> NewRelicMetricsLogger:
newrelic_logger: Final = NewRelicMetricsLogger(
newrelic_api_key=credentials.get("newrelic_api_key") or "",
newrelic_region=credentials.get("newrelic_region"),
)
in_memory_dynamic_logger_cache.set_cache(
credentials=credentials,
service_name="newrelic",
logging_obj=newrelic_logger,
)
verbose_logger.debug("New Relic: Created and cached new NewRelicMetricsLogger for team-scoped credentials")
return newrelic_logger
@staticmethod
def get_dynamic_newrelic_logging_config(
standard_callback_dynamic_params: StandardCallbackDynamicParams,
) -> NewRelicLoggingConfig:
return NewRelicLoggingConfig(
newrelic_api_key=standard_callback_dynamic_params.get("newrelic_api_key"),
newrelic_region=standard_callback_dynamic_params.get("newrelic_region"),
)
@staticmethod
def _dynamic_newrelic_credentials_are_passed(
standard_callback_dynamic_params: StandardCallbackDynamicParams,
) -> bool:
return standard_callback_dynamic_params.get("newrelic_api_key") is not None

74
litellm/litellm_core_utils/litellm_logging.py
View file

@ -613,37 +613,60 @@ class Logging(LiteLLMLoggingBaseClass):
processed_list: Final[list[str | Callable | CustomLogger]] = []
for callback in callback_list:
if isinstance(callback, str) and callback in litellm._known_custom_logger_compatible_callbacks:
# For callbacks that support team-scoped credentials (e.g. datadog),
# pass only the relevant dynamic params as custom_logger_init_args.
_custom_logger_init_args: dict | None = None
if callback == "datadog":
# dd_* params are blocked from standard_callback_dynamic_params
# (request-level security); only the proxy-stamped team/key
# callback vars are admin-configured and trusted.
_custom_logger_init_args = {k: v for k, v in self._trusted_callback_vars if k.startswith("dd_")}
callback_class = _init_custom_logger_compatible_class(
callback,
internal_usage_cache=None,
llm_router=None,
custom_logger_init_args=_custom_logger_init_args,
)
if callback_class is not None:
processed_list.append(callback_class)
for callback_instance in self._resolve_dynamic_callback_string(callback):
processed_list.append(callback_instance)
# If processing dynamic_success_callbacks, add to dynamic_async_success_callbacks
if dynamic_callbacks_type == "success":
if self.dynamic_async_success_callbacks is None:
self.dynamic_async_success_callbacks = []
self.dynamic_async_success_callbacks.append(callback_class)
self.dynamic_async_success_callbacks.append(callback_instance)
elif dynamic_callbacks_type == "failure":
if self.dynamic_async_failure_callbacks is None:
self.dynamic_async_failure_callbacks = []
self.dynamic_async_failure_callbacks.append(callback_class)
self.dynamic_async_failure_callbacks.append(callback_instance)
else:
processed_list.append(callback)
return processed_list
def _resolve_dynamic_callback_string(self, callback: str) -> "tuple[CustomLogger, ...]":
"""
Resolve a known callback name to the logger instance(s) it dispatches to.
For callbacks that support team-scoped credentials (datadog, newrelic),
only the proxy-stamped team/key callback vars are passed as
custom_logger_init_args: dd_*/newrelic_* params are blocked from
standard_callback_dynamic_params (request-level security), so the
trusted-vars channel is the only way credentials reach a per-team logger.
"""
_trusted_var_prefix: Final = "dd_" if callback == "datadog" else "newrelic_" if callback == "newrelic" else None
_custom_logger_init_args: Final[dict | None] = (
{k: v for k, v in self._trusted_callback_vars if k.startswith(_trusted_var_prefix)}
if _trusted_var_prefix is not None
else None
)
callback_class: Final = _init_custom_logger_compatible_class(
callback,
internal_usage_cache=None,
llm_router=None,
custom_logger_init_args=_custom_logger_init_args,
)
if callback_class is None:
return ()
# With team creds, "newrelic" resolves to the per-team METRICS logger;
# resolve the name again without creds so the trace logger (OTel v2 /
# legacy agent) keeps receiving this request.
_newrelic_trace_class: Final = (
_init_custom_logger_compatible_class(callback, internal_usage_cache=None, llm_router=None)
if callback == "newrelic" and _custom_logger_init_args and _custom_logger_init_args.get("newrelic_api_key")
else None
)
if _newrelic_trace_class is not None and _newrelic_trace_class is not callback_class:
return (callback_class, _newrelic_trace_class)
return (callback_class,)
def initialize_standard_callback_dynamic_params(self, kwargs: dict | None = None) -> StandardCallbackDynamicParams:
"""
Initialize the standard callback dynamic params from the kwargs
@ -4642,6 +4665,19 @@ def _init_custom_logger_compatible_class(
_in_memory_loggers.append(gitlab_logger)
return gitlab_logger
elif logging_integration == "newrelic":
if custom_logger_init_args.get("newrelic_api_key"):
# Team-scoped credentials: per-team METRICS logger, isolated per
# credential set via DynamicLoggingCache. The trace logger for
# this name stays on the global path below.
from litellm.integrations.newrelic.newrelic_team_handler import (
NewRelicHandler,
)
return NewRelicHandler.get_newrelic_logger_for_request(
standard_callback_dynamic_params=custom_logger_init_args,
in_memory_dynamic_logger_cache=in_memory_dynamic_logger_cache,
)
_v2 = _maybe_construct_otel_v2("newrelic", _in_memory_loggers)
if _v2 is not None:
return _v2

10
litellm/litellm_core_utils/specialty_caches/dynamic_logging_cache.py
View file

@ -13,6 +13,7 @@ import json
from typing import Any, Final
import litellm
from litellm._logging import verbose_logger
from litellm.constants import _DEFAULT_TTL_FOR_HTTPX_CLIENTS
from ...caching import InMemoryCache
@ -46,6 +47,15 @@ class LangfuseInMemoryCache(InMemoryCache):
_created_langfuse_logger.Langfuse.flush()
_created_langfuse_logger.Langfuse.shutdown()
# Loggers with a periodic flush task (e.g. NewRelicMetricsLogger) expose
# stop() so eviction actually ends the task instead of leaking it.
_evicted_stop: Final = getattr(self.cache_dict[key], "stop", None)
if callable(_evicted_stop):
try:
_evicted_stop()
except Exception: # noqa: BLE001 # a failing stop() must not block eviction
verbose_logger.debug("DynamicLoggingCache: stop() raised during eviction", exc_info=True)
#########################################################
# Call parent class to remove key from cache
#########################################################

114
litellm/types/integrations/newrelic.py
View file

@ -1,3 +1,10 @@
from collections.abc import Mapping, Sequence
from dataclasses import dataclass
from types import MappingProxyType
from typing import Final, Literal
from typing_extensions import ReadOnly, TypedDict
from litellm.types.integrations.custom_logger import StandardCustomLoggerInitParams
@ -5,3 +12,110 @@ class NewRelicInitParams(StandardCustomLoggerInitParams):
"""
Params for initializing a New Relic logger on litellm
"""
#: Region -> Metric API endpoint. A fixed table by design: team config picks a
#: region enum rather than a free-form endpoint, so callback vars can never
#: redirect metrics to an arbitrary host.
NEWRELIC_METRIC_ENDPOINT_BY_REGION: Final[Mapping[str, str]] = MappingProxyType(
{
"us": "https://metric-api.newrelic.com/metric/v1",
"eu": "https://metric-api.eu.newrelic.com/metric/v1",
}
)
NEWRELIC_DEFAULT_REGION: Final = "us"
#: Metric API caps a payload at 2000 data points / 1MB compressed; each queued
#: record expands to at most 6 metrics, so cap the per-flush record count well
#: below that.
NEWRELIC_METRICS_MAX_BATCH_SIZE: Final = 250
#: Hard cap on records retained across failed flushes (5xx/network requeue).
#: Beyond this the oldest records are dropped.
NEWRELIC_METRICS_MAX_RETRY_QUEUE_SIZE: Final = 10_000
# Outer passes over a stopped logger's queue: each pass retries the whole
# queue, so records that arrive mid-drain still get attempts before the bounded
# terminal drop. Serialized by a per-logger drain lock, so this bounds work.
NEWRELIC_METRICS_MAX_DRAIN_PASSES: Final = 3
# Metric API caps attribute values; 255 keeps caller-controlled model strings
# from inflating the shared batch payload into a 413
NEWRELIC_METRIC_ATTRIBUTE_MAX_LEN: Final = 255
NEWRELIC_METRIC_REQUESTS: Final = "litellm.requests"
NEWRELIC_METRIC_COST_USD: Final = "litellm.cost.usd"
NEWRELIC_METRIC_PROMPT_TOKENS: Final = "litellm.tokens.prompt"
NEWRELIC_METRIC_COMPLETION_TOKENS: Final = "litellm.tokens.completion"
NEWRELIC_METRIC_TOTAL_TOKENS: Final = "litellm.tokens.total"
NEWRELIC_METRIC_REQUEST_DURATION_MS: Final = "litellm.request.duration_ms"
class NewRelicSummaryValue(TypedDict):
"""Value shape of a Metric API ``summary`` data point."""
count: ReadOnly[int]
sum: ReadOnly[float]
min: ReadOnly[float]
max: ReadOnly[float]
class NewRelicCountMetric(TypedDict):
name: ReadOnly[str]
type: ReadOnly[Literal["count"]]
value: ReadOnly[float]
attributes: ReadOnly[Mapping[str, str]]
class NewRelicSummaryMetric(TypedDict):
name: ReadOnly[str]
type: ReadOnly[Literal["summary"]]
value: ReadOnly[NewRelicSummaryValue]
attributes: ReadOnly[Mapping[str, str]]
NewRelicMetric = NewRelicCountMetric | NewRelicSummaryMetric
#: ``interval.ms`` has a dot in it, so the functional TypedDict form is required.
NewRelicMetricCommon = TypedDict(
"NewRelicMetricCommon",
{ # mutable-ok: functional TypedDict requires a dict-literal fields argument ("interval.ms" key)
"timestamp": ReadOnly[int],
"interval.ms": ReadOnly[int],
},
)
class NewRelicMetricEnvelope(TypedDict):
"""One element of the Metric API request body (``[{common, metrics}]``)."""
common: ReadOnly[NewRelicMetricCommon]
metrics: ReadOnly[Sequence[NewRelicMetric]]
@dataclass(frozen=True, slots=True)
class NewRelicMetricRecord:
"""One request's contribution to the per-flush aggregation."""
team_id: str
team_alias: str
model_group: str
model: str
custom_llm_provider: str
status: str
response_cost: float
prompt_tokens: int
completion_tokens: int
total_tokens: int
duration_ms: float
@property
def bucket_key(self) -> tuple[str, str, str, str, str, str]:
return (
self.team_id,
self.team_alias,
self.model_group,
self.model,
self.custom_llm_provider,
self.status,
)

825
tests/test_litellm/integrations/newrelic/test_newrelic_metrics.py Normal file
View file

@ -0,0 +1,825 @@
"""
Batching tests for NewRelicMetricsLogger: flush-window interval computation,
dimension-bucket aggregation, the 4xx-drop vs 5xx/network-requeue policy, the
retry-queue cap, and the stop flag that ends the periodic flush loop.
"""
import asyncio
import gzip
import json
from unittest.mock import AsyncMock, patch
import pytest
from httpx import HTTPStatusError, Request, Response
from litellm.integrations.newrelic.newrelic_metrics import (
NewRelicMetricsLogger,
_bucket_metrics,
build_metric_payload,
)
from litellm.types.integrations.newrelic import (
NEWRELIC_METRIC_COMPLETION_TOKENS,
NEWRELIC_METRIC_COST_USD,
NEWRELIC_METRIC_ENDPOINT_BY_REGION,
NEWRELIC_METRIC_PROMPT_TOKENS,
NEWRELIC_METRIC_REQUEST_DURATION_MS,
NEWRELIC_METRIC_REQUESTS,
NEWRELIC_METRIC_TOTAL_TOKENS,
NewRelicMetricRecord,
)
def _record(
team_id="team-a",
team_alias=None,
model="gpt-4o",
model_group=None,
status="success",
response_cost=0.5,
prompt_tokens=10,
completion_tokens=20,
total_tokens=30,
duration_ms=100.0,
) -> NewRelicMetricRecord:
return NewRelicMetricRecord(
team_id=team_id,
team_alias=team_alias if team_alias is not None else f"{team_id}-alias",
model_group=model_group if model_group is not None else f"{model}-group",
model=model,
custom_llm_provider="openai",
status=status,
response_cost=response_cost,
prompt_tokens=prompt_tokens,
completion_tokens=completion_tokens,
total_tokens=total_tokens,
duration_ms=duration_ms,
)
def _standard_logging_object(team_id="team-a", response_cost=0.25) -> dict:
return {
"metadata": {"user_api_key_team_id": team_id, "user_api_key_team_alias": f"{team_id}-alias"},
"model_group": "gpt-4o-group",
"model": "gpt-4o",
"custom_llm_provider": "openai",
"status": "success",
"response_cost": response_cost,
"prompt_tokens": 10,
"completion_tokens": 20,
"total_tokens": 30,
"response_time": 0.1,
}
def _make_logger(**kwargs) -> NewRelicMetricsLogger:
with patch("asyncio.create_task"):
return NewRelicMetricsLogger(newrelic_api_key="test-key", **kwargs)
def _response(status_code: int, text: str = "") -> Response:
return Response(status_code, request=Request("POST", "https://example.com"), text=text)
def _raises(status_code: int):
"""Mock the way AsyncHTTPHandler.post really behaves: raise_for_status() turns
every non-2xx into an HTTPStatusError rather than returning the response."""
resp = _response(status_code)
return AsyncMock(side_effect=HTTPStatusError("err", request=resp.request, response=resp))
def _metrics_by_name(payload, name):
return [m for m in payload[0]["metrics"] if m["name"] == name]
class TestBuildMetricPayload:
def test_interval_and_timestamp_reflect_flush_window(self):
payload = build_metric_payload((_record(),), window_start=1_000.0, now=1_007.5)
assert payload[0]["common"]["timestamp"] == 1_000_000
assert payload[0]["common"]["interval.ms"] == 7_500
def test_interval_is_at_least_one_ms(self):
payload = build_metric_payload((_record(),), window_start=1_000.0, now=1_000.0)
assert payload[0]["common"]["interval.ms"] == 1
def test_single_record_metric_values(self):
payload = build_metric_payload(
(_record(response_cost=0.5, prompt_tokens=10, completion_tokens=20, total_tokens=30, duration_ms=100.0),),
window_start=1_000.0,
now=1_005.0,
)
by_name = {m["name"]: m for m in payload[0]["metrics"]}
assert by_name[NEWRELIC_METRIC_REQUESTS]["value"] == 1.0
assert by_name[NEWRELIC_METRIC_REQUESTS]["type"] == "count"
assert by_name[NEWRELIC_METRIC_COST_USD]["value"] == 0.5
assert by_name[NEWRELIC_METRIC_PROMPT_TOKENS]["value"] == 10.0
assert by_name[NEWRELIC_METRIC_COMPLETION_TOKENS]["value"] == 20.0
assert by_name[NEWRELIC_METRIC_TOTAL_TOKENS]["value"] == 30.0
duration = by_name[NEWRELIC_METRIC_REQUEST_DURATION_MS]
assert duration["type"] == "summary"
assert duration["value"] == {"count": 1, "sum": 100.0, "min": 100.0, "max": 100.0}
assert by_name[NEWRELIC_METRIC_REQUESTS]["attributes"] == {
"team_id": "team-a",
"team_alias": "team-a-alias",
"model_group": "gpt-4o-group",
"model": "gpt-4o",
"custom_llm_provider": "openai",
"status": "success",
}
def test_aggregates_across_dimension_buckets(self):
"""Two teams x two models in one queue land in the right bucket sums.
team_alias and model_group are held constant so bucketing provably keys on
team_id and model themselves, not on correlated fields.
"""
shared = {"team_alias": "shared-alias", "model_group": "shared-group"}
records = (
_record(team_id="team-a", model="gpt-4o", response_cost=0.1, total_tokens=10, duration_ms=50.0, **shared),
_record(team_id="team-a", model="gpt-4o", response_cost=0.2, total_tokens=20, duration_ms=150.0, **shared),
_record(
team_id="team-a", model="claude-4", response_cost=0.4, total_tokens=40, duration_ms=200.0, **shared
),
_record(team_id="team-b", model="gpt-4o", response_cost=0.8, total_tokens=80, duration_ms=300.0, **shared),
)
payload = build_metric_payload(records, window_start=1_000.0, now=1_005.0)
cost_by_bucket = {
(m["attributes"]["team_id"], m["attributes"]["model"]): m["value"]
for m in _metrics_by_name(payload, NEWRELIC_METRIC_COST_USD)
}
assert cost_by_bucket == {
("team-a", "gpt-4o"): pytest.approx(0.3),
("team-a", "claude-4"): pytest.approx(0.4),
("team-b", "gpt-4o"): pytest.approx(0.8),
}
requests_by_bucket = {
(m["attributes"]["team_id"], m["attributes"]["model"]): m["value"]
for m in _metrics_by_name(payload, NEWRELIC_METRIC_REQUESTS)
}
assert requests_by_bucket == {
("team-a", "gpt-4o"): 2.0,
("team-a", "claude-4"): 1.0,
("team-b", "gpt-4o"): 1.0,
}
duration_by_bucket = {
(m["attributes"]["team_id"], m["attributes"]["model"]): m["value"]
for m in _metrics_by_name(payload, NEWRELIC_METRIC_REQUEST_DURATION_MS)
}
assert duration_by_bucket[("team-a", "gpt-4o")] == {"count": 2, "sum": 200.0, "min": 50.0, "max": 150.0}
def test_status_is_a_bucket_dimension(self):
records = (
_record(status="success", response_cost=0.1),
_record(status="failure", response_cost=0.0),
)
payload = build_metric_payload(records, window_start=1_000.0, now=1_005.0)
statuses = {m["attributes"]["status"] for m in _metrics_by_name(payload, NEWRELIC_METRIC_REQUESTS)}
assert statuses == {"success", "failure"}
def test_empty_attribute_values_are_omitted(self):
record = NewRelicMetricRecord(
team_id="",
team_alias="",
model_group="",
model="gpt-4o",
custom_llm_provider="openai",
status="success",
response_cost=0.0,
prompt_tokens=0,
completion_tokens=0,
total_tokens=0,
duration_ms=0.0,
)
payload = build_metric_payload((record,), window_start=1_000.0, now=1_005.0)
attributes = payload[0]["metrics"][0]["attributes"]
assert "team_id" not in attributes
assert "team_alias" not in attributes
assert "model_group" not in attributes
class TestQueueAndFlush:
@pytest.mark.asyncio
async def test_log_event_queues_record_from_standard_logging_object(self):
logger = _make_logger()
await logger.async_log_success_event(
kwargs={"standard_logging_object": _standard_logging_object()},
response_obj={},
start_time=None,
end_time=None,
)
assert len(logger.log_queue) == 1
record = logger.log_queue[0]
assert record.team_id == "team-a"
assert record.response_cost == 0.25
assert record.duration_ms == pytest.approx(100.0)
@pytest.mark.asyncio
async def test_failure_event_queues_record(self):
logger = _make_logger()
slo = _standard_logging_object()
slo["status"] = "failure"
await logger.async_log_failure_event(
kwargs={"standard_logging_object": slo},
response_obj={},
start_time=None,
end_time=None,
)
assert len(logger.log_queue) == 1
assert logger.log_queue[0].status == "failure"
@pytest.mark.asyncio
async def test_threshold_flush_uses_flush_queue(self):
logger = _make_logger()
logger.batch_size = 1
logger.flush_queue = AsyncMock()
await logger.async_log_success_event(
kwargs={"standard_logging_object": _standard_logging_object()},
response_obj={},
start_time=None,
end_time=None,
)
logger.flush_queue.assert_awaited_once()
@pytest.mark.asyncio
async def test_flush_queue_updates_last_flush_time_on_success(self):
logger = _make_logger()
logger.log_queue = [_record()]
logger.last_flush_time = 0
logger.async_client.post = AsyncMock(return_value=_response(202))
await logger.flush_queue()
assert logger.log_queue == []
assert logger.last_flush_time > 0
@pytest.mark.asyncio
async def test_flush_advances_window_even_on_requeue(self):
# The window start advances every flush cycle so requeued records report
# in the next window instead of freezing interval.ms under sustained
# failure, and an idle gap never inflates the next batch's window
logger = _make_logger()
logger.log_queue = [_record()]
logger.last_flush_time = 123.0
logger.async_client.post = _raises(500)
await logger.flush_queue()
assert logger.last_flush_time > 123.0
assert len(logger.log_queue) == 1
@pytest.mark.asyncio
async def test_sent_payload_window_starts_at_last_flush_time(self):
logger = _make_logger()
logger.log_queue = [_record()]
logger.last_flush_time = 2_000.0
logger.async_client.post = AsyncMock(return_value=_response(202))
with patch("litellm.integrations.newrelic.newrelic_metrics.time.time", return_value=2_010.0):
await logger.async_send_batch()
sent = logger.async_client.post.await_args.kwargs
body = json.loads(gzip.decompress(sent["data"]).decode("utf-8"))
assert body[0]["common"]["timestamp"] == 2_000_000
assert body[0]["common"]["interval.ms"] == 10_000
assert sent["headers"]["Api-Key"] == "test-key"
assert sent["headers"]["Content-Encoding"] == "gzip"
assert sent["url"] == NEWRELIC_METRIC_ENDPOINT_BY_REGION["us"]
class TestBatchSizeCap:
@pytest.mark.asyncio
async def test_flush_sends_at_most_batch_size_records_per_request(self):
"""A queue grown past the batch size by requeues must go out in chunks:
one oversized request would breach the Metric API data point cap and get
the whole retry backlog dropped as a 4xx."""
logger = _make_logger()
logger.batch_size = 2
logger.log_queue = [_record(model=f"model-{i}") for i in range(5)]
logger.async_client.post = AsyncMock(return_value=_response(202))
await logger.flush_queue()
sent_counts = [
sum(
metric["value"]
for metric in json.loads(gzip.decompress(call.kwargs["data"]).decode("utf-8"))[0]["metrics"]
if metric["name"] == NEWRELIC_METRIC_REQUESTS
)
for call in logger.async_client.post.await_args_list
]
assert sent_counts == [2.0, 2.0, 1.0]
assert logger.log_queue == []
@pytest.mark.asyncio
async def test_failed_chunk_stops_the_flush_and_keeps_order(self):
"""A 5xx on the first chunk ends the flush instead of hammering the same
failing endpoint with the rest of the backlog, and the requeue keeps the
records in chronological order."""
logger = _make_logger()
logger.batch_size = 2
records = [_record(model=f"model-{i}") for i in range(5)]
logger.log_queue = list(records)
logger.async_client.post = _raises(500)
await logger.flush_queue()
assert logger.async_client.post.await_count == 1
assert logger.log_queue == records
class TestFlushConcurrency:
@pytest.mark.asyncio
async def test_records_appended_during_flush_await_survive(self):
"""A record appended by a concurrent request while the POST is in flight
must survive the flush, not be clobbered by a queue replacement."""
logger = _make_logger()
logger.log_queue = [_record(team_id="team-a")]
interleaved = _record(team_id="team-interleaved")
async def _post_appending_mid_flight(**kwargs):
logger.log_queue.append(interleaved)
return _response(202)
logger.async_client.post = AsyncMock(side_effect=_post_appending_mid_flight)
await logger.async_send_batch()
assert logger.log_queue == [interleaved]
body = json.loads(gzip.decompress(logger.async_client.post.await_args.kwargs["data"]).decode("utf-8"))
team_ids = {m["attributes"]["team_id"] for m in body[0]["metrics"]}
assert team_ids == {"team-a"}
@pytest.mark.asyncio
async def test_records_appended_during_failed_flush_await_survive_requeue(self):
"""The requeue path must also preserve interleaved records: batch is
prepended in place, never assigned over the live queue."""
logger = _make_logger()
original = _record(team_id="team-a")
logger.log_queue = [original]
interleaved = _record(team_id="team-interleaved")
async def _post_appending_mid_flight(**kwargs):
logger.log_queue.append(interleaved)
raise HTTPStatusError('e', request=_response(500).request, response=_response(500))
logger.async_client.post = AsyncMock(side_effect=_post_appending_mid_flight)
await logger.async_send_batch()
assert logger.log_queue == [original, interleaved]
class TestErrorPolicy:
@pytest.mark.asyncio
async def test_4xx_drops_batch(self):
logger = _make_logger()
logger.log_queue = [_record(), _record(team_id="team-b")]
logger.async_client.post = AsyncMock(return_value=_response(400, text="bad request"))
await logger.async_send_batch()
assert logger.log_queue == []
assert logger.async_client.post.await_count == 1
@pytest.mark.asyncio
async def test_403_drops_batch_and_names_permanent_credential_failure(self):
logger = _make_logger()
logger.log_queue = [_record()]
logger.async_client.post = _raises(403)
with patch("litellm.integrations.newrelic.newrelic_metrics.verbose_logger") as mock_logger:
await logger.async_send_batch()
assert logger.log_queue == []
warning_text = " ".join(str(arg) for call in mock_logger.warning.call_args_list for arg in call.args)
assert "permanent credential failure" in warning_text
@pytest.mark.asyncio
async def test_5xx_requeues_batch(self):
records = [_record(), _record(team_id="team-b")]
logger = _make_logger()
logger.log_queue = list(records)
logger.async_client.post = _raises(500)
await logger.async_send_batch()
assert logger.log_queue == records
@pytest.mark.asyncio
async def test_network_error_requeues_batch(self):
records = [_record()]
logger = _make_logger()
logger.log_queue = list(records)
logger.async_client.post = AsyncMock(side_effect=ConnectionError("boom"))
await logger.async_send_batch()
assert logger.log_queue == records
@pytest.mark.asyncio
async def test_requeue_is_capped_dropping_oldest(self):
logger = _make_logger()
logger.max_queue_size = 3
oldest = _record(team_id="oldest")
rest = [_record(team_id=f"team-{i}") for i in range(3)]
logger.log_queue = [oldest, *rest]
logger.async_client.post = _raises(500)
await logger.async_send_batch()
assert logger.log_queue == rest
@pytest.mark.asyncio
async def test_requeued_records_are_resent_with_new_records(self):
logger = _make_logger()
logger.log_queue = [_record()]
logger.async_client.post = _raises(500)
await logger.async_send_batch()
logger.log_queue.append(_record(team_id="team-b"))
logger.async_client.post = AsyncMock(return_value=_response(202))
await logger.async_send_batch()
sent = logger.async_client.post.await_args.kwargs
body = json.loads(gzip.decompress(sent["data"]).decode("utf-8"))
team_ids = {m["attributes"]["team_id"] for m in body[0]["metrics"]}
assert team_ids == {"team-a", "team-b"}
assert logger.log_queue == []
class TestStopFlag:
@pytest.mark.asyncio
async def test_stop_ends_periodic_flush_loop(self):
logger = _make_logger()
logger.flush_interval = 0.01
logger.flush_queue = AsyncMock()
task = asyncio.create_task(logger.periodic_flush())
await asyncio.sleep(0.05)
assert not task.done()
logger.stop()
await asyncio.wait_for(task, timeout=1.0)
assert task.done()
@pytest.mark.asyncio
async def test_stopped_logger_exits_after_one_final_drain(self):
logger = _make_logger()
logger.flush_interval = 0.01
logger._final_drain = AsyncMock()
logger._stopped = True
await asyncio.wait_for(logger.periodic_flush(), timeout=1.0)
logger._final_drain.assert_awaited_once()
@pytest.mark.asyncio
async def test_eviction_drains_queued_records(self):
"""Eviction must post what is already queued, not silently discard it."""
from litellm.litellm_core_utils.specialty_caches.dynamic_logging_cache import (
DynamicLoggingCache,
)
cache = DynamicLoggingCache()
logger = _make_logger()
logger.log_queue = [_record(), _record(team_id="team-b")]
logger.async_client.post = AsyncMock(return_value=_response(202))
credentials = {"newrelic_api_key": "test-key", "newrelic_region": None}
cache.set_cache(credentials=credentials, service_name="newrelic", logging_obj=logger)
key = cache.get_cache_key(args={**credentials, "service_name": "newrelic"})
cache.cache._remove_key(key)
for _ in range(10):
await asyncio.sleep(0)
logger.async_client.post.assert_awaited_once()
body = json.loads(gzip.decompress(logger.async_client.post.await_args.kwargs["data"]).decode("utf-8"))
team_ids = {m["attributes"]["team_id"] for m in body[0]["metrics"]}
assert team_ids == {"team-a", "team-b"}
assert logger.log_queue == []
@pytest.mark.asyncio
async def test_dynamic_logging_cache_eviction_calls_stop(self):
from litellm.litellm_core_utils.specialty_caches.dynamic_logging_cache import (
DynamicLoggingCache,
)
cache = DynamicLoggingCache()
logger = _make_logger()
credentials = {"newrelic_api_key": "test-key", "newrelic_region": None}
cache.set_cache(credentials=credentials, service_name="newrelic", logging_obj=logger)
key = cache.get_cache_key(args={**credentials, "service_name": "newrelic"})
cache.cache._remove_key(key)
assert logger._stopped is True
assert cache.get_cache(credentials=credentials, service_name="newrelic") is None
@pytest.mark.asyncio
async def test_append_after_eviction_drain_self_flushes():
"""An in-flight callback holding an evicted (stopped) logger still delivers
its record: with no periodic loop left, the append itself drains."""
logger = _make_logger()
with patch.object(
logger.async_client, "post", new=AsyncMock(return_value=_response(202))
) as mock_post:
logger.stop()
await logger.async_log_success_event(
{"standard_logging_object": _standard_logging_object()}, None, None, None
)
assert mock_post.await_count >= 1, "record appended after stop() must be flushed, not stranded"
assert logger.log_queue == []
@pytest.mark.asyncio
async def test_final_drain_retries_transient_failure_then_delivers():
"""A transient 5xx during the eviction drain must not strand the last
batch: the final drain retries on its own (no periodic loop is left)."""
logger = _make_logger()
err = _response(500)
responses = [HTTPStatusError('e', request=err.request, response=err), HTTPStatusError('e', request=err.request, response=err), _response(202)]
post_mock = AsyncMock(side_effect=responses)
with patch.object(logger, "async_client") as client, patch("asyncio.sleep", new=AsyncMock()):
client.post = post_mock
await logger._log_async_event(standard_logging_object=_standard_logging_object())
await logger._final_drain()
assert post_mock.await_count == 3
assert logger.log_queue == []
@pytest.mark.asyncio
async def test_final_drain_drops_after_bounded_passes_under_lock():
"""A permanently failing destination is retried across bounded passes, then
the remainder is dropped under flush_lock and logged, never stranded. A
second drain over the now-empty queue is a no-op."""
logger = _make_logger()
post_mock = _raises(500)
with patch.object(logger, "async_client") as client, patch("asyncio.sleep", new=AsyncMock()):
client.post = post_mock
await logger._log_async_event(standard_logging_object=_standard_logging_object())
await logger._final_drain()
after_first = post_mock.await_count
await logger._final_drain()
assert after_first >= 1, "the failing destination was retried before the drop"
assert post_mock.await_count == after_first, "second drain over an empty queue is a no-op"
assert logger.log_queue == [], "exhausted retries end in a logged drop, not a stranded queue"
def test_attribute_values_bounded_against_payload_bombs():
"""A caller-controlled high-entropy model string is truncated in metric
attributes so one record cannot inflate the shared batch past the Metric
API payload cap and take out other users' metrics."""
record = _record(model="m" * 5000)
metrics = _bucket_metrics((record,))
for metric in metrics:
assert len(metric["attributes"]["model"]) == 255
@pytest.mark.asyncio
async def test_idle_gap_does_not_inflate_next_window():
"""Empty flush cycles advance the window start, so a burst after idling
reports an interval close to the flush cadence, not the whole idle gap."""
logger = _make_logger()
logger.last_flush_time = 100.0
with patch.object(logger, "async_client") as client:
client.post = AsyncMock(return_value=_response(202))
await logger.flush_queue()
assert logger.last_flush_time > 100.0
@pytest.mark.asyncio
async def test_mid_drain_append_delivered_against_healthy_destination():
"""A record a callback appends while a drain is running is picked up by a
later pass and delivered when the destination is healthy; nothing stranded."""
logger = _make_logger()
logger.stop()
late_record = _record(model="late-model")
injected = {"done": False}
posted = []
async def _capture(url, headers=None, content=None, **kw):
posted.append(content)
if not injected["done"]:
injected["done"] = True
logger.log_queue.append(late_record)
return _response(202)
with patch.object(logger, "async_client") as client, patch("asyncio.sleep", new=AsyncMock()):
client.post = _capture
logger.log_queue.append(_record(model="first"))
await logger._drain_with_retry()
assert logger.log_queue == [], "the mid-drain append was drained too, nothing stranded"
assert len(posted) >= 2, "both the original and the mid-drain record were sent"
@pytest.mark.asyncio
async def test_drain_attempts_every_chunk_not_just_the_head_under_failure():
"""Regression: with more than batch_size records queued on a stopped logger
and a persistently failing destination, every record must be attempted before
the bounded terminal drop. The periodic path stops at the first failing chunk,
so a drain that reused it would drop the un-sent tail (records past the head
chunk) as if it had tried them, silently undercounting the team's usage."""
logger = _make_logger()
logger.stop()
logger.batch_size = 2
logger.log_queue = [_record(model=f"m{i}") for i in range(5)]
sent_models = []
async def _capture_then_fail(url, data=None, headers=None, **kw):
body = json.loads(gzip.decompress(data).decode("utf-8"))
sent_models.extend(
m["attributes"]["model"] for m in body[0]["metrics"] if m["name"] == NEWRELIC_METRIC_REQUESTS
)
resp = _response(503)
raise HTTPStatusError("err", request=resp.request, response=resp)
with patch("asyncio.sleep", new=AsyncMock()):
logger.async_client.post = _capture_then_fail
await logger._drain_with_retry()
assert set(sent_models) == {"m0", "m1", "m2", "m3", "m4"}, "every chunk, including the tail, was attempted"
assert logger.log_queue == [], "the exhausted batch is dropped after bounded passes, nothing stranded"
@pytest.mark.asyncio
async def test_drain_delivers_the_tail_once_the_destination_recovers():
"""The tail beyond the head chunk must be delivered, not stranded, once a
transiently failing destination recovers within the drain's passes."""
logger = _make_logger()
logger.stop()
logger.batch_size = 2
logger.log_queue = [_record(model=f"m{i}") for i in range(5)]
delivered_models = []
posts = {"n": 0}
async def _fail_first_pass_then_recover(url, data=None, headers=None, **kw):
posts["n"] += 1
if posts["n"] <= 3: # the first pass's three chunks all fail
resp = _response(503)
raise HTTPStatusError("err", request=resp.request, response=resp)
body = json.loads(gzip.decompress(data).decode("utf-8"))
delivered_models.extend(
m["attributes"]["model"] for m in body[0]["metrics"] if m["name"] == NEWRELIC_METRIC_REQUESTS
)
return _response(202)
with patch("asyncio.sleep", new=AsyncMock()):
logger.async_client.post = _fail_first_pass_then_recover
await logger._drain_with_retry()
assert set(delivered_models) == {"m0", "m1", "m2", "m3", "m4"}, "all chunks delivered after recovery"
assert logger.log_queue == [], "nothing left stranded once the destination recovered"
@pytest.mark.asyncio
async def test_terminal_drop_leaves_untried_late_arrival_for_next_drain():
"""Against a permanently failing destination, the terminal drop clears only
the records this drain actually tried; a record a callback appends during the
final pass, after that pass's snapshot, is left in the queue for its own
serialized drain, never wiped un-tried."""
logger = _make_logger()
logger.stop()
from litellm.types.integrations.newrelic import NEWRELIC_METRICS_MAX_DRAIN_PASSES
late_record = _record(model="late-arrival")
posts = {"n": 0}
async def _fail_and_append_on_final_pass(url, data=None, headers=None, **kw):
posts["n"] += 1
# One record means one post per pass, so the final pass's post is the
# Nth; append then, after the drain has already snapshotted the queue.
if posts["n"] == NEWRELIC_METRICS_MAX_DRAIN_PASSES:
logger.log_queue.append(late_record)
resp = _response(503)
raise HTTPStatusError("err", request=resp.request, response=resp)
with patch("asyncio.sleep", new=AsyncMock()):
logger.async_client.post = _fail_and_append_on_final_pass
logger.log_queue.append(_record(model="doomed"))
await logger._drain_with_retry()
assert logger.log_queue == [late_record], "the un-tried late arrival is left for its own drain, not dropped"
@pytest.mark.asyncio
async def test_record_appended_on_an_early_pass_is_not_dropped_short_of_the_retry_budget():
"""A record a callback appends during an early drain pass entered the queue
after this drain's snapshot, so it has not seen the full retry budget. The
terminal drop must clear only records queued when the drain began, leaving
the early-pass arrival for its own serialized drain instead of dropping it
after fewer than the configured attempts."""
logger = _make_logger()
logger.stop()
early_record = _record(model="early-pass-arrival")
posts = {"n": 0}
async def _fail_and_append_on_first_pass(url, data=None, headers=None, **kw):
posts["n"] += 1
# One record queued at start means the first pass's post is the 1st;
# append during it, before this drain's later passes.
if posts["n"] == 1:
logger.log_queue.append(early_record)
resp = _response(503)
raise HTTPStatusError("err", request=resp.request, response=resp)
with patch("asyncio.sleep", new=AsyncMock()):
logger.async_client.post = _fail_and_append_on_first_pass
logger.log_queue.append(_record(model="doomed"))
await logger._drain_with_retry()
assert logger.log_queue == [early_record], "the early-pass arrival is left for its own drain, not dropped short"
@pytest.mark.asyncio
async def test_post_stop_drains_are_serialized():
"""A callback that appends to a stopped logger and starts its own drain must
queue behind an already-running drain, not race it: otherwise one drain's
terminal clear could wipe a record the other is still responsible for.
Proven by holding the first drain inside its flush and asserting the second
has not entered its own flush until the first releases."""
logger = _make_logger()
logger._stopped = True # stopped without scheduling a background drain
logger.log_queue.append(_record(model="r1"))
entered = []
release = asyncio.Event()
async def blocking_flush():
entered.append(len(entered) + 1)
if len(entered) == 1:
await release.wait()
logger.log_queue.clear()
logger._drain_flush_once = blocking_flush
t1 = asyncio.create_task(logger._drain_with_retry())
await asyncio.sleep(0.02) # let t1 acquire the drain lock and enter flush
assert entered == [1], f"first drain did not enter flush: {entered}"
t2 = asyncio.create_task(logger._drain_with_retry())
await asyncio.sleep(0.02) # t2 must block on the drain lock, not enter flush
assert entered == [1], f"second drain raced the first: {entered}"
release.set()
await asyncio.gather(t1, t2)
assert logger.log_queue == []
@pytest.mark.asyncio
async def test_raised_403_is_dropped_not_requeued():
"""AsyncHTTPHandler.post raises HTTPStatusError on 4xx, so a 403 (permanent
bad key) arrives as an exception, not a response. It must be dropped, never
requeued, or a revoked key retries forever."""
logger = _make_logger()
logger.log_queue.append(_record())
logger.async_client.post = _raises(403)
await logger.async_send_batch()
assert logger.log_queue == [], "a permanent 403 must drop, not requeue"
@pytest.mark.asyncio
async def test_raised_500_is_requeued():
"""A raised 5xx is transient and must be requeued for retry."""
logger = _make_logger()
record = _record()
logger.log_queue.append(record)
logger.async_client.post = _raises(503)
await logger.async_send_batch()
assert logger.log_queue == [record], "a transient 5xx must requeue"
@pytest.mark.asyncio
@pytest.mark.parametrize("status", [429, 408])
async def test_transient_4xx_is_requeued_not_dropped(status):
"""The Metric API returns 429 when it throttles (and 408 on a request
timeout); both are transient and expect a retry, so the batch must be
requeued rather than permanently dropped like a 400/403."""
logger = _make_logger()
record = _record()
logger.log_queue.append(record)
logger.async_client.post = _raises(status)
await logger.async_send_batch()
assert logger.log_queue == [record], f"a transient {status} must requeue, not drop"
@pytest.mark.asyncio
@pytest.mark.parametrize("status", [200, 201, 204])
async def test_any_2xx_is_treated_as_delivered_not_requeued(status):
"""The Metric API answers 202, but any 2xx means the destination accepted the
batch. Treating a non-202 2xx as a failure would re-queue and re-send data
New Relic already stored, duplicating the team's metrics until the cap drops."""
logger = _make_logger()
logger.log_queue.append(_record())
logger.async_client.post = AsyncMock(return_value=_response(status))
await logger.async_send_batch()
assert logger.log_queue == [], f"a {status} success must drop, not requeue and duplicate"

274
tests/test_litellm/integrations/newrelic/test_newrelic_team_handler.py Normal file
View file

@ -0,0 +1,274 @@
"""
Tests for team-scoped New Relic metrics callback support.
Verifies that NewRelicMetricsLogger is instantiated with per-team credentials
(newrelic_api_key, newrelic_region) with no environment fallback, and that
NewRelicHandler correctly resolves and caches per-team loggers.
"""
import copy
from unittest.mock import patch
import pytest
from litellm.integrations.newrelic.newrelic_metrics import NewRelicMetricsLogger
from litellm.integrations.newrelic.newrelic_team_handler import NewRelicHandler
from litellm.litellm_core_utils.initialize_dynamic_callback_params import (
TRUSTED_CALLBACK_VARS_FIELD,
)
from litellm.litellm_core_utils.specialty_caches.dynamic_logging_cache import (
DynamicLoggingCache,
)
from litellm.types.integrations.newrelic import NEWRELIC_METRIC_ENDPOINT_BY_REGION
from litellm.types.utils import StandardCallbackDynamicParams
US_ENDPOINT = NEWRELIC_METRIC_ENDPOINT_BY_REGION["us"]
EU_ENDPOINT = NEWRELIC_METRIC_ENDPOINT_BY_REGION["eu"]
class TestNewRelicMetricsLoggerCredentialKwargs:
"""The logger takes credentials by injection only; env vars never leak in."""
def test_init_with_explicit_credentials(self):
with patch("asyncio.create_task"):
logger = NewRelicMetricsLogger(newrelic_api_key="team_key", newrelic_region="eu")
assert logger.newrelic_api_key == "team_key"
assert logger.metric_api_url == EU_ENDPOINT
def test_init_defaults_to_us_region(self):
with patch("asyncio.create_task"):
logger = NewRelicMetricsLogger(newrelic_api_key="team_key")
assert logger.metric_api_url == US_ENDPOINT
def test_unknown_region_falls_back_to_us(self):
with patch("asyncio.create_task"):
logger = NewRelicMetricsLogger(newrelic_api_key="team_key", newrelic_region="mars")
assert logger.metric_api_url == US_ENDPOINT
def test_region_is_case_insensitive(self):
with patch("asyncio.create_task"):
logger = NewRelicMetricsLogger(newrelic_api_key="team_key", newrelic_region="EU")
assert logger.metric_api_url == EU_ENDPOINT
def test_init_raises_without_api_key(self):
with pytest.raises(ValueError, match="newrelic_api_key"):
with patch("asyncio.create_task"):
NewRelicMetricsLogger(newrelic_api_key="")
def test_init_never_falls_back_to_env_license_key(self, monkeypatch):
"""A missing team key must fail, never silently reuse the operator's key."""
monkeypatch.setenv("NEW_RELIC_LICENSE_KEY", "operator-license-key")
with pytest.raises(ValueError, match="newrelic_api_key"):
with patch("asyncio.create_task"):
NewRelicMetricsLogger(newrelic_api_key="")
class TestNewRelicHandler:
"""The handler resolves the correct logger per team."""
def test_creates_team_logger_with_dynamic_credentials(self):
cache = DynamicLoggingCache()
params = StandardCallbackDynamicParams(newrelic_api_key="team_a_key", newrelic_region="eu")
with patch("asyncio.create_task"):
result = NewRelicHandler.get_newrelic_logger_for_request(
standard_callback_dynamic_params=params,
in_memory_dynamic_logger_cache=cache,
)
assert result.newrelic_api_key == "team_a_key"
assert result.metric_api_url == EU_ENDPOINT
def test_caches_team_logger(self):
cache = DynamicLoggingCache()
params = StandardCallbackDynamicParams(newrelic_api_key="team_b_key")
with patch("asyncio.create_task"):
result1 = NewRelicHandler.get_newrelic_logger_for_request(
standard_callback_dynamic_params=params,
in_memory_dynamic_logger_cache=cache,
)
result2 = NewRelicHandler.get_newrelic_logger_for_request(
standard_callback_dynamic_params=params,
in_memory_dynamic_logger_cache=cache,
)
assert result1 is result2
def test_different_teams_get_different_loggers(self):
cache = DynamicLoggingCache()
params_a = StandardCallbackDynamicParams(newrelic_api_key="team_a_key")
params_b = StandardCallbackDynamicParams(newrelic_api_key="team_b_key")
with patch("asyncio.create_task"):
result_a = NewRelicHandler.get_newrelic_logger_for_request(
standard_callback_dynamic_params=params_a,
in_memory_dynamic_logger_cache=cache,
)
result_b = NewRelicHandler.get_newrelic_logger_for_request(
standard_callback_dynamic_params=params_b,
in_memory_dynamic_logger_cache=cache,
)
assert result_a is not result_b
assert result_a.newrelic_api_key == "team_a_key"
assert result_b.newrelic_api_key == "team_b_key"
def test_region_is_part_of_cache_key(self):
"""Same key, different region must not share a logger (different endpoints)."""
cache = DynamicLoggingCache()
with patch("asyncio.create_task"):
result_us = NewRelicHandler.get_newrelic_logger_for_request(
standard_callback_dynamic_params=StandardCallbackDynamicParams(newrelic_api_key="key"),
in_memory_dynamic_logger_cache=cache,
)
result_eu = NewRelicHandler.get_newrelic_logger_for_request(
standard_callback_dynamic_params=StandardCallbackDynamicParams(
newrelic_api_key="key", newrelic_region="eu"
),
in_memory_dynamic_logger_cache=cache,
)
assert result_us is not result_eu
assert result_us.metric_api_url == US_ENDPOINT
assert result_eu.metric_api_url == EU_ENDPOINT
def test_request_blocked_callback_params_includes_newrelic(self):
from litellm.litellm_core_utils.initialize_dynamic_callback_params import (
_request_blocked_callback_params,
)
assert "newrelic_api_key" in _request_blocked_callback_params
assert "newrelic_region" in _request_blocked_callback_params
class TestDynamicCredentialDetection:
def test_no_credentials(self):
params = StandardCallbackDynamicParams()
assert NewRelicHandler._dynamic_newrelic_credentials_are_passed(params) is False
def test_region_only_is_not_credentials(self):
params = StandardCallbackDynamicParams(newrelic_region="eu")
assert NewRelicHandler._dynamic_newrelic_credentials_are_passed(params) is False
def test_api_key_is_credentials(self):
params = StandardCallbackDynamicParams(newrelic_api_key="key")
assert NewRelicHandler._dynamic_newrelic_credentials_are_passed(params) is True
class TestStandardCallbackDynamicParamsIncludesNewRelic:
def test_newrelic_params_in_annotations(self):
annotations = StandardCallbackDynamicParams.__annotations__
assert "newrelic_api_key" in annotations
assert "newrelic_region" in annotations
def _build_logging_obj(kwargs: dict, *, with_newrelic_callback: bool = True):
from litellm.litellm_core_utils.litellm_logging import Logging
with patch("asyncio.create_task"):
return Logging(
model="gpt-4",
messages=[{"role": "user", "content": "hi"}],
stream=False,
call_type="completion",
start_time="2026-01-01",
litellm_call_id="test-call-id",
function_id="test-func",
dynamic_success_callbacks=["newrelic"] if with_newrelic_callback else None,
kwargs=kwargs,
)
def _metrics_loggers(logging_obj) -> list[NewRelicMetricsLogger]:
return [cb for cb in (logging_obj.dynamic_success_callbacks or []) if isinstance(cb, NewRelicMetricsLogger)]
class TestTeamCallbackFlowPassesNewRelicCredentials:
"""
newrelic_* credentials reach NewRelicHandler only from the proxy-stamped trusted
field. Anything the caller put in the request body must not, or a caller could
pair its own newrelic_region with the team's ingest key.
"""
def test_trusted_callback_vars_reach_newrelic_handler(self):
logging_obj = _build_logging_obj(
{
TRUSTED_CALLBACK_VARS_FIELD: {"newrelic_api_key": "team-nr-key-123", "newrelic_region": "eu"},
"model": "gpt-4",
"litellm_params": {"metadata": {}},
}
)
metrics_loggers = _metrics_loggers(logging_obj)
assert len(metrics_loggers) == 1, "NewRelicMetricsLogger should be initialized from team callback_vars"
assert metrics_loggers[0].newrelic_api_key == "team-nr-key-123"
assert metrics_loggers[0].metric_api_url == EU_ENDPOINT
def test_trace_logger_still_dispatched_alongside_metrics(self):
"""The metrics logger must not displace the trace logger for the same name."""
logging_obj = _build_logging_obj(
{
TRUSTED_CALLBACK_VARS_FIELD: {"newrelic_api_key": "team-nr-key-123"},
"model": "gpt-4",
"litellm_params": {"metadata": {}},
}
)
non_metrics = [
cb for cb in (logging_obj.dynamic_success_callbacks or []) if not isinstance(cb, NewRelicMetricsLogger)
]
assert len(non_metrics) == 1, "trace logger (OTel v2 or legacy agent) must remain in the dynamic list"
assert len(_metrics_loggers(logging_obj)) == 1
async_non_metrics = [
cb
for cb in (logging_obj.dynamic_async_success_callbacks or [])
if not isinstance(cb, NewRelicMetricsLogger)
]
assert len(async_non_metrics) == 1
def test_request_kwargs_newrelic_params_are_ignored(self):
logging_obj = _build_logging_obj(
{
"newrelic_api_key": "caller-nr-key",
"newrelic_region": "eu",
"model": "gpt-4",
"litellm_params": {"metadata": {}},
}
)
assert _metrics_loggers(logging_obj) == []
def test_logging_object_stays_deepcopyable(self):
logging_obj = _build_logging_obj(
{
TRUSTED_CALLBACK_VARS_FIELD: {"newrelic_api_key": "team-nr-key-123"},
"model": "gpt-4",
"litellm_params": {"metadata": {}},
},
with_newrelic_callback=False,
)
assert copy.deepcopy(logging_obj)._trusted_callback_vars == logging_obj._trusted_callback_vars
def test_caller_cannot_redirect_team_credentials(self):
"""The exfil shape: caller's newrelic_region paired with the team's key."""
logging_obj = _build_logging_obj(
{
TRUSTED_CALLBACK_VARS_FIELD: {"newrelic_api_key": "team-nr-key-123"},
"newrelic_region": "eu",
"model": "gpt-4",
"litellm_params": {"metadata": {}},
}
)
metrics_loggers = _metrics_loggers(logging_obj)
assert len(metrics_loggers) == 1
assert metrics_loggers[0].newrelic_api_key == "team-nr-key-123"
assert metrics_loggers[0].metric_api_url == US_ENDPOINT