fehmikaya/rag_agent_langgraph

1import streamlit as st2 3import time4import shutil5import os6 7from ragagent import RAGAgent8 9from langchain_community.document_loaders import PyPDFDirectoryLoader10 11icons = {"assistant": "robot.png", "user": "man-kddi.png"}12DATA_DIR = "data"13# Ensure data directory exists14os.makedirs(DATA_DIR, exist_ok=True)15 16def init_agent_with_docs():17 18    docs=[]19    if os.path.exists(DATA_DIR):20        try:21            pdf_loader = PyPDFDirectoryLoader(DATA_DIR)22            pdf_docs = pdf_loader.load()23            if pdf_docs:24                docs.append(pdf_docs)25                st.session_state["console_out"] += "Pdf's loaded\n"26        except Exception as e:27            st.error("PyPDFLoader Exception: " + e)28    return RAGAgent(docs)29 30def remove_old_files():31    if len(os.listdir(DATA_DIR)) !=0:32        st.session_state["console_out"] += "remove_old_files\n"33        shutil.rmtree(DATA_DIR)34        os.makedirs(DATA_DIR)35 36def streamer(text):37    for i in text:38        yield i39        time.sleep(0.02)40 41if "console_out" not in st.session_state:42    st.session_state["console_out"] = ""43 44# Streamlit app initialization45st.title("RAG AGENT")46st.markdown("Multi PDF and Web Search - Llama 3")47st.markdown("Routing retrieval, Fallback to web search, Fix hallucinations and check answers")48 49if 'messages' not in st.session_state:50    st.session_state.messages = [{'role': 'assistant', "content": "Hello! Upload PDF's and ask me anything about the content."}]51 52for message in st.session_state.messages:53    with st.chat_message(message['role'], avatar=icons[message['role']]):54        st.write(message['content'])55 56with st.sidebar:57    uploaded_files = st.file_uploader("Upload your PDF Files and Click Submit & Process", type="pdf", accept_multiple_files=True)58 59    if st.button("Submit & Process"):60        with st.spinner("Processing..."):61            st.session_state["console_out"] = ""62            remove_old_files()  63            for index, file in enumerate(uploaded_files):64                # filepath = os.path.join(DATA_DIR, f"saved_pdf_{index}.pdf")65                filepath = os.path.join(DATA_DIR, file.name)66                with open(filepath, "wb") as f:67                    f.write(file.getbuffer())68            st.session_state.agent = init_agent_with_docs()69            remove_old_files()  70            st.success("Done")71    st.text_area("Console", st.session_state["console_out"], height=250)72 73user_prompt = st.chat_input("Ask me anything about the content of the PDF:")74 75if user_prompt and uploaded_files:76    st.session_state.messages.append({'role': 'user', "content": user_prompt})77    response = "Could not find an answer."78    with st.chat_message("user", avatar="man-kddi.png"):79        st.write(user_prompt)80 81    # Trigger assistant's response retrieval and update UI82    with st.spinner("Thinking..."):83        inputs = {"question": user_prompt}84        for output in st.session_state.agent.app.stream(inputs):85            for key, value in output.items():86                if "generation" in value:87                    response = value["generation"]88        st.session_state["console_out"] = st.session_state.agent.logs89    with st.chat_message("user", avatar="robot.png"):90        st.write_stream(streamer(response))91    st.session_state.messages.append({'role': 'assistant', "content": response})92 93    st.rerun()