Update app.py

app.py

@@ -581,17 +581,27 @@ def merge_pdfs(paths, output_path):
         writer.write(f) 
     return output_path 
 
-def gemini_generate_content(prompt_text, file_upload_obj=None, image_obj=None, model_name="gemini-2.5-pro", fallback_model="gemini-2.5-flash"): 
+def gemini_generate_content(prompt_text, file_upload_obj=None, image_obj=None, model_name="gemini-2.5-pro", fallback_model="gemini-2.5-flash", file_path=None): 
     """ 
     Send prompt_text and optionally an uploaded file (or an image object/list) to the model using NEW SDK.
     Automatically rotates through available API keys on RESOURCE_EXHAUSTED errors.
-    Note: Files uploaded with one API key cannot be accessed by another API key.
+    When rotating keys with file uploads, re-uploads the file with the new API key.
+    
+    Args:
+        prompt_text: The prompt to send
+        file_upload_obj: Previously uploaded file object (optional)
+        image_obj: Image or list of images (optional)
+        model_name: Primary model to use
+        fallback_model: Fallback model if primary fails
+        file_path: Local file path (needed for re-upload when rotating keys)
+    
     Returns textual response and prints progress. 
     """ 
     contents = [prompt_text] 
+    current_file_obj = file_upload_obj
     
-    if file_upload_obj: 
-        contents.append(file_upload_obj) 
+    if current_file_obj: 
+        contents.append(current_file_obj) 
         
     if image_obj: 
         if isinstance(image_obj, list):
