feat - allow accessing data post success call

litellm/proxy/hooks/azure_content_safety.py

@@ -1,11 +1,16 @@
-from litellm.integrations.custom_logger import CustomLogger
-from litellm.caching import DualCache
-from litellm.proxy._types import UserAPIKeyAuth
-import litellm, traceback, sys, uuid
-from fastapi import HTTPException
-from litellm._logging import verbose_proxy_logger
+import sys
+import traceback
+import uuid
 from typing import Optional
 
+from fastapi import HTTPException
+
+import litellm
+from litellm._logging import verbose_proxy_logger
+from litellm.caching import DualCache
+from litellm.integrations.custom_logger import CustomLogger
+from litellm.proxy._types import UserAPIKeyAuth
+
 
 class _PROXY_AzureContentSafety(
     CustomLogger
@@ -15,12 +20,12 @@ class _PROXY_AzureContentSafety(
     def __init__(self, endpoint, api_key, thresholds=None):
         try:
             from azure.ai.contentsafety.aio import ContentSafetyClient
-            from azure.core.credentials import AzureKeyCredential
             from azure.ai.contentsafety.models import (
-                TextCategory,
                 AnalyzeTextOptions,
                 AnalyzeTextOutputType,
+                TextCategory,
             )
+            from azure.core.credentials import AzureKeyCredential
             from azure.core.exceptions import HttpResponseError
         except Exception as e:
             raise Exception(
@@ -132,6 +137,7 @@ class _PROXY_AzureContentSafety(
 
     async def async_post_call_success_hook(
         self,
+        data: dict,
         user_api_key_dict: UserAPIKeyAuth,
         response,
     ):

litellm/proxy/hooks/dynamic_rate_limiter.py

@@ -254,7 +254,7 @@ class _PROXY_DynamicRateLimitHandler(CustomLogger):
         return None
 
     async def async_post_call_success_hook(
-        self, user_api_key_dict: UserAPIKeyAuth, response
+        self, data: dict, user_api_key_dict: UserAPIKeyAuth, response
     ):
         try:
             if isinstance(response, ModelResponse):
@@ -287,7 +287,9 @@ class _PROXY_DynamicRateLimitHandler(CustomLogger):
 
                 return response
             return await super().async_post_call_success_hook(
-                user_api_key_dict, response
+                data=data,
+                user_api_key_dict=user_api_key_dict,
+                response=response,
             )
         except Exception as e:
             verbose_proxy_logger.exception(

litellm/proxy/hooks/presidio_pii_masking.py

@@ -322,6 +322,7 @@ class _OPTIONAL_PresidioPIIMasking(CustomLogger):
 
     async def async_post_call_success_hook(
         self,
+        data: dict,
         user_api_key_dict: UserAPIKeyAuth,
         response: Union[ModelResponse, EmbeddingResponse, ImageResponse],
     ):