diff --git a/litellm/integrations/prometheus_helpers/prometheus_api.py b/litellm/integrations/prometheus_helpers/prometheus_api.py index 472c3b61bc..c599390198 100644 --- a/litellm/integrations/prometheus_helpers/prometheus_api.py +++ b/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_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 diff --git a/litellm/proxy/spend_tracking/spend_management_endpoints.py b/litellm/proxy/spend_tracking/spend_management_endpoints.py index 295546a3d4..ec09159b0e 100644 --- a/litellm/proxy/spend_tracking/spend_management_endpoints.py +++ b/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: - sql_query = """SELECT * FROM "MonthlyGlobalSpend" ORDER BY "date";""" - - response = await prisma_client.db.query_raw(query=sql_query) + prometheus_api_enabled = is_prometheus_connected() + if prometheus_api_enabled: + response = await get_daily_spend_from_prometheus(api_key=api_key) return response else: - sql_query = """ - SELECT * FROM "MonthlyGlobalSpendPerKey" - WHERE "api_key" = $1 - ORDER BY "date"; - """ + if api_key is None: + sql_query = """SELECT * FROM "MonthlyGlobalSpend" 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()