cut

logger.py

@@ -1,27 +1,26 @@
-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
+from io import StringIO
+import contextlib
+import sys
 
-log_enabled = False
+# Redirect stdout and stderr
+original_stdout = sys.stdout
+original_stderr = sys.stderr
+sys.stdout = StringIO()
+sys.stderr = StringIO()
 
-def log_response(response):
-    if log_enabled:
-        with st.chat_message("ai"):
-            st.markdown("Agent Response\n {}".format(response))
-    print(response)
+@st.cache
+def capture_logs(func):
+    with contextlib.redirect_stdout(sys.stdout), contextlib.redirect_stderr(sys.stderr):
+        func()
 
-# Custom logging handler to append log messages to the chat
-class ChatHandler(logging.Handler):
-    def __init__(self):
-        super().__init__()
+# Get captured log messages
+def get_captured_logs():
+    captured_stdout = sys.stdout.getvalue()
+    captured_stderr = sys.stderr.getvalue()
+    return captured_stdout, captured_stderr
 
-    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)
+# Restore original stdout and stderr
+def restore_original_streams():
+    sys.stdout = original_stdout
+    sys.stderr = original_stderr