Update app.py

app.py

@@ -1,16 +1,15 @@
 import streamlit as st
-from langchain_openai import ChatOpenAI
+from langchain_openai import ChatOpenAI, OpenAIEmbeddings
 import os
 import dotenv
-from langchain_community.document_loaders import WebBaseLoader
+from langchain_community.document_loaders import TextLoader, PyPDFLoader, CSVLoader, UnstructuredPowerPointLoader, UnstructuredWordDocumentLoader, UnstructuredExcelLoader
 from langchain_text_splitters import RecursiveCharacterTextSplitter
 from langchain_chroma import Chroma
-from langchain_openai import OpenAIEmbeddings
 from langchain.chains.combine_documents import create_stuff_documents_chain
 from langchain_core.prompts import ChatPromptTemplate, MessagesPlaceholder
 from langchain_core.messages import HumanMessage, AIMessage
 from langchain.memory import ConversationBufferMemory
-from langchain.document_loaders import PyPDFLoader
+import tempfile
 
 # Set page config
 st.set_page_config(page_title="Enterprise document search + chat", layout="wide")
@@ -27,153 +26,156 @@ with st.sidebar:
 
 # Main app logic
 if "OPENAI_API_KEY" in os.environ:
-    # Initialize components
-    @st.cache_resource
-    def initialize_components():
-        dotenv.load_dotenv()
-        chat = ChatOpenAI(model="gpt-3.5-turbo-1106", temperature=0.2)
-
-        # #loader1 = WebBaseLoader("https://www.tbankltd.com/")
-        # loader1 = PyPDFLoader("Tbank resources.pdf")
-        # loader2 = PyPDFLoader("International Banking Services.pdf")
-        # data1 = loader1.load()
-        # data2 = loader2.load()
-        # data = data1 + data2
-        st.header('Multiple File Upload')
-        uploaded_files = st.file_uploader('Upload your files',accept_multiple_files=True, type=['txt', 'pdf', 'csv', 'ppt', 'doc', 'xls', 'pptx', 'xlsx'])  
-        def load_file(file):
-            file_extension = os.path.splitext(file.name)[1].lower()
-            with tempfile.NamedTemporaryFile(delete=False, suffix=file_extension) as temp_file:
-                temp_file.write(file.getvalue())
-                temp_file_path = temp_file.name
-        
-            if file_extension == '.txt':
-                loader = TextLoader(temp_file_path)
-            elif file_extension == '.pdf':
-                loader = PyPDFLoader(temp_file_path)
-            elif file_extension == '.csv':
-                loader = CSVLoader(temp_file_path)
-            elif file_extension in ['.ppt', '.pptx']:
-                loader = UnstructuredPowerPointLoader(temp_file_path)
-            elif file_extension in ['.doc', '.docx']:
-                loader = UnstructuredWordDocumentLoader(temp_file_path)
-            elif file_extension in ['.xls', '.xlsx']:
-                loader = UnstructuredExcelLoader(temp_file_path)
-            else:
-                os.unlink(temp_file_path)
-                raise ValueError(f"Unsupported file type: {file_extension}")
-        
-            documents = loader.load()
+    st.header('Multiple File Upload')
+    uploaded_files = st.file_uploader('Upload your files', accept_multiple_files=True, type=['txt', 'pdf', 'csv', 'ppt', 'doc', 'xls', 'pptx', 'xlsx'])
+
+    def load_file(file):
+        file_extension = os.path.splitext(file.name)[1].lower()
+        with tempfile.NamedTemporaryFile(delete=False, suffix=file_extension) as temp_file:
+            temp_file.write(file.getvalue())
+            temp_file_path = temp_file.name
+
+        if file_extension == '.txt':
+            loader = TextLoader(temp_file_path)
+        elif file_extension == '.pdf':
+            loader = PyPDFLoader(temp_file_path)
+        elif file_extension == '.csv':
+            loader = CSVLoader(temp_file_path)
+        elif file_extension in ['.ppt', '.pptx']:
+            loader = UnstructuredPowerPointLoader(temp_file_path)
+        elif file_extension in ['.doc', '.docx']:
+            loader = UnstructuredWordDocumentLoader(temp_file_path)
+        elif file_extension in ['.xls', '.xlsx']:
+            loader = UnstructuredExcelLoader(temp_file_path)
+        else:
             os.unlink(temp_file_path)
