Update logger.py

logger.py

@@ -1,26 +1,32 @@
+import logging #### todo add transformer logs
 import streamlit as st
-from io import StringIO
-import contextlib
-import sys
-
-# Redirect stdout and stderr
-original_stdout = sys.stdout
-original_stderr = sys.stderr
-sys.stdout = StringIO()
-sys.stderr = StringIO()
-
-@st.cache
-def capture_logs(func):
-    with contextlib.redirect_stdout(sys.stdout), contextlib.redirect_stderr(sys.stderr):
-        func()
-
-# Get captured log messages
-def get_captured_logs():
-    captured_stdout = sys.stdout.getvalue()
-    captured_stderr = sys.stderr.getvalue()
-    return captured_stdout, captured_stderr
-
-# Restore original stdout and stderr
-def restore_original_streams():
-    sys.stdout = original_stdout
-    sys.stderr = original_stderr
+# 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 = True
+
+
+
+
+
+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}")
+
+# Add the custom handler to the transformers_logger
+chat_handler = ChatHandler()
+transformers_logger.addHandler(chat_handler)