Implemented stdout redirect method - logger widget in progress

app.py

@@ -2,8 +2,9 @@ import os
 
 import streamlit as st
 
-from chatbot.utils import download_test_data
-from chatbot.utils import load_data
+from chatbot import redirect as rd
+from chatbot.data import download_test_data
+from chatbot.data import load_data
 
 # add OpenAI API key to environemntal variables
 os.environ["OPENAI_API_KEY"] = st.secrets["OPENAI_API_KEY"]
@@ -15,6 +16,12 @@ if "messages" not in st.session_state.keys():  # Initialize the chat message his
     st.session_state.messages = [{"role": "assistant", "content": "Ask me a question about André's research!"}]
 
 def main():
+    # setup logger sidebar
+    #st.sidebar.text("Standard output log:")
+    #_sidebar_out = st.sidebar.empty()
+    #with rd.stdout(to=_sidebar_out, format='text'):
+    #    print("test")
+
     # setup dataset
     download_test_data()
     index = load_data()

chatbot/redirect.py

@@ -0,0 +1,133 @@
+import streamlit as st
+import io
+import contextlib
+import sys
+import re
+
+
+class _Redirect:
+    """
+    Based on: https://gist.github.com/schaumb/037f139035d93cff3ad9f4f7e5f739ce
+    Also see: https://github.com/streamlit/streamlit/issues/268#issuecomment-810478208
+    """
+    class IOStuff(io.StringIO):
+        def __init__(self, trigger, max_buffer, buffer_separator, regex, dup=None):
+            super().__init__()
+            self._trigger = trigger
+            self._max_buffer = max_buffer
+            self._buffer_separator = buffer_separator
+            self._regex = regex and re.compile(regex)
+            self._dup = dup
+
+        def write(self, __s: str) -> int:
+            if self._max_buffer:
+                concatenated_len = super().tell() + len(__s)
+                if concatenated_len > self._max_buffer:
+                    rest = self.get_filtered_output()[concatenated_len - self._max_buffer:]
+                    if self._buffer_separator is not None:
+                        rest = rest.split(self._buffer_separator, 1)[-1]
+                    super().seek(0)
+                    super().write(rest)
+                    super().truncate(super().tell() + len(__s))
+            res = super().write(__s)
+            if self._dup is not None:
+                self._dup.write(__s)
+            self._trigger(self.get_filtered_output())
+            return res
+
+        def get_filtered_output(self):
+            if self._regex is None or self._buffer_separator is None:
+                return self.getvalue()
+
+            return self._buffer_separator.join(filter(self._regex.search, self.getvalue().split(self._buffer_separator)))
+
+        def print_at_end(self):
+            self._trigger(self.get_filtered_output())
+
+    def __init__(self, stdout=None, stderr=False, format=None, to=None, max_buffer=None, buffer_separator='\n',
+                 regex=None, duplicate_out=False):
+        self.io_args = {'trigger': self._write, 'max_buffer': max_buffer, 'buffer_separator': buffer_separator,
+                        'regex': regex}
+        self.redirections = []
+        self.st = None
+        self.stderr = stderr is True
+        self.stdout = stdout is True or (stdout is None and not self.stderr)
+        self.format = format or 'code'
+        self.to = to
+        self.fun = None
+        self.duplicate_out = duplicate_out or None
+        self.active_nested = None
+
+        if not self.stdout and not self.stderr:
+            raise ValueError("one of stdout or stderr must be True")
+
+        if self.format not in ['text', 'markdown', 'latex', 'code', 'write']:
+            raise ValueError(
+                f"format need oneof the following: {', '.join(['text', 'markdown', 'latex', 'code', 'write'])}")
+
+        if self.to and (not hasattr(self.to, 'text') or not hasattr(self.to, 'empty')):
+            raise ValueError(f"'to' is not a streamlit container object")
+
+    def __enter__(self):
+        if self.st is not None:
+            if self.to is None:
+                if self.active_nested is None:
+                    self.active_nested = self(format=self.format, max_buffer=self.io_args['max_buffer'],
+                                              buffer_separator=self.io_args['buffer_separator'],
+                                              regex=self.io_args['regex'], duplicate_out=self.duplicate_out)
+                return self.active_nested.__enter__()
+            else:
+                raise Exception("Already entered")
+        to = self.to or st
+
+        to.text(f"Redirected output from "
+                f"{'stdout and stderr' if self.stdout and self.stderr else 'stdout' if self.stdout else 'stderr'}"
+                f"{' [' + self.io_args['regex'] + ']' if self.io_args['regex'] else ''}"
+                f":")
+        self.st = to.empty()
+        self.fun = getattr(self.st, self.format)
+
+        io_obj = None
+
+        def redirect(to_duplicate):
+            nonlocal io_obj
+            io_obj = _Redirect.IOStuff(dup=self.duplicate_out and to_duplicate, **self.io_args)
+            redirection = contextlib.redirect_stdout(io_obj)
+            self.redirections.append((redirection, io_obj))
+            redirection.__enter__()
+
+        if self.stderr:
+            redirect(sys.stderr)
+        if self.stdout:
+            redirect(sys.stdout)
+
+        return io_obj
+
+    def __call__(self, to=None, format=None, max_buffer=None, buffer_separator='\n', regex=None, duplicate_out=False):
+        return _Redirect(self.stdout, self.stderr, format=format, to=to, max_buffer=max_buffer,
+                         buffer_separator=buffer_separator, regex=regex, duplicate_out=duplicate_out)
+
+    def __exit__(self, *exc):
+        if self.active_nested is not None:
+            nested = self.active_nested
+            if nested.active_nested is None:
+                self.active_nested = None
+            return nested.__exit__(*exc)
+
+        res = None
+        for redirection, io_obj in reversed(self.redirections):
+            res = redirection.__exit__(*exc)
+            io_obj.print_at_end()
+
+        self.redirections = []
+        self.st = None
+        self.fun = None
+        return res
+
+    def _write(self, data):
+        self.fun(data)
+
+
+stdout = _Redirect(max_buffer=1, buffer_separator="\n")
+stderr = _Redirect(stderr=True)
+stdouterr = _Redirect(stdout=True, stderr=True)