@@ -618,6 +628,25 @@ def gemini_generate_content(prompt_text, file_upload_obj=None, image_obj=None, m
         current_client = client_manager.get_current_client()
         current_key_num = client_manager.current_key_index + 1
         
+        # Update contents with current file object
+        contents = [prompt_text]
+        if current_file_obj:
+            contents.append(current_file_obj)
+        if image_obj:
+            if isinstance(image_obj, list):
+                for img_path in image_obj:
+                    if isinstance(img_path, str):
+                        pil_img = Image.open(img_path)
+                        contents.append(pil_img)
+                    else:
+                        contents.append(img_path)
+            else:
+                if isinstance(image_obj, str):
+                    pil_img = Image.open(image_obj)
+                    contents.append(pil_img)
+                else:
+                    contents.append(image_obj)
+        
         # Try primary model first
         try:
             print(f"🔑 Using API key #{current_key_num} with model {model_name}")
@@ -661,17 +690,33 @@ def gemini_generate_content(prompt_text, file_upload_obj=None, image_obj=None, m
                         
                         # Now try next API key if available
                         if attempt < max_attempts - 1:
-                            # Check if we have file uploads - they won't work with different keys
-                            if file_upload_obj:
-                                print("⚠️ WARNING: File uploads cannot be shared across API keys!")
-                                print("   You need to re-upload files with each new API key.")
-                                print("   Stopping rotation to avoid PERMISSION_DENIED errors.")
-                                raise Exception(f"All models exhausted for API key #{current_key_num}. Cannot rotate keys with file uploads.")
-                            
-                            client_manager.rotate_to_next_key()
-                            attempt += 1
-                            print(f"🔄 Trying next API key (attempt {attempt + 1}/{max_attempts})...")
-                            continue
+                            # Check if we have file uploads and can re-upload
+                            if file_upload_obj and file_path:
+                                print(f"🔄 Rotating to next API key and re-uploading file...")
+                                client_manager.rotate_to_next_key()
+                                
+                                # Re-upload file with new API key
+                                try:
+                                    print(f"📤 Re-uploading file with API key #{client_manager.current_key_index + 1}...")
+                                    current_file_obj = upload_to_gemini(file_path)
+                                    print(f"✅ File re-uploaded successfully")
+                                except Exception as upload_error:
+                                    print(f"❌ Failed to re-upload file: {upload_error}")
+                                    raise Exception(f"Failed to re-upload file with new API key: {upload_error}")
+                                
+                                attempt += 1
+                                print(f"🔄 Retrying with next API key (attempt {attempt + 1}/{max_attempts})...")
+                                continue
+                            elif file_upload_obj and not file_path:
+                                print("⚠️ WARNING: Cannot rotate API keys - file_path not provided for re-upload!")
+                                print("   To enable API key rotation with file uploads, pass file_path parameter.")
+                                raise Exception(f"All models exhausted for API key #{current_key_num}. Cannot rotate without file_path.")
+                            else:
+                                # No file uploads, safe to rotate
+                                client_manager.rotate_to_next_key()
+                                attempt += 1
+                                print(f"🔄 Trying next API key (attempt {attempt + 1}/{max_attempts})...")
+                                continue
                         else:
                             raise Exception(f"All {max_attempts} API key(s) exhausted with both models.")
                     else:
@@ -681,7 +726,24 @@ def gemini_generate_content(prompt_text, file_upload_obj=None, image_obj=None, m
             elif "403" in error_str or "PERMISSION_DENIED" in error_str:
                 # This happens when trying to access a file uploaded with a different API key
                 print(f"⚠️ Permission denied - likely due to file uploaded with different API key")
-                raise Exception(f"File access denied. Files uploaded with one API key cannot be accessed by another. Error: {e}")
+                
+                # Try to re-upload if we have the file path
+                if file_path and attempt < max_attempts - 1:
+                    print(f"🔄 Attempting to re-upload file with next API key...")
+                    client_manager.rotate_to_next_key()
+                    
+                    try:
+                        print(f"📤 Re-uploading file with API key #{client_manager.current_key_index + 1}...")
+                        current_file_obj = upload_to_gemini(file_path)
+                        print(f"✅ File re-uploaded successfully")
+                        attempt += 1
+                        print(f"🔄 Retrying with next API key (attempt {attempt + 1}/{max_attempts})...")
+                        continue
+                    except Exception as upload_error:
+                        print(f"❌ Failed to re-upload file: {upload_error}")
+                        raise Exception(f"Failed to re-upload file with new API key: {upload_error}")
+                else:
+                    raise Exception(f"File access denied. Cannot re-upload without file_path. Error: {e}")
             
             else:
                 # Other error - try fallback model with same key
@@ -698,23 +760,39 @@ def gemini_generate_content(prompt_text, file_upload_obj=None, image_obj=None, m
                 except Exception as e2:
                     print(f"❌ Fallback also failed: {e2}")
                     
-                    # If we have more keys and no file uploads, try them
-                    if attempt < max_attempts - 1 and not file_upload_obj:
-                        client_manager.rotate_to_next_key()
-                        attempt += 1
-                        print(f"🔄 Trying next API key (attempt {attempt + 1}/{max_attempts})...")
-                        continue
-                    else:
-                        if file_upload_obj:
-                            raise Exception(f"All models failed. Cannot rotate keys with file uploads. Last error: {e2}")
+                    # If we have more keys, try them
+                    if attempt < max_attempts - 1:
+                        if file_upload_obj and file_path:
+                            print(f"🔄 Rotating to next API key and re-uploading file...")
+                            client_manager.rotate_to_next_key()
+                            
+                            try:
+                                print(f"📤 Re-uploading file with API key #{client_manager.current_key_index + 1}...")
+                                current_file_obj = upload_to_gemini(file_path)
+                                print(f"✅ File re-uploaded successfully")
+                            except Exception as upload_error:
+                                print(f"❌ Failed to re-upload file: {upload_error}")
+                                raise Exception(f"Failed to re-upload file with new API key: {upload_error}")
+                            
+                            attempt += 1
+                            print(f"🔄 Retrying with next API key (attempt {attempt + 1}/{max_attempts})...")
+                            continue
+                        elif file_upload_obj and not file_path:
+                            raise Exception(f"All models failed. Cannot rotate keys without file_path. Last error: {e2}")
                         else:
-                            raise Exception(f"All attempts failed. Last error: {e2}")
+                            client_manager.rotate_to_next_key()
+                            attempt += 1
+                            print(f"🔄 Trying next API key (attempt {attempt + 1}/{max_attempts})...")
+                            continue
+                    else:
+                        raise Exception(f"All attempts failed. Last error: {e2}")
     
     # If we exhausted all attempts
     raise Exception(f"❌ All {max_attempts} API key(s) exhausted. Please check your quota or try again later.")
 
 
 
+
 # ---------------- PARSERS ---------------- 
 def extract_question_ids_from_qpms(text: str):
     """Extract question IDs from QP+MS transcript."""
@@ -1270,7 +1348,7 @@ def align_and_grade_pipeline(qp_path, ms_path, ans_path, subject="Maths", imprin
 
         print("1.i) Transcribing QP+MS (questions first, then full markscheme, with graph detection)...")
         qpms_prompt = QP_MS_TRANSCRIPTION_PROMPT["content"] + "\nAt the end, also list all questions in the markscheme where a graph is expected, in the format:\nGraph expected in:\n- Question <number> → Page <number>\n(One per line, after ==== MARKSCHEME END ====)"
-        qpms_text = gemini_generate_content(qpms_prompt, file_upload_obj=merged_uploaded, model_name="gemini-2.5-flash", fallback_model="gemini-2.5-flash-preview-09-2025")
+        qpms_text = gemini_generate_content(qpms_prompt, file_upload_obj=merged_uploaded, model_name="gemini-2.5-flash", fallback_model="gemini-2.5-flash-preview-09-2025", file_path=merged_qpms_path)
         print("📄 QP+MS transcription received. Saving debug file: debug_qpms_transcript.txt")
         with open("debug_qpms_transcript.txt", "w", encoding="utf-8") as f:
             f.write(qpms_text)
@@ -1288,7 +1366,7 @@ def align_and_grade_pipeline(qp_path, ms_path, ans_path, subject="Maths", imprin
 
         print("1.ii) Building AS transcription prompt with expected question IDs and graph detection, sending to Gemini...")
         as_prompt = build_as_cot_prompt_with_expected_ids(extracted_ids, qpms_text) + "\nAt the end, also list all answers where a graph is found, in the format:\nGraph found in:\n- Answer <number> → Page <number>\n(One per line, after all answers)"
-        as_text = gemini_generate_content(as_prompt, file_upload_obj=ans_uploaded, model_name="gemini-2.5-flash", fallback_model="gemini-2.5-flash-preview-09-2025")
+        as_text = gemini_generate_content(as_prompt, file_upload_obj=ans_uploaded, model_name="gemini-2.5-flash", fallback_model="gemini-2.5-flash-preview-09-2025", file_path=ans_path)
         print("📝 AS transcription received. Saving debug file: debug_as_transcript.txt")
         with open("debug_as_transcript.txt", "w", encoding="utf-8") as f:
             f.write(as_text)