Update app.py via AI Editor

app.py

@@ -162,7 +162,7 @@ def process_document(action, selected_filename=None, chat_input=None):
         if chat_input:
             prompt += f"User additional instructions: {chat_input}\n"
         prompt += f"\nFile Name: {selected_filename}\n\n"
-        result_holder = {"text": None, "docx_bytes": None, "docx_name": None}
+        result_holder = {"text": None, "docx_name": None}
         def thread_shred():
             global shredded_document
             shredded_document = ""
@@ -173,8 +173,8 @@ def process_document(action, selected_filename=None, chat_input=None):
                 logging.info("Document shredded successfully.")
                 docx_bytes = save_shredded_as_docx(result, selected_filename)
                 generated_docx_name = f"{os.path.splitext(selected_filename)[0]}_shredded.docx"
+                shredded_documents[generated_docx_name] = docx_bytes
                 result_holder["text"] = result
-                result_holder["docx_bytes"] = docx_bytes
                 result_holder["docx_name"] = generated_docx_name
             except Exception as e:
                 shredded_document = f"Error during shredding: {e}"
@@ -184,13 +184,10 @@ def process_document(action, selected_filename=None, chat_input=None):
         t = Thread(target=thread_shred)
         t.start()
         t.join()
-        if result_holder["docx_bytes"] and result_holder["docx_name"]:
-            docx_name = result_holder["docx_name"]
-            try:
-                uploaded_documents[docx_name] = result_holder["text"]
-            except Exception as e:
-                logging.error(f"Error adding shredded docx to uploaded_documents: {e}")
-        return result_holder["text"], result_holder["docx_bytes"], result_holder["docx_name"], result_holder["text"]
+        if result_holder["docx_name"] and result_holder["text"]:
+            if result_holder["docx_name"] not in uploaded_documents:
+                uploaded_documents[result_holder["docx_name"]] = result_holder["text"]
+        return result_holder["text"], None, result_holder["docx_name"], result_holder["text"]
 
     elif action == 'generate':
         if not shredded_document:
@@ -204,7 +201,7 @@ def process_document(action, selected_filename=None, chat_input=None):
         if chat_input:
             prompt += f"User additional instructions: {chat_input}\n"
         prompt += f"\nFile Name: {selected_filename}\n\n{shredded_document}"
-        result_holder = {"text": None, "docx_bytes": None, "docx_name": None}
+        result_holder = {"text": None, "docx_name": None}
         def thread_generate():
             global generated_response
             generated_response = ""
@@ -215,8 +212,8 @@ def process_document(action, selected_filename=None, chat_input=None):
                 logging.info("Proposal response generated successfully.")
                 docx_bytes = save_proposal_as_docx(result, selected_filename)
                 generated_docx_name = f"{os.path.splitext(selected_filename)[0]}_proposal.docx"
+                generated_documents[generated_docx_name] = docx_bytes
                 result_holder["text"] = result
-                result_holder["docx_bytes"] = docx_bytes
                 result_holder["docx_name"] = generated_docx_name
             except Exception as e:
                 generated_response = f"Error during generation: {e}"
@@ -226,7 +223,7 @@ def process_document(action, selected_filename=None, chat_input=None):
         t = Thread(target=thread_generate)
         t.start()
         t.join()
-        return result_holder["text"], result_holder["docx_bytes"], result_holder["docx_name"], result_holder["text"]
+        return result_holder["text"], None, result_holder["docx_name"], result_holder["text"]
 
     elif action == 'proposal':
         if not doc_content:
@@ -244,7 +241,7 @@ def process_document(action, selected_filename=None, chat_input=None):
             prompt += f"User additional instructions: {chat_input}\n"
         prompt += f"\n---\nRFP/SOW/PWS/RFI ({rfp_filename}):\n{doc_content}\n"
         logging.info(f"Sending proposal prompt to Gemini. RFP: {rfp_filename}")
-        result_holder = {"text": None, "docx_bytes": None, "docx_name": None}
+        result_holder = {"text": None, "docx_name": None}
         def thread_proposal():
             try:
                 logging.info("Connecting to Gemini for proposal.")
@@ -252,8 +249,8 @@ def process_document(action, selected_filename=None, chat_input=None):
                 logging.info("Received proposal results from Gemini.")
                 docx_bytes = save_proposal_as_docx(response, rfp_filename)
                 generated_docx_name = f"{os.path.splitext(rfp_filename)[0]}_proposal.docx"
+                generated_documents[generated_docx_name] = docx_bytes
                 result_holder["text"] = response
