AIProdAndInnov/Chat-with-Docs

1from llama_index import GPTVectorStoreIndex, SimpleDirectoryReader2from llama_index import download_loader3from pandasai.llm.openai import OpenAI4from matplotlib import pyplot as plt5import streamlit as st6import pandas as pd7import os8 9 10documents_folder = "./documents"11 12# Load PandasAI loader, Which is a wrapper over PandasAI library13PandasAIReader = download_loader("PandasAIReader")14 15st.title("Welcome to `ChatwithDocs`")16st.header("Interact with Documents such as `PDFs/CSV/Docs` using the power of LLMs\nPowered by `LlamaIndex🦙` \nCheckout the [GITHUB Repo Here](https://github.com/anoopshrma/Chat-with-Docs) and Leave a star⭐")17 18 19def get_csv_result(df, query):20  reader = PandasAIReader(llm=csv_llm)21  response = reader.run_pandas_ai(22    df, 23    query, 24    is_conversational_answer=False25    )26  return response27 28def save_file(doc):        29    fn = os.path.basename(doc.name)30    # open read and write the file into the server31    open(documents_folder+'/'+fn, 'wb').write(doc.read())32    # Check for the current filename, If new filename33    # clear the previous cached vectors and update the filename 34    # with current name     35    if st.session_state.get('file_name'):36        if st.session_state.file_name != fn:37            st.cache_resource.clear()38            st.session_state['file_name'] = fn39    else:40        st.session_state['file_name'] = fn41 42    return fn43 44def remove_file(file_path):45    # Remove the file from the Document folder once 46    # vectors are created47    if os.path.isfile(documents_folder+'/'+file_path):48        os.remove(documents_folder+'/'+file_path)49 50    51 52@st.cache_resource53def create_index():54    # Create vectors for the file stored under Document folder. 55    # NOTE: You can create vectors for multiple files at once.56    documents = SimpleDirectoryReader(documents_folder).load_data()57    index = GPTVectorStoreIndex.from_documents(documents)58    return index59 60 61 62def query_doc(vector_index, query):63    # Applies Similarity Algo, Finds the nearest match and 64    # take the match and user query to OpenAI for rich response65    query_engine = vector_index.as_query_engine()66    response = query_engine.query(query)67    return response68 69 70api_key = st.text_input("Enter your OpenAI API key here:", type="password")71if api_key:72    os.environ['OPENAI_API_KEY'] = api_key73    csv_llm = OpenAI(api_token=api_key)74 75 76tab1, tab2= st.tabs(["CSV", "PDFs/Docs"])77 78with tab1:79   80   st.write("Chat with CSV files using PandasAI loader with LlamaIndex")81   input_csv = st.file_uploader("Upload your CSV file", type=['csv'])82   83   if input_csv is not None:      84    st.info("CSV Uploaded Successfully")85    df = pd.read_csv(input_csv)86    st.dataframe(df, use_container_width=True)87    88 89   st.write("---")90   91   input_text = st.text_area("Ask your query")92   93   if input_text is not None:94    if st.button("Send"):95      st.info("Your query: "+ input_text)96      with st.spinner('Processing your query...'):97        response = get_csv_result(df, input_text)98      if plt.get_fignums():99          st.pyplot(plt.gcf())100      else:101        st.success(response)102 103    104with tab2:105   st.write("Chat with PDFs/Docs")106   input_doc = st.file_uploader("Upload your Docs")107   108   if input_doc is not None: 109       st.info("Doc Uploaded Successfully")110       file_name = save_file(input_doc)111       index = create_index()112       remove_file(file_name)113       114   115   st.write("---")116   input_text = st.text_area("Ask your question")117   118   if input_text is not None:119    if st.button("Ask"):120        st.info("Your query: \n" +input_text)121        with st.spinner("Processing your query.."):122            response = query_doc(index, input_text)123            print(response)124        125        st.success(response)126 127        st.write("---")128        # Shows the source documents context which 129        # has been used to prepare the response130        st.write("Source Documents")131        st.write(response.get_formatted_sources())