Update app.py

app.py

@@ -10,15 +10,18 @@ def save_docs(docs):
     import shutil
     import os
 
-    destination_dir = "/home/user/app/docs/"
-    os.makedirs(destination_dir, exist_ok=True)
+    output_dir="/content/docs/"
 
-    output_dir="/home/user/app/docs/"
+    if os.path.exists(output_dir):
+        shutil.rmtree(output_dir)
+
+    if not os.path.exists(output_dir):
+        os.makedirs(output_dir)
 
     for doc in docs:
-      shutil.copy(doc.name, output_dir)
+        shutil.copy(doc.name, output_dir)
 
-    return "File(s) saved successfully!"
+    return "Successful!"
 
 def process_docs():
 
@@ -26,21 +29,31 @@ def process_docs():
     from langchain.document_loaders import DirectoryLoader
     from langchain.document_loaders import TextLoader
     from langchain.document_loaders import Docx2txtLoader
+    from langchain.document_loaders.csv_loader import CSVLoader
+    from langchain.document_loaders import UnstructuredExcelLoader
     from langchain.vectorstores import FAISS
     from langchain.embeddings.openai import OpenAIEmbeddings
     from langchain.text_splitter import RecursiveCharacterTextSplitter
 
-    loader1 = DirectoryLoader('/home/user/app/docs/', glob="./*.pdf", loader_cls=PyPDFLoader)
+    loader1 = DirectoryLoader('/content/docs/', glob="./*.pdf", loader_cls=PyPDFLoader)
     document1 = loader1.load()
 
-    loader2 = DirectoryLoader('/home/user/app/docs/', glob="./*.txt", loader_cls=TextLoader)
+    loader2 = DirectoryLoader('/content/docs/', glob="./*.txt", loader_cls=TextLoader)
     document2 = loader2.load()
 
-    loader3 = DirectoryLoader('/home/user/app/docs/', glob="./*.docx", loader_cls=Docx2txtLoader)
+    loader3 = DirectoryLoader('/content/docs/', glob="./*.docx", loader_cls=Docx2txtLoader)
     document3 = loader3.load()
 
+    loader4 = DirectoryLoader('/content/docs/', glob="./*.csv", loader_cls=CSVLoader)
+    document4 = loader4.load()
+
+    loader5 = DirectoryLoader('/content/docs/', glob="./*.xlsx", loader_cls=UnstructuredExcelLoader)
+    document5 = loader5.load()
+
     document1.extend(document2)
     document1.extend(document3)
