Update app

This reverts commit 48881b203e602383d6a33df91bec1cf5aeae0a90.

app.py

@@ -2,48 +2,52 @@ 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
 
-from chat_profile import ChatProfileRoleEnum
-from stream_handler import PrintRetrievalHandler, StreamHandler
+# TODO: refactor
+# TODO: extract class
+# TODO: modularize
+# TODO: hide side bar
+# TODO: make the page attactive
 
-# Configuration
+# configs
 LLM_MODEL_NAME = "gpt-3.5-turbo"
 EMBEDDING_MODEL_NAME = "all-MiniLM-L6-v2"
 
-
-# Set up Streamlit app
-def setup_streamlit_app():
-    st.set_page_config(
-        page_title=":books: InkChatGPT: Chat with Documents",
-        page_icon="📚",
-        initial_sidebar_state="collapsed",
-        menu_items={
-            "Get Help": "https://x.com/vinhnx",
-            "Report a bug": "https://github.com/vinhnx/InkChatGPT/issues",
-            "About": "InkChatGPT is a Streamlit application that allows users to upload PDF documents and engage in a conversational Q&A with a language model (LLM) based on the content of those documents.",
-        },
-    )
-
-    st.image("./assets/icon.jpg", width=100)
-    st.header(":gray[:books: InkChatGPT]", divider="blue")
-    st.write("**Chat** with Documents")
-
+st.set_page_config(
+    page_title=":books: InkChatGPT: Chat with Documents",
+    page_icon="📚",
+    initial_sidebar_state="collapsed",
+    menu_items={
+        "Get Help": "https://x.com/vinhnx",
+        "Report a bug": "https://github.com/vinhnx/InkChatGPT/issues",
+        "About": "InkChatGPT is a Streamlit application that allows users to upload PDF documents and engage in a conversational Q&A with a language model (LLM) based on the content of those documents.",
+    },
+)
+
+st.image("./assets/icon.jpg", width=100)
+st.header(
+    ":gray[:books: InkChatGPT]",
+    divider="blue",
+)
+st.write("**Chat** with Documents")
 
 # Setup memory for contextual conversation
 msgs = StreamlitChatMessageHistory()
 
 
-# Load and process documents
+@st.cache_resource(ttl="1h")
 def configure_retriever(uploaded_files):
     # Read documents
     docs = []
@@ -66,6 +70,7 @@ def configure_retriever(uploaded_files):
             st.write("This document format is not supported!")
             return None
 
+        # loader = PyPDFLoader(temp_filepath)
         docs.extend(loader.load())
 
     # Split documents
@@ -84,68 +89,91 @@ def configure_retriever(uploaded_files):
     return retriever
 
 
-# Main function
-def main():
-    setup_streamlit_app()
+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")
 
-    with st.sidebar.expander("Documents"):
-        st.subheader("Files")
-        uploaded_files = st.file_uploader(
-            label="Select files",
-            type=["pdf", "txt", "docx"],
-            accept_multiple_files=True,
-        )
+    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)
 
-    with st.sidebar.expander("Setup"):
-        st.subheader("API Key")
-        openai_api_key = st.text_input("OpenAI API Key", type="password")
 
-        is_empty_chat_messages = len(msgs.messages) == 0
-        if is_empty_chat_messages or st.button("Clear message history"):
-            msgs.clear()
-            msgs.add_ai_message("How can I help you?")
+class PrintRetrievalHandler(BaseCallbackHandler):
+    def __init__(self, container):
+        self.status = container.status("**Thinking...**")
+        self.container = container
 
-    if not openai_api_key:
-        st.info("Please add your OpenAI API key in the sidebar to continue.")
-        st.stop()
+    def on_retriever_start(self, serialized: dict, query: str, **kwargs):
+        self.status.write(f"**Checking document for query:** `{query}`. Please wait...")
 
-    if uploaded_files:
-        retriever = configure_retriever(uploaded_files)
+    def on_retriever_end(self, documents, **kwargs):
+        self.container.empty()
 
-        memory = ConversationBufferMemory(
-            memory_key="chat_history", chat_memory=msgs, return_messages=True
-        )
 
-        # Setup LLM and QA chain
-        llm = ChatOpenAI(
-            model_name=LLM_MODEL_NAME,
-            openai_api_key=openai_api_key,
-            temperature=0,
-            streaming=True,
-        )
+with st.sidebar.expander("Documents"):
+    st.subheader("Files")
+    uploaded_files = st.file_uploader(
+        label="Select files", type=["pdf", "txt", "docx"], accept_multiple_files=True
+    )
+
+
+with st.sidebar.expander("Setup"):
+    st.subheader("API Key")
+    openai_api_key = st.text_input("OpenAI API Key", type="password")
 
-        chain = ConversationalRetrievalChain.from_llm(
-            llm, retriever=retriever, memory=memory, verbose=False
-        )
+    is_empty_chat_messages = len(msgs.messages) == 0
+    if is_empty_chat_messages or st.button("Clear message history"):
+        msgs.clear()
+        msgs.add_ai_message("How can I help you?")
 
-        avatars = {
-            ChatProfileRoleEnum.Human: "user",
-            ChatProfileRoleEnum.AI: "assistant",
-        }
+if not openai_api_key:
+    st.info("Please add your OpenAI API key in the sidebar to continue.")
+    st.stop()
 
-        for msg in msgs.messages:
-            st.chat_message(avatars[msg.type]).write(msg.content)
+if uploaded_files:
+    retriever = configure_retriever(uploaded_files)
+
+    memory = ConversationBufferMemory(
+        memory_key="chat_history", chat_memory=msgs, return_messages=True
+    )
+
+    # Setup LLM and QA chain
+    llm = ChatOpenAI(
+        model_name=LLM_MODEL_NAME,
+        openai_api_key=openai_api_key,
+        temperature=0,
+        streaming=True,
+    )
+
+    chain = ConversationalRetrievalChain.from_llm(
+        llm, retriever=retriever, memory=memory, verbose=False
+    )
 
-        if user_query := st.chat_input(placeholder="Ask me anything!"):
-            st.chat_message("user").write(user_query)
+    avatars = {
+        ChatProfileRoleEnum.Human: "user",
+        ChatProfileRoleEnum.AI: "assistant",
+    }
 
-            with st.chat_message("assistant"):
-                retrieval_handler = PrintRetrievalHandler(st.empty())
-                stream_handler = StreamHandler(st.empty())
-                response = chain.run(
-                    user_query, callbacks=[retrieval_handler, stream_handler]
-                )
+    for msg in msgs.messages:
+        st.chat_message(avatars[msg.type]).write(msg.content)
 
+    if user_query := st.chat_input(placeholder="Ask me anything!"):
+        st.chat_message("user").write(user_query)
 
-if __name__ == "__main__":
-    main()
+        with st.chat_message("assistant"):
+            retrieval_handler = PrintRetrievalHandler(st.empty())
+            stream_handler = StreamHandler(st.empty())
+            response = chain.run(
+                user_query, callbacks=[retrieval_handler, stream_handler]
+            )

stream_handler.py

@@ -1,37 +0,0 @@
-import os
-import streamlit as st
-
-from langchain.callbacks.base import BaseCallbackHandler
-
-
-# Callback handlers
-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()

token_stream_handler.py

@@ -0,0 +1,13 @@
+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)