User metrics for promethus

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

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
 

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(

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(

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,

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():
     """

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"

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