-            return documents
+            raise ValueError(f"Unsupported file type: {file_extension}")
+
+        documents = loader.load()
+        os.unlink(temp_file_path)
+        return documents
+
+    # Process uploaded files
+    if uploaded_files:
+        all_documents = []
+        for file in uploaded_files:
+            all_documents.extend(load_file(file))
+
         text_splitter = RecursiveCharacterTextSplitter(chunk_size=1000, chunk_overlap=200)
-        all_splits = text_splitter.split_documents(documents)
-        embeddings = OpenAIEmbeddings(model="text-embedding-ada-002")
-        vectorstore = Chroma.from_documents(documents=all_splits, embedding=embeddings)
-        retriever = vectorstore.as_retriever(k=4)
+        all_splits = text_splitter.split_documents(all_documents)
 
-        SYSTEM_TEMPLATE = """
-        You are an advanced AI assistant designed for document search and chatbot functionality. Your primary functions are:
+        # Initialize components
+        @st.cache_resource
+        def initialize_components(_all_splits):
+            dotenv.load_dotenv()
+            chat = ChatOpenAI(model="gpt-3.5-turbo-1106", temperature=0.2)
+            embeddings = OpenAIEmbeddings(model="text-embedding-ada-002")
+            vectorstore = Chroma.from_documents(documents=_all_splits, embedding=embeddings)
+            retriever = vectorstore.as_retriever(k=4)
 
-        1. Process and structure multiple documents in various formats, including:
-            .txt, .pdf, .csv, .ppt, .doc, .xls, .pptx, and .xlsx
+            SYSTEM_TEMPLATE = """
+            You are an advanced AI assistant designed for document search and chatbot functionality. Your primary functions are:
 
-        2. Extract and organize information from these unstructured documents into a coherent, searchable format.
+            1. Process and structure multiple documents in various formats, including:
+                .txt, .pdf, .csv, .ppt, .doc, .xls, .pptx, and .xlsx
 
-        3. Retrieve relevant information from the processed documents based on user queries.
+            2. Extract and organize information from these unstructured documents into a coherent, searchable format.
 
-        4. Act as a chatbot, engaging in conversations about the content of the documents.
+            3. Retrieve relevant information from the processed documents based on user queries.
 
-        5. Provide accurate and contextual responses to user questions, drawing solely from the information contained within the processed documents.
+            4. Act as a chatbot, engaging in conversations about the content of the documents.
 
-        6. If a user's question is not related to the content of the provided documents, politely inform them that you can only answer questions based on the information in the given documents.
+            5. Provide accurate and contextual responses to user questions, drawing solely from the information contained within the processed documents.
 
-        7. When answering, cite the specific document or section where the information was found, if possible.
+            6. If a user's question is not related to the content of the provided documents, politely inform them that you can only answer questions based on the information in the given documents.
 
-        8. If there's ambiguity in a query, ask for clarification to ensure you provide the most relevant information.
+            7. When answering, cite the specific document or section where the information was found, if possible.
 
-        9. Maintain confidentiality and do not share or discuss information from one user's documents with other users.
+            8. If there's ambiguity in a query, ask for clarification to ensure you provide the most relevant information.
 
-        Remember, your knowledge is limited to the content of the documents you've been given to process. Do not provide information or answer questions that are outside the scope of these documents. Always strive for accuracy and relevance in your responses.
+            9. Maintain confidentiality and do not share or discuss information from one user's documents with other users.
 
-        <context>
-        {context}
-        </context>
+            Remember, your knowledge is limited to the content of the documents you've been given to process. Do not provide information or answer questions that are outside the scope of these documents. Always strive for accuracy and relevance in your responses.
 
-        Chat History:
-        {chat_history}
-        """
+            <context>
+            {context}
+            </context>
 
-        question_answering_prompt = ChatPromptTemplate.from_messages(
-            [
-                (
-                    "system",
-                    SYSTEM_TEMPLATE,
-                ),
-                MessagesPlaceholder(variable_name="chat_history"),
-                MessagesPlaceholder(variable_name="messages"),
-            ]
-        )
+            Chat History:
+            {chat_history}
+            """
 
