Update app.py

app.py

@@ -1,50 +1,131 @@
-# Q&A Chatbot
-from langchain.llms import OpenAI
-import os
-
-from dotenv import load_dotenv, find_dotenv
-_ = load_dotenv(find_dotenv()) # read local .env file
-
-from langchain_community.document_loaders import PyPDFLoader
-from langchain.indexes import VectorstoreIndexCreator
-from langchain.vectorstores import DocArrayInMemorySearch
+from llama_index import GPTVectorStoreIndex, SimpleDirectoryReader
+from llama_index import download_loader
+from pandasai.llm.openai import OpenAI
+from matplotlib import pyplot as plt
 import streamlit as st
+import pandas as pd
+import os
 
 
-os.getenv("OPENAI_API_KEY")
-
-## Function to load OpenAI model and get respones
-
-loader = PyPDFLoader("Firm Data.pdf")
-index = VectorstoreIndexCreator(
-    vectorstore_cls=DocArrayInMemorySearch
-).from_loaders([loader])
+documents_folder = "./documents"
 
+# Load PandasAI loader, Which is a wrapper over PandasAI library
+PandasAIReader = download_loader("PandasAIReader")
 
-def get_openai_response(index, question):
-    llm_replacement_model = OpenAI(temperature=0,
-                               model='gpt-3.5-turbo-instruct')
-    response = index.query(question,
-                       llm = llm_replacement_model)
-    return response
+st.title("Welcome to `ChatwithDocs`")
+st.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⭐")
 
-st.set_page_config(page_title="Ask Questions on S&P 500 companies financials")
 
-st.header("Langchain Application")
+def get_csv_result(df, query):
+  reader = PandasAIReader(llm=csv_llm)
+  response = reader.run_pandas_ai(
+    df, 
+    query, 
+    is_conversational_answer=False
+    )
+  return response
 
-input=st.text_input("Input: ",key="input")
-response=get_openai_response(index, input)
+def save_file(doc):        
+    fn = os.path.basename(doc.name)
+    # open read and write the file into the server
+    open(documents_folder+'/'+fn, 'wb').write(doc.read())
+    # Check for the current filename, If new filename
+    # clear the previous cached vectors and update the filename 
+    # with current name     
+    if st.session_state.get('file_name'):
+        if st.session_state.file_name != fn:
+            st.cache_resource.clear()
+            st.session_state['file_name'] = fn
+    else:
+        st.session_state['file_name'] = fn
 
-submit=st.button("Ask the question")
+    return fn
 
-## If ask button is clicked
+def remove_file(file_path):
+    # Remove the file from the Document folder once 
+    # vectors are created
+    if os.path.isfile(documents_folder+'/'+file_path):
+        os.remove(documents_folder+'/'+file_path)
 
-if submit:
-    st.subheader("The Response is")
-    st.write(response)
+    
 
+@st.cache_resource
+def create_index():
+    # Create vectors for the file stored under Document folder. 
+    # NOTE: You can create vectors for multiple files at once.
+    documents = SimpleDirectoryReader(documents_folder).load_data()
+    index = GPTVectorStoreIndex.from_documents(documents)
+    return index
 
 
 
+def query_doc(vector_index, query):
+    # Applies Similarity Algo, Finds the nearest match and 
+    # take the match and user query to OpenAI for rich response
+    query_engine = vector_index.as_query_engine()
+    response = query_engine.query(query)
+    return response
 
 
+api_key = st.text_input("Enter your OpenAI API key here:", type="password")
+if api_key:
+    os.environ['OPENAI_API_KEY'] = api_key
+    csv_llm = OpenAI(api_token=api_key)
+
+
+tab1, tab2= st.tabs(["CSV", "PDFs/Docs"])
+
+with tab1:
+   
+   st.write("Chat with CSV files using PandasAI loader with LlamaIndex")
+   input_csv = st.file_uploader("Upload your CSV file", type=['csv'])
+   
+   if input_csv is not None:      
+    st.info("CSV Uploaded Successfully")
+    df = pd.read_csv(input_csv)
+    st.dataframe(df, use_container_width=True)
+    
+
+   st.write("---")
+   
+   input_text = st.text_area("Ask your query")
+   
+   if input_text is not None:
+    if st.button("Send"):
+      st.info("Your query: "+ input_text)
+      with st.spinner('Processing your query...'):
+        response = get_csv_result(df, input_text)
+      if plt.get_fignums():
+          st.pyplot(plt.gcf())
+      else:
+        st.success(response)
+
+    
+with tab2:
+   st.write("Chat with PDFs/Docs")
+   input_doc = st.file_uploader("Upload your Docs")
+   
+   if input_doc is not None: 
+       st.info("Doc Uploaded Successfully")
+       file_name = save_file(input_doc)
+       index = create_index()
+       remove_file(file_name)
+       
+   
+   st.write("---")
+   input_text = st.text_area("Ask your question")
+   
+   if input_text is not None:
+    if st.button("Ask"):
+        st.info("Your query: \n" +input_text)
+        with st.spinner("Processing your query.."):
+            response = query_doc(index, input_text)
+            print(response)
+        
+        st.success(response)
+
+        st.write("---")
+        # Shows the source documents context which 
+        # has been used to prepare the response
+        st.write("Source Documents")
+        st.write(response.get_formatted_sources())