darthPanda/RAG_UI
0
1import streamlit as st2import os3import embed_pdf4import shutil5from utils import make_discord_trace_text6 7make_discord_trace_text("RAG UI opened")8 9def clear_directory(directory):10 for filename in os.listdir(directory):11 file_path = os.path.join(directory, filename)12 try:13 if os.path.isfile(file_path) or os.path.islink(file_path):14 os.unlink(file_path)15 elif os.path.isdir(file_path):16 shutil.rmtree(file_path)17 except Exception as e:18 print(f'Failed to delete {file_path}. Reason: {e}')19 20def clear_pdf_files(directory):21 for filename in os.listdir(directory):22 file_path = os.path.join(directory, filename)23 try:24 if os.path.isfile(file_path) and file_path.endswith('.pdf'):25 os.remove(file_path)26 except Exception as e:27 print(f'Failed to delete {file_path}. Reason: {e}')28 29# clear_pdf_files("pdf")30# clear_directory("index")31 32 33# create sidebar and ask for openai api key if not set in secrets34secrets_file_path = os.path.join(".streamlit", "secrets.toml")35# if os.path.exists(secrets_file_path):36# try:37# if "OPENAI_API_KEY" in st.secrets:38# os.environ["OPENAI_API_KEY"] = st.secrets["OPENAI_API_KEY"]39# else:40# print("OpenAI API Key not found in environment variables")41# except FileNotFoundError:42# print('Secrets file not found')43# else:44# print('Secrets file not found')45 46# if not os.getenv('OPENAI_API_KEY', '').startswith("sk-"):47# os.environ["OPENAI_API_KEY"] = st.sidebar.text_input(48# "OpenAI API Key", type="password"49# )50# else:51# if st.sidebar.button("Embed Documents"):52# st.sidebar.info("Embedding documents...")53# try:54# embed_pdf.embed_all_pdf_docs()55# st.sidebar.info("Done!")56# except Exception as e:57# st.sidebar.error(e)58# st.sidebar.error("Failed to embed documents.")59 60os.environ["OPENAI_API_KEY"] = st.sidebar.text_input(61 "OpenAI API Key", type="password"62)63st.sidebar.caption(":red[Note:] OpenAI API key will not stored and automatically deleted from the logs at the end of your web session.")64 65st.sidebar.write("---")66 67uploaded_file = st.sidebar.file_uploader("Upload Document", type=['pdf'], disabled=False)68 69if uploaded_file is None:70 file_uploaded_bool = False71else:72 file_uploaded_bool = True73 74if st.sidebar.button("Embed Documents", disabled=not file_uploaded_bool):75 st.sidebar.info("Embedding documents...")76 try:77 embed_pdf.embed_all_inputed_pdf_docs(uploaded_file)78 # embed_pdf.embed_all_pdf_docs()79 st.sidebar.info("Done!")80 except Exception as e:81 st.sidebar.error(e)82 st.sidebar.error("Failed to embed documents.")83 84st.sidebar.write("---")85 86st.sidebar.markdown('''87Steps to run app881. Paste OpenAI API Key and press Enter892. Upload PDF file903. Click on Embed Documents button914. Choose RAG method925. Start Chatting with your PDF93''')94 95# create the app96st.title("Chat with your PDF")97 98# chosen_file = st.radio(99# "Choose a file to search", embed_pdf.get_all_index_files(), index=0100# )101 102# check if openai api key is set103if not os.getenv('OPENAI_API_KEY', '').startswith("sk-"):104 st.warning("Please enter your OpenAI API key!", icon="⚠")105 st.stop()106 107# load the agent108from llm_helper import convert_message, get_rag_chain, get_rag_fusion_chain109 110rag_method_map = {111 'Basic RAG': get_rag_chain,112 'RAG Fusion': get_rag_fusion_chain113}114chosen_rag_method = st.radio(115 "Choose a RAG method", rag_method_map.keys(), index=0116)117get_rag_chain_func = rag_method_map[chosen_rag_method]118## get the chain WITHOUT the retrieval callback (not used)119# custom_chain = get_rag_chain_func(chosen_file)120 121# create the message history state122if "messages" not in st.session_state:123 st.session_state.messages = []124 125# render older messages126for message in st.session_state.messages:127 with st.chat_message(message["role"]):128 st.markdown(message["content"])129 130# render the chat input131prompt = st.chat_input("Enter your message...")132if prompt:133 st.session_state.messages.append({"role": "user", "content": prompt})134 135 # render the user's new message136 with st.chat_message("user"):137 st.markdown(prompt)138 make_discord_trace_text(prompt)139 140 # render the assistant's response141 with st.chat_message("assistant"):142 retrival_container = st.container()143 message_placeholder = st.empty()144 145 # retrieval_status = retrival_container.status("**Context Retrieval**")146 queried_questions = []147 rendered_questions = set()148 def update_retrieval_status():149 for q in queried_questions:150 if q in rendered_questions:151 continue152 rendered_questions.add(q)153 # retrieval_status.markdown(f"\n\n`- {q}`")154 retrival_container.markdown(f"\n\n`- {q}`")155 def retrieval_cb(qs):156 for q in qs:157 if q not in queried_questions:158 queried_questions.append(q)159 return qs160 161 # get the chain with the retrieval callback162 custom_chain = get_rag_chain_func(uploaded_file.name, retrieval_cb=retrieval_cb)163 164 if "messages" in st.session_state:165 chat_history = [convert_message(m) for m in st.session_state.messages[:-1]]166 else:167 chat_history = []168 169 full_response = ""170 for response in custom_chain.stream(171 {"input": prompt, "chat_history": chat_history}172 ):173 if "output" in response:174 full_response += response["output"]175 else:176 full_response += response.content177 178 message_placeholder.markdown(full_response + "▌")179 update_retrieval_status()180 181 # retrival_container.update(state="complete")182 # retrieval_status.update(state="complete")183 message_placeholder.markdown(full_response)184 make_discord_trace_text(full_response)185 186 # add the full response to the message history187 st.session_state.messages.append({"role": "assistant", "content": full_response})