feat(newrelic): export team max and remaining budget gauges to the Metric API (#40542)

Co-authored-by: yassin <yassin@berri.ai>
Co-authored-by: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>

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

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

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