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