diff --git a/litellm/integrations/custom_batch_logger.py b/litellm/integrations/custom_batch_logger.py index c9e24913900..bfc78b93715 100644 --- a/litellm/integrations/custom_batch_logger.py +++ b/litellm/integrations/custom_batch_logger.py @@ -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) diff --git a/litellm/integrations/newrelic/newrelic_metrics.py b/litellm/integrations/newrelic/newrelic_metrics.py new file mode 100644 index 00000000000..25dbfc2bdb2 --- /dev/null +++ b/litellm/integrations/newrelic/newrelic_metrics.py @@ -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, + ) diff --git a/litellm/integrations/newrelic/newrelic_team_handler.py b/litellm/integrations/newrelic/newrelic_team_handler.py new file mode 100644 index 00000000000..ae52a6d4efb --- /dev/null +++ b/litellm/integrations/newrelic/newrelic_team_handler.py @@ -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 diff --git a/litellm/litellm_core_utils/litellm_logging.py b/litellm/litellm_core_utils/litellm_logging.py index c0750bb94e7..3018f0c4d24 100644 --- a/litellm/litellm_core_utils/litellm_logging.py +++ b/litellm/litellm_core_utils/litellm_logging.py @@ -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 diff --git a/litellm/litellm_core_utils/specialty_caches/dynamic_logging_cache.py b/litellm/litellm_core_utils/specialty_caches/dynamic_logging_cache.py index f63c60dd430..da3ac366bfd 100644 --- a/litellm/litellm_core_utils/specialty_caches/dynamic_logging_cache.py +++ b/litellm/litellm_core_utils/specialty_caches/dynamic_logging_cache.py @@ -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 ######################################################### diff --git a/litellm/types/integrations/newrelic.py b/litellm/types/integrations/newrelic.py index 96d9a201ad7..36e4d02c2a8 100644 --- a/litellm/types/integrations/newrelic.py +++ b/litellm/types/integrations/newrelic.py @@ -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, + ) diff --git a/tests/test_litellm/integrations/newrelic/test_newrelic_metrics.py b/tests/test_litellm/integrations/newrelic/test_newrelic_metrics.py new file mode 100644 index 00000000000..9c75e0b0a47 --- /dev/null +++ b/tests/test_litellm/integrations/newrelic/test_newrelic_metrics.py @@ -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" diff --git a/tests/test_litellm/integrations/newrelic/test_newrelic_team_handler.py b/tests/test_litellm/integrations/newrelic/test_newrelic_team_handler.py new file mode 100644 index 00000000000..f4460a615df --- /dev/null +++ b/tests/test_litellm/integrations/newrelic/test_newrelic_team_handler.py @@ -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