Initial Push of small files

app.py

@@ -0,0 +1,244 @@
+import streamlit as st
+from pypdf import PdfReader
+# import replicate
+import os
+from pathlib import Path
+from dotenv import load_dotenv
+import pickle
+import timeit
+from PIL import Image
+import datetime
+import base64
+
+from langchain.embeddings import HuggingFaceEmbeddings
+from langchain.vectorstores import FAISS
+from langchain.document_loaders import PyPDFLoader
+from langchain.text_splitter import RecursiveCharacterTextSplitter
+from langchain.document_loaders import PyPDFLoader, DirectoryLoader
+from langchain.memory import ConversationBufferMemory
+from langchain.chains import ConversationalRetrievalChain
+from langchain.prompts.prompt import PromptTemplate
+from langchain.llms import LlamaCpp
+from langchain.callbacks.manager import CallbackManager
+from langchain.callbacks.streaming_stdout import StreamingStdOutCallbackHandler
+from langchain.vectorstores import Chroma
+from langchain.document_loaders import PyPDFDirectoryLoader
+from langchain.retrievers import BM25Retriever, EnsembleRetriever
+from langchain.chat_models import ChatOpenAI
+from langchain.agents.agent_toolkits import create_retriever_tool
+from langchain.agents.agent_toolkits import create_conversational_retrieval_agent
+from langchain.utilities import SerpAPIWrapper
+
+from utils import build_embedding_model, build_llm
+from utils import load_ensemble_retriver, load_text_chunks, load_vectorstore, load_conversational_retrievel_agent
+
+load_dotenv()
+# Getting current timestamp to keep track of historical conversations
+current_timestamp = datetime.datetime.now()
+timestamp_string = current_timestamp.strftime("%Y-%m-%d %H:%M:%S")
+
+#Directories path
+persist_directory= "Database/PDF_HTML_CHROMA_DB"
+all_docs_pkl_directory= 'Database/text_chunks_html_pdf.pkl'
+
+# Initliazing sesstion states in Streamlit to cache different stuffs like model iniitialization and there by avoid re-running of alredy initialized stuffs over and again.
+if "llm" not in st.session_state:
+    st.session_state["llm"] = build_llm()
+
+if "embeddings" not in st.session_state:
+    st.session_state["embeddings"] = build_embedding_model()
+    
+if "vector_db" not in st.session_state:
+    st.session_state["vector_db"] = load_vectorstore(persist_directory=persist_directory, embeddings=st.session_state["embeddings"])
+ 
+if "text_chunks" not in st.session_state:
+    st.session_state["text_chunks"] = load_text_chunks(text_chunks_pkl_dir=all_docs_pkl_directory)
+
+if "ensemble_retriver" not in st.session_state:
+    st.session_state["ensemble_retriver"] = load_ensemble_retriver(text_chunks=st.session_state["text_chunks"], embeddings=st.session_state["embeddings"], chroma_vectorstore=st.session_state["vector_db"] )
+
+if "agent_executor" not in st.session_state:
+    st.session_state["agent_executor"] = load_conversational_retrievel_agent(retriever=st.session_state["ensemble_retriver"], llm=st.session_state["llm"])    
+
+
+
+# App title
+st.set_page_config(
+    page_title="OMP Search Bot",
+    layout="wide",
+    initial_sidebar_state="expanded",
+)
+
+st.markdown("""
+        <style>
+               .block-container {
+                    padding-top: 2.2rem}
+        </style>
+        """, unsafe_allow_html=True)
+# To get header in the App
+col1, col2= st.columns(2)
+
+title1 = """
+<p style="font-size: 26px;text-align: right; color: #0C3453; font-weight: bold">OPM Retirement Services Assistant</p>
+"""
+
+def clear_chat_history():
+        st.session_state.messages = [{"role": "assistant", "content": "How may I assist you today?"}]
+
+file_ = open("opm_logo.png", "rb")
+contents = file_.read()
+data_url = base64.b64encode(contents).decode("utf-8")
+file_.close()
+
+st.markdown(
+    f"""
+    <div style="background-color: white; padding: 15px; border-radius: 10px;">
+        <div style="display: flex; justify-content: space-between;">
+            <div>
+                <img src="data:image/png;base64,{data_url}" style="max-width: 100%;" alt="OPM Logo" />
+            </div>
+            <div style="flex: 1; padding: 15px;">
+                {title1}
+    """,
+    unsafe_allow_html=True
+)
+st.write("")
+
+
+st.write('<p style="color: #B0B0B0;margin: 0;">OPM is here to help you transition from serving the American people to enjoying your retirement. This retirement services assistant shows our commitment to supporting new and existing retirees throughout the retirement journey.  Our assistant is trained on 1500+ documents related to OPM retirement services and can answer your questions in conversational style.  Just ask away..</p>', unsafe_allow_html=True)
+
+st.markdown("""---""")
+
+text_html = """
+    <p style="font-size: 14px; text-align: center; color: #727477; margin: 0;">
+        Type your question in conversational style
+    </p>
+    <p style="font-size: 14px; text-align: center; color: #727477; margin: 0;">
+        Example: What are interim benefits?
+    </p>
+"""
+
+st.write(text_html, unsafe_allow_html=True)
+
+
+with st.sidebar:
+    st.subheader("")       
+        
+if st.session_state["vector_db"] and st.session_state["llm"]:
+        # Store LLM generated responses
+    if "messages" not in st.session_state.keys():
+        st.session_state.messages = [{"role": "assistant", "content": "How may I assist you today?", "Source":""}]
+
+    # Display or clear chat messages
+    for message in st.session_state.messages:
+        with st.chat_message(message["role"]):
+            st.write(message["content"])
+            if message["Source"]=="":
+                st.write("")
+            else:
+                with st.expander("source"):
+                    for idx, item in enumerate(message["Source"]):
+                        st.markdown(item["Page"])
+                        st.markdown(item["Source"])
+                        st.markdown(item["page_content"])
+                        st.write("---")
+
+
+    # Initialize the session state to store chat history
+    if "stored_session" not in st.session_state:
+        st.session_state["stored_session"] = []
+
+    # Create a list to store expanders
+    if "expanders" not in st.session_state:
+        st.session_state["expanders"] = []
+    
+        # Define a function to add a new chat expander
+    def add_chat_expander(chat_history):
+        current_timestamp = datetime.datetime.now()
+        timestamp_string = current_timestamp.strftime("%Y-%m-%d %H:%M:%S")
+        st.session_state["expanders"].append({"timestamp": timestamp_string, "chat_history": chat_history})
+                
+    def clear_chat_history():
+        """
+        To remove existing chat history and start new conversation
+        """
+        stored_session = []
+        for dict_message in st.session_state.messages:
+            if dict_message["role"] == "user":
+                string_dialogue = "User: " + dict_message["content"] + "\n\n"
+                st.session_state["stored_session"].append(string_dialogue)
+
+            else:
+                string_dialogue = "Assistant: " + dict_message["content"] + "\n\n"
+                st.session_state["stored_session"].append(string_dialogue)
+            stored_session.append(string_dialogue)
+        
+            # Add a new chat expander
+        add_chat_expander(stored_session)
+        st.session_state.messages = [{"role": "assistant", "content": "How may I assist you today?", "Source":""}]
+        
+    st.sidebar.button('New chat', on_click=clear_chat_history, use_container_width=True)
+    st.sidebar.text("")
+    st.sidebar.write('<p style="font-size: 16px;text-align: center; color: #727477; font-weight: bold">Chat history</p>', unsafe_allow_html=True)
+    # Display existing chat expanders
+    for expander_info in st.session_state["expanders"]:
+        with st.sidebar.expander("Conversation ended at:"+"\n\n"+expander_info["timestamp"]):
+            for message in expander_info["chat_history"]:
+                if message.startswith("User:"):
+                    st.write(f'<span style="color: #EF6A6A;">{message}</span>', unsafe_allow_html=True)
+                elif message.startswith("Assistant:"):
+                    st.write(f'<span style="color: #F7BD45;">{message}</span>', unsafe_allow_html=True)
+                else:
+                    st.write(message)
+
+
+    def generate_llm_response(agent_executor, prompt_input):
+        result = agent_executor({"input": prompt_input})
+        return [result['output'], result['intermediate_steps']]
+
+
+    # User-provided prompt
+    if prompt := st.chat_input(disabled= not st.session_state["vector_db"]):
+        st.session_state.messages.append({"role": "user", "content": prompt, "Source":""})
+        with st.chat_message("user"):
+            st.write(prompt)
+
+    # Generate a new response if last message is not from assistant
+    if st.session_state.messages[-1]["role"] != "assistant":
+        with st.chat_message("assistant"):
+            with st.spinner("Searching..."):
+                start = timeit.default_timer()
+                response = generate_llm_response(agent_executor=st.session_state["agent_executor"], prompt_input=prompt)
+                placeholder = st.empty()
+                full_response = ''
+                for item in response[0]:
+                    full_response += item
+                placeholder.markdown(full_response)
+                # The following logic will work in the way given below.
+                # -- Check if intermediary steps are present in the output of the given prompt.
+                # -- If not, we can conclude that, agent has used internet search as tool.
+                # -- Check if intermediary steps are present in the output of the prompt.
+                # -- If intermediary steps are present, it means agent has used exising custom knowledge base for iformation retrival and therefore we need to give souce docs as output along with LLM's reponse.
+                if len(response[1])>0:
+                    st.text("-------------------------------------")
+                    docs= st.session_state["ensemble_retriver"].get_relevant_documents(prompt)   
+                    source_doc_list= []  
+                    for doc in docs:
+                        source_doc_list.append(doc.dict())                   
+                    merged_source_doc= []   
+                    with st.expander("source"): 
+                        for idx, item in enumerate(source_doc_list):
+                            source_doc = {"Page": f"Source {idx + 1}", "Source": f"**Source:** {item['metadata']['source'].split('/')[-1]}",
+                                        "page_content":item["page_content"]}
+                            merged_source_doc.append(source_doc)
+                            st.markdown(f"Source {idx + 1}")
+                            st.markdown(f"**Source:** {item['metadata']['source'].split('/')[-1]}")
+                            st.markdown(item["page_content"])
+                            st.write("---")  # Add a separator between entries
+                    message = {"role": "assistant", "content": full_response, "Source":merged_source_doc}
+                    st.session_state.messages.append(message)
+                else:
+                    message = {"role": "assistant", "content": full_response, "Source":""}
+                    st.session_state.messages.append(message)
+        end = timeit.default_timer()
+        print(f"Time to retrieve response: {end - start}")

