Update app.py via AI Editor

app.py

@@ -8,8 +8,7 @@ import pandas as pd
 import logging
 from docx import Document
 import mimetypes
-from threading import Thread, Lock, Event
-import threading
+from threading import Lock
 import time
 
 import google.generativeai as genai
@@ -38,8 +37,6 @@ shredded_document = None
 generated_response = None
 
 gemini_lock = Lock()
-stream_buffer = {"preview": ""}
-stream_event = Event()
 
 def decode_document(decoded_bytes):
     try:
@@ -89,7 +86,7 @@ def upload_to_gemini_file(decoded_bytes, filename):
         logging.error(f"Exception during file upload to Gemini: {e}")
         return None
 
-def gemini_generate_content_stream(prompt, file_id=None, chat_input=None, cancel_event=None):
+def gemini_generate_content(prompt, file_id=None, chat_input=None):
     try:
         files = []
         if file_id:
@@ -111,19 +108,10 @@ def gemini_generate_content_stream(prompt, file_id=None, chat_input=None, cancel
             )
         )
         result = response.text if hasattr(response, "text") else str(response)
-        chunk_size = 400
-        for i in range(0, len(result), chunk_size):
-            if cancel_event is not None and cancel_event.is_set():
-                logging.info("Gemini stream cancelled by user.")
-                yield "[Cancelled by user]\n"
-                break
-            yield result[:i+chunk_size]
-            time.sleep(0.08)
-        if cancel_event is not None and cancel_event.is_set():
-            yield "[Cancelled by user]\n"
+        return result
     except Exception as e:
         logging.error("Error during Gemini generate_content: %s", e)
-        yield f"Error during Gemini completion: {e}"
+        return f"Error during Gemini completion: {e}"
 
 def save_shredded_as_docx(shredded_text, rfp_filename):
     doc = Document()
@@ -155,8 +143,8 @@ def save_compliance_as_docx(compliance_text, rfp_filename):
     memf.seek(0)
     return memf.read()
 
-def process_document(action, selected_filename=None, chat_input=None, rfp_decoded_bytes=None, cancel_event=None, stream=True, selected_generated_filename=None):
-    global shredded_document, generated_response, stream_buffer
+def process_document(action, selected_filename=None, chat_input=None, rfp_decoded_bytes=None, selected_generated_filename=None):
+    global shredded_document, generated_response
 
     logging.info(f"Process document called with action: {action}")
 