-                result_holder["docx_bytes"] = docx_bytes
                 result_holder["docx_name"] = generated_docx_name
             except Exception as e:
                 logging.error("Error during Gemini proposal request: %s", e)
@@ -261,7 +258,7 @@ def process_document(action, selected_filename=None, chat_input=None):
         t = Thread(target=thread_proposal)
         t.start()
         t.join()
-        return result_holder["text"], result_holder["docx_bytes"], result_holder["docx_name"], result_holder["text"]
+        return result_holder["text"], None, result_holder["docx_name"], result_holder["text"]
 
     elif action == 'compliance':
         return "Compliance checking not implemented yet.", None, None, None
@@ -347,8 +344,8 @@ def get_ctx_input_n_clicks(ctx, idx):
         return []
 
 app.layout = dbc.Container([
-    dcc.Store(id='shred-store', data={'text': None, 'docx_bytes': None, 'docx_name': None}),
-    dcc.Store(id='proposal-store', data={'text': None, 'docx_bytes': None, 'docx_name': None}),
+    dcc.Store(id='shred-store', data={'text': None, 'docx_name': None}),
+    dcc.Store(id='proposal-store', data={'text': None, 'docx_name': None}),
     dbc.Row([
         dbc.Col([
             dbc.Card([
@@ -514,8 +511,8 @@ def master_callback(
 
     upload_triggered = False
 
-    shred_store = {'text': None, 'docx_bytes': None, 'docx_name': None}
-    proposal_store = {'text': None, 'docx_bytes': None, 'docx_name': None}
+    shred_store = {'text': None, 'docx_name': None}
+    proposal_store = {'text': None, 'docx_name': None}
 
     rfp_delete_clicks = safe_get_n_clicks(ctx, 4)
     proposal_delete_clicks = safe_get_n_clicks(ctx, 8)
@@ -628,14 +625,10 @@ def master_callback(
 
     if triggered_id == 'shred-action-btn':
         output_data_upload = dcc.Loading(type="default", children=html.Div("Shredding document...", style={"wordWrap": "break-word"}))
-        shred_text, shredded_docx_bytes, shredded_docx_name, shredded_text = process_document('shred', doc_value, chat_input)
-        shred_store = {'text': shred_text, 'docx_bytes': shredded_docx_bytes, 'docx_name': shredded_docx_name}
-
-        if shredded_docx_bytes and shredded_docx_name:
-            shredded_documents[shredded_docx_name] = shredded_docx_bytes
-            logging.info(f"Shredded docx saved: {shredded_docx_name}")
-            if shredded_docx_name not in uploaded_documents:
-                uploaded_documents[shredded_docx_name] = shredded_text
+        shred_text, _, shredded_docx_name, shredded_text = process_document('shred', doc_value, chat_input)
+        shred_store = {'text': shred_text, 'docx_name': shredded_docx_name}
+
+        if shredded_docx_name:
             doc_options = [{'label': fn, 'value': fn} for fn in uploaded_documents.keys()]
             doc_value = shredded_docx_name if shredded_docx_name in uploaded_documents else (next(iter(uploaded_documents), None) if uploaded_documents else None)
             uploaded_doc_list = get_uploaded_doc_list(uploaded_documents)
@@ -650,11 +643,10 @@ def master_callback(
 
     if triggered_id == 'proposal-action-btn':
         output_data_upload = dcc.Loading(type="default", children=html.Div("Generating proposal...", style={"wordWrap": "break-word"}))
-        proposal_text, proposal_docx_bytes, proposal_docx_name, proposal_text_preview = process_document('proposal', doc_value, chat_input)
-        proposal_store = {'text': proposal_text, 'docx_bytes': proposal_docx_bytes, 'docx_name': proposal_docx_name}
+        proposal_text, _, proposal_docx_name, proposal_text_preview = process_document('proposal', doc_value, chat_input)
+        proposal_store = {'text': proposal_text, 'docx_name': proposal_docx_name}
         new_generated_doc_value = generated_doc_value
-        if proposal_docx_bytes and proposal_docx_name:
-            generated_documents[proposal_docx_name] = proposal_docx_bytes
+        if proposal_docx_name:
             generated_doc_options = [{'label': fn, 'value': fn} for fn in generated_documents.keys()]
             new_generated_doc_value = proposal_docx_name if proposal_docx_name in generated_documents else (next(iter(generated_documents), None) if generated_documents else None)
             generated_doc_list = get_generated_doc_list(generated_documents)