Merge pull request #1 from andreped/dev

Added OpenAI key to secrets; improved verbose; added Redirect class for stdout logging

.gitignore

@@ -2,3 +2,5 @@ venv/
 data/
 .DS_Store
 config.json
+.streamlit/
+secrets.toml

app.py

@@ -1,9 +1,12 @@
-import json
+import os
 
 import streamlit as st
 
-from chatbot.utils import download_test_data
-from chatbot.utils import load_data
+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"]
 
 # Initialize message history
 st.header("Chat with André's research 💬 📚")
@@ -11,15 +14,17 @@ st.header("Chat with André's research 💬 📚")
 if "messages" not in st.session_state.keys():  # Initialize the chat message history
     st.session_state.messages = [{"role": "assistant", "content": "Ask me a question about André's research!"}]
 
-# Load config values
-with open(r"config.json") as config_file:
-    config_details = json.load(config_file)
-
 
 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(config_details)
+    index = load_data()
     chat_engine = index.as_chat_engine(chat_mode="condense_question", verbose=True)
 
     if prompt := st.chat_input("Your question"):  # Prompt for user input and save to chat history

chatbot/{utils.py → data.py}

@@ -14,37 +14,46 @@ from llama_index.llms import AzureOpenAI
 def download_test_data():
     # url = f"https://drive.google.com/drive/folders/uc?export=download&confirm=pbef&id={file_id}"
     url = "https://drive.google.com/drive/folders/1uDSAWtLvp1YPzfXUsK_v6DeWta16pq6y"
-    with st.spinner(text="Downloading test data. Might take a few seconds."):
+    with st.spinner(text="Downloading test data. This might take a minute."):
+        # @TODO: replace gown solution with a custom solution compatible with GitHub and
+        # use st.progress to get more verbose during download
         download_folder(url=url, quiet=False, use_cookies=False, output="./data/")
 
 
 @st.cache_resource(show_spinner=False)
-def load_data(config_details):
+def load_data():
     with st.spinner(text="Loading and indexing the provided dataset – hang tight! This may take a few seconds."):
         documents = SimpleDirectoryReader(input_dir="./data", recursive=True).load_data()
+
+    with st.spinner(text="Setting up Azure OpenAI..."):
         llm = AzureOpenAI(
             model="gpt-3.5-turbo",
-            engine=config_details["ENGINE"],
+            engine=st.secrets["ENGINE"],
             temperature=0.5,
-            api_key=os.getenv("OPENAI_API_KEY"),
-            api_base=config_details["OPENAI_API_BASE"],
+            api_key=os.environ["OPENAI_API_KEY"],
+            api_base=st.secrets["OPENAI_API_BASE"],
             api_type="azure",
-            api_version=config_details["OPENAI_API_VERSION"],
+            api_version=st.secrets["OPENAI_API_VERSION"],
             system_prompt="You are an expert on André's research and your job is to answer"
             "technical questions. Assume that all questions are related to"
-            "André's research. Keep your answers technical and based on facts"
-            " – do not hallucinate features.",
+            "André's research. Keep your answers technical and based on facts;"
+            " do not hallucinate features.",
         )
+
+    with st.spinner(text="Setting up OpenAI Embedding..."):
         # You need to deploy your own embedding model as well as your own chat completion model
         embed_model = OpenAIEmbedding(
             model="text-embedding-ada-002",
-            deployment_name=config_details["ENGINE_EMBEDDING"],
-            api_key=os.getenv("OPENAI_API_KEY"),
-            api_base=config_details["OPENAI_API_BASE"],
+            deployment_name=st.secrets["ENGINE_EMBEDDING"],
+            api_key=os.environ["OPENAI_API_KEY"],
+            api_base=st.secrets["OPENAI_API_BASE"],
             api_type="azure",
-            api_version=config_details["OPENAI_API_VERSION"],
+            api_version=st.secrets["OPENAI_API_VERSION"],
+            embed_batch_size=10,  # set to low value to reduce rate limit -> may degrade response runtime
         )
-        service_context = ServiceContext.from_defaults(llm=llm, embed_model=embed_model)
+
+    with st.spinner(text="Setting up Vector Store Index..."):
+        service_context = ServiceContext.from_defaults(llm=llm, embed_model=embed_model)  # , chunk_size=512)
         set_global_service_context(service_context)
         index = VectorStoreIndex.from_documents(documents)  # , service_context=service_context)
         return index

chatbot/redirect.py

@@ -0,0 +1,165 @@
+import contextlib
+import io
+import re
+import sys
+
+import streamlit as st
+
+
+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: {self.to}")
+
+    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)