utils/logger.py

import logging #### todo add transformer logs
import streamlit as st
# Configure the logging settings for transformers
transformers_logger = logging.getLogger("transformers.file_utils")
transformers_logger.setLevel(logging.INFO)  # Set the desired logging level

log_enabled = False


##################################### logger #################################

 


def log_response(response):
    if log_enabled:
        with st.chat_message("ai"):
            st.markdown("Agent Response\n {}".format(response))
    print(response)

# Custom logging handler to append log messages to the chat
class ChatHandler(logging.Handler):
    def __init__(self):
        super().__init__()


    def emit(self, record):
        log_message = self.format(record)
        #with st.chat_message("ai"):
        st.markdown(f"Log: {log_message}")
        with st.chat_message("ai"):
            st.markdown("Agent Response\n {}".format(record))

# Add the custom handler to the transformers_logger
chat_handler = ChatHandler()
transformers_logger.addHandler(chat_handler)