requirements.txt

@@ -0,0 +1,15 @@
+chromadb==0.4.6
+langchain==0.0.278
+openai==0.27.8
+numpy==1.25.2
+pandas==2.0.3
+Pillow==9.5.0
+pypdf==3.15.1
+PyPDF2==3.0.1
+python-dotenv==1.0.0
+sentence-transformers==2.2.2
+streamlit==1.25.0
+streamlit-chat==0.1.1
+rank-bm25==0.2.2
+google-search-results==2.4.2
+tiktoken

utils.py

@@ -0,0 +1,279 @@
+import streamlit as st
+from pypdf import PdfReader
+import os
+from pathlib import Path
+from dotenv import load_dotenv
+import pickle
+import timeit
+from PIL import Image
+import zipfile
+import datetime
+import shutil
+from collections import defaultdict
+import pandas as pd
+
+from langchain.embeddings import HuggingFaceEmbeddings
+from langchain.document_loaders import PyPDFLoader
+from langchain.text_splitter import RecursiveCharacterTextSplitter
+from langchain.document_loaders import PyPDFLoader, DirectoryLoader
+from langchain.memory import ConversationBufferMemory
+from langchain.chains import ConversationalRetrievalChain
+from langchain.prompts.prompt import PromptTemplate
+from langchain.vectorstores import Chroma
+from langchain.document_loaders import PyPDFDirectoryLoader
+from langchain.retrievers import BM25Retriever, EnsembleRetriever
+from langchain.document_loaders import UnstructuredHTMLLoader
+from langchain.llms import OpenAI
+from langchain.chat_models import ChatOpenAI
+from langchain.agents.agent_toolkits import create_retriever_tool
+from langchain.agents.agent_toolkits import create_conversational_retrieval_agent
+from langchain.utilities import SerpAPIWrapper
+from langchain.agents import Tool
+from langchain.agents import load_tools
+
+load_dotenv()
+
+
+current_timestamp = datetime.datetime.now()
+timestamp_string = current_timestamp.strftime("%Y-%m-%d %H:%M:%S")
+
+
+def build_llm():
+    '''
+    Loading OpenAI model
+    '''
+    # llm= OpenAI(temperature=0.2)
+    llm= ChatOpenAI(temperature = 0, max_tokens=256)
+    return llm
+
+def build_embedding_model():
+    '''
+    Loading Sentence transformer model for text embedding
+    '''
+    embeddings = HuggingFaceEmbeddings(model_name='sentence-transformers/all-MiniLM-L6-v2',
+                                       model_kwargs={'device': 'cpu'})
+    return embeddings
+
+def unzip_opm():
+    # Specify the path to your ZIP file
+    zip_file_path = r'OPM_Files/OPM_Retirement_backup-20230902T130906Z-001.zip'
+
+    # Get the directory where the ZIP file is located
+    extract_path = os.path.dirname(zip_file_path)
+
+    # Create a folder with the same name as the ZIP file (without the .zip extension)
+    extract_folder = os.path.splitext(os.path.basename(zip_file_path))[0]
+    extract_folder_path = os.path.join(extract_path, extract_folder)
+
+    # Create the folder if it doesn't exist
+    if not os.path.exists(extract_folder_path):
+        os.makedirs(extract_folder_path)
+
+    # Open the ZIP file for reading
+    with zipfile.ZipFile(zip_file_path, 'r') as zip_ref:
+        # Extract all the contents into the created folder
+        zip_ref.extractall(extract_folder_path)
+
+    print(f'Unzipped {zip_file_path} to {extract_folder_path}')
+    return extract_folder_path
+
+
+
+
+
+    return 
+
+def count_files_by_type(folder_path):
+    '''
+    Counting files by file type in the specified folder
+    '''
+    file_count_by_type = defaultdict(int)
+    
+    for root, _, files in os.walk(folder_path):
+        for file in files:
+            _, extension = os.path.splitext(file)
+            file_count_by_type[extension] += 1
+    
+    return file_count_by_type
+
+def generate_file_count_table(file_count_by_type):
+    '''
+    Generate a table files count file type
+    '''
+    data = {"File Type": [], "Number of Files": []}
+    for extension, count in file_count_by_type.items():
+        data["File Type"].append(extension)
+        data["Number of Files"].append(count)
+    
+    df = pd.DataFrame(data)
+    df = df.sort_values(by="Number of Files", ascending=False)  # Sort by number of files
+    return df
+
+def move_files_to_folders(folder_path):
+    '''
+    Move files to respective folder. Example, PDF docs to PDFs folder, HTML docs to HTMLs folder.
+    '''
+    for root, _, files in os.walk(folder_path):
+        for file in files:
+            _, extension = os.path.splitext(file)
+            source_path = os.path.join(root, file)
+            
+            if extension == '.pdf':
+                dest_folder = "PDFs"
+            elif extension == '.html':
+                dest_folder = "HTMLs"
+            else:
+                continue
+            
+            dest_path = os.path.join(dest_folder, file)
+            os.makedirs(dest_folder, exist_ok=True)
+            shutil.copy(source_path, dest_path)
+
+
+
+def load_vectorstore(persist_directory, embeddings):
+    '''
+    This function will try first to load chroma database from the disk. If it does exist,
+    It will do the following,
+        1) Load the pdfs
+        2) create text chunks
+        3) Index it and store it in a Chroma DB
+        4) Peform the same for HTML files
+        5) Store the final chroma db in the disk
+    '''
+    if os.path.exists(persist_directory):
+        print("Using existing vectore store for these documents.")
+        vectorstore = Chroma(persist_directory=persist_directory, embedding_function=embeddings)
+        print("Chroma DB loaded from the disk")
+        return vectorstore
+    else:
+        folder_path= unzip_opm()
+        print("Vector store is not available. Creating new one.")
+        file_count_by_type = count_files_by_type(folder_path)
+        file_count_table = generate_file_count_table(file_count_by_type)
+        print("File Count Table:")
+        print(file_count_table)
+        #move files into respective folders
+        move_files_to_folders(folder_path)
+        print("PDF and HTML files copied to separate folders.")
+        
+        # Load the pdf files from the pdffolder in order to create new chroma db
+        pdf_folder_path= f"{folder_path}/PDFs" #pdf folder
+        html_folder_path= f"{folder_path}/HTMLs" #html folder
+        pdf_dir_loader = PyPDFDirectoryLoader(pdf_folder_path)
+        pdf_pages = pdf_dir_loader.load()
+        print("PDF files are loaded from the folder.")
+
+        
+        #Loading HTML files from the html folder in order to create new chroma db 
+        HTML_docs_path_list = [os.path.join(html_folder_path, f) for f in os.listdir(html_folder_path) if os.path.isfile(os.path.join(html_folder_path, f))]
+
+        html_loaders= []
+        for html_file in HTML_docs_path_list:
+            loader = UnstructuredHTMLLoader(html_file)
+            html_loaders.append(loader)
+
+        html_pages = []
+        docs_cannot_load= []
+        for loader in html_loaders:
+            try:
+                html_pages.extend(loader.load())
+            except:
+                print("Cannot load the file:", loader)
+                docs_cannot_load.append(loader)
+        print("HTML files are loaded from the folder.")
+        # Create text chunks from the PDF docs
+        text_splitter = RecursiveCharacterTextSplitter(
+            # Set a really small chunk size, just to show.
+            chunk_size = 1000,
+            chunk_overlap  = 200,
+            length_function = len,
+            is_separator_regex = False,
+        )
+
+        pdf_texts = text_splitter.transform_documents(pdf_pages)
+        # Create text chunks from the HTML docs
+        html_texts = text_splitter.transform_documents(html_pages)
+        # Merging all the text chunks (HTML + PDF)
+        all_texts= pdf_texts+html_texts
+        print("PDF and HTML docs are split into chunks and created a final list representing all the chunks.")
+
+        # Create embeddings for all the text chunks and store it in a Chroma DB
+        vectorstore = Chroma.from_documents(all_texts,
+                                            embeddings,
+                                            persist_directory=persist_directory)
+        vectorstore.persist()
+        print("Chroma DB created and loaded")
+        return vectorstore
+
+
+def load_text_chunks(text_chunks_pkl_dir):
+    '''
+    Loading the pickle file that holds all the documents from the disk.
+    If it does not exist, create new one.
+    Text documents are required to create BM25 Retriever. But loading all the documents in
+    every session will be a time consuming process. So we are storing all the docs in a pickle file
+    and load the pickle file from the disk to overcome this problem.
+    '''
+    try:
+        print("Text chunks are loading from the disk")
+        with open(text_chunks_pkl_dir, 'rb') as file:
+            cached_text_chunks = pickle.load(file)
+        # Now, `cached_text_chunks` contains your cached data
+        print("Text chunks are loaded from the disk")
+        return cached_text_chunks
+    except:
+        print("Creating text chunks from the docs and caching it.")
+        folder_path= unzip_opm()
+        pdf_folder_path= f"{folder_path}/PDFs" #pdf folder
+        html_folder_path= f"{folder_path}/HTMLs" #html folder
+        pdf_dir_loader = PyPDFDirectoryLoader(pdf_folder_path)
+        pdf_pages = pdf_dir_loader.load()
+        HTML_docs_path_list = [os.path.join(html_folder_path, f) for f in os.listdir(html_folder_path) if os.path.isfile(os.path.join(html_folder_path, f))]
+
+        html_loaders= []
+        for html_file in HTML_docs_path_list:
+            loader = UnstructuredHTMLLoader(html_file)
+            html_loaders.append(loader)
+
+        html_pages = []
+        for loader in html_loaders:
+            try:
+                html_pages.extend(loader.load())
+            except:
+                print("Cannot load the file:", loader)
+        all_texts= pdf_pages+html_pages
+       # Cache the list to a file
+        with open('text_chunks.pkl', 'wb') as file:
+            pickle.dump(all_texts, file)
+        print("Text chunks are created and cached")
+        
+def load_ensemble_retriver(text_chunks, embeddings, chroma_vectorstore):
+    """Load ensemble retiriever with BM25 and Chroma as individual retrievers"""
+    bm25_retriever = BM25Retriever.from_documents(text_chunks)
+    bm25_retriever.k = 1
+    chroma_retriever = chroma_vectorstore.as_retriever(search_kwargs={"k": 1})  
+    ensemble_retriever = EnsembleRetriever(retrievers=[bm25_retriever, chroma_retriever], weights=[0.3, 0.7])
+    return ensemble_retriever
+
+
+def load_conversational_retrievel_agent(retriever, llm):
+    '''Load Conversational Retrievel agent with following tasks as tools,
+    1) OPM Knowledge base query
+    2) INternet search with SerpAPI
+    This agent combines RAG, chat interfaces, agents.
+    '''
+    retriever_tool = create_retriever_tool(
+    retriever, 
+    "Search_US_Office_of_Personnel_Management_Document",
+    "Searches and returns documents regarding the U.S. Office of Personnel Management (OPM).")
+    search_api = SerpAPIWrapper()
+    search_api_tool = Tool(
+            name = "Current_Search",
+            func=search_api.run,
+            description="useful for when you need to answer questions about current events or the current state of the world"
+        )
+    tools = [retriever_tool]
+    agent_executor = create_conversational_retrieval_agent(llm, tools, verbose=True, max_token_limit=512)
+    return agent_executor
+