AjiNiktech/Document_search

1import streamlit as st2from langchain_openai import ChatOpenAI, OpenAIEmbeddings3import os4import dotenv5from langchain_community.document_loaders import TextLoader, PyPDFLoader, CSVLoader, UnstructuredPowerPointLoader, UnstructuredWordDocumentLoader, UnstructuredExcelLoader6from langchain_text_splitters import RecursiveCharacterTextSplitter7from langchain_chroma import Chroma8from langchain.chains.combine_documents import create_stuff_documents_chain9from langchain_core.prompts import ChatPromptTemplate, MessagesPlaceholder10from langchain_core.messages import HumanMessage, AIMessage11from langchain.memory import ConversationBufferMemory12import tempfile13 14# Set page config15st.set_page_config(page_title="Enterprise document search + chat", layout="wide")16 17# Streamlit app header18st.title("Enterprise document helpdesk")19 20# Initialize session state21if 'api_key_entered' not in st.session_state:22    st.session_state.api_key_entered = False23 24# Sidebar25with st.sidebar:26    st.header("Configuration")27    api_key = st.text_input("Enter your OpenAI API Key:", type="password")28    if api_key:29        os.environ["OPENAI_API_KEY"] = api_key30        st.session_state.api_key_entered = True31 32    if st.session_state.api_key_entered:33        st.header('Document Upload and Processing')34        uploaded_files = st.file_uploader('Upload your files', accept_multiple_files=True, type=['txt', 'pdf', 'csv', 'ppt', 'doc', 'xls', 'pptx', 'xlsx'])35 36        def load_file(file):37            file_extension = os.path.splitext(file.name)[1].lower()38            with tempfile.NamedTemporaryFile(delete=False, suffix=file_extension) as temp_file:39                temp_file.write(file.getvalue())40                temp_file_path = temp_file.name41        42            if file_extension == '.txt':43                loader = TextLoader(temp_file_path)44            elif file_extension == '.pdf':45                loader = PyPDFLoader(temp_file_path)46            elif file_extension == '.csv':47                loader = CSVLoader(temp_file_path)48            elif file_extension in ['.ppt', '.pptx']:49                loader = UnstructuredPowerPointLoader(temp_file_path)50            elif file_extension in ['.doc', '.docx']:51                loader = UnstructuredWordDocumentLoader(temp_file_path)52            elif file_extension in ['.xls', '.xlsx']:53                loader = UnstructuredExcelLoader(temp_file_path)54            else:55                os.unlink(temp_file_path)56                raise ValueError(f"Unsupported file type: {file_extension}")57        58            documents = loader.load()59            os.unlink(temp_file_path)60            return documents61 62        def summarize_documents(documents):63            chat = ChatOpenAI(model="gpt-3.5-turbo-1106", temperature=0.2)64            65            combined_text = " ".join([doc.page_content for doc in documents])66            67            prompt = f"""Summarize the following document in a concise manner, highlighting the key points:68 69            {combined_text}70 71            Summary:"""72            73            response = chat.invoke(prompt)74            return response.content75 76        # Process uploaded files77        if uploaded_files:78            if st.button("Process Documents"):79                with st.spinner("Processing documents..."):80                    all_documents = []81                    for file in uploaded_files:82                        all_documents.extend(load_file(file))83 84                    text_splitter = RecursiveCharacterTextSplitter(chunk_size=1000, chunk_overlap=200)85                    all_splits = text_splitter.split_documents(all_documents)86 87                    # Store processed documents in session state88                    st.session_state.processed_documents = all_splits89                st.success("Documents processed successfully!")90                91                # Add a button for summarization92                if st.button("Generate Summary"):93                    with st.spinner("Generating summary..."):94                        summary = summarize_documents(st.session_state.processed_documents)95                        st.session_state.document_summary = summary96                    st.success("Summary generated successfully!")97 98        # Display the summary if it exists99        if 'document_summary' in st.session_state:100            st.subheader("Document Summary")101            st.write(st.session_state.document_summary)102 103# Main app logic104if st.session_state.api_key_entered:105    # Initialize components106    @st.cache_resource107    def initialize_components():108        dotenv.load_dotenv()109        chat = ChatOpenAI(model="gpt-3.5-turbo-1106", temperature=0.2)110        embeddings = OpenAIEmbeddings(model="text-embedding-ada-002")111        return chat, embeddings112 113    # Load components114    chat, embeddings = initialize_components()115 116    # Create vectorstore and retriever only if documents are processed117    if 'processed_documents' in st.session_state:118        vectorstore = Chroma.from_documents(documents=st.session_state.processed_documents, embedding=embeddings)119        retriever = vectorstore.as_retriever(k=4)120 121        SYSTEM_TEMPLATE = """122        You are an advanced AI assistant designed for document search and chatbot functionality. Your primary functions are:123 124        1. Process and structure multiple documents in various formats, including:125            .txt, .pdf, .csv, .ppt, .doc, .xls, .pptx, and .xlsx126 127        2. Extract and organize information from these unstructured documents into a coherent, searchable format.128 129        3. Retrieve relevant information from the processed documents based on user queries.130 131        4. Act as a chatbot, engaging in conversations about the content of the documents.132 133        5. Provide accurate and contextual responses to user questions, drawing solely from the information contained within the processed documents.134 135        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.136 137        7. When answering, cite the specific document or section where the information was found, if possible.138 139        8. If there's ambiguity in a query, ask for clarification to ensure you provide the most relevant information.140 141        9. Maintain confidentiality and do not share or discuss information from one user's documents with other users.142 143        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.144 145        <context>146        {context}147        </context>148 149        Chat History:150        {chat_history}151        """152 153        question_answering_prompt = ChatPromptTemplate.from_messages(154            [155                (156                    "system",157                    SYSTEM_TEMPLATE,158                ),159                MessagesPlaceholder(variable_name="chat_history"),160                MessagesPlaceholder(variable_name="messages"),161            ]162        )163 164        document_chain = create_stuff_documents_chain(chat, question_answering_prompt)165 166        # Initialize memory for each session167        if "memory" not in st.session_state:168            st.session_state.memory = ConversationBufferMemory(memory_key="chat_history", return_messages=True)169 170        # Chat interface171        st.subheader("Chat with Assistant")172 173        # Initialize chat history174        if "messages" not in st.session_state:175            st.session_state.messages = []176 177        # Display chat messages from history on app rerun178        for message in st.session_state.messages:179            with st.chat_message(message["role"]):180                st.markdown(message["content"])181 182        # React to user input183        if prompt := st.chat_input("What would you like to know about Document?"):184            # Display user message in chat message container185            st.chat_message("user").markdown(prompt)186            # Add user message to chat history187            st.session_state.messages.append({"role": "user", "content": prompt})188 189            with st.chat_message("assistant"):190                message_placeholder = st.empty()191 192                # Retrieve relevant documents193                docs = retriever.get_relevant_documents(prompt)194 195                # Generate response196                response = document_chain.invoke(197                    {198                        "context": docs,199                        "chat_history": st.session_state.memory.load_memory_variables({})["chat_history"],200                        "messages": [201                            HumanMessage(content=prompt)202                        ],203                    }204                )205 206                # The response is already a string, so we can use it directly207                full_response = response208                message_placeholder.markdown(full_response)209 210            # Add assistant response to chat history211            st.session_state.messages.append({"role": "assistant", "content": full_response})212 213            # Update memory214            st.session_state.memory.save_context({"input": prompt}, {"output": full_response})215 216    else:217        st.info("Please upload and process documents to start chatting.")218 219else:220    st.info("Please enter your OpenAI API Key in the sidebar to start.")221 222# Add a footer223st.markdown("---")224st.markdown("By AI Planet")