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
This commit is contained in:
Krish Dholakia 2024-09-18 22:39:15 -07:00 committed by GitHub
parent 0bbee5e286
commit c23d1e6ddc
2 changed files with 90 additions and 15 deletions

View file

@ -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_SELECTED_INSTANCE = litellm.get_secret("PROMETHEUS_SELECTED_INSTANCE")
PROMETHEUS_URL: Optional[str] = get_secret("PROMETHEUS_URL") # type: ignore
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

View file

@ -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,6 +2011,12 @@ async def global_spend_logs(
return response
prometheus_api_enabled = is_prometheus_connected()
if prometheus_api_enabled:
response = await get_daily_spend_from_prometheus(api_key=api_key)
return response
else:
if api_key is None:
sql_query = """SELECT * FROM "MonthlyGlobalSpend" ORDER BY "date";"""