Custom Handler not logging response

[prahaladd]

I am attempting to write a simple custom handler based on the documentation provided here .
I am able to intercept requests, however if I am unable to intercept the response and log it to a file. Here is my implementation based on the documentation and this bug report. What am I doing wrong here?

I am using the Ollama provider and send a request to the Ollama server as below

curl -i -sSL --location 'http://0.0.0.0:4000/api/chat' \
    --header 'Authorization: Bearer sk-1234' \
    --header 'Content-Type: application/json' \
    --data '{
      "model": "llama3.2",
      "messages": [{"role": "user", "content": "what tools are available?"}]
    }'

This is my config file:

model_list:
  - model_name: llama3.2
    litellm_params:
      model: litellm_proxy/llama3.2
      api_base: http://localhost:11434
      #api_key: "os.environ/AZURE_API_KEY"
      #api_version: "2024-07-01-preview" # [OPTIONAL] litellm uses the latest azure api_version by default

litellm_settings:      
  turn_off_message_logging: false
  callbacks: /app/plugins/custom_handler.proxy_handler_instance
  json_logs: false
general_settings:
  master_key: sk-1234
  store_prompts_in_spend_logs: true
  pass_through_endpoints:
    - path: "/api/chat"                                  # route you want to add to LiteLLM Proxy Server
      target: "http://host.docker.internal:11434/api/chat"          # URL this route should forward requests to
      headers:                                            # headers to forward to this URL
        #Authorization: "bearer os.environ/COHERE_API_KEY" # (Optional) Auth Header to forward to your Endpoint
        content-type: application/json                    # (Optional) Extra Headers to pass to this endpoint 
        accept: application/json
      forward_headers: True     

from litellm.integrations.custom_logger import CustomLogger
#from litellm.types.utils import ModelResponseStream
from typing import Any, AsyncGenerator, Optional, Literal
from pprint import pprint

# This file includes the custom callbacks for LiteLLM Proxy
# Once defined, these can be passed in proxy_config.yaml
class MyCustomHandler(CustomLogger):
    def __init__(self):
        pass

    #### CALL HOOKS - proxy only ####

    async def async_pre_call_hook(self, user_api_key_dict: Any, cache: Any, data: dict, call_type: Literal[
            "completion",
            "text_completion",
            "embeddings",
            "image_generation",
            "moderation",
            "audio_transcription",
        ]):
        # Pretty print the data dictionary
        print("Pre-call hook...")
        pprint(data)
        logger.info({"event": "pre_call_hook", "data": data})
        # Example: modify the model before making the LLM API call
        #data["model"] = "my-new-model"
        
        return data

    async def async_post_call_failure_hook(
        self,
        request_data: dict,
        original_exception: Exception,
        user_api_key_dict: Any,
        traceback_str: Optional[str] = None,
    ):
        # Custom logic for handling failures
        pass

    async def async_post_call_success_hook(
        self,
        data: dict,
        user_api_key_dict: Any,
        response,
    ):
        # Custom logic for handling successful calls
        # Log entry to identify the method
        print("Post-call success hook...")
        logger.info({"event": "post_call_success_hook", "data": data, "response": str(response)})
        # Optionally, pretty print the response for debugging
        pprint(response)

    async def async_moderation_hook(
        self,
        data: dict,
        user_api_key_dict: Any,
        call_type: Literal["completion", "embeddings", "image_generation", "moderation", "audio_transcription"],
    ):
        # Custom moderation logic
        pass

    async def async_post_call_streaming_hook(
        self,
        user_api_key_dict: Any,
        response: str,
    ):
        # Custom logic for streaming responses
        print("Post-call streaming hook...")
        pprint(response)
        logger.info({"event": "post_call_streaming_hook", "response": response})

    async def async_post_call_streaming_iterator_hook(
        self,
        user_api_key_dict: Any,
        response: Any,
        request_data: dict,
    ) -> AsyncGenerator[Any, None]:
        """
        Passes the entire stream to the guardrail
        This is useful for plugins that need to see the entire stream.
        """
        async for item in response:
            print(item)
            logger.info({"event": "post_call_streaming_iterator_hook", "item": item})
            yield item

# Instance to be referenced in proxy config
proxy_handler_instance = MyCustomHandler() 

import logging
import json

class JSONLogFormatter(logging.Formatter):
    def format(self, record):
        log_message = {
            "timestamp": self.formatTime(record, "%Y-%m-%dT%H:%M:%S"),
            "level": record.levelname,
            "message": record.getMessage(),
        }
        return json.dumps(log_message)
   
logger = logging.getLogger("litellm_logger")
logger.setLevel(logging.INFO)

logger.propagate = False

console_handler = logging.StreamHandler()
console_handler.setFormatter(JSONLogFormatter())

logger.addHandler(console_handler)