Update app.py

app.py

@@ -22,6 +22,7 @@ from duckduckgo_search import DDGS
 import requests
 import tempfile
 
+
 tokenizer = tiktoken.encoding_for_model("gpt-3.5-turbo")
 
 # create the length function
@@ -149,14 +150,15 @@ def add_files_to_zip(session_id):
                 arcname = os.path.relpath(file_path, session_id)
                 zipObj.write(file_path, arcname)
 
-
 ## Search files functions ##
 
 def search_docs(topic, max_references):
+  print(f"SEARCH PDF : {topic}")
   doc_list = []
   with DDGS() as ddgs:
     i=0
     for r in ddgs.text('{} filetype:pdf'.format(topic), region='wt-wt', safesearch='On', timelimit='n'):
+      #doc_list.append(str(r))
       if i>=max_references:
         break
       doc_list.append("TITLE : " + r['title'] + " -- BODY : " + r['body'] + " -- URL : " + r['href'])
@@ -164,7 +166,7 @@ def search_docs(topic, max_references):
   return doc_list
 
 
-def store_files(references):
+def store_files(references, ret_names=False):
     url_list=[]
     temp_files = []
     for ref in references:
@@ -182,11 +184,13 @@ def store_files(references):
                 temp_file = tempfile.NamedTemporaryFile(delete=False,prefix=filename, suffix='.pdf')
             temp_file.write(response.content)
             temp_file.close()
-            temp_files.append(temp_file)
+            if ret_names:
+                temp_files.append(temp_file.name)
+            else:
+                temp_files.append(temp_file)
     
     return temp_files
-
-                
+    
 ## Summary functions ##
 
 ## Load each doc from the vector store
@@ -289,7 +293,7 @@ def embed_files(files,ui_session_id,progress=gr.Progress(),progress_step=0.05):
     
     print("EMBEDDED, before embeddeding: ",session_id,len(db.index_to_docstore_id))
     for file_id,file in enumerate(files):
-        print("ID : ", file_id,"FILE : ", file)
+        print("ID : ", file_id, "FILE : ", file)
         file_type = file.name.split('.')[-1].lower()
         source = file.name.split('/')[-1]
         print(f"current file: {source}")
@@ -330,19 +334,25 @@ def embed_files(files,ui_session_id,progress=gr.Progress(),progress_step=0.05):
     progress(progress_step, desc = 'db zipped')
     return f"{session_id}.zip",ui_session_id
 
-def display_docs(docs):
-  output_str = ''
-  for i, doc in enumerate(docs):
-      source = doc.metadata['source'].split('/')[-1]
-      output_str += f"Ref: {i+1}\n{repr(doc.page_content)}\nSource: {source}\n\n"
-  return output_str
 
 
 def add_to_db(references,ui_session_id):
     files = store_files(references)
     return embed_files(files,ui_session_id)
+
+def export_files(references):
+    files = store_files(references, ret_names=True)
+    #paths = [file.name for file in files]
+    return files
     
 
+def display_docs(docs):
+  output_str = ''
+  for i, doc in enumerate(docs):
+      source = doc.metadata['source'].split('/')[-1]
+      output_str += f"Ref: {i+1}\n{repr(doc.page_content)}\nSource: {source}\n\n"
+  return output_str
+
 def ask_gpt(query, apikey,history,ui_session_id):
     session_id = f"PDFAISS-{ui_session_id}"
     try:
@@ -368,7 +378,10 @@ with gr.Blocks() as demo:
     gr.Markdown("Upload your documents and question them.")
     with gr.Accordion("Open to enter your API key", open=False):
         apikey_input = gr.Textbox(placeholder="Type here your OpenAI API key to use Summarization and Q&A", label="OpenAI API Key",type='password')
-    with gr.Tab("Upload PDF & TXT"):
+
+
+        
+    with gr.Tab("Upload PDF & TXT"): 
         with gr.Accordion("Get files from the web", open=False):
             with gr.Column():
                 topic_input = gr.Textbox(placeholder="Type your research", label="Research")
@@ -379,6 +392,8 @@ with gr.Blocks() as demo:
                 dd_documents.style(container=True)
                 with gr.Row():
                     btn_dl = gr.Button("Add these files to the Database")
+                    btn_export = gr.Button("Export selected files ⬇⬇")
+                
         tb_session_id = gr.Textbox(label='session id')
         docs_input = gr.File(file_count="multiple", file_types=[".txt", ".pdf",".zip",".docx"])
         db_output = gr.outputs.File(label="Download zipped database")
@@ -406,10 +421,12 @@ with gr.Blocks() as demo:
 
     btn_search.click(search_docs, inputs=[topic_input, max_files], outputs=dd_documents)
     btn_dl.click(add_to_db, inputs=[dd_documents,tb_session_id], outputs=[db_output,tb_session_id])
+    topic_input.submit(export_files, inputs=dd_documents, outputs=docs_input)
+    btn_export.click(export_files, inputs=dd_documents, outputs=docs_input)
     btn_generate_db.click(embed_files, inputs=[docs_input,tb_session_id], outputs=[db_output,tb_session_id])
     btn_reset_db.click(reset_database,inputs=[tb_session_id],outputs=[db_output])
     btn_summary.click(summarize_docs, inputs=[apikey_input,tb_session_id], outputs=summary_output)
     btn_askGPT.click(ask_gpt, inputs=[query_input,apikey_input,history,tb_session_id], outputs=[answer_output,sources,history])
-
+#
 demo.queue(concurrency_count=10)
 demo.launch(debug=False,share=False)