+    document1.extend(document4)
+    document1.extend(document5)
 
     text_splitter = RecursiveCharacterTextSplitter(
         chunk_size=1000,
@@ -52,11 +65,26 @@ def process_docs():
     embeddings = OpenAIEmbeddings()
 
     docs_db = FAISS.from_documents(docs, embeddings)
-    docs_db.save_local("/home/user/app/docs_db/")
+    docs_db.save_local("/content/docs_db/")
 
-    return "File(s) processed successfully!"
+    return "Successful!"
 
-def formatted_response(docs, response):
+global agent
+
+def create_agent():
+
+    from langchain.chat_models import ChatOpenAI
+    from langchain.chains.conversation.memory import ConversationSummaryBufferMemory
+    from langchain.chains import ConversationChain
+    global agent
+
+    llm = ChatOpenAI(model_name='gpt-3.5-turbo-16k')
+    memory = ConversationSummaryBufferMemory(llm=llm, max_token_limit=1000)
+    agent = ConversationChain(llm=llm, memory=memory, verbose=True)
+
+    return "Successful!"
+
+def formatted_response(docs, question, response, state):
 
     formatted_output = response + "\n\nSources"
 
@@ -71,49 +99,39 @@ def formatted_response(docs, response):
         else:
             formatted_output += f"\n{doc_name}"
 
-    return formatted_output
+    state.append((question, formatted_output))
+    return state, state
 
-def search_docs(question):
+def search_docs(prompt, question, state):
 
     from langchain.embeddings.openai import OpenAIEmbeddings
     from langchain.vectorstores import FAISS
-    from langchain.chains.question_answering import load_qa_chain
     from langchain.callbacks import get_openai_callback
-    from langchain.chat_models import ChatOpenAI
+    global agent
+    agent = agent
+
+    state = state or []
 
     embeddings = OpenAIEmbeddings()
-    docs_db = FAISS.load_local("/home/user/app/docs_db/", embeddings)
+    docs_db = FAISS.load_local("/content/docs_db/", embeddings)
     docs = docs_db.similarity_search(question)
 
-    llm = ChatOpenAI(model_name='gpt-3.5-turbo')
-    chain = load_qa_chain(llm, chain_type="stuff")
+    prompt += "\n\n"
+    prompt += question
+    prompt += "\n\n"
+    prompt += str(docs)
 
     with get_openai_callback() as cb:
-        response = chain.run(input_documents=docs, question=question)
+        response = agent.predict(input=prompt)
         print(cb)
 
-    return formatted_response(docs, response)
-
-def delete_docs():
-
-    import shutil
-
-    path1 = "/home/user/app/docs/"
-    path2 = "/home/user/app/docs_db/"
-
-    try:
-        shutil.rmtree(path1)
-        shutil.rmtree(path2)
-        return "Deleted Successfully"
-
-    except:
-        return "Already Deleted"
+    return formatted_response(docs, question, response, state)
 
 import gradio as gr
 
 css = """
 .col{
-    max-width: 50%;
+    max-width: 75%;
     margin: 0 auto;
     display: flex;
     flex-direction: column;
@@ -123,9 +141,9 @@ css = """
 """
 
 with gr.Blocks(css=css) as demo:
-    gr.Markdown("## <center>Lawyer GPT</center>")
+    gr.Markdown("## <center>All in One Document Chatting App</center>")
 
-    with gr.Tab("Your AI Legal Assistant"):
+    with gr.Tab("Chat With Your Documents"):
       with gr.Column(elem_classes="col"):
 
         with gr.Tab("Upload and Process Documents"):
@@ -138,28 +156,33 @@ with gr.Blocks(css=css) as demo:
             docs_process_button = gr.Button("Process")
             docs_process_output = gr.Textbox(label="Output")
 
-            gr.ClearButton([docs_upload_input, docs_upload_output, docs_process_output])
+            create_agent_button = gr.Button("Create Agent")
+            create_agent_output = gr.Textbox(label="Output")
+
+            gr.ClearButton([docs_upload_input, docs_upload_output, docs_process_output, create_agent_output])
 
         with gr.Tab("Query Documents"):
           with gr.Column():
 
-            docs_search_input = gr.Textbox(label="Enter Question")
-            docs_search_button = gr.Button("Search")
-            docs_search_output = gr.Textbox(label="Output")
+            docs_prompt_input = gr.Textbox(label="Custom Prompt")
 
-            docs_delete_button = gr.Button("Delete")
-            docs_delete_output = gr.Textbox(label="Output")
+            docs_chatbot = gr.Chatbot(label="Chats")
+            docs_state = gr.State()
 
-            gr.ClearButton([docs_search_input, docs_search_output, docs_delete_output])
+            docs_search_input = gr.Textbox(label="Question")
+            docs_search_button = gr.Button("Search")
+
+            gr.ClearButton([docs_prompt_input, docs_search_input])
 
     #########################################################################################################
+
     docs_upload_button.click(save_docs, inputs=docs_upload_input, outputs=docs_upload_output)
     docs_process_button.click(process_docs, inputs=None, outputs=docs_process_output)
+    create_agent_button.click(create_agent, inputs=None, outputs=create_agent_output)
 
-    docs_search_button.click(search_docs, inputs=docs_search_input, outputs=docs_search_output)
+    docs_search_button.click(search_docs, inputs=[docs_prompt_input, docs_search_input, docs_state], outputs=[docs_chatbot, docs_state])
 
-    docs_delete_button.click(delete_docs, inputs=None, outputs=docs_delete_output)
     #########################################################################################################
 
 demo.queue()
-demo.launch()
+demo.launch(debug=True, share=True)