-        document_chain = create_stuff_documents_chain(chat, question_answering_prompt)
+            question_answering_prompt = ChatPromptTemplate.from_messages(
+                [
+                    (
+                        "system",
+                        SYSTEM_TEMPLATE,
+                    ),
+                    MessagesPlaceholder(variable_name="chat_history"),
+                    MessagesPlaceholder(variable_name="messages"),
+                ]
+            )
 
-        return retriever, document_chain
-        # else:
-        #   st.warning("Please Upload File to Continue")
+            document_chain = create_stuff_documents_chain(chat, question_answering_prompt)
 
-    # Load components
-    with st.spinner("Initializing Assistant..."):
-        retriever, document_chain = initialize_components()
+            return retriever, document_chain
 
-    # Initialize memory for each session
-    if "memory" not in st.session_state:
-        st.session_state.memory = ConversationBufferMemory(memory_key="chat_history", return_messages=True)
+        # Load components
+        with st.spinner("Initializing Assistant..."):
+            retriever, document_chain = initialize_components(all_splits)
 
-    # Chat interface
-    st.subheader("Chat with Assistant")
+        # Initialize memory for each session
+        if "memory" not in st.session_state:
+            st.session_state.memory = ConversationBufferMemory(memory_key="chat_history", return_messages=True)
 
-    # Initialize chat history
-    if "messages" not in st.session_state:
-        st.session_state.messages = []
+        # Chat interface
+        st.subheader("Chat with Assistant")
 
-    # Display chat messages from history on app rerun
-    for message in st.session_state.messages:
-        with st.chat_message(message["role"]):
-            st.markdown(message["content"])
+        # Initialize chat history
+        if "messages" not in st.session_state:
+            st.session_state.messages = []
 
-    # React to user input
-    if prompt := st.chat_input("What would you like to know about Document?"):
-        # Display user message in chat message container
-        st.chat_message("user").markdown(prompt)
-        # Add user message to chat history
-        st.session_state.messages.append({"role": "user", "content": prompt})
+        # Display chat messages from history on app rerun
+        for message in st.session_state.messages:
+            with st.chat_message(message["role"]):
+                st.markdown(message["content"])
 
-        with st.chat_message("assistant"):
-            message_placeholder = st.empty()
+        # React to user input
+        if prompt := st.chat_input("What would you like to know about Document?"):
+            # Display user message in chat message container
+            st.chat_message("user").markdown(prompt)
+            # Add user message to chat history
+            st.session_state.messages.append({"role": "user", "content": prompt})
 
-            # Retrieve relevant documents
-            docs = retriever.get_relevant_documents(prompt)
+            with st.chat_message("assistant"):
+                message_placeholder = st.empty()
 
-            # Generate response
-            response = document_chain.invoke(
-                {
-                    "context": docs,
-                    "chat_history": st.session_state.memory.load_memory_variables({})["chat_history"],
-                    "messages": [
-                        HumanMessage(content=prompt)
-                    ],
-                }
-            )
+                # Retrieve relevant documents
+                docs = retriever.get_relevant_documents(prompt)
+
+                # Generate response
+                response = document_chain.invoke(
+                    {
+                        "context": docs,
+                        "chat_history": st.session_state.memory.load_memory_variables({})["chat_history"],
+                        "messages": [
+                            HumanMessage(content=prompt)
+                        ],
+                    }
+                )
+
+                # The response is already a string, so we can use it directly
+                full_response = response
+                message_placeholder.markdown(full_response)
 
-            # The response is already a string, so we can use it directly
-            full_response = response
-            message_placeholder.markdown(full_response)
+            # Add assistant response to chat history
+            st.session_state.messages.append({"role": "assistant", "content": full_response})
 
-        # Add assistant response to chat history
-        st.session_state.messages.append({"role": "assistant", "content": full_response})
+            # Update memory
+            st.session_state.memory.save_context({"input": prompt}, {"output": full_response})
 
-        # Update memory
-        st.session_state.memory.save_context({"input": prompt}, {"output": full_response})
+    else:
+        st.warning("Please upload files to continue.")
 
 else:
     st.warning("Please enter your OpenAI API Key in the sidebar to start the chatbot.")