feat - allow accessing data post success call

litellm/proxy/custom_callbacks1.py

@@ -40,6 +40,7 @@ class MyCustomHandler(
 
     async def async_post_call_success_hook(
         self,
+        data: dict,
         user_api_key_dict: UserAPIKeyAuth,
         response,
     ):

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],
     ):

litellm/proxy/proxy_server.py

@@ -3136,7 +3136,7 @@ async def chat_completion(
 
         ### CALL HOOKS ### - modify outgoing data
         response = await proxy_logging_obj.post_call_success_hook(
-            user_api_key_dict=user_api_key_dict, response=response
+            data=data, user_api_key_dict=user_api_key_dict, response=response
         )
 
         hidden_params = (
@@ -3350,6 +3350,11 @@ async def completion(
                 media_type="text/event-stream",
                 headers=custom_headers,
             )
+        ### CALL HOOKS ### - modify outgoing data
+        response = await proxy_logging_obj.post_call_success_hook(
+            data=data, user_api_key_dict=user_api_key_dict, response=response
+        )
+
         fastapi_response.headers.update(
             get_custom_headers(
                 user_api_key_dict=user_api_key_dict,

litellm/proxy/utils.py

@@ -717,6 +717,7 @@ class ProxyLogging:
 
     async def post_call_success_hook(
         self,
+        data: dict,
         response: Union[ModelResponse, EmbeddingResponse, ImageResponse],
         user_api_key_dict: UserAPIKeyAuth,
     ):
@@ -738,7 +739,9 @@ class ProxyLogging:
                     _callback = callback  # type: ignore
                 if _callback is not None and isinstance(_callback, CustomLogger):
                     await _callback.async_post_call_success_hook(
-                        user_api_key_dict=user_api_key_dict, response=response
+                        user_api_key_dict=user_api_key_dict,
+                        data=data,
+                        response=response,
                     )
             except Exception as e:
                 raise e