Update app.py via AI Editor

app.py

@@ -38,9 +38,8 @@ 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_buffer = {"preview": ""}
 stream_event = Event()
 
 def decode_document(decoded_bytes):
@@ -92,7 +91,6 @@ def upload_to_gemini_file(decoded_bytes, filename):
         return 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:
@@ -107,7 +105,6 @@ def gemini_generate_content_stream(prompt, file_id=None, chat_input=None, cancel
             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(
@@ -115,7 +112,7 @@ 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  # chars per chunk
+        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.")
@@ -151,6 +148,7 @@ def save_proposal_as_docx(proposal_text, base_filename):
 
 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
@@ -224,7 +222,6 @@ def process_document(action, selected_filename=None, chat_input=None, rfp_decode
             return "No RFP/SOW/PWS/RFI document selected.", None, None, None
         rfp_filename = selected_filename
         rfp_fileid = uploaded_documents_fileid.get(selected_filename)
-        # Upload file to Gemini if not already uploaded (if rfp_decoded_bytes provided and no fileid)
         if not rfp_fileid and rfp_decoded_bytes is not None:
             try:
                 fileid = upload_to_gemini_file(rfp_decoded_bytes, rfp_filename)
@@ -569,7 +566,9 @@ 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')
+        Input('cancel-action-btn', 'n_clicks'),
+        Input('stream-interval', 'n_intervals'),
+        State('stream-status', 'data')
     ],
     prevent_initial_call=True
 )
@@ -581,7 +580,8 @@ def master_callback(
     shredded_delete_clicks, shredded_doc_children,
     select_generated_value,
     chat_input, selected_filename, selected_proposal_dropdown, selected_generated_dropdown_state,
-    cancel_clicks
+    cancel_clicks,
+    stream_n_intervals, stream_status
 ):
     ctx = callback_context
     triggered_id = ctx.triggered[0]['prop_id'].split('.')[0] if ctx.triggered else None
@@ -607,6 +607,33 @@ def master_callback(
 
     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, 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)
+            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,
+                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, dash.no_update, stream_status, False
+
     if triggered_id == 'cancel-action-btn':
         stream_event.set()
         streaming = False
@@ -735,7 +762,6 @@ def master_callback(
 
     output_data_upload = html.Div("No action taken yet.", style={"wordWrap": "break-word"})
 
-    # 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:
@@ -752,7 +778,6 @@ def master_callback(
         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)
@@ -812,26 +837,6 @@ def master_callback(
         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)