Update app.py via AI Editor

app.py

@@ -8,7 +8,9 @@ import pandas as pd
 import logging
 from docx import Document
 import mimetypes
-from threading import Thread
+from threading import Thread, Lock, Event
+import threading
+import time
 
 import google.generativeai as genai
 
@@ -36,6 +38,11 @@ shredded_documents = {}
 shredded_document = None
 generated_response = None
 
+# Streaming and cancel globals
+gemini_lock = Lock()
+stream_buffer = {"preview": ""}  # Only one operation at a time
+stream_event = Event()
+
 def decode_document(decoded_bytes):
     try:
         content = decoded_bytes.decode('utf-8')
@@ -84,7 +91,8 @@ def upload_to_gemini_file(decoded_bytes, filename):
         logging.error(f"Exception during file upload to Gemini: {e}")
         return None
 
-def gemini_generate_content(prompt, file_id=None, chat_input=None):
+def gemini_generate_content_stream(prompt, file_id=None, chat_input=None, cancel_event=None):
+    # Streaming Gemini not officially supported, so simulate by splitting result into manageable chunks
     try:
         files = []
         if file_id:
@@ -99,6 +107,7 @@ def gemini_generate_content(prompt, file_id=None, chat_input=None):
             content_list.append("\n\n")
         content_list.append(prompt)
         model = genai.GenerativeModel(GEMINI_MODEL)
+        # No streaming in SDK, so emulate
         response = model.generate_content(
             contents=content_list,
             generation_config=genai.types.GenerationConfig(
@@ -106,11 +115,19 @@ def gemini_generate_content(prompt, file_id=None, chat_input=None):
             )
         )
         result = response.text if hasattr(response, "text") else str(response)
-        logging.info("Received response from Gemini.")
-        return result
+        chunk_size = 400  # chars per chunk
+        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"
     except Exception as e:
         logging.error("Error during Gemini generate_content: %s", e)
-        return f"Error during Gemini completion: {e}"
+        yield f"Error during Gemini completion: {e}"
 
 def save_shredded_as_docx(shredded_text, rfp_filename):
     doc = Document()
@@ -132,8 +149,8 @@ def save_proposal_as_docx(proposal_text, base_filename):
     memf.seek(0)
     return memf.read()
 
-def process_document(action, selected_filename=None, chat_input=None, rfp_decoded_bytes=None):
-    global shredded_document, generated_response
+def process_document(action, selected_filename=None, chat_input=None, rfp_decoded_bytes=None, cancel_event=None, stream=True):
+    global shredded_document, generated_response, stream_buffer
     logging.info(f"Process document called with action: {action}")
 
     doc_content = None
