feat(prometheus_api.py): support querying prometheus metrics for all-up + key-level spend on UI (#5782)

enables getting aggregated view from prometheus api

Makes proxy UI reliable in prod

litellm/integrations/prometheus_helpers/prometheus_api.py

@@ -5,16 +5,19 @@ Helper functions to query prometheus API
 import asyncio
 import os
 import time
+from datetime import datetime, timedelta
+from typing import Optional
 
 import litellm
+from litellm import get_secret
 from litellm._logging import verbose_logger
 from litellm.llms.custom_httpx.http_handler import (
     get_async_httpx_client,
     httpxSpecialProvider,
 )
 
-PROMETHEUS_URL = litellm.get_secret("PROMETHEUS_URL")
+PROMETHEUS_URL: Optional[str] = get_secret("PROMETHEUS_URL")  # type: ignore
-PROMETHEUS_SELECTED_INSTANCE = litellm.get_secret("PROMETHEUS_SELECTED_INSTANCE")
+PROMETHEUS_SELECTED_INSTANCE: Optional[str] = get_secret("PROMETHEUS_SELECTED_INSTANCE")  # type: ignore
 async_http_handler = get_async_httpx_client(
     llm_provider=httpxSpecialProvider.LoggingCallback
 )
@@ -73,3 +76,65 @@ async def get_fallback_metric_from_prometheus():
                 response_message += "\n"
         verbose_logger.debug("response message %s", response_message)
     return response_message
+
+
+def is_prometheus_connected() -> bool:
+    if PROMETHEUS_URL is not None:
+        return True
+    return False
+
+
+async def get_daily_spend_from_prometheus(api_key: Optional[str]):
+    """
+    Expected Response Format:
+    [
+    {
+        "date": "2024-08-18T00:00:00+00:00",
+        "spend": 1.001818099998933
+    },
+    ...]
+    """
+    if PROMETHEUS_URL is None:
+        raise ValueError(
+            "PROMETHEUS_URL not set please set 'PROMETHEUS_URL=<>' in .env"
+        )
+
+    # Calculate the start and end dates for the last 30 days
+    end_date = datetime.utcnow()
+    start_date = end_date - timedelta(days=30)
+
+    # Format dates as ISO 8601 strings with UTC offset
+    start_str = start_date.isoformat() + "+00:00"
+    end_str = end_date.isoformat() + "+00:00"
+
+    url = f"{PROMETHEUS_URL}/api/v1/query_range"
+
+    if api_key is None:
+        query = "sum(delta(litellm_spend_metric_total[1d]))"
+    else:
+        query = (
+            f'sum(delta(litellm_spend_metric_total{{hashed_api_key="{api_key}"}}[1d]))'
+        )
+
+    params = {
+        "query": query,
+        "start": start_str,
+        "end": end_str,
+        "step": "86400",  # Step size of 1 day in seconds
+    }
+
+    response = await async_http_handler.get(url, params=params)
+    _json_response = response.json()
+    verbose_logger.debug("json response from prometheus /query api %s", _json_response)
+    results = response.json()["data"]["result"]
+    formatted_results = []
+
+    for result in results:
+        metric_data = result["values"]
+        for timestamp, value in metric_data:
+            # Convert timestamp to ISO 8601 string with UTC offset
+            date = datetime.fromtimestamp(float(timestamp)).isoformat() + "+00:00"
+            spend = float(value)
+            formatted_results.append({"date": date, "spend": spend})
+
+    return formatted_results

litellm/proxy/spend_tracking/spend_management_endpoints.py

@@ -1942,7 +1942,7 @@ async def global_spend_for_internal_user(
 
         user_id = user_api_key_dict.user_id
         if user_id is None:
-            raise ValueError(f"/global/spend/logs Error: User ID is None")
+            raise ValueError("/global/spend/logs Error: User ID is None")
         if api_key is not None:
             sql_query = """
                 SELECT * FROM "MonthlyGlobalSpendPerUserPerKey"
@@ -1971,7 +1971,7 @@ async def global_spend_for_internal_user(
     include_in_schema=False,
 )
 async def global_spend_logs(
-    api_key: str = fastapi.Query(
+    api_key: Optional[str] = fastapi.Query(
         default=None,
         description="API Key to get global spend (spend per day for last 30d). Admin-only endpoint",
     ),
@@ -1986,6 +1986,10 @@ async def global_spend_logs(
     """
     import traceback
 
+    from litellm.integrations.prometheus_helpers.prometheus_api import (
+        get_daily_spend_from_prometheus,
+        is_prometheus_connected,
+    )
     from litellm.proxy.proxy_server import prisma_client
 
     try:
@@ -2007,22 +2011,28 @@ async def global_spend_logs(
 
             return response
 
-        if api_key is None:
+        prometheus_api_enabled = is_prometheus_connected()
-            sql_query = """SELECT * FROM "MonthlyGlobalSpend" ORDER BY "date";"""
-
-            response = await prisma_client.db.query_raw(query=sql_query)
 
+        if prometheus_api_enabled:
+            response = await get_daily_spend_from_prometheus(api_key=api_key)
             return response
         else:
-            sql_query = """
+            if api_key is None:
-                SELECT * FROM "MonthlyGlobalSpendPerKey"
+                sql_query = """SELECT * FROM "MonthlyGlobalSpend" ORDER BY "date";"""
-                WHERE "api_key" = $1
-                ORDER BY "date";
-                """
 
-            response = await prisma_client.db.query_raw(sql_query, api_key)
+                response = await prisma_client.db.query_raw(query=sql_query)
 
-            return response
+                return response
+            else:
+                sql_query = """
+                    SELECT * FROM "MonthlyGlobalSpendPerKey"
+                    WHERE "api_key" = $1
+                    ORDER BY "date";
+                    """
+
+                response = await prisma_client.db.query_raw(sql_query, api_key)
+
+                return response
 
     except Exception as e:
         error_trace = traceback.format_exc()