msaifee/Multi-Loader-RAG

1import streamlit as st2from langchain_community.document_loaders import TextLoader, PyPDFLoader, WebBaseLoader3from langchain_text_splitters import RecursiveCharacterTextSplitter4from langchain_openai import OpenAIEmbeddings5from langchain_community.vectorstores import FAISS6from langchain.llms import OpenAI7from langchain.chains import RetrievalQA8import os9import tempfile10from dotenv import load_dotenv11 12# Load environment variables13load_dotenv()14 15os.environ["LANGCHAIN_TRACING_V2"] = "true"16os.environ["LANGCHAIN_PROJECT"]="Multi Loader RAG"17 18 19# Streamlit app title20st.title("Multi Loader RAG")21 22# File upload and web link input23st.header("Upload Documents")24text_file = st.file_uploader("Upload a Text File", type=["txt"])25pdf_file = st.file_uploader("Upload a PDF File", type=["pdf"])26web_link = st.text_input("Enter a Web URL")27 28# Load documents function29def load_documents(text_file, pdf_file, web_link):30    docs = []31    32    # Load text file33    if text_file is not None:34        with tempfile.NamedTemporaryFile(delete=False, suffix=".txt") as tmp_file:35            tmp_file.write(text_file.getvalue())36            tmp_file_path = tmp_file.name37        text_loader = TextLoader(tmp_file_path)38        docs.extend(text_loader.load())39        os.remove(tmp_file_path)40    41    # Load PDF file42    if pdf_file is not None:43        with tempfile.NamedTemporaryFile(delete=False, suffix=".pdf") as tmp_file:44            tmp_file.write(pdf_file.getvalue())45            tmp_file_path = tmp_file.name46        pdf_loader = PyPDFLoader(tmp_file_path)47        docs.extend(pdf_loader.load())48        os.remove(tmp_file_path)49    50    # Load web content51    if web_link:52        web_loader = WebBaseLoader([web_link])53        docs.extend(web_loader.load())54    55    return docs56 57# Split documents function58def split_documents(docs, chunk_size, chunk_overlap):59    text_splitter = RecursiveCharacterTextSplitter(60        chunk_size=chunk_size,61        chunk_overlap=chunk_overlap62    )63    return text_splitter.split_documents(docs)64 65# Create FAISS vector store function66def create_vector_store(splits):67    embeddings = OpenAIEmbeddings()68    vectorstore = FAISS.from_documents(splits, embeddings)69    return vectorstore70 71# Main app logic72if st.button("Process Documents"):73    if not (text_file or pdf_file or web_link):74        st.error("Please upload at least one document or provide a web link.")75    else:76        with st.spinner("Processing documents..."):77            # Load documents78            documents = load_documents(text_file, pdf_file, web_link)79            80            # Split documents81            splits = split_documents(documents, 1000, 300)82            83            # Create FAISS vector store84            st.session_state.vector_store = create_vector_store(splits)85            86            st.success("Documents processed and FAISS vector store created!")87 88st.header("Get Summary/Answer")89query = st.text_input("Enter your query")90 91if st.button("Search"):92    if st.session_state.vector_store is None:93        st.error("Please process documents first.")94    elif not query:95        st.error("Please enter a query.")96    else:97        with st.spinner("Searching..."): 98            # Create retriever and chain99            retriever = st.session_state.vector_store.as_retriever(100                search_type="similarity",101                search_kwargs={"k": 5}102            )103            llm = OpenAI(temperature=0.6)104            qa_chain = RetrievalQA.from_chain_type(105                llm=llm,106                chain_type="stuff",107                retriever=retriever,108                return_source_documents=True109            )110 111            # Execute query112            result = qa_chain({"query": query})113 114            # Display the result115            st.markdown("### Answer:")116            st.write(result["result"])