diff --git a/litellm/integrations/newrelic/newrelic_metrics.py b/litellm/integrations/newrelic/newrelic_metrics.py index 25dbfc2bdb2..da952b78d3f 100644 --- a/litellm/integrations/newrelic/newrelic_metrics.py +++ b/litellm/integrations/newrelic/newrelic_metrics.py @@ -5,8 +5,9 @@ NR Reference API: https://docs.newrelic.com/docs/data-apis/ingest-apis/metric-ap `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. +into count/summary metrics, plus one max/remaining budget gauge pair per team +taken from the team's latest record. `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 @@ -47,11 +48,14 @@ from litellm.types.integrations.newrelic import ( NEWRELIC_METRIC_PROMPT_TOKENS, NEWRELIC_METRIC_REQUEST_DURATION_MS, NEWRELIC_METRIC_REQUESTS, + NEWRELIC_METRIC_TEAM_MAX_BUDGET, + NEWRELIC_METRIC_TEAM_REMAINING_BUDGET, NEWRELIC_METRIC_TOTAL_TOKENS, NEWRELIC_METRICS_MAX_BATCH_SIZE, NEWRELIC_METRICS_MAX_DRAIN_PASSES, NEWRELIC_METRICS_MAX_RETRY_QUEUE_SIZE, NewRelicCountMetric, + NewRelicGaugeMetric, NewRelicMetric, NewRelicMetricCommon, NewRelicMetricEnvelope, @@ -98,6 +102,8 @@ def _metric_record_from_payload(standard_logging_object: StandardLoggingPayload) 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, + team_max_budget=metadata.get("user_api_key_team_max_budget") if metadata else None, + team_spend=metadata.get("user_api_key_team_spend") if metadata else None, ) @@ -140,6 +146,33 @@ def _bucket_metrics(bucket_records: tuple[NewRelicMetricRecord, ...]) -> tuple[N return (*count_metrics, summary_metric) +def _team_budget_gauges(record: NewRelicMetricRecord) -> tuple[NewRelicMetric, ...]: + team_max_budget: Final = record.team_max_budget + if team_max_budget is None: + return () + 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", record.team_id), ("team_alias", record.team_alias)) + if value + } + remaining_budget: Final = team_max_budget - (record.team_spend or 0.0) - record.response_cost + return ( + NewRelicGaugeMetric( + name=NEWRELIC_METRIC_TEAM_MAX_BUDGET, type="gauge", value=team_max_budget, attributes=attributes + ), + NewRelicGaugeMetric( + name=NEWRELIC_METRIC_TEAM_REMAINING_BUDGET, type="gauge", value=remaining_budget, attributes=attributes + ), + ) + + +def _team_budget_metrics(records: tuple[NewRelicMetricRecord, ...]) -> tuple[NewRelicMetric, ...]: + latest_by_team: Final[Mapping[str, NewRelicMetricRecord]] = MappingProxyType( + {record.team_id: record for record in records if record.team_id} + ) + return tuple(gauge for record in latest_by_team.values() for gauge in _team_budget_gauges(record)) + + def build_metric_payload( records: tuple[NewRelicMetricRecord, ...], *, @@ -158,7 +191,7 @@ def build_metric_payload( "timestamp": int(window_start * 1000), "interval.ms": interval_ms, } - return (NewRelicMetricEnvelope(common=common, metrics=metrics),) + return (NewRelicMetricEnvelope(common=common, metrics=(*metrics, *_team_budget_metrics(records))),) class NewRelicMetricsLogger(CustomBatchLogger): diff --git a/litellm/types/integrations/newrelic.py b/litellm/types/integrations/newrelic.py index 36e4d02c2a8..b5905ad0b93 100644 --- a/litellm/types/integrations/newrelic.py +++ b/litellm/types/integrations/newrelic.py @@ -27,9 +27,9 @@ NEWRELIC_METRIC_ENDPOINT_BY_REGION: Final[Mapping[str, str]] = MappingProxyType( 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 +#: record expands to at most 6 bucket metrics plus 2 team budget gauges (8), so cap +#: the per-flush record count well below 2000 / 8. +NEWRELIC_METRICS_MAX_BATCH_SIZE: Final = 200 #: Hard cap on records retained across failed flushes (5xx/network requeue). #: Beyond this the oldest records are dropped. @@ -48,6 +48,8 @@ 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" +NEWRELIC_METRIC_TEAM_MAX_BUDGET: Final = "litellm.team.max_budget" +NEWRELIC_METRIC_TEAM_REMAINING_BUDGET: Final = "litellm.team.remaining_budget" class NewRelicSummaryValue(TypedDict): @@ -66,6 +68,13 @@ class NewRelicCountMetric(TypedDict): attributes: ReadOnly[Mapping[str, str]] +class NewRelicGaugeMetric(TypedDict): + name: ReadOnly[str] + type: ReadOnly[Literal["gauge"]] + value: ReadOnly[float] + attributes: ReadOnly[Mapping[str, str]] + + class NewRelicSummaryMetric(TypedDict): name: ReadOnly[str] type: ReadOnly[Literal["summary"]] @@ -73,7 +82,7 @@ class NewRelicSummaryMetric(TypedDict): attributes: ReadOnly[Mapping[str, str]] -NewRelicMetric = NewRelicCountMetric | NewRelicSummaryMetric +NewRelicMetric = NewRelicCountMetric | NewRelicGaugeMetric | NewRelicSummaryMetric #: ``interval.ms`` has a dot in it, so the functional TypedDict form is required. @@ -108,6 +117,8 @@ class NewRelicMetricRecord: completion_tokens: int total_tokens: int duration_ms: float + team_max_budget: float | None = None + team_spend: float | None = None @property def bucket_key(self) -> tuple[str, str, str, str, str, str]: diff --git a/tests/test_litellm/integrations/newrelic/test_newrelic_metrics.py b/tests/test_litellm/integrations/newrelic/test_newrelic_metrics.py index 9c75e0b0a47..a9074e89806 100644 --- a/tests/test_litellm/integrations/newrelic/test_newrelic_metrics.py +++ b/tests/test_litellm/integrations/newrelic/test_newrelic_metrics.py @@ -24,6 +24,8 @@ from litellm.types.integrations.newrelic import ( NEWRELIC_METRIC_PROMPT_TOKENS, NEWRELIC_METRIC_REQUEST_DURATION_MS, NEWRELIC_METRIC_REQUESTS, + NEWRELIC_METRIC_TEAM_MAX_BUDGET, + NEWRELIC_METRIC_TEAM_REMAINING_BUDGET, NEWRELIC_METRIC_TOTAL_TOKENS, NewRelicMetricRecord, ) @@ -40,6 +42,8 @@ def _record( completion_tokens=20, total_tokens=30, duration_ms=100.0, + team_max_budget=None, + team_spend=None, ) -> NewRelicMetricRecord: return NewRelicMetricRecord( team_id=team_id, @@ -53,12 +57,28 @@ def _record( completion_tokens=completion_tokens, total_tokens=total_tokens, duration_ms=duration_ms, + team_max_budget=team_max_budget, + team_spend=team_spend, ) -def _standard_logging_object(team_id="team-a", response_cost=0.25) -> dict: +def _standard_logging_object( + team_id="team-a", response_cost=0.25, team_max_budget: float | None = None, team_spend: float | None = None +) -> dict: + budget_metadata = { + key: value + for key, value in ( + ("user_api_key_team_max_budget", team_max_budget), + ("user_api_key_team_spend", team_spend), + ) + if value is not None + } return { - "metadata": {"user_api_key_team_id": team_id, "user_api_key_team_alias": f"{team_id}-alias"}, + "metadata": { + "user_api_key_team_id": team_id, + "user_api_key_team_alias": f"{team_id}-alias", + **budget_metadata, + }, "model_group": "gpt-4o-group", "model": "gpt-4o", "custom_llm_provider": "openai", @@ -204,6 +224,74 @@ class TestBuildMetricPayload: assert "model_group" not in attributes +class TestTeamBudgetGauges: + def test_latest_record_per_team_drives_one_gauge_pair(self): + records = ( + _record(team_id="team-a", model="gpt-4o", response_cost=0.5, team_max_budget=100.0, team_spend=10.0), + _record(team_id="team-a", model="claude-4", response_cost=2.0, team_max_budget=100.0, team_spend=10.5), + _record(team_id="team-b", response_cost=1.0, team_max_budget=None, team_spend=3.0), + _record(team_id="", team_alias="", response_cost=1.0, team_max_budget=50.0, team_spend=1.0), + ) + payload = build_metric_payload(records, window_start=1_000.0, now=1_005.0) + + max_budget_gauges = _metrics_by_name(payload, NEWRELIC_METRIC_TEAM_MAX_BUDGET) + remaining_gauges = _metrics_by_name(payload, NEWRELIC_METRIC_TEAM_REMAINING_BUDGET) + assert [(m["type"], m["value"], m["attributes"]) for m in max_budget_gauges] == [ + ("gauge", 100.0, {"team_id": "team-a", "team_alias": "team-a-alias"}) + ] + assert [(m["type"], m["attributes"]) for m in remaining_gauges] == [ + ("gauge", {"team_id": "team-a", "team_alias": "team-a-alias"}) + ] + assert remaining_gauges[0]["value"] == pytest.approx(100.0 - 10.5 - 2.0) + assert len(_metrics_by_name(payload, NEWRELIC_METRIC_COST_USD)) == 4 + + def test_missing_team_spend_counts_only_this_request(self): + payload = build_metric_payload( + (_record(response_cost=0.25, team_max_budget=10.0, team_spend=None),), window_start=1_000.0, now=1_005.0 + ) + + assert _metrics_by_name(payload, NEWRELIC_METRIC_TEAM_REMAINING_BUDGET)[0]["value"] == pytest.approx(9.75) + + @pytest.mark.asyncio + async def test_budget_gauges_reach_the_metric_api_from_standard_logging_metadata(self): + logger = _make_logger() + logger.async_client.post = AsyncMock(return_value=_response(202)) + slo = _standard_logging_object(response_cost=0.25, team_max_budget=20.0, team_spend=4.5) + + await logger.async_log_success_event( + kwargs={"standard_logging_object": slo}, response_obj={}, start_time=None, end_time=None + ) + await logger.flush_queue() + + body = json.loads(gzip.decompress(logger.async_client.post.await_args.kwargs["data"]).decode("utf-8")) + by_name = {m["name"]: m for m in body[0]["metrics"]} + assert by_name[NEWRELIC_METRIC_TEAM_MAX_BUDGET] == { + "name": NEWRELIC_METRIC_TEAM_MAX_BUDGET, + "type": "gauge", + "value": 20.0, + "attributes": {"team_id": "team-a", "team_alias": "team-a-alias"}, + } + assert by_name[NEWRELIC_METRIC_TEAM_REMAINING_BUDGET]["type"] == "gauge" + assert by_name[NEWRELIC_METRIC_TEAM_REMAINING_BUDGET]["value"] == pytest.approx(15.25) + assert by_name[NEWRELIC_METRIC_COST_USD]["value"] == 0.25 + + @pytest.mark.asyncio + async def test_no_budget_metadata_sends_no_gauges(self): + logger = _make_logger() + logger.async_client.post = AsyncMock(return_value=_response(202)) + + await logger.async_log_success_event( + kwargs={"standard_logging_object": _standard_logging_object()}, + response_obj={}, + start_time=None, + end_time=None, + ) + await logger.flush_queue() + + body = json.loads(gzip.decompress(logger.async_client.post.await_args.kwargs["data"]).decode("utf-8")) + assert {m["type"] for m in body[0]["metrics"]} == {"count", "summary"} + + class TestQueueAndFlush: @pytest.mark.asyncio async def test_log_event_queues_record_from_standard_logging_object(self):