diff --git a/litellm/integrations/prometheus.py b/litellm/integrations/prometheus.py index c01f7481277..d5146dc37b2 100644 --- a/litellm/integrations/prometheus.py +++ b/litellm/integrations/prometheus.py @@ -20,7 +20,7 @@ from typing import ( import litellm from litellm._logging import print_verbose, verbose_logger from litellm.integrations.custom_logger import CustomLogger -from litellm.proxy._types import LiteLLM_TeamTable, UserAPIKeyAuth +from litellm.proxy._types import LiteLLM_TeamTable, LiteLLM_UserTable, UserAPIKeyAuth from litellm.types.integrations.prometheus import * from litellm.types.integrations.prometheus import _sanitize_prometheus_label_name from litellm.types.utils import StandardLoggingPayload @@ -191,6 +191,30 @@ class PrometheusLogger(CustomLogger): ), ) + # Remaining Budget for User + self.litellm_remaining_user_budget_metric = self._gauge_factory( + "litellm_remaining_user_budget_metric", + "Remaining budget for user", + labelnames=self.get_labels_for_metric( + "litellm_remaining_user_budget_metric" + ), + ) + + # Max Budget for User + self.litellm_user_max_budget_metric = self._gauge_factory( + "litellm_user_max_budget_metric", + "Maximum budget set for user", + labelnames=self.get_labels_for_metric("litellm_user_max_budget_metric"), + ) + + self.litellm_user_budget_remaining_hours_metric = self._gauge_factory( + "litellm_user_budget_remaining_hours_metric", + "Remaining hours for user budget to be reset", + labelnames=self.get_labels_for_metric( + "litellm_user_budget_remaining_hours_metric" + ), + ) + ######################################## # LiteLLM Virtual API KEY metrics ######################################## @@ -916,6 +940,7 @@ class PrometheusLogger(CustomLogger): user_api_key_alias=user_api_key_alias, litellm_params=litellm_params, response_cost=response_cost, + user_id=user_id, ) # set proxy virtual key rpm/tpm metrics @@ -1022,6 +1047,7 @@ class PrometheusLogger(CustomLogger): user_api_key_alias: Optional[str], litellm_params: dict, response_cost: float, + user_id: Optional[str] = None, ): _team_spend = litellm_params.get("metadata", {}).get( "user_api_key_team_spend", None @@ -1036,6 +1062,14 @@ class PrometheusLogger(CustomLogger): _api_key_max_budget = litellm_params.get("metadata", {}).get( "user_api_key_max_budget", None ) + + _user_spend = litellm_params.get("metadata", {}).get( + "user_api_key_user_spend", None + ) + _user_max_budget = litellm_params.get("metadata", {}).get( + "user_api_key_user_max_budget", None + ) + await self._set_api_key_budget_metrics_after_api_request( user_api_key=user_api_key, user_api_key_alias=user_api_key_alias, @@ -1052,6 +1086,13 @@ class PrometheusLogger(CustomLogger): response_cost=response_cost, ) + await self._set_user_budget_metrics_after_api_request( + user_id=user_id, + user_spend=_user_spend, + user_max_budget=_user_max_budget, + response_cost=response_cost, + ) + def _increment_top_level_request_and_spend_metrics( self, end_user_id: Optional[str], @@ -1907,6 +1948,37 @@ class PrometheusLogger(CustomLogger): data_type="keys", ) + async def _initialize_user_budget_metrics(self): + """ + Initialize user budget metrics by reusing the generic pagination logic. + """ + from litellm.proxy._types import LiteLLM_UserTable + from litellm.proxy.proxy_server import prisma_client + + if prisma_client is None: + verbose_logger.debug( + "Prometheus: skipping user metrics initialization, DB not initialized" + ) + return + + async def fetch_users( + page_size: int, page: int + ) -> Tuple[List[LiteLLM_UserTable], Optional[int]]: + skip = (page - 1) * page_size + users = await prisma_client.db.litellm_usertable.find_many( + skip=skip, + take=page_size, + order={"created_at": "desc"}, + ) + total_count = await prisma_client.db.litellm_usertable.count() + return users, total_count + + await self._initialize_budget_metrics( + data_fetch_function=fetch_users, + set_metrics_function=self._set_user_list_budget_metrics, + data_type="users", + ) + async def initialize_remaining_budget_metrics(self): """ Handler for initializing remaining budget metrics for all teams to avoid metric discrepancies. @@ -1939,11 +2011,12 @@ class PrometheusLogger(CustomLogger): async def _initialize_remaining_budget_metrics(self): """ - Helper to initialize remaining budget metrics for all teams and API keys. + Helper to initialize remaining budget metrics for all teams, API keys, and users. """ - verbose_logger.debug("Emitting key, team budget metrics....") + verbose_logger.debug("Emitting key, team, user budget metrics....") await self._initialize_team_budget_metrics() await self._initialize_api_key_budget_metrics() + await self._initialize_user_budget_metrics() async def _set_key_list_budget_metrics( self, keys: List[Union[str, UserAPIKeyAuth]] @@ -1958,6 +2031,11 @@ class PrometheusLogger(CustomLogger): for team in teams: self._set_team_budget_metrics(team) + async def _set_user_list_budget_metrics(self, users: List[LiteLLM_UserTable]): + """Helper function to set budget metrics for a list of users""" + for user in users: + self._set_user_budget_metrics(user) + async def _set_team_budget_metrics_after_api_request( self, user_api_team: Optional[str], @@ -2175,6 +2253,120 @@ class PrometheusLogger(CustomLogger): return user_api_key_dict + async def _set_user_budget_metrics_after_api_request( + self, + user_id: Optional[str], + user_spend: Optional[float], + user_max_budget: Optional[float], + response_cost: float, + ): + """ + Set user budget metrics after an LLM API request + + - Assemble a LiteLLM_UserTable object + - looks up user info from db if not available in metadata + - Set user budget metrics + """ + if user_id: + user_object = await self._assemble_user_object( + user_id=user_id, + spend=user_spend, + max_budget=user_max_budget, + response_cost=response_cost, + ) + + self._set_user_budget_metrics(user_object) + + async def _assemble_user_object( + self, + user_id: str, + spend: Optional[float], + max_budget: Optional[float], + response_cost: float, + ) -> LiteLLM_UserTable: + """ + Assemble a LiteLLM_UserTable object + + for fields not available in metadata, we fetch from db + Fields not available in metadata: + - `budget_reset_at` + """ + from litellm.proxy.auth.auth_checks import get_user_object + from litellm.proxy.proxy_server import prisma_client, user_api_key_cache + + _total_user_spend = (spend or 0) + response_cost + user_object = LiteLLM_UserTable( + user_id=user_id, + spend=_total_user_spend, + max_budget=max_budget, + ) + try: + user_info = await get_user_object( + user_id=user_id, + prisma_client=prisma_client, + user_api_key_cache=user_api_key_cache, + ) + except Exception as e: + verbose_logger.debug( + f"[Non-Blocking] Prometheus: Error getting user info: {str(e)}" + ) + return user_object + + if user_info: + user_object.budget_reset_at = user_info.budget_reset_at + + return user_object + + def _set_user_budget_metrics( + self, + user: LiteLLM_UserTable, + ): + """ + Set user budget metrics for a single user + + - Remaining Budget + - Max Budget + - Budget Reset At + """ + enum_values = UserAPIKeyLabelValues( + user=user.user_id, + ) + + _labels = prometheus_label_factory( + supported_enum_labels=self.get_labels_for_metric( + metric_name="litellm_remaining_user_budget_metric" + ), + enum_values=enum_values, + ) + self.litellm_remaining_user_budget_metric.labels(**_labels).set( + self._safe_get_remaining_budget( + max_budget=user.max_budget, + spend=user.spend, + ) + ) + + if user.max_budget is not None: + _labels = prometheus_label_factory( + supported_enum_labels=self.get_labels_for_metric( + metric_name="litellm_user_max_budget_metric" + ), + enum_values=enum_values, + ) + self.litellm_user_max_budget_metric.labels(**_labels).set(user.max_budget) + + if user.budget_reset_at is not None: + _labels = prometheus_label_factory( + supported_enum_labels=self.get_labels_for_metric( + metric_name="litellm_user_budget_remaining_hours_metric" + ), + enum_values=enum_values, + ) + self.litellm_user_budget_remaining_hours_metric.labels(**_labels).set( + self._get_remaining_hours_for_budget_reset( + budget_reset_at=user.budget_reset_at + ) + ) + def _get_remaining_hours_for_budget_reset(self, budget_reset_at: datetime) -> float: """ Get remaining hours for budget reset diff --git a/litellm/proxy/_types.py b/litellm/proxy/_types.py index 7a01f4db6f7..d1d84257ff7 100644 --- a/litellm/proxy/_types.py +++ b/litellm/proxy/_types.py @@ -2175,6 +2175,8 @@ class UserAPIKeyAuth( user_tpm_limit: Optional[int] = None user_rpm_limit: Optional[int] = None user_email: Optional[str] = None + user_spend: Optional[float] = None + user_max_budget: Optional[float] = None request_route: Optional[str] = None user: Optional[Any] = None # Expanded user object when expand=user is used diff --git a/litellm/proxy/auth/user_api_key_auth.py b/litellm/proxy/auth/user_api_key_auth.py index 9b53d9a3a80..44c2ec0b61a 100644 --- a/litellm/proxy/auth/user_api_key_auth.py +++ b/litellm/proxy/auth/user_api_key_auth.py @@ -1308,6 +1308,8 @@ async def _return_user_api_key_auth_obj( user_tpm_limit=user_obj.tpm_limit, user_rpm_limit=user_obj.rpm_limit, user_email=user_obj.user_email, + user_spend=getattr(user_obj, "spend", None), + user_max_budget=getattr(user_obj, "max_budget", None), ) if user_obj is not None and _is_user_proxy_admin(user_obj=user_obj): user_api_key_kwargs.update( diff --git a/litellm/proxy/litellm_pre_call_utils.py b/litellm/proxy/litellm_pre_call_utils.py index 5b5723efc3d..7dc6741f023 100644 --- a/litellm/proxy/litellm_pre_call_utils.py +++ b/litellm/proxy/litellm_pre_call_utils.py @@ -999,6 +999,13 @@ async def add_litellm_data_to_request( # noqa: PLR0915 "user_api_key_model_max_budget" ] = user_api_key_dict.model_max_budget + # User spend, budget - used by prometheus.py + # Follow same pattern as team and API key budgets + data[_metadata_variable_name]["user_api_key_user_spend"] = user_api_key_dict.user_spend + data[_metadata_variable_name][ + "user_api_key_user_max_budget" + ] = user_api_key_dict.user_max_budget + data[_metadata_variable_name]["user_api_key_metadata"] = user_api_key_dict.metadata _headers = dict(request.headers) _headers.pop( diff --git a/litellm/types/integrations/prometheus.py b/litellm/types/integrations/prometheus.py index 6a254fc8252..bc8f06dcc40 100644 --- a/litellm/types/integrations/prometheus.py +++ b/litellm/types/integrations/prometheus.py @@ -175,6 +175,9 @@ DEFINED_PROMETHEUS_METRICS = Literal[ "litellm_remaining_api_key_budget_metric", "litellm_api_key_max_budget_metric", "litellm_api_key_budget_remaining_hours_metric", + "litellm_remaining_user_budget_metric", + "litellm_user_max_budget_metric", + "litellm_user_budget_remaining_hours_metric", "litellm_deployment_state", "litellm_deployment_failure_responses", "litellm_deployment_total_requests", @@ -396,6 +399,18 @@ class PrometheusMetricLabels: litellm_remaining_api_key_budget_metric ) + litellm_remaining_user_budget_metric = [ + UserAPIKeyLabelNames.USER.value, + ] + + litellm_user_max_budget_metric = [ + UserAPIKeyLabelNames.USER.value, + ] + + litellm_user_budget_remaining_hours_metric = [ + UserAPIKeyLabelNames.USER.value, + ] + # Add deployment metrics litellm_deployment_failure_responses = [ UserAPIKeyLabelNames.REQUESTED_MODEL.value, diff --git a/tests/otel_tests/test_prometheus.py b/tests/otel_tests/test_prometheus.py index 883562e8820..ce3031b5141 100644 --- a/tests/otel_tests/test_prometheus.py +++ b/tests/otel_tests/test_prometheus.py @@ -442,6 +442,24 @@ async def get_key_info(session: aiohttp.ClientSession, key: str) -> Dict[str, An return await response.json() +async def get_user_info(session: aiohttp.ClientSession, user_id: str) -> Dict[str, Any]: + """Fetch user info and return the response""" + from urllib.parse import quote + + # URL encode user_id to handle special characters + encoded_user_id = quote(user_id, safe="") + url = f"http://0.0.0.0:4000/user/info?user_id={encoded_user_id}" + headers = { + "Authorization": "Bearer sk-1234", + } + + async with session.get(url, headers=headers) as response: + assert ( + response.status == 200 + ), f"Failed to get user info. Status: {response.status}" + return await response.json() + + def extract_key_budget_metrics(metrics_text: str, key_id: str) -> Dict[str, float]: """Extract budget-related metrics for a specific key""" import re @@ -466,6 +484,33 @@ def extract_key_budget_metrics(metrics_text: str, key_id: str) -> Dict[str, floa return metrics +def extract_user_budget_metrics(metrics_text: str, user_id: str) -> Dict[str, float]: + """Extract budget-related metrics for a specific user""" + import re + + metrics = {} + + # Escape user_id for regex pattern matching + escaped_user_id = re.escape(user_id) + + # Get remaining budget + remaining_pattern = f'litellm_remaining_user_budget_metric{{user="{escaped_user_id}"}} ([0-9.]+)' + remaining_match = re.search(remaining_pattern, metrics_text) + metrics["remaining"] = float(remaining_match.group(1)) if remaining_match else None + + # Get total budget + total_pattern = f'litellm_user_max_budget_metric{{user="{escaped_user_id}"}} ([0-9.]+)' + total_match = re.search(total_pattern, metrics_text) + metrics["total"] = float(total_match.group(1)) if total_match else None + + # Get remaining hours + hours_pattern = f'litellm_user_budget_remaining_hours_metric{{user="{escaped_user_id}"}} ([0-9.]+)' + hours_match = re.search(hours_pattern, metrics_text) + metrics["remaining_hours"] = float(hours_match.group(1)) if hours_match else None + + return metrics + + @pytest.mark.asyncio async def test_key_budget_metrics(): """ @@ -476,6 +521,8 @@ async def test_key_budget_metrics(): 4. Verify request costs are being tracked correctly 5. Verify prometheus metrics match /key/info spend data """ + from datetime import datetime, timedelta, timezone + async with aiohttp.ClientSession() as session: # Setup test key with unique alias unique_alias = f"budget_test_key_{uuid.uuid4()}" @@ -483,6 +530,7 @@ async def test_key_budget_metrics(): "key_alias": unique_alias, "max_budget": 10, "budget_duration": "7d", + "budget_reset_at": (datetime.now(timezone.utc) + timedelta(days=7)).isoformat(), } key = await create_test_key_with_budget(session, key_data) @@ -543,6 +591,94 @@ async def test_key_budget_metrics(): ), f"Spend mismatch: Prometheus={key_info_remaining_budget}, Key Info={first_budget['remaining']}" +@pytest.mark.asyncio +async def test_user_budget_metrics(): + """ + Test user budget tracking metrics: + 1. Create a user with max_budget + 2. Make chat completion requests using OpenAI SDK with the user's key + 3. Verify budget decreases over time + 4. Verify request costs are being tracked correctly + 5. Verify prometheus metrics match /user/info spend data + """ + from datetime import datetime, timedelta, timezone + + async with aiohttp.ClientSession() as session: + # Setup test user with unique user_id + unique_user_id = f"budget_test_user_{uuid.uuid4()}" + user_data = { + "user_id": unique_user_id, + "max_budget": 10, + "budget_duration": "7d", + "budget_reset_at": (datetime.now(timezone.utc) + timedelta(days=7)).isoformat(), + } + user_info = await create_test_user(session, user_data) + print("user_info", user_info) + user_id = user_info["user_id"] + print("user_id", user_id) + # Get the key that was created with the user + key = user_info["key"] + + # Initialize OpenAI client with the user's key + client = AsyncOpenAI(base_url="http://0.0.0.0:4000", api_key=key) + + # Make initial request and check budget + await client.chat.completions.create( + model="fake-openai-endpoint", + messages=[{"role": "user", "content": f"Hello {uuid.uuid4()}"}], + ) + + await asyncio.sleep(11) # Wait for metrics to update + + # Get metrics after request + metrics_after_first = await get_prometheus_metrics(session) + print("metrics_after_first request", metrics_after_first) + first_budget = extract_user_budget_metrics(metrics_after_first, user_id) + + print(f"Budget after 1 request: {first_budget}") + assert ( + first_budget["remaining"] is not None + ), "remaining budget metric should be present" + assert ( + first_budget["total"] is not None + ), "total budget metric should be present" + assert ( + first_budget["remaining"] < 10.0 + ), "remaining budget should be less than 10.0 after first request" + assert first_budget["total"] == 10.0, "Total budget metric is incorrect" + print("first_budget['remaining_hours']", first_budget["remaining_hours"]) + # The budget reset time is now standardized - for "7d" it resets on Monday at midnight + # So we'll check if it's within a reasonable range (0-7 days depending on current day of week) + assert ( + first_budget["remaining_hours"] is not None + ), "remaining hours metric should be present" + assert ( + 0 <= first_budget["remaining_hours"] <= 168 + ), "Budget remaining hours should be within a reasonable range (0-7 days depending on day of week)" + + # Get user info and verify spend matches prometheus metrics + user_info_response = await get_user_info(session, user_id) + print("user_info_response", user_info_response) + _user_info_data = user_info_response["user_info"] + + # Calculate spend from prometheus (total - remaining) + user_info_spend = float(_user_info_data["spend"]) + user_info_max_budget = float(_user_info_data["max_budget"]) + user_info_remaining_budget = user_info_max_budget - user_info_spend + print("\n\n\n###### Final budget metrics ######\n\n\n") + print("user_info_remaining_budget", user_info_remaining_budget) + print("prometheus_remaining_budget", first_budget["remaining"]) + print( + "diff between user_info_remaining_budget and prometheus_remaining_budget", + user_info_remaining_budget - first_budget["remaining"], + ) + + # Verify spends match within a small delta (floating point comparison) + assert ( + abs(user_info_remaining_budget - first_budget["remaining"]) <= 0.001 + ), f"Spend mismatch: Prometheus={user_info_remaining_budget}, User Info={first_budget['remaining']}" + + @pytest.mark.asyncio async def test_user_email_metrics(): """ diff --git a/tests/test_litellm/proxy/auth/test_user_api_key_auth.py b/tests/test_litellm/proxy/auth/test_user_api_key_auth.py index fcc8c1f0f2e..46a7b0ff435 100644 --- a/tests/test_litellm/proxy/auth/test_user_api_key_auth.py +++ b/tests/test_litellm/proxy/auth/test_user_api_key_auth.py @@ -347,3 +347,60 @@ async def test_proxy_admin_expired_key_from_cache(): finally: # Clean up - restore original values if needed pass + + +@pytest.mark.asyncio +async def test_return_user_api_key_auth_obj_user_spend_and_budget(): + """ + Test that _return_user_api_key_auth_obj correctly sets user_spend and user_max_budget + from user_obj attributes. + """ + from datetime import datetime + + from litellm.proxy._types import UserAPIKeyAuth + from litellm.proxy.auth.user_api_key_auth import _return_user_api_key_auth_obj + + user_obj = type( + "LiteLLM_UserTable", + (), + { + "tpm_limit": 1000, + "rpm_limit": 100, + "user_email": "test@example.com", + "spend": 250.0, + "max_budget": 1000.0, + "user_role": "internal_user", + }, + ) + + api_key = "sk-test-key" + valid_token_dict = { + "user_id": "test-user", + "org_id": "test-org", + } + route = "/chat/completions" + start_time = datetime.now() + + mock_service_logger = MagicMock() + mock_service_logger.async_service_success_hook = AsyncMock() + + with patch( + "litellm.proxy.auth.user_api_key_auth.user_api_key_service_logger_obj", + new=mock_service_logger, + ): + result = await _return_user_api_key_auth_obj( + user_obj=user_obj, + api_key=api_key, + parent_otel_span=None, + valid_token_dict=valid_token_dict, + route=route, + start_time=start_time, + user_role=None, + ) + + assert isinstance(result, UserAPIKeyAuth) + assert result.user_spend == 250.0 + assert result.user_max_budget == 1000.0 + assert result.user_tpm_limit == 1000 + assert result.user_rpm_limit == 100 + assert result.user_email == "test@example.com" diff --git a/tests/test_litellm/proxy/test_litellm_pre_call_utils.py b/tests/test_litellm/proxy/test_litellm_pre_call_utils.py index fd39b308a7a..ebc74af1d1d 100644 --- a/tests/test_litellm/proxy/test_litellm_pre_call_utils.py +++ b/tests/test_litellm/proxy/test_litellm_pre_call_utils.py @@ -160,6 +160,44 @@ async def test_add_litellm_data_to_request_parses_string_metadata(): assert updated_data["metadata"]["generation_name"] == "gen123" +@pytest.mark.asyncio +async def test_add_litellm_data_to_request_user_spend_and_budget(): + from litellm.proxy.litellm_pre_call_utils import add_litellm_data_to_request + + request_mock = MagicMock(spec=Request) + request_mock.url.path = "/v1/completions" + request_mock.url = MagicMock() + request_mock.url.__str__.return_value = "http://localhost/v1/completions" + request_mock.method = "POST" + request_mock.query_params = {} + request_mock.headers = {"Content-Type": "application/json"} + request_mock.client = MagicMock() + request_mock.client.host = "127.0.0.1" + + data = {"model": "gpt-3.5-turbo", "messages": [{"role": "user", "content": "Hello"}]} + + user_api_key_dict = UserAPIKeyAuth( + api_key="hashed-key", + metadata={}, + team_metadata={}, + user_spend=150.0, + user_max_budget=500.0, + ) + + updated_data = await add_litellm_data_to_request( + data=data, + request=request_mock, + user_api_key_dict=user_api_key_dict, + proxy_config=MagicMock(), + general_settings={}, + version="test-version", + ) + + metadata = updated_data.get("metadata", {}) + assert metadata["user_api_key_user_spend"] == 150.0 + assert metadata["user_api_key_user_max_budget"] == 500.0 + + @pytest.mark.asyncio async def test_add_litellm_data_to_request_audio_transcription_multipart(): from litellm.proxy.litellm_pre_call_utils import add_litellm_data_to_request