deletion

app.py

@@ -61,6 +61,28 @@ class DocumentManager:
     def set_provide_citation(self, value):
         self.provide_citation = value
 
+    def delete_files(self):
+        self.docs = []
+        self.files = []
+        self.source_documents = []
+        self.db = Chroma(embedding_function=OpenAIEmbeddings(openai_api_key=self.api_key))
+
+        self.db._client_settings.allow_reset = True
+        self.db._client.reset()
+
+        self.retreiver = self.db.as_retriever(search_type="mmr", search_kwargs={'fetch_k': 30}, return_source_documents=True)
+        self.llm = ChatOpenAI(model_name="gpt-4", temperature=0, streaming=True, openai_api_key=self.api_key)
+
+        self.qa = RetrievalQA.from_chain_type(
+            llm=self.llm,
+            chain_type="stuff", 
+            retriever=Chroma(
+                embedding_function=OpenAIEmbeddings(openai_api_key=self.api_key))
+                    .as_retriever(search_type="mmr", search_kwargs={'fetch_k': 30}, return_source_documents=True),
+            return_source_documents=True)
+
+        return gr.Markdown(self.generate_file_markdown(), label="Uploaded files")
+
     def reset_qa(self):
         documents = self.text_splitter.split_documents(self.docs)
         self.retriever = Chroma.from_documents(documents, OpenAIEmbeddings(openai_api_key=self.api_key)).as_retriever(search_type="mmr", search_kwargs={'fetch_k': 30}, return_source_documents=True)
@@ -100,7 +122,9 @@ class DocumentManager:
 
         summed_history = " ".join(sum(history, []))
         question = "You have access to these documents:" + self.generate_file_markdown() + ". Do not make things up, only say what you have a primary source document for. ---- CHAT HISTORY : " + summed_history + " --- SYSTEM PROMPT: " + self.user_prompt + " -- Answer this question: " + message
+
         result = self.qa({"query": question})
+
         self.source_documents = result["source_documents"]
         self.result = result["result"]
         self.question = question
@@ -130,6 +154,12 @@ class DocumentManager:
 
         return gr.Markdown(self.generate_file_markdown(), label="Uploaded files") 
 
+    def create_delete_button(self, value):
+        if value and self.api_key:
+            return gr.Button("Delete files", scale=4, interactive=True)
+        else:
+            return gr.Button("Delete files", scale=4, interactive=False)
+
 def create_demo():
     doc_manager = DocumentManager()
 
@@ -163,6 +193,12 @@ def create_demo():
             prompt_textbox = gr.Textbox(label="Prompt", value=doc_manager.get_user_prompt())
             prompt_textbox.change(doc_manager.set_user_prompt, prompt_textbox)
 
+            with gr.Row():
+                allow_delete_checkbox = gr.Checkbox(value=False, label="Allow deletion of files") 
+                delete_button = doc_manager.create_delete_button(False)
+                delete_button.click(doc_manager.delete_files, outputs=file_output)
+                allow_delete_checkbox.select(doc_manager.create_delete_button, outputs=delete_button, inputs=allow_delete_checkbox)
+
             settings_tab.select(doc_manager.create_api_key_status_display, outputs=api_key_status)
 
     return demo