@@ -163,24 +180,36 @@ def process_document(action, selected_filename=None, chat_input=None, rfp_decode
             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
+            global shredded_document, stream_buffer
             shredded_document = ""
+            stream_buffer["preview"] = ""
             try:
                 logging.info("Starting Gemini generate_content for shredding.")
-                result = gemini_generate_content(prompt, file_id=doc_fileid, chat_input=chat_input)
+                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.")
-                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
             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()
@@ -216,23 +245,35 @@ def process_document(action, selected_filename=None, chat_input=None, rfp_decode
         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
+            global generated_response, stream_buffer
             generated_response = ""
+            stream_buffer["preview"] = ""
             try:
                 logging.info("Connecting to Gemini for proposal.")
-                response = gemini_generate_content(prompt, file_id=rfp_fileid, chat_input=chat_input)
+                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"
+                    generated_documents[generated_docx_name] = docx_bytes
+                    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.")
-                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_name"] = generated_docx_name
             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()
@@ -254,20 +295,16 @@ def get_uploaded_doc_list(docdict):
     doc_list = []
     for filename in docdict:
         file_content = docdict[filename]
-        # Always encode as docx for download, but if not docx, encode as txt
         if filename.lower().endswith('.docx'):
             mime = "application/vnd.openxmlformats-officedocument.wordprocessingml.document"
             b64 = None
-            # docx as text, not as file: re-encode to docx for download
             try:
-                # If it's a docx, but only text available, save as docx
                 docx_bytes = save_shredded_as_docx(file_content, filename)
                 b64 = base64.b64encode(docx_bytes).decode('utf-8')
             except Exception:
                 b64 = base64.b64encode(file_content.encode('utf-8')).decode('utf-8')
         else:
             mime = "text/plain"
-            # Always encode as utf-8 text for download
             b64 = base64.b64encode(file_content.encode('utf-8')).decode('utf-8')
         download_link = html.A(
             filename,
@@ -380,6 +417,8 @@ def get_generated_rfp_download_section(selected_value):
 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=500, n_intervals=0, disabled=True),
     dbc.Row([
         dbc.Col([
             dbc.Card([
@@ -482,6 +521,7 @@ app.layout = dbc.Container([
                         dbc.Button("Recover", id="recover-action-btn", className="me-3 mb-2 btn-tertiary"),
                         dbc.Button("Virtual Board", id="board-action-btn", className="me-3 mb-2 btn-tertiary"),
                         dbc.Button("LOE", id="loe-action-btn", className="mb-2 btn-tertiary"),
+                        dbc.Button("Cancel", id="cancel-action-btn", className="ms-3 mb-2 btn-danger", color="danger"),
                     ], className="mt-3 mb-3 d-flex flex-wrap"),
                     dcc.Loading(
                         id="loading",
@@ -507,6 +547,8 @@ app.layout = dbc.Container([
     Output('select-generated-dropdown', 'options'),
     Output('select-generated-dropdown', 'value'),
     Output('generated-rfp-section', 'children'),
+    Output('stream-status', 'data'),
+    Output('stream-interval', 'disabled'),
     [
         Input('shred-action-btn', 'n_clicks'),
         Input('proposal-action-btn', 'n_clicks'),
@@ -527,6 +569,7 @@ app.layout = dbc.Container([
         State('select-document-dropdown', 'value'),
         State('select-proposal-dropdown', 'value'),
         State('select-generated-dropdown', 'value'),
+        Input('cancel-action-btn', 'n_clicks')
     ],
     prevent_initial_call=True
 )
@@ -537,7 +580,8 @@ def master_callback(
     generated_delete_clicks, selected_generated,
     shredded_delete_clicks, shredded_doc_children,
     select_generated_value,
-    chat_input, selected_filename, selected_proposal_dropdown, selected_generated_dropdown_state
+    chat_input, selected_filename, selected_proposal_dropdown, selected_generated_dropdown_state,
+    cancel_clicks
 ):
     ctx = callback_context
     triggered_id = ctx.triggered[0]['prop_id'].split('.')[0] if ctx.triggered else None
@@ -552,6 +596,7 @@ def master_callback(
 
     shred_store = {'text': None, 'docx_name': None}
     proposal_store = {'text': None, 'docx_name': None}
+    streaming = False
 
     rfp_delete_clicks = safe_get_n_clicks(ctx, 4)
     proposal_delete_clicks = safe_get_n_clicks(ctx, 8)
@@ -560,6 +605,30 @@ def master_callback(
 
     uploaded_rfp_decoded_bytes = None
 
+    global gemini_lock, stream_buffer, stream_event
+
+    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)
+        shredded_doc_list_items = get_shredded_doc_list(shredded_documents)
+        uploaded_doc_list = get_uploaded_doc_list(uploaded_documents)
+        generated_doc_list = get_generated_doc_list(generated_documents)
+        generated_doc_options = [{'label': fn, 'value': fn} for fn in generated_documents.keys()]
+        generated_doc_value = select_generated_value if select_generated_value in generated_documents else (next(iter(generated_documents), None) if generated_documents else None)
+        generated_rfp_section = get_generated_rfp_download_section(generated_doc_value)
+        return (
+            shred_store, proposal_store, output_data_upload,
+            uploaded_doc_list, doc_options, doc_value,
+            shredded_doc_list_items,
+            generated_doc_list, generated_doc_options, generated_doc_value,
+            generated_rfp_section,
+            {'streaming': False},
+            True
+        )
+
     if triggered_id == 'upload-document' and rfp_content is not None and rfp_filename:
         content_type, content_string = rfp_content.split(',')
         decoded = base64.b64decode(content_string)
@@ -666,54 +735,41 @@ def master_callback(
 
     output_data_upload = html.Div("No action taken yet.", style={"wordWrap": "break-word"})
 
-    if triggered_id == 'shred-action-btn':
-        output_data_upload = dcc.Loading(type="default", children=html.Div("Shredding document...", style={"wordWrap": "break-word"}), style={"textAlign": "center"})
-        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)
-            shredded_doc_list_items = get_shredded_doc_list(shredded_documents)
-        output_data_upload = dcc.Markdown(shred_text, style={"whiteSpace": "pre-wrap", "wordWrap": "break-word"})
-        generated_rfp_section = get_generated_rfp_download_section(generated_doc_value)
+    # Stream logic: only one at a time
+    if triggered_id in ['shred-action-btn', 'proposal-action-btn']:
+        got_lock = gemini_lock.acquire(blocking=False)
+        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,
+                uploaded_doc_list, doc_options, doc_value,
+                shredded_doc_list_items,
+                generated_doc_list, generated_doc_options, generated_doc_value,
+                generated_rfp_section,
+                {'streaming': False},
+                True
+            )
+        stream_event.clear()
+        stream_buffer["preview"] = ""
+        streaming = True
+        # Launch thread and let interval update the preview
+        def stream_gemini_thread(action, doc_value, chat_input, rfp_decoded_bytes):
+            try:
+                process_document(action, doc_value, chat_input, rfp_decoded_bytes, cancel_event=stream_event, stream=True)
+            finally:
+                gemini_lock.release()
+        t = Thread(target=stream_gemini_thread, args=(('shred' if triggered_id=='shred-action-btn' else 'proposal'), doc_value, chat_input, uploaded_rfp_decoded_bytes))
+        t.daemon = True
+        t.start()
+        output_data_upload = dcc.Markdown("Starting Gemini operation...", style={"whiteSpace": "pre-wrap", "wordWrap": "break-word"})
         return (
             shred_store, proposal_store, output_data_upload,
             uploaded_doc_list, doc_options, doc_value,
             shredded_doc_list_items,
             generated_doc_list, generated_doc_options, generated_doc_value,
-            generated_rfp_section
-        )
-
-    if triggered_id == 'proposal-action-btn':
-        output_data_upload = dcc.Loading(type="default", children=html.Div("Generating proposal...", style={"wordWrap": "break-word"}), style={"textAlign": "center"})
-        # get decoded bytes for the selected RFP doc for upload if needed
-        rfp_decoded_bytes = None
-        if doc_value and doc_value in uploaded_documents:
-            # Try to get from upload-document if that was just uploaded
-            if rfp_content is not None and rfp_filename == doc_value:
-                content_type, content_string = rfp_content.split(',')
-                rfp_decoded_bytes = base64.b64decode(content_string)
-            # otherwise, not available, unless user re-uploads
-        proposal_text, _, proposal_docx_name, proposal_text_preview = process_document('proposal', doc_value, chat_input, rfp_decoded_bytes)
-        proposal_store = {'text': proposal_text, 'docx_name': proposal_docx_name}
-        new_generated_doc_value = generated_doc_value
-        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)
-            generated_rfp_section = get_generated_rfp_download_section(new_generated_doc_value)
-            logging.info(f"Generated proposal docx saved: {proposal_docx_name}")
-        else:
-            generated_rfp_section = get_generated_rfp_download_section(generated_doc_value)
-        output_data_upload = dcc.Markdown(proposal_text, style={"whiteSpace": "pre-wrap", "wordWrap": "break-word"})
-        return (
-            shred_store, proposal_store, output_data_upload,
-            uploaded_doc_list, doc_options, doc_value,
-            shredded_doc_list_items,
-            generated_doc_list, generated_doc_options, new_generated_doc_value,
-            generated_rfp_section
+            generated_rfp_section,
+            {'streaming': True},
+            False
         )
 
     if triggered_id == 'select-generated-dropdown':
@@ -723,7 +779,9 @@ def master_callback(
             uploaded_doc_list, doc_options, doc_value,
             shredded_doc_list_items,
             generated_doc_list, generated_doc_options, generated_doc_value,
-            generated_rfp_section
+            generated_rfp_section,
+            {'streaming': False},
+            True
         )
 
     if upload_triggered:
@@ -736,7 +794,9 @@ def master_callback(
             uploaded_doc_list, doc_options, doc_value,
             shredded_doc_list_items,
             generated_doc_list, generated_doc_options, generated_doc_value,
-            generated_rfp_section
+            generated_rfp_section,
+            {'streaming': False},
+            True
         )
 
     doc_value = doc_value if doc_value in uploaded_documents else (next(iter(uploaded_documents), None) if uploaded_documents else None)
@@ -747,9 +807,31 @@ def master_callback(
         uploaded_doc_list, doc_options, doc_value,
         shredded_doc_list_items,
         generated_doc_list, generated_doc_options, generated_doc_value,
-        generated_rfp_section
+        generated_rfp_section,
+        {'streaming': False},
+        True
     )
 
+@app.callback(
+    Output('output-data-upload', 'children'),
+    Output('stream-status', 'data'),
+    Output('stream-interval', 'disabled'),
+    Input('stream-interval', 'n_intervals'),
+    State('stream-status', 'data'),
+    prevent_initial_call=True
+)
+def stream_output_callback(n_intervals, stream_status):
+    global stream_buffer, gemini_lock
+    if not stream_status.get('streaming'):
+        return dash.no_update, stream_status, True
+    preview = stream_buffer.get("preview", "")
+    # If lock is free, means operation done
+    still_streaming = gemini_lock.locked()
+    if not still_streaming:
+        stream_status['streaming'] = False
+        return dcc.Markdown(preview, style={"whiteSpace": "pre-wrap", "wordWrap": "break-word"}), stream_status, True
+    return dcc.Markdown(preview, style={"whiteSpace": "pre-wrap", "wordWrap": "break-word"}), stream_status, False
+
 if __name__ == '__main__':
     print("Starting the Dash application...")
     app.run(debug=True, host='0.0.0.0', port=7860, threaded=True)