Update app.py via AI Editor

app.py

@@ -5,11 +5,12 @@ import dash
 from dash import dcc, html, Input, Output, State, callback_context
 import dash_bootstrap_components as dbc
 import pandas as pd
-import openai
 import logging
 from docx import Document
-import requests
 import mimetypes
+from threading import Thread
+
+import google.generativeai as genai
 
 logging.basicConfig(
     level=logging.INFO,
@@ -19,9 +20,12 @@ logging.basicConfig(
 app = dash.Dash(__name__, external_stylesheets=[dbc.themes.BOOTSTRAP])
 server = app.server
 
-OPENAI_API_KEY = os.environ.get("OPENAI_API_KEY", "")
-openai.api_key = OPENAI_API_KEY
-OPENAI_MODEL = "gpt-4.1"
+GOOGLE_API_KEY = os.environ.get("GOOGLE_API_KEY", "")
+genai.configure(api_key=GOOGLE_API_KEY)
+GEMINI_MODEL = "models/gemini-2.5-pro-preview-03-25"
+
+MAX_INPUT_TOKENS = 1048576
+MAX_OUTPUT_TOKENS = 65536
 
 uploaded_documents = {}
 uploaded_documents_fileid = {}
@@ -64,61 +68,49 @@ def guess_mime_type(filename):
             mime_type = "application/octet-stream"
     return mime_type
 
-def upload_to_openai_file(decoded_bytes, filename):
+def upload_to_gemini_file(decoded_bytes, filename):
     try:
-        url = "https://api.openai.com/v1/files"
         mime_type = guess_mime_type(filename)
-        headers = {
-            "Authorization": f"Bearer {OPENAI_API_KEY}"
-        }
-        files = {
-            "file": (filename, io.BytesIO(decoded_bytes), mime_type)
-        }
-        data = {
-            "purpose": "assistants"
-        }
-        logging.info(f"Uploading file to OpenAI /v1/files: {filename}, size={len(decoded_bytes)}, mime_type={mime_type}")
-        resp = requests.post(url, headers=headers, files=files, data=data)
-        if resp.status_code == 200:
-            resp_json = resp.json()
-            file_id = resp_json.get("id")
-            if file_id:
-                logging.info(f"File uploaded to OpenAI /v1/files: {filename}, file_id: {file_id}")
-                return file_id
-            else:
-                logging.error(f"OpenAI /v1/files did not return file id. Response: {resp_json}")
-                return None
+        file_obj = io.BytesIO(decoded_bytes)
+        file_obj.name = filename
+        myfile = genai.upload_file(file_obj, display_name=filename)
+        if hasattr(myfile, "name"):
+            logging.info(f"File uploaded to Gemini: {filename}, file_id: {myfile.name}")
+            return myfile.name
         else:
-            logging.error(f"Failed uploading file to OpenAI API /v1/files: {resp.status_code}, {resp.text}")
+            logging.error(f"Gemini file upload did not return file name/id. Response: {myfile}")
             return None
     except Exception as e:
-        logging.error(f"Exception during file upload to OpenAI /v1/files: {e}")
+        logging.error(f"Exception during file upload to Gemini: {e}")
         return None
 
-def openai_stream_generate(prompt, file_id=None):
-    result_text = []
+def gemini_generate_content(prompt, file_id=None, chat_input=None):
     try:
-        logging.info("Connecting to OpenAI gpt-4.1 for streaming completion...")
-        messages = [{"role": "user", "content": prompt}]
-        # The OpenAI ChatCompletion API does NOT accept a file_id/files argument for completions.
-        # For now, just ignore file_id and rely on prompt content/context.
-        stream = openai.ChatCompletion.create(
-            model=OPENAI_MODEL,
-            messages=messages,
-            max_tokens=32768,
-            temperature=0.2,
-            stream=True
+        files = []
+        if file_id:
+            try:
+                gemini_file = genai.get_file(file_id)
+                files.append(gemini_file)
+            except Exception as e:
+                logging.error(f"Could not fetch Gemini file for id {file_id}: {e}")
+        content_list = []
+        if files:
+            content_list.extend(files)
+            content_list.append("\n\n")
+        content_list.append(prompt)
+        model = genai.GenerativeModel(GEMINI_MODEL)
+        response = model.generate_content(
+            contents=content_list,
+            generation_config=genai.types.GenerationConfig(
+                max_output_tokens=MAX_OUTPUT_TOKENS
+            )
         )
-        for chunk in stream:
-            if 'choices' in chunk and chunk['choices'] and 'delta' in chunk['choices'][0] and 'content' in chunk['choices'][0]['delta']:
-                piece = chunk['choices'][0]['delta']['content']
-                if piece:
-                    result_text.append(piece)
-        logging.info("Received response from OpenAI streaming.")
-        return ''.join(result_text)
+        result = response.text if hasattr(response, "text") else str(response)
+        logging.info("Received response from Gemini.")
+        return result
     except Exception as e:
-        logging.error("Error during OpenAI streaming request: %s", e)
-        return f"Error during streaming: {e}"
+        logging.error("Error during Gemini generate_content: %s", e)
+        return f"Error during Gemini completion: {e}"
 
 def save_shredded_as_docx(shredded_text, rfp_filename):
     doc = Document()
@@ -177,8 +169,8 @@ def process_document(action, selected_filename=None, chat_input=None, selected_p
             global shredded_document
             shredded_document = ""
             try:
-                logging.info("Starting streaming shredding operation with OpenAI.")
-                result = openai_stream_generate(prompt, file_id=doc_fileid)
+                logging.info("Starting Gemini generate_content for shredding.")
+                result = gemini_generate_content(prompt, file_id=doc_fileid, chat_input=chat_input)
                 shredded_document = result
                 logging.info("Document shredded successfully.")
                 docx_bytes = save_shredded_as_docx(result, selected_filename)
@@ -191,7 +183,6 @@ def process_document(action, selected_filename=None, chat_input=None, selected_p
                 logging.error("Error in thread_shred: %s", e)
                 result_holder["text"] = shredded_document
         shredded_document = "Shredding in progress..."
-        from threading import Thread
         t = Thread(target=thread_shred)
         t.start()
         t.join()
@@ -214,8 +205,8 @@ def process_document(action, selected_filename=None, chat_input=None, selected_p
             global generated_response
             generated_response = ""
             try:
-                logging.info("Starting streaming generation operation with OpenAI.")
-                result = openai_stream_generate(prompt, file_id=doc_fileid)
+                logging.info("Starting Gemini generate_content for generation.")
+                result = gemini_generate_content(prompt, file_id=doc_fileid, chat_input=chat_input)
                 generated_response = result
                 logging.info("Proposal response generated successfully.")
                 docx_bytes = save_proposal_as_docx(result, selected_filename)
@@ -228,7 +219,6 @@ def process_document(action, selected_filename=None, chat_input=None, selected_p
                 logging.error("Error in thread_generate: %s", e)
                 result_holder["text"] = generated_response
         generated_response = "Generating response..."
-        from threading import Thread
         t = Thread(target=thread_generate)
         t.start()
         t.join()
@@ -273,23 +263,21 @@ def process_document(action, selected_filename=None, chat_input=None, selected_p
             prompt += f"User additional instructions: {chat_input}\n"
         prompt += f"\n---\nRFP/SOW/PWS/RFI ({rfp_filename}):\n"
         prompt += f"\n---\nGenerated Document ({generated_docname}):\n{generated_doc_content}\n"
-        logging.info(f"Sending proposal prompt to OpenAI. RFP: {rfp_filename}, Generated Doc: {generated_docname}")
+        logging.info(f"Sending proposal prompt to Gemini. RFP: {rfp_filename}, Generated Doc: {generated_docname}")
         result_holder = {"text": None, "docx_bytes": None, "docx_name": None}
         def thread_proposal():
             try:
-                logging.info("Connecting to OpenAI gpt-4.1 for proposal streaming...")
-                # No file_id support for ChatCompletion, so ignore files here.
-                result = openai_stream_generate(prompt)
-                logging.info("Received proposal results from OpenAI.")
-                docx_bytes = save_proposal_as_docx(result, rfp_filename)
+                logging.info("Connecting to Gemini for proposal.")
+                response = gemini_generate_content(prompt, file_id=rfp_fileid, chat_input=chat_input)
+                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]}_{os.path.splitext(generated_docname)[0]}_proposal.docx"
-                result_holder["text"] = result
+                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 streaming proposal request: %s", e)
-                result_holder["text"] = f"Error during streaming: {e}"
-        from threading import Thread
+                logging.error("Error during Gemini proposal request: %s", e)
+                result_holder["text"] = f"Error during Gemini completion: {e}"
         t = Thread(target=thread_proposal)
         t.start()
         t.join()
@@ -538,7 +526,7 @@ def master_callback(
         text = decode_document(decoded)
         fileid = None
         if rfp_filename.lower().endswith(('.pdf', '.docx', '.xlsx', '.xls')):
-            fileid = upload_to_openai_file(decoded, rfp_filename)
+            fileid = upload_to_gemini_file(decoded, rfp_filename)
         if text is not None:
             uploaded_documents[rfp_filename] = text
             if fileid:
@@ -554,7 +542,7 @@ def master_callback(
         text = decode_document(decoded)
         fileid = None
         if proposal_filename.lower().endswith(('.pdf', '.docx', '.xlsx', '.xls')):
-            fileid = upload_to_openai_file(decoded, proposal_filename)
+            fileid = upload_to_gemini_file(decoded, proposal_filename)
         if text is not None:
             uploaded_proposals[proposal_filename] = text
             if fileid:
@@ -572,6 +560,10 @@ def master_callback(
                 if del_filename in uploaded_documents:
                     del uploaded_documents[del_filename]
                     if del_filename in uploaded_documents_fileid:
+                        try:
+                            genai.delete_file(uploaded_documents_fileid[del_filename])
+                        except Exception as e:
+                            logging.warning(f"Failed to delete Gemini file {del_filename}: {e}")
                         del uploaded_documents_fileid[del_filename]
                     logging.info(f"Document deleted: {del_filename}")
                     if selected_doc == del_filename:
@@ -587,6 +579,10 @@ def master_callback(
                 if del_filename in uploaded_proposals:
                     del uploaded_proposals[del_filename]
                     if del_filename in uploaded_proposals_fileid:
+                        try:
+                            genai.delete_file(uploaded_proposals_fileid[del_filename])
+                        except Exception as e:
+                            logging.warning(f"Failed to delete Gemini proposal file {del_filename}: {e}")
                         del uploaded_proposals_fileid[del_filename]
                     logging.info(f"Proposal deleted: {del_filename}")
                     if selected_proposal == del_filename: