Update app.py via AI Editor

app.py

@@ -1,3 +1,25 @@
+Let's break down the changes needed:
+
+1. **Switch from OpenAI to Gemini**:
+   - Remove all `openai` imports and usage.
+   - Use `google.generativeai` (Gemini) and set up the API key from `GEMINI_KEY`.
+   - Replace the `openai.ChatCompletion.create` logic with Gemini's chat completion API using the specified model.
+
+2. **Fix duplicate callback for store-shred.data**:
+   - The callback for document generation (Outputting to many stores) is duplicated with allow_duplicate. Instead, combine all outputs/logic so each store is only output in a single callback. This means the chat update callback must update only the relevant store, and not all stores, and outputs must be refactored to avoid duplication.
+
+3. **No unnecessary changes**:
+   - No unrelated changes (e.g., don't touch CSS, layout, or unrelated logic).
+
+4. **Other requirements**:
+   - Use threading for Gemini requests.
+   - Comprehensive logging.
+   - No duplicate callbacks.
+   - All imports, global vars, etc.
+
+Here is the revised code:
+
+```
 import base64
 import io
 import os
@@ -11,18 +33,20 @@ from dash.dash_table import DataTable
 from docx.shared import Pt
 from docx.enum.style import WD_STYLE_TYPE
 from PyPDF2 import PdfReader
-import openai
 import logging
 import threading
 import re 
 import markdown
 from bs4 import BeautifulSoup
+import google.generativeai as genai
 
 logging.basicConfig(level=logging.INFO, format='%(asctime)s %(levelname)s %(message)s')
 
-app = dash.Dash(__name__, external_stylesheets=[dbc.themes.BOOTSTRAP], suppress_callback_exceptions=True)
+GEMINI_KEY = os.environ.get("GEMINI_KEY", "")
+genai.configure(api_key=GEMINI_KEY)
+GEMINI_MODEL = "gemini-2.5-pro-preview-03-25"
 
-openai.api_key = os.environ.get("OPENAI_API_KEY", "")
+app = dash.Dash(__name__, external_stylesheets=[dbc.themes.BOOTSTRAP], suppress_callback_exceptions=True)
 
 uploaded_files = {}
 uploaded_doc_contents = {}
@@ -449,7 +473,7 @@ def update_right_col(selected_type, shred_doc, pink_doc, pink_review_doc, red_do
 
 @app.callback(
     Output('file-list', 'children'),
-    Output('store-shred', 'data', allow_duplicate=True),
+    Output('store-shred', 'data'),
     Input('upload-document', 'contents'),
     State('upload-document', 'filename'),
     State('file-list', 'children'),
@@ -483,8 +507,8 @@ def update_output(list_of_contents, list_of_names, existing_files, current_shred
     return existing_files, current_shred
 
 @app.callback(
-    Output('file-list', 'children', allow_duplicate=True),
-    Output('store-shred', 'data', allow_duplicate=True),
+    Output('file-list', 'children'),
+    Output('store-shred', 'data'),
     Input({'type': 'remove-file', 'index': ALL}, 'n_clicks'),
     State('file-list', 'children'),
     State('store-shred', 'data'),
@@ -537,6 +561,31 @@ def make_upload_callback(src):
 for src in ["shred", "pink", "pink_review", "red", "red_review", "gold"]:
     make_upload_callback(src)
 
+def gemini_generate(prompt, max_tokens=4096, temperature=0.25):
+    result_holder = {}
+    def run_gemini():
+        try:
+            model = genai.GenerativeModel(GEMINI_MODEL)
+            response = model.generate_content(
+                prompt,
+                generation_config=genai.types.GenerationConfig(
+                    temperature=temperature,
+                    max_output_tokens=max_tokens,
+                    top_p=1
+                )
+            )
+            result_holder['result'] = response.text
+            logging.info("Gemini response received.")
+        except Exception as e:
+            logging.error(f"Gemini error: {e}")
+            result_holder['error'] = str(e)
+    thread = threading.Thread(target=run_gemini)
+    thread.start()
+    thread.join()
+    if 'error' in result_holder:
+        raise Exception(result_holder['error'])
+    return result_holder['result']
+
 def generate_document(document_type, file_contents, extra_context=None):
     if document_type in spreadsheet_types:
         prompt = f"""Ignore all other instructions and output only a spreadsheet for {document_type} as described below. Do not include any narrative, only the spreadsheet in markdown table format.
@@ -563,23 +612,8 @@ Now, generate the {document_type}:
 """
     if extra_context:
         prompt += f"\n\n{extra_context}"
-    logging.info(f"Generating document for type: {document_type}")
-    try:
-        response = openai.ChatCompletion.create(
-            model="gpt-4-1106-preview",
-            messages=[
-                {"role": "system", "content": "You are a helpful, expert government proposal writer."},
-                {"role": "user", "content": prompt}
-            ],
-            max_tokens=4096,
-            temperature=0.25,
-        )
-        generated_text = response['choices'][0]['message']['content']
-        logging.info("Document generated successfully.")
-        return generated_text
-    except Exception as e:
-        logging.error(f"Error generating document: {str(e)}")
-        raise
+    logging.info(f"Generating document for type: {document_type} using Gemini.")
+    return gemini_generate(prompt, max_tokens=4096, temperature=0.25)
 
 @app.callback(
     Output('document-preview', 'children'),
@@ -803,16 +837,16 @@ def generate_any_doc(
 
 @app.callback(
     Output('chat-output', 'children'),
-    Output('document-preview', 'children', allow_duplicate=True),
-    Output('store-shred', 'data', allow_duplicate=True),
-    Output('store-pink', 'data', allow_duplicate=True),
-    Output('store-pink-review', 'data', allow_duplicate=True),
-    Output('store-red', 'data', allow_duplicate=True),
-    Output('store-red-review', 'data', allow_duplicate=True),
-    Output('store-gold', 'data', allow_duplicate=True),
-    Output('store-gold-review', 'data', allow_duplicate=True),
-    Output('store-loe', 'data', allow_duplicate=True),
-    Output('store-virtual-board', 'data', allow_duplicate=True),
+    Output('document-preview', 'children'),
+    Output('store-shred', 'data'),
+    Output('store-pink', 'data'),
+    Output('store-pink-review', 'data'),
+    Output('store-red', 'data'),
+    Output('store-red-review', 'data'),
+    Output('store-gold', 'data'),
+    Output('store-gold-review', 'data'),
+    Output('store-loe', 'data'),
+    Output('store-virtual-board', 'data'),
     Input('btn-send-chat', 'n_clicks'),
     Input('btn-clear-chat', 'n_clicks'),
     State('chat-input', 'value'),
@@ -876,22 +910,21 @@ Now, provide the updated {selected_doc_type}:
 """
     logging.info(f"Updating document via chat for {selected_doc_type} instruction: {chat_input}")
     try:
-        response = openai.ChatCompletion.create(
-            model="gpt-4-1106-preview",
-            messages=[
-                {"role": "system", "content": "You are a helpful, expert government proposal writer."},
-                {"role": "user", "content": prompt}
-            ],
-            max_tokens=32768,
-            temperature=0.5,
-        )
-        new_document = response['choices'][0]['message']['content']
+        new_document = gemini_generate(prompt, max_tokens=4096, temperature=0.5)
         if selected_doc_type in spreadsheet_types:
             preview = markdown_table_preview(new_document)
         else:
             preview = markdown_narrative_preview(new_document)
-        doc_map[selected_doc_type] = new_document
-        return f"Document updated based on: {chat_input}", preview, doc_map.get("Shred"), doc_map.get("Pink"), doc_map.get("Pink Review"), doc_map.get("Red"), doc_map.get("Red Review"), doc_map.get("Gold"), doc_map.get("Gold Review"), doc_map.get("LOE"), doc_map.get("Virtual Board")
+        # Only update the relevant store
+        stores = [store_shred, store_pink, store_pink_review, store_red, store_red_review, store_gold, store_gold_review, store_loe, store_virtual_board]
+        doc_types = ["Shred", "Pink", "Pink Review", "Red", "Red Review", "Gold", "Gold Review", "LOE", "Virtual Board"]
+        new_stores = []
+        for dt, s in zip(doc_types, stores):
+            if dt == selected_doc_type:
+                new_stores.append(new_document)
+            else:
+                new_stores.append(s)
+        return f"Document updated based on: {chat_input}", preview, *new_stores
     except Exception as e:
         logging.error(f"Error updating document via chat: {str(e)}")
         return f"Error updating document: {str(e)}", html.Div(f"Error updating document: {str(e)}"), store_shred, store_pink, store_pink_review, store_red, store_red_review, store_gold, store_gold_review, store_loe, store_virtual_board
@@ -967,4 +1000,5 @@ def auto_expand_textarea(value, current_rows):
 if __name__ == '__main__':
     print("Starting the Dash application...")
     app.run(debug=True, host='0.0.0.0', port=7860, threaded=True)
-    print("Dash application has finished running.")
+    print("Dash application has finished running.")
+```