Merge pull request #24432 from BerriAI/krrishdholakia/project-id-tracking

feat(proxy): add project_alias tracking in callbacks

Author: Krish Dholakia

litellm/litellm_core_utils/litellm_logging.py

@@ -2416,6 +2416,7 @@ class LiteLLM_VerificationTokenView(LiteLLM_VerificationToken):
     organization_metadata: Optional[dict] = None
 
     # Project Params
+    project_alias: Optional[str] = None
     project_metadata: Optional[dict] = None
 
     # Time stamps
@@ -3228,6 +3229,7 @@ class SpendLogsMetadata(TypedDict):
     user_api_key_alias: Optional[str]
     user_api_key_team_id: Optional[str]
     user_api_key_project_id: Optional[str]
+    user_api_key_project_alias: Optional[str]
     user_api_key_org_id: Optional[str]
     user_api_key_user_id: Optional[str]
     user_api_key_team_alias: Optional[str]

litellm/proxy/_types.py

@@ -836,6 +836,7 @@ async def _user_api_key_auth_builder(  # noqa: PLR0915
                         )
                         if _jwt_project_obj is not None:
                             valid_token.project_metadata = _jwt_project_obj.metadata
+                            valid_token.project_alias = _jwt_project_obj.project_alias
 
                     # run through common checks
                     _ = await common_checks(
@@ -1431,6 +1432,7 @@ async def _user_api_key_auth_builder(  # noqa: PLR0915
                 )
                 if _project_obj is not None:
                     valid_token.project_metadata = _project_obj.metadata
+                    valid_token.project_alias = _project_obj.project_alias
 
             global_proxy_spend = None
             if (
@@ -1888,6 +1890,7 @@ async def _run_post_custom_auth_checks(
         )
         if _project_obj is not None:
             valid_token.project_metadata = _project_obj.metadata
+            valid_token.project_alias = _project_obj.project_alias
 
     if general_settings.get("custom_auth_run_common_checks", False):
         _ = await common_checks(

litellm/proxy/auth/user_api_key_auth.py

@@ -27,14 +27,16 @@ async def create_missing_views(db: _db):  # noqa: PLR0915
         await db.execute_raw(
             """
                 CREATE VIEW "LiteLLM_VerificationTokenView" AS
-                SELECT 
-                v.*, 
-                t.spend AS team_spend, 
-                t.max_budget AS team_max_budget, 
-                t.tpm_limit AS team_tpm_limit, 
-                t.rpm_limit AS team_rpm_limit
+                SELECT
+                v.*,
+                t.spend AS team_spend,
+                t.max_budget AS team_max_budget,
+                t.tpm_limit AS team_tpm_limit,
+                t.rpm_limit AS team_rpm_limit,
+                p.project_alias AS project_alias
                 FROM "LiteLLM_VerificationToken" v
-                LEFT JOIN "LiteLLM_TeamTable" t ON v.team_id = t.team_id;
+                LEFT JOIN "LiteLLM_TeamTable" t ON v.team_id = t.team_id
+                LEFT JOIN "LiteLLM_ProjectTable" p ON v.project_id = p.project_id;
             """
         )
 

litellm/proxy/db/create_views.py

@@ -18,11 +18,9 @@
     log_db_metrics,
 )
 from litellm.proxy.auth.route_checks import RouteChecks
+from litellm.proxy.litellm_pre_call_utils import LiteLLMProxyRequestSetup
 from litellm.proxy.utils import ProxyUpdateSpend
-from litellm.types.utils import (
-    StandardLoggingPayload,
-    StandardLoggingUserAPIKeyMetadata,
-)
+from litellm.types.utils import StandardLoggingPayload
 from litellm.utils import get_end_user_id_for_cost_tracking
 
 
@@ -51,25 +49,8 @@ async def async_post_call_failure_hook(
         from litellm.proxy.proxy_server import proxy_logging_obj
 
         _metadata = dict(
-            StandardLoggingUserAPIKeyMetadata(
-                user_api_key_hash=user_api_key_dict.api_key,
-                user_api_key_alias=user_api_key_dict.key_alias,
-                user_api_key_spend=user_api_key_dict.spend,
-                user_api_key_max_budget=user_api_key_dict.max_budget,
-                user_api_key_budget_reset_at=(
-                    user_api_key_dict.budget_reset_at.isoformat()
-                    if user_api_key_dict.budget_reset_at
-                    else None
-                ),
-                user_api_key_user_email=user_api_key_dict.user_email,
-                user_api_key_user_id=user_api_key_dict.user_id,
-                user_api_key_team_id=user_api_key_dict.team_id,
-                user_api_key_org_id=user_api_key_dict.org_id,
-                user_api_key_project_id=user_api_key_dict.project_id,
-                user_api_key_team_alias=user_api_key_dict.team_alias,
-                user_api_key_end_user_id=user_api_key_dict.end_user_id,
-                user_api_key_request_route=user_api_key_dict.request_route,
-                user_api_key_auth_metadata=user_api_key_dict.metadata,
+            LiteLLMProxyRequestSetup.get_sanitized_user_information_from_key(
+                user_api_key_dict=user_api_key_dict
             )
         )
         _metadata["user_api_key"] = user_api_key_dict.api_key

litellm/proxy/hooks/proxy_track_cost_callback.py

@@ -658,6 +658,7 @@ def get_sanitized_user_information_from_key(
             user_api_key_max_budget=user_api_key_dict.max_budget,
             user_api_key_team_id=user_api_key_dict.team_id,
             user_api_key_project_id=user_api_key_dict.project_id,
+            user_api_key_project_alias=user_api_key_dict.project_alias,
             user_api_key_user_id=user_api_key_dict.user_id,
             user_api_key_org_id=user_api_key_dict.org_id,
             user_api_key_team_alias=user_api_key_dict.team_alias,

litellm/proxy/litellm_pre_call_utils.py

@@ -55,14 +55,14 @@
     _read_request_body,
     _safe_get_request_headers,
 )
+from litellm.proxy.litellm_pre_call_utils import LiteLLMProxyRequestSetup
 from litellm.proxy.utils import get_server_root_path, normalize_route_for_root_path
 from litellm.secret_managers.main import get_secret_str
 from litellm.types.llms.custom_http import httpxSpecialProvider
 from litellm.types.passthrough_endpoints.pass_through_endpoints import (
     EndpointType,
     PassthroughStandardLoggingPayload,
 )
-from litellm.types.utils import StandardLoggingUserAPIKeyMetadata
 
 from .streaming_handler import PassThroughStreamingHandler
 from .success_handler import PassThroughEndpointLogging
@@ -502,25 +502,8 @@ def _init_kwargs_for_pass_through_endpoint(
                 litellm_params_in_body[k] = _parsed_body.pop(k, None)
 
         _metadata = dict(
-            StandardLoggingUserAPIKeyMetadata(
-                user_api_key_hash=user_api_key_dict.api_key,
-                user_api_key_alias=user_api_key_dict.key_alias,
-                user_api_key_user_email=user_api_key_dict.user_email,
-                user_api_key_user_id=user_api_key_dict.user_id,
-                user_api_key_team_id=user_api_key_dict.team_id,
-                user_api_key_org_id=user_api_key_dict.org_id,
-                user_api_key_project_id=user_api_key_dict.project_id,
-                user_api_key_team_alias=user_api_key_dict.team_alias,
-                user_api_key_end_user_id=user_api_key_dict.end_user_id,
-                user_api_key_request_route=user_api_key_dict.request_route,
-                user_api_key_spend=user_api_key_dict.spend,
-                user_api_key_max_budget=user_api_key_dict.max_budget,
-                user_api_key_budget_reset_at=(
-                    user_api_key_dict.budget_reset_at.isoformat()
-                    if user_api_key_dict.budget_reset_at
-                    else None
-                ),
-                user_api_key_auth_metadata=user_api_key_dict.metadata,
+            LiteLLMProxyRequestSetup.get_sanitized_user_information_from_key(
+                user_api_key_dict=user_api_key_dict
             )
         )
 

litellm/proxy/pass_through_endpoints/pass_through_endpoints.py

@@ -1898,9 +1898,9 @@ async def _handle_logging_proxy_only_error(
                     normalized_call_type = CallTypes.aembedding.value
             if normalized_call_type is not None:
                 litellm_logging_obj.call_type = normalized_call_type
-                litellm_logging_obj.model_call_details["call_type"] = (
-                    normalized_call_type
-                )
+                litellm_logging_obj.model_call_details[
+                    "call_type"
+                ] = normalized_call_type
             # Pass-through endpoints are logged via the callback loop's
             # async_post_call_failure_hook — skip pre_call and failure handlers.
             if litellm_logging_obj.call_type == CallTypes.pass_through.value:
@@ -2498,7 +2498,8 @@ async def check_view_exists(self):
             required_view = "LiteLLM_VerificationTokenView"
             expected_views_str = ", ".join(f"'{view}'" for view in expected_views)
             pg_schema = os.getenv("DATABASE_SCHEMA", "public")
-            ret = await self.db.query_raw(f"""
+            ret = await self.db.query_raw(
+                f"""
                 WITH existing_views AS (
                     SELECT viewname
                     FROM pg_views
@@ -2510,7 +2511,8 @@ async def check_view_exists(self):
                     (SELECT COUNT(*) FROM existing_views) AS view_count,
                     ARRAY_AGG(viewname) AS view_names
                 FROM existing_views
-                """)
+                """
+            )
             expected_total_views = len(expected_views)
             if ret[0]["view_count"] == expected_total_views:
                 verbose_proxy_logger.info("All necessary views exist!")
@@ -2519,7 +2521,8 @@ async def check_view_exists(self):
                 ## check if required view exists ##
                 if ret[0]["view_names"] and required_view not in ret[0]["view_names"]:
                     await self.health_check()  # make sure we can connect to db
-                    await self.db.execute_raw("""
+                    await self.db.execute_raw(
+                        """
                             CREATE VIEW "LiteLLM_VerificationTokenView" AS
                             SELECT
                             v.*,
@@ -2529,7 +2532,8 @@ async def check_view_exists(self):
                             t.rpm_limit AS team_rpm_limit
                             FROM "LiteLLM_VerificationToken" v
                             LEFT JOIN "LiteLLM_TeamTable" t ON v.team_id = t.team_id;
-                        """)
+                        """
+                    )
 
                     verbose_proxy_logger.info(
                         "LiteLLM_VerificationTokenView Created in DB!"
@@ -2964,6 +2968,7 @@ async def get_data(  # noqa: PLR0915
                             t.members_with_roles AS team_members_with_roles,
                             t.object_permission_id AS team_object_permission_id,
                             t.organization_id as org_id,
+                            p.project_alias AS project_alias,
                             tm.spend AS team_member_spend,
                             m.aliases AS team_model_aliases,
                             -- Added comma to separate b.* columns
@@ -2981,6 +2986,7 @@ async def get_data(  # noqa: PLR0915
                         LEFT JOIN "LiteLLM_TeamMembership" AS tm ON v.team_id = tm.team_id AND tm.user_id = v.user_id
                         LEFT JOIN "LiteLLM_ModelTable" m ON t.model_id = m.id
                         LEFT JOIN "LiteLLM_BudgetTable" AS b ON v.budget_id = b.budget_id
+                        LEFT JOIN "LiteLLM_ProjectTable" AS p ON v.project_id = p.project_id
                         LEFT JOIN "LiteLLM_OrganizationTable" AS o ON v.organization_id = o.organization_id
                         LEFT JOIN "LiteLLM_BudgetTable" AS b2 ON o.budget_id = b2.budget_id
                         WHERE v.token = '{token}'

litellm/proxy/utils.py

@@ -2488,6 +2488,7 @@ class StandardLoggingUserAPIKeyMetadata(TypedDict):
     user_api_key_org_id: Optional[str]
     user_api_key_team_id: Optional[str]
     user_api_key_project_id: Optional[str]
+    user_api_key_project_alias: Optional[str]
     user_api_key_user_id: Optional[str]
     user_api_key_user_email: Optional[str]
     user_api_key_team_alias: Optional[str]