@@ -186,44 +174,17 @@ def process_document(action, selected_filename=None, chat_input=None, rfp_decode
         if chat_input:
             prompt += f"User additional instructions: {chat_input}\n"
         prompt += f"\nFile Name: {selected_filename}\n\n"
-        result_holder = {"text": None, "docx_name": None}
-
-        def thread_shred():
-            global shredded_document, stream_buffer
-            shredded_document = ""
-            stream_buffer["preview"] = ""
-            try:
-                logging.info("Starting Gemini generate_content for shredding.")
-                for partial in gemini_generate_content_stream(prompt, file_id=doc_fileid, chat_input=chat_input, cancel_event=cancel_event):
-                    stream_buffer["preview"] = partial
-                    if cancel_event is not None and cancel_event.is_set():
-                        break
-                result = stream_buffer["preview"]
-                shredded_document = result
-                if not cancel_event or not cancel_event.is_set():
-                    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_name"] = generated_docx_name
-                else:
-                    result_holder["text"] = "[Cancelled by user]\n"
-                    result_holder["docx_name"] = None
-                logging.info("Document shredded successfully.")
-            except Exception as e:
-                shredded_document = f"Error during shredding: {e}"
-                logging.error("Error in thread_shred: %s", e)
-                result_holder["text"] = shredded_document
-
-        shredded_document = "Shredding in progress..."
-        stream_buffer["preview"] = ""
-        t = Thread(target=thread_shred)
-        t.start()
-        t.join()
-        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"]
+        result = gemini_generate_content(prompt, file_id=doc_fileid, chat_input=chat_input)
+        shredded_document = result
+        if result and not result.startswith("Error"):
+            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
+            if generated_docx_name not in uploaded_documents:
+                uploaded_documents[generated_docx_name] = result
+            return result, None, generated_docx_name, result
+        else:
+            return result, None, None, result
 
     elif action == 'proposal':
         if not doc_content:
@@ -249,42 +210,16 @@ def process_document(action, selected_filename=None, chat_input=None, rfp_decode
         if chat_input:
             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_name": None}
-
-        def thread_proposal():
-            global generated_response, stream_buffer
-            generated_response = ""
-            stream_buffer["preview"] = ""
-            try:
-                logging.info("Connecting to Gemini for proposal.")
-                for partial in gemini_generate_content_stream(prompt, file_id=rfp_fileid, chat_input=chat_input, cancel_event=cancel_event):
-                    stream_buffer["preview"] = partial
-                    if cancel_event is not None and cancel_event.is_set():
-                        break
-                response = stream_buffer["preview"]
-                generated_response = response
-                if not cancel_event or not cancel_event.is_set():
-                    docx_bytes = save_proposal_as_docx(response, rfp_filename)
-                    generated_docx_name = f"{os.path.splitext(rfp_filename)[0]}_proposal.docx"
-                    uploaded_proposals[generated_docx_name] = response
-                    uploaded_proposals_fileid[generated_docx_name] = None
-                    result_holder["text"] = response
-                    result_holder["docx_name"] = generated_docx_name
-                else:
-                    result_holder["text"] = "[Cancelled by user]\n"
-                    result_holder["docx_name"] = None
-                logging.info("Received proposal results from Gemini.")
-            except Exception as e:
-                generated_response = f"Error during Gemini completion: {e}"
-                logging.error("Error during Gemini proposal request: %s", e)
-                result_holder["text"] = generated_response
-
-        stream_buffer["preview"] = ""
-        t = Thread(target=thread_proposal)
-        t.start()
-        t.join()
-        return result_holder["text"], None, result_holder["docx_name"], result_holder["text"]
+        result = gemini_generate_content(prompt, file_id=rfp_fileid, chat_input=chat_input)
+        generated_response = result
+        if result and not result.startswith("Error"):
+            docx_bytes = save_proposal_as_docx(result, rfp_filename)
+            generated_docx_name = f"{os.path.splitext(rfp_filename)[0]}_proposal.docx"
+            uploaded_proposals[generated_docx_name] = result
+            uploaded_proposals_fileid[generated_docx_name] = None
+            return result, None, generated_docx_name, result
+        else:
+            return result, None, None, result
 
     elif action == 'compliance':
         if not selected_generated_filename or selected_generated_filename not in uploaded_proposals:
@@ -308,37 +243,15 @@ def process_document(action, selected_filename=None, chat_input=None, rfp_decode
             "---\nGenerated Proposal Document:\n"
             f"{proposal_text}\n"
         )
-        result_holder = {"text": None, "docx_name": None}
-
-        def thread_compliance():
-            global stream_buffer
-            stream_buffer["preview"] = ""
-            try:
-                logging.info("Starting Gemini generate_content for compliance check.")
-                for partial in gemini_generate_content_stream(prompt, file_id=None, chat_input=None, cancel_event=cancel_event):
-                    stream_buffer["preview"] = partial
-                    if cancel_event is not None and cancel_event.is_set():
-                        break
-                result = stream_buffer["preview"]
-                result_holder["text"] = result
-                if not cancel_event or not cancel_event.is_set():
-                    docx_bytes = save_compliance_as_docx(result, selected_filename)
-                    compliance_docx_name = f"{os.path.splitext(selected_filename)[0]}_compliance_check.docx"
-                    uploaded_documents[compliance_docx_name] = result
-                    shredded_documents[compliance_docx_name] = docx_bytes
-                    result_holder["docx_name"] = compliance_docx_name
-                else:
-                    result_holder["docx_name"] = None
-                logging.info("Compliance check completed.")
-            except Exception as e:
-                logging.error("Error during compliance check: %s", e)
-                result_holder["text"] = f"Error during compliance check: {e}"
-
-        stream_buffer["preview"] = ""
-        t = Thread(target=thread_compliance)
-        t.start()
-        t.join()
-        return result_holder["text"], None, result_holder["docx_name"], result_holder["text"]
+        result = gemini_generate_content(prompt, file_id=None, chat_input=None)
+        if result and not result.startswith("Error"):
+            docx_bytes = save_compliance_as_docx(result, selected_filename)
+            compliance_docx_name = f"{os.path.splitext(selected_filename)[0]}_compliance_check.docx"
+            uploaded_documents[compliance_docx_name] = result
+            shredded_documents[compliance_docx_name] = docx_bytes
+            return result, None, compliance_docx_name, result
+        else:
+            return result, None, None, result
 
     elif action == 'recover':
         return "Recovery not implemented yet.", None, None, None
@@ -427,10 +340,6 @@ def get_uploaded_proposal_list(docdict):
     return dbc.ListGroup(doc_list, flush=True)
 
 app.layout = dbc.Container([
-    dcc.Store(id='shred-store', data={'text': None, 'docx_name': None}),
-    dcc.Store(id='proposal-store', data={'text': None, 'docx_name': None}),
-    dcc.Store(id='stream-status', data={'streaming': False}),
-    dcc.Interval(id='stream-interval', interval=5000, n_intervals=0, disabled=True),
     dbc.Row([
         dbc.Col([
             dbc.Card([
@@ -531,8 +440,6 @@ app.layout = dbc.Container([
 ], fluid=True)
 
 @app.callback(
-    Output('shred-store', 'data'),
-    Output('proposal-store', 'data'),
     Output('output-data-upload', 'children'),
     Output('uploaded-doc-list', 'children'),
     Output('select-document-dropdown', 'options'),
@@ -541,8 +448,6 @@ app.layout = dbc.Container([
     Output('uploaded-proposal-list', 'children'),
     Output('select-proposal-dropdown', 'options'),
     Output('select-proposal-dropdown', 'value'),
-    Output('stream-status', 'data'),
-    Output('stream-interval', 'disabled'),
     [
         Input('shred-action-btn', 'n_clicks'),
         Input('proposal-action-btn', 'n_clicks'),
@@ -561,9 +466,7 @@ app.layout = dbc.Container([
         State('chat-input', 'value'),
         State('select-document-dropdown', 'value'),
         State('select-proposal-dropdown', 'value'),
-        Input('cancel-action-btn', 'n_clicks'),
-        Input('stream-interval', 'n_intervals'),
-        State('stream-status', 'data')   # <-- Moved this to the last State, matching the last argument
+        Input('cancel-action-btn', 'n_clicks')
     ],
     prevent_initial_call=True
 )
@@ -574,8 +477,7 @@ def master_callback(
     shredded_delete_clicks, shredded_doc_children,
     select_proposal_dropdown_value,
     chat_input, selected_filename, selected_proposal_dropdown, selected_proposal_dropdown_state,
-    cancel_clicks,
-    stream_n_intervals, stream_status
+    cancel_clicks
 ):
     ctx = callback_context
     triggered_id = ctx.triggered[0]['prop_id'].split('.')[0] if ctx.triggered else None
@@ -588,46 +490,15 @@ def master_callback(
 
     upload_triggered = False
 
-    shred_store = {'text': None, 'docx_name': None}
-    proposal_store = {'text': None, 'docx_name': None}
-    streaming = False
-
     rfp_delete_clicks = safe_get_n_clicks(ctx, 5)
     proposal_delete_clicks = safe_get_n_clicks(ctx, 9)
     shredded_delete_clicks = safe_get_n_clicks(ctx, 11)
 
     uploaded_rfp_decoded_bytes = None
 
-    global gemini_lock, stream_buffer, stream_event
-
-    if triggered_id == 'stream-interval':
-        if not stream_status or not stream_status.get('streaming'):
-            return dash.no_update, dash.no_update, dash.no_update, dash.no_update, dash.no_update, dash.no_update, dash.no_update, dash.no_update, dash.no_update, dash.no_update, stream_status, True
-        preview = stream_buffer.get("preview", "")
-        still_streaming = gemini_lock.locked()
-        if not still_streaming:
-            stream_status['streaming'] = False
-            output_data_upload = dcc.Markdown(preview, style={"whiteSpace": "pre-wrap", "wordWrap": "break-word"})
-            doc_options = [{'label': fn, 'value': fn} for fn in uploaded_documents.keys()]
-            doc_value = selected_doc if selected_doc in uploaded_documents else (next(iter(uploaded_documents), None) if uploaded_documents else None)
-            shredded_doc_list_items = get_shredded_doc_list(shredded_documents)
-            uploaded_doc_list = get_uploaded_doc_list(uploaded_documents)
-            uploaded_proposal_list = get_uploaded_proposal_list(uploaded_proposals)
-            proposal_options = [{'label': fn, 'value': fn} for fn in uploaded_proposals.keys()]
-            proposal_value = select_proposal_dropdown_value if select_proposal_dropdown_value in uploaded_proposals else (next(iter(uploaded_proposals), None) if uploaded_proposals else None)
-            return (
-                shred_store, proposal_store, output_data_upload,
-                uploaded_doc_list, doc_options, doc_value,
-                shredded_doc_list_items,
-                uploaded_proposal_list, proposal_options, proposal_value,
-                stream_status,
-                True
-            )
-        return dash.no_update, dash.no_update, dcc.Markdown(preview, style={"whiteSpace": "pre-wrap", "wordWrap": "break-word"}), dash.no_update, dash.no_update, dash.no_update, dash.no_update, dash.no_update, dash.no_update, dash.no_update, stream_status, False
+    global gemini_lock
 
     if triggered_id == 'cancel-action-btn':
-        stream_event.set()
-        streaming = False
         output_data_upload = html.Div("[Cancelled by user]\n", style={"wordWrap": "break-word"})
         doc_options = [{'label': fn, 'value': fn} for fn in uploaded_documents.keys()]
         doc_value = selected_doc if selected_doc in uploaded_documents else (next(iter(uploaded_documents), None) if uploaded_documents else None)
@@ -637,12 +508,10 @@ def master_callback(
         proposal_options = [{'label': fn, 'value': fn} for fn in uploaded_proposals.keys()]
         proposal_value = select_proposal_dropdown_value if select_proposal_dropdown_value in uploaded_proposals else (next(iter(uploaded_proposals), None) if uploaded_proposals else None)
         return (
-            shred_store, proposal_store, output_data_upload,
+            output_data_upload,
             uploaded_doc_list, doc_options, doc_value,
             shredded_doc_list_items,
-            uploaded_proposal_list, proposal_options, proposal_value,
-            {'streaming': False},
-            True
+            uploaded_proposal_list, proposal_options, proposal_value
         )
 
     if triggered_id == 'upload-document' and rfp_content is not None and rfp_filename:
@@ -742,33 +611,29 @@ def master_callback(
         if not got_lock:
             output_data_upload = html.Div("Another Gemini operation is in progress. Please wait or cancel.", style={"wordWrap": "break-word"})
             return (
-                shred_store, proposal_store, output_data_upload,
+                output_data_upload,
                 uploaded_doc_list, doc_options, doc_value,
                 shredded_doc_list_items,
-                uploaded_proposal_list, proposal_options, proposal_value,
-                {'streaming': False},
-                True
+                uploaded_proposal_list, proposal_options, proposal_value
             )
-        stream_event.clear()
-        stream_buffer["preview"] = ""
-        streaming = True
-        def stream_gemini_thread(action, doc_value, chat_input, rfp_decoded_bytes, selected_generated_filename):
-            try:
-                process_document(action, doc_value, chat_input, rfp_decoded_bytes, cancel_event=stream_event, stream=True, selected_generated_filename=selected_generated_filename)
-            finally:
-                gemini_lock.release()
-        action_name = "shred" if triggered_id=="shred-action-btn" else ("proposal" if triggered_id=="proposal-action-btn" else "compliance")
-        t = Thread(target=stream_gemini_thread, args=(action_name, doc_value, chat_input, uploaded_rfp_decoded_bytes, select_proposal_dropdown_value))
-        t.daemon = True
-        t.start()
-        output_data_upload = dcc.Markdown("Starting Gemini operation...", style={"whiteSpace": "pre-wrap", "wordWrap": "break-word"})
+        try:
+            action_name = "shred" if triggered_id=="shred-action-btn" else ("proposal" if triggered_id=="proposal-action-btn" else "compliance")
+            result, _, _, _ = process_document(action_name, doc_value, chat_input, uploaded_rfp_decoded_bytes, select_proposal_dropdown_value)
+            output_data_upload = dcc.Markdown(result, style={"whiteSpace": "pre-wrap", "wordWrap": "break-word"})
+        finally:
+            gemini_lock.release()
+        doc_options = [{'label': fn, 'value': fn} for fn in uploaded_documents.keys()]
+        doc_value = doc_value if doc_value in uploaded_documents else (next(iter(uploaded_documents), None) if uploaded_documents else None)
+        proposal_options = [{'label': fn, 'value': fn} for fn in uploaded_proposals.keys()]
+        proposal_value = proposal_value if proposal_value in uploaded_proposals else (next(iter(uploaded_proposals), None) if uploaded_proposals else None)
+        shredded_doc_list_items = get_shredded_doc_list(shredded_documents)
+        uploaded_doc_list = get_uploaded_doc_list(uploaded_documents)
+        uploaded_proposal_list = get_uploaded_proposal_list(uploaded_proposals)
         return (
-            shred_store, proposal_store, output_data_upload,
+            output_data_upload,
             uploaded_doc_list, doc_options, doc_value,
             shredded_doc_list_items,
-            uploaded_proposal_list, proposal_options, proposal_value,
-            {'streaming': True},
-            False
+            uploaded_proposal_list, proposal_options, proposal_value
         )
 
     if upload_triggered:
@@ -776,23 +641,19 @@ def master_callback(
         proposal_value = proposal_value if proposal_value in uploaded_proposals else (next(iter(uploaded_proposals), None) if uploaded_proposals else None)
         output_data_upload = html.Div("Upload/Delete completed.", style={"wordWrap": "break-word"})
         return (
-            shred_store, proposal_store, output_data_upload,
+            output_data_upload,
             uploaded_doc_list, doc_options, doc_value,
             shredded_doc_list_items,
-            uploaded_proposal_list, proposal_options, proposal_value,
-            {'streaming': False},
-            True
+            uploaded_proposal_list, proposal_options, proposal_value
         )
 
     doc_value = doc_value if doc_value in uploaded_documents else (next(iter(uploaded_documents), None) if uploaded_documents else None)
     proposal_value = proposal_value if proposal_value in uploaded_proposals else (next(iter(uploaded_proposals), None) if uploaded_proposals else None)
     return (
-        shred_store, proposal_store, output_data_upload,
+        output_data_upload,
         uploaded_doc_list, doc_options, doc_value,
         shredded_doc_list_items,
-        uploaded_proposal_list, proposal_options, proposal_value,
-        {'streaming': False},
-        True
+        uploaded_proposal_list, proposal_options, proposal_value
     )
 
 if __name__ == '__main__':