Spaces:
Sleeping
Sleeping
| import streamlit as st | |
| import time | |
| import shutil | |
| import os | |
| from ragagent import RAGAgent | |
| from langchain_community.document_loaders import PyPDFDirectoryLoader | |
| icons = {"assistant": "robot.png", "user": "man-kddi.png"} | |
| DATA_DIR = "data" | |
| # Ensure data directory exists | |
| os.makedirs(DATA_DIR, exist_ok=True) | |
| def init_agent_with_docs(): | |
| docs=[] | |
| if os.path.exists(DATA_DIR): | |
| try: | |
| pdf_loader = PyPDFDirectoryLoader(DATA_DIR) | |
| pdf_docs = pdf_loader.load() | |
| if pdf_docs: | |
| docs.append(pdf_docs) | |
| st.session_state["console_out"] += "Pdf's loaded\n" | |
| except Exception as e: | |
| st.error("PyPDFLoader Exception: " + e) | |
| return RAGAgent(docs) | |
| def remove_old_files(): | |
| if len(os.listdir(DATA_DIR)) !=0: | |
| st.session_state["console_out"] += "remove_old_files\n" | |
| shutil.rmtree(DATA_DIR) | |
| os.makedirs(DATA_DIR) | |
| def streamer(text): | |
| for i in text: | |
| yield i | |
| time.sleep(0.02) | |
| if "console_out" not in st.session_state: | |
| st.session_state["console_out"] = "" | |
| # Streamlit app initialization | |
| st.title("RAG AGENT") | |
| st.markdown("Multi PDF and Web Search - Llama 3") | |
| st.markdown("Routing retrieval, Fallback to web search, Fix hallucinations and check answers") | |
| if 'messages' not in st.session_state: | |
| st.session_state.messages = [{'role': 'assistant', "content": "Hello! Upload PDF's and ask me anything about the content."}] | |
| for message in st.session_state.messages: | |
| with st.chat_message(message['role'], avatar=icons[message['role']]): | |
| st.write(message['content']) | |
| with st.sidebar: | |
| uploaded_files = st.file_uploader("Upload your PDF Files and Click Submit & Process", type="pdf", accept_multiple_files=True) | |
| if st.button("Submit & Process"): | |
| with st.spinner("Processing..."): | |
| st.session_state["console_out"] = "" | |
| remove_old_files() | |
| for index, file in enumerate(uploaded_files): | |
| # filepath = os.path.join(DATA_DIR, f"saved_pdf_{index}.pdf") | |
| filepath = os.path.join(DATA_DIR, file.name) | |
| with open(filepath, "wb") as f: | |
| f.write(file.getbuffer()) | |
| st.session_state.agent = init_agent_with_docs() | |
| remove_old_files() | |
| st.success("Done") | |
| st.text_area("Console", st.session_state["console_out"], height=250) | |
| user_prompt = st.chat_input("Ask me anything about the content of the PDF:") | |
| if user_prompt and uploaded_files: | |
| st.session_state.messages.append({'role': 'user', "content": user_prompt}) | |
| response = "Could not find an answer." | |
| with st.chat_message("user", avatar="man-kddi.png"): | |
| st.write(user_prompt) | |
| # Trigger assistant's response retrieval and update UI | |
| with st.spinner("Thinking..."): | |
| inputs = {"question": user_prompt} | |
| for output in st.session_state.agent.app.stream(inputs): | |
| for key, value in output.items(): | |
| if "generation" in value: | |
| response = value["generation"] | |
| st.session_state["console_out"] = st.session_state.agent.logs | |
| with st.chat_message("user", avatar="robot.png"): | |
| st.write_stream(streamer(response)) | |
| st.session_state.messages.append({'role': 'assistant', "content": response}) | |
| st.rerun() |