Refactor

app.py

@@ -2,24 +2,18 @@ import os
 import tempfile
 
 import streamlit as st
-from chat_profile import ChatProfileRoleEnum
-
 from langchain.callbacks.base import BaseCallbackHandler
 from langchain.chains import ConversationalRetrievalChain
 from langchain.chat_models import ChatOpenAI
-from langchain_community.document_loaders import Docx2txtLoader, PyPDFLoader, TextLoader
 from langchain.embeddings import HuggingFaceEmbeddings
 from langchain.memory import ConversationBufferMemory
 from langchain.memory.chat_message_histories import StreamlitChatMessageHistory
 from langchain.text_splitter import RecursiveCharacterTextSplitter
+from langchain_community.document_loaders import Docx2txtLoader, PyPDFLoader, TextLoader
 from langchain_community.vectorstores import DocArrayInMemorySearch
-from streamlit_extras.add_vertical_space import add_vertical_space
 
-# TODO: refactor
-# TODO: extract class
-# TODO: modularize
-# TODO: hide side bar
-# TODO: make the page attactive
+from chat_profile import ChatProfileRoleEnum
+from calback_handler import StreamHandler, PrintRetrievalHandler
 
 # configs
 LLM_MODEL_NAME = "gpt-3.5-turbo"
@@ -89,38 +83,6 @@ def configure_retriever(uploaded_files):
     return retriever
 
 
-class StreamHandler(BaseCallbackHandler):
-    def __init__(
-        self, container: st.delta_generator.DeltaGenerator, initial_text: str = ""
-    ):
-        self.container = container
-        self.text = initial_text
-        self.run_id_ignore_token = None
-
-    def on_llm_start(self, serialized: dict, prompts: list, **kwargs):
-        # Workaround to prevent showing the rephrased question as output
-        if prompts[0].startswith("Human"):
-            self.run_id_ignore_token = kwargs.get("run_id")
-
-    def on_llm_new_token(self, token: str, **kwargs) -> None:
-        if self.run_id_ignore_token == kwargs.get("run_id", False):
-            return
-        self.text += token
-        self.container.markdown(self.text)
-
-
-class PrintRetrievalHandler(BaseCallbackHandler):
-    def __init__(self, container):
-        self.status = container.status("**Thinking...**")
-        self.container = container
-
-    def on_retriever_start(self, serialized: dict, query: str, **kwargs):
-        self.status.write(f"**Checking document for query:** `{query}`. Please wait...")
-
-    def on_retriever_end(self, documents, **kwargs):
-        self.container.empty()
-
-
 with st.sidebar.expander("Documents"):
     st.subheader("Files")
     uploaded_files = st.file_uploader(

calback_handler.py

@@ -0,0 +1,31 @@
+from langchain.callbacks.base import BaseCallbackHandler
+
+
+class StreamHandler(BaseCallbackHandler):
+    def __init__(self, container, initial_text: str = ""):
+        self.container = container
+        self.text = initial_text
+        self.run_id_ignore_token = None
+
+    def on_llm_start(self, serialized: dict, prompts: list, **kwargs):
+        # Workaround to prevent showing the rephrased question as output
+        if prompts[0].startswith("Human"):
+            self.run_id_ignore_token = kwargs.get("run_id")
+
+    def on_llm_new_token(self, token: str, **kwargs) -> None:
+        if self.run_id_ignore_token == kwargs.get("run_id", False):
+            return
+        self.text += token
+        self.container.markdown(self.text)
+
+
+class PrintRetrievalHandler(BaseCallbackHandler):
+    def __init__(self, container):
+        self.status = container.status("**Thinking...**")
+        self.container = container
+
+    def on_retriever_start(self, serialized: dict, query: str, **kwargs):
+        self.status.write(f"**Checking document for query:** `{query}`. Please wait...")
+
+    def on_retriever_end(self, documents, **kwargs):
+        self.container.empty()

requirements.txt

@@ -5,4 +5,5 @@ langchain
 streamlit
 streamlit_chat
 streamlit-extras
-pypdf
+pypdf
+docx2txt

token_stream_handler.py

@@ -1,13 +0,0 @@
-import os
-
-from langchain.callbacks.base import BaseCallbackHandler
-
-
-class StreamHandler(BaseCallbackHandler):
-    def __init__(self, container, initial_text=""):
-        self.container = container
-        self.text = initial_text
-
-    def on_llm_new_token(self, token: str, **kwargs) -> None:
-        self.text += token
-        self.container.markdown(self.text)