Spaces:
Sleeping
Sleeping
File size: 3,338 Bytes
ff86c9f e536d7d 4ef548e a86a7c9 a56fdd5 a86a7c9 0ce97ce 4ef548e 9b8317e 0ce97ce 9b8317e 0ce97ce 9b8317e e1bc7be 62bdde9 4ef548e 1ba663b 2e9fddc f1ac089 2e9fddc a1edf53 e1bc7be 4ef548e 0ce97ce a86a7c9 ff86c9f ebda27a ff86c9f d9e74a8 bf74a90 fe75a15 ff86c9f 3f664a4 ff86c9f f1ac089 ff86c9f 334a17b f1ac089 ff86c9f 7138439 0ce97ce b2a1a43 4e87b58 4ef548e f1ac089 013c3ab 0ce97ce ff86c9f 28677c7 0959dbd 334a17b ff86c9f 6444235 ff86c9f e1bc7be ff86c9f e1bc7be 013c3ab e1bc7be ac08c88 ff86c9f bf74a90 e378c05 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 |
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() |