LangChainDemo/OPM_Retirement_Assistant
0
1import streamlit as st2from pypdf import PdfReader3# import replicate4import os5from pathlib import Path6from dotenv import load_dotenv7import pickle8import timeit9from PIL import Image10import datetime11import base6412 13from langchain.embeddings import HuggingFaceEmbeddings14from langchain.vectorstores import FAISS15from langchain.document_loaders import PyPDFLoader16from langchain.text_splitter import RecursiveCharacterTextSplitter17from langchain.document_loaders import PyPDFLoader, DirectoryLoader18from langchain.memory import ConversationBufferMemory19from langchain.chains import ConversationalRetrievalChain20from langchain.prompts.prompt import PromptTemplate21from langchain.llms import LlamaCpp22from langchain.callbacks.manager import CallbackManager23from langchain.callbacks.streaming_stdout import StreamingStdOutCallbackHandler24from langchain.vectorstores import Chroma25from langchain.document_loaders import PyPDFDirectoryLoader26from langchain.retrievers import BM25Retriever, EnsembleRetriever27from langchain.chat_models import ChatOpenAI28from langchain.agents.agent_toolkits import create_retriever_tool29from langchain.agents.agent_toolkits import create_conversational_retrieval_agent30from langchain.utilities import SerpAPIWrapper31 32from utils import build_embedding_model, build_llm33from utils import load_retriver,load_vectorstore, load_conversational_retrievel_chain34 35load_dotenv()36# Getting current timestamp to keep track of historical conversations37current_timestamp = datetime.datetime.now()38timestamp_string = current_timestamp.strftime("%Y-%m-%d %H:%M:%S")39 40#Directories path41persist_directory= "Database/PDF_HTML_CHROMA_DB"42all_docs_pkl_directory= 'Database/text_chunks_html_pdf.pkl'43 44# 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.45if "llm" not in st.session_state:46 st.session_state["llm"] = build_llm()47 48if "embeddings" not in st.session_state:49 st.session_state["embeddings"] = build_embedding_model()50 51if "vector_db" not in st.session_state:52 st.session_state["vector_db"] = load_vectorstore(persist_directory=persist_directory, embeddings=st.session_state["embeddings"])53 54# if "text_chunks" not in st.session_state:55# st.session_state["text_chunks"] = load_text_chunks(text_chunks_pkl_dir=all_docs_pkl_directory)56 57if "retriever" not in st.session_state:58 st.session_state["retriever"] = load_retriver(chroma_vectorstore=st.session_state["vector_db"])59 60if "conversation_chain" not in st.session_state:61 st.session_state["conversation_chain"] = load_conversational_retrievel_chain(retriever=st.session_state["retriever"], llm=st.session_state["llm"]) 62 63 64 65# App title66st.set_page_config(67 page_title="OMP Search Bot",68 layout="wide",69 initial_sidebar_state="expanded",70)71 72st.markdown("""73 <style>74 .block-container {75 padding-top: 2.2rem}76 </style>77 """, unsafe_allow_html=True)78# To get header in the App79col1, col2= st.columns(2)80 81title1 = """82<p style="font-size: 26px;text-align: right; color: #0C3453; font-weight: bold">OPM Retirement Services Assistant</p>83"""84 85def clear_chat_history():86 """87 Clear chat and start new chat88 """89 st.session_state.messages = [{"role": "assistant", "content": "How may I assist you today?"}]90 91#loading OPM logo92file_ = open("opm_logo.png", "rb")93contents = file_.read()94data_url = base64.b64encode(contents).decode("utf-8")95file_.close()96 97st.markdown(98 f"""99 <div style="background-color: white; padding: 15px; border-radius: 10px;">100 <div style="display: flex; justify-content: space-between;">101 <div>102 <img src="data:image/png;base64,{data_url}" style="max-width: 100%;" alt="OPM Logo" />103 </div>104 <div style="flex: 1; padding: 15px;">105 {title1}106 """,107 unsafe_allow_html=True108)109st.write("")110 111 112st.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)113 114st.markdown("""---""")115 116text_html = """117 <p style="font-size: 24px; text-align: center; color:blue; margin: 0;">118 Type your question below in conversational style language.119 </p>120 <p style="font-size: 18px; text-align: center; color: blue; margin: 0;">121 Sample Questions:<br>122 can I work part-time and get retirement benefits <br>123 will I get paid for my unused annual leave <br>124 how do I report the death of a federal employee <br>125 what are interim benefits 126 </p>127"""128 129st.write(text_html, unsafe_allow_html=True)130 131 132with st.sidebar:133 st.subheader("") 134 135if st.session_state["vector_db"] and st.session_state["llm"]:136 # Store LLM generated responses137 if "messages" not in st.session_state.keys():138 st.session_state.messages = [{"role": "assistant", "content": "How may I assist you today?", "Source":""}]139 140 # Display or clear chat messages141 for message in st.session_state.messages:142 with st.chat_message(message["role"]):143 st.write(message["content"])144 if message["Source"]=="":145 st.write("")146 else:147 with st.expander("source"):148 for idx, item in enumerate(message["Source"]):149 st.markdown(item["Page"])150 st.markdown(item["Source"])151 st.markdown(item["page_content"])152 st.write("---")153 154 155 # Initialize the session state to store chat history156 if "stored_session" not in st.session_state:157 st.session_state["stored_session"] = []158 159 # Create a list to store expanders160 if "expanders" not in st.session_state:161 st.session_state["expanders"] = []162 163 # Define a function to add a new chat expander164 def add_chat_expander(chat_history):165 current_timestamp = datetime.datetime.now()166 timestamp_string = current_timestamp.strftime("%Y-%m-%d %H:%M:%S")167 st.session_state["expanders"].append({"timestamp": timestamp_string, "chat_history": chat_history})168 169 def clear_chat_history():170 """171 To remove existing chat history and start new conversation172 """173 stored_session = []174 for dict_message in st.session_state.messages:175 if dict_message["role"] == "user":176 string_dialogue = "User: " + dict_message["content"] + "\n\n"177 st.session_state["stored_session"].append(string_dialogue)178 179 else:180 string_dialogue = "Assistant: " + dict_message["content"] + "\n\n"181 st.session_state["stored_session"].append(string_dialogue)182 stored_session.append(string_dialogue)183 184 # Add a new chat expander185 add_chat_expander(stored_session)186 st.session_state.messages = [{"role": "assistant", "content": "How may I assist you today?", "Source":""}]187 188 st.sidebar.button('New chat', on_click=clear_chat_history, use_container_width=True)189 st.sidebar.text("")190 st.sidebar.write('<p style="font-size: 16px;text-align: center; color: #727477; font-weight: bold">Chat history</p>', unsafe_allow_html=True)191 # Display existing chat expanders192 for expander_info in st.session_state["expanders"]:193 with st.sidebar.expander("Conversation ended at:"+"\n\n"+expander_info["timestamp"]):194 for message in expander_info["chat_history"]:195 if message.startswith("User:"):196 st.write(f'<span style="color: #EF6A6A;">{message}</span>', unsafe_allow_html=True)197 elif message.startswith("Assistant:"):198 st.write(f'<span style="color: #F7BD45;">{message}</span>', unsafe_allow_html=True)199 else:200 st.write(message)201 202 203 def generate_llm_response(conversation_chain, prompt_input):204 # output= conversation_chain({'question': prompt_input})205 res = conversation_chain(prompt_input)206 return res['result']207 208 209 # User-provided prompt210 if prompt := st.chat_input(disabled= not st.session_state["vector_db"]):211 st.session_state.messages.append({"role": "user", "content": prompt, "Source":""})212 with st.chat_message("user"):213 st.write(prompt)214 215 # Generate a new response if last message is not from assistant216 if st.session_state.messages[-1]["role"] != "assistant":217 with st.chat_message("assistant"):218 with st.spinner("Searching..."):219 start = timeit.default_timer()220 response = generate_llm_response(conversation_chain=st.session_state["conversation_chain"], prompt_input=prompt)221 placeholder = st.empty()222 full_response = ''223 for item in response:224 full_response += item225 placeholder.markdown(full_response)226 if response:227 st.text("-------------------------------------")228 docs= st.session_state["retriever"].get_relevant_documents(prompt) 229 source_doc_list= [] 230 for doc in docs:231 source_doc_list.append(doc.dict()) 232 merged_source_doc= [] 233 with st.expander("source"): 234 for idx, item in enumerate(source_doc_list):235 source_doc = {"Page": f"Source {idx + 1}", "Source": f"**Source:** {item['metadata']['source'].split('/')[-1]}",236 "page_content":item["page_content"]}237 merged_source_doc.append(source_doc)238 st.markdown(f"Source {idx + 1}")239 st.markdown(f"**Source:** {item['metadata']['source'].split('/')[-1]}")240 st.markdown(item["page_content"])241 st.write("---") # Add a separator between entries242 message = {"role": "assistant", "content": full_response, "Source":merged_source_doc}243 st.session_state.messages.append(message)244 st.markdown("👍 👎 Create Ticket")245 # else:246 # with st.expander("source"):247 # message = {"role": "assistant", "content": full_response, "Source":""}248 # st.session_state.messages.append(message)249 end = timeit.default_timer()250 print(f"Time to retrieve response: {end - start}")251