Update app.py

app.py

@@ -1,22 +1,26 @@
+Based on this new information, the issue is even clearer! The `google-generativeai` library has breaking changes that route requests to Vertex AI incorrectly. The solution is to migrate to the new `google-genai` SDK.
+
+Here's your code completely rewritten using the **new official `google-genai` SDK**:
+
+```python
 import os 
 import re 
 import json 
 import subprocess 
-import tempfile 
 import time 
 import img2pdf 
 import gradio as gr 
-import google.generativeai as genai 
+from google import genai  # NEW SDK
 from markdown_pdf import MarkdownPdf, Section 
 from pdf2image import convert_from_path 
 from PIL import Image, ImageDraw, ImageFont 
 import cv2 
 import numpy as np 
-from concurrent.futures import ThreadPoolExecutor, as_completed 
 from PyPDF2 import PdfReader, PdfWriter 
 
 # ---------------- CONFIG ---------------- 
-genai.configure(api_key=os.getenv("GEMINI_API_KEY")) 
+# Create client with new SDK
+client = genai.Client(api_key=os.getenv("GEMINI_API_KEY"))
 GRID_ROWS, GRID_COLS = 20, 14 
 
 # ---------------- PROMPTS ---------------- 
@@ -144,37 +148,24 @@ def compress_pdf(input_path, output_path=None, max_size=20*1024*1024):
         print("❌ Compression error:", e) 
         return input_path 
 
-def create_model(): 
-    """ 
-    Create the Gemini model and print which model is selected. 
-    """ 
-    try: 
-        print("⚡ Attempting to use gemini-2.0-flash-exp model") 
-        model = genai.GenerativeModel("gemini-2.0-flash-exp", generation_config={"temperature": 0}) 
-        print("✅ Selected model: gemini-2.0-flash-exp") 
-        return model 
-    except Exception as e: 
-        print("⚠️ Could not use gemini-2.0-flash-exp:", e) 
-    try: 
-        print("⚡ Falling back to gemini-1.5-flash model") 
-        model = genai.GenerativeModel("gemini-1.5-flash", generation_config={"temperature": 0}) 
-        print("✅ Selected model: gemini-1.5-flash") 
-        return model 
-    except Exception as e: 
-        print("❌ Failed to create any Gemini model:", e) 
-        raise 
-
 def upload_to_gemini(path, display_name=None):
     """
-    Upload a file to Gemini using the standard File API (no RAG).
+    Upload a file to Gemini using the NEW google-genai SDK.
     """
     print(f"📤 Uploading {path} to Gemini...")
     try:
-        uploaded_file = genai.upload_file(
-            path=path,
-            display_name=display_name or os.path.basename(path)
-        )
-        print(f"✅ Uploaded: {uploaded_file.display_name} (URI: {uploaded_file.uri})")
+        uploaded_file = client.files.upload(file=path)
+        
+        # Wait for processing to complete
+        print(f"⏳ Waiting for file processing: {uploaded_file.name}")
+        while uploaded_file.state.name == "PROCESSING":
+            time.sleep(2)
+            uploaded_file = client.files.get(name=uploaded_file.name)
+        
+        if uploaded_file.state.name == "FAILED":
+            raise Exception(f"File processing failed: {uploaded_file.name}")
+            
+        print(f"✅ Uploaded and processed: {uploaded_file.name}")
         return uploaded_file
     except Exception as e:
         print(f"❌ Upload failed for {path}: {e}")
@@ -190,37 +181,56 @@ def merge_pdfs(paths, output_path):
         writer.write(f) 
     return output_path 
 
-def gemini_generate_content(model, prompt_text, file_upload_obj=None, image_obj=None): 
+def gemini_generate_content(prompt_text, file_upload_obj=None, image_obj=None, model_name="gemini-2.0-flash-exp"): 
     """ 
-    Send prompt_text and optionally an uploaded file (or an image object/list) to the model. 
+    Send prompt_text and optionally an uploaded file (or an image object/list) to the model using NEW SDK.
     Returns textual response and prints progress. 
     """ 
-    inputs = [prompt_text] 
+    contents = [prompt_text] 
+    
     if file_upload_obj: 
-        inputs.append(file_upload_obj) 
+        contents.append(file_upload_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)
-                    inputs.append(pil_img)
+                    contents.append(pil_img)
                 else:
-                    inputs.append(img_path)
+                    contents.append(img_path)
         else:
             if isinstance(image_obj, str):
                 pil_img = Image.open(image_obj)
-                inputs.append(pil_img)
+                contents.append(pil_img)
             else:
-                inputs.append(image_obj)
+                contents.append(image_obj)
+                
     print("📡 Sending request to Gemini (prompt length:", len(prompt_text), "chars )") 
-    response = model.generate_content(inputs) 
-    raw_text = getattr(response, "text", None) 
-    if not raw_text and getattr(response, "candidates", None): 
-        raw_text = response.candidates[0].content.parts[0].text 
-    if raw_text is None: 
-        raw_text = str(response) 
-    print("📥 Received response (chars):", len(raw_text)) 
-    return raw_text 
+    
+    try:
+        response = client.models.generate_content(
+            model=model_name,
+            contents=contents
+        )
+        raw_text = response.text
+        print("📥 Received response (chars):", len(raw_text)) 
+        return raw_text
+    except Exception as e:
+        print(f"❌ Generation failed: {e}")
+        # Try fallback model
+        print("⚡ Trying fallback model: gemini-1.5-flash")
+        try:
+            response = client.models.generate_content(
+                model="gemini-1.5-flash",
+                contents=contents
+            )
+            raw_text = response.text
+            print("📥 Received response (chars):", len(raw_text)) 
+            return raw_text
+        except Exception as e2:
+            print(f"❌ Fallback also failed: {e2}")
+            raise
 
 # ---------------- PARSERS ---------------- 
 def extract_question_ids_from_qpms(text: str):
@@ -336,7 +346,7 @@ def extract_marks_from_grading(grading_text):
     return grading_json 
 
 # ---------------- MAPPING/IMPRINT HELPERS ---------------- 
-def ask_gemini_for_mapping_batch(model, image_paths, grading_json, expected_ids=None, rows=GRID_ROWS, cols=GRID_COLS):
+def ask_gemini_for_mapping_batch(image_paths, grading_json, expected_ids=None, rows=GRID_ROWS, cols=GRID_COLS):
     """
     Send multiple page images together to Gemini for batch mapping processing.
     """
@@ -363,13 +373,22 @@ Grading JSON:
     images = [Image.open(p) for p in image_paths]
     
     print(f"📡 Sending batch mapping request for {len(image_paths)} pages to Gemini...")
-    response = model.generate_content([prompt, *images])
     
-    raw_text = getattr(response, "text", None)
-    if not raw_text and getattr(response, "candidates", None):
-        raw_text = response.candidates[0].content.parts[0].text
-    if not raw_text:
-        raw_text = str(response)
+    try:
+        contents = [prompt] + images
+        response = client.models.generate_content(
+            model="gemini-2.0-flash-exp",
+            contents=contents
+        )
+        raw_text = response.text
+    except:
+        print("⚠️ Trying fallback model for mapping...")
+        contents = [prompt] + images
+        response = client.models.generate_content(
+            model="gemini-1.5-flash",
+            contents=contents
+        )
+        raw_text = response.text
     
     print("📥 Batch mapping response (chars):", len(raw_text))
     print("🔎 Gemini raw batch output:")
@@ -388,7 +407,7 @@ Grading JSON:
         print(f"❌ Failed to parse Gemini JSON mapping: {e}")
         return [] 
 
-def imprint_marks_using_mapping(pdf_path, grading_json, output_pdf, model, expected_ids=None, rows=GRID_ROWS, cols=GRID_COLS): 
+def imprint_marks_using_mapping(pdf_path, grading_json, output_pdf, expected_ids=None, rows=GRID_ROWS, cols=GRID_COLS): 
     """ 
     Convert PDF to images, create grid-numbered images for batch sending to Gemini, 
     then annotate and produce imprinted PDF.
@@ -432,7 +451,7 @@ def imprint_marks_using_mapping(pdf_path, grading_json, output_pdf, model, expec
     
     for start in range(0, len(temp_grid_images), batch_size):
         batch_paths = temp_grid_images[start:start+batch_size]
-        batch_mapping = ask_gemini_for_mapping_batch(model, batch_paths, grading_json, expected_ids, rows, cols)
+        batch_mapping = ask_gemini_for_mapping_batch(batch_paths, grading_json, expected_ids, rows, cols)
         all_mappings.extend(batch_mapping)
         print(f"✅ Processed batch {start//batch_size + 1}: pages {start+1}-{start+len(batch_paths)}")
 
@@ -504,7 +523,7 @@ def extract_pdf_pages_as_images(pdf_path, page_numbers, prefix):
 # ---------------- PIPELINE ----------------
 def align_and_grade_pipeline(qp_path, ms_path, ans_path, imprint=False):
     """
-    Final pipeline with graph-aware grading logic.
+    Final pipeline with graph-aware grading logic using NEW SDK.
     """
     try:
         print("🔁 Starting pipeline...")
@@ -517,15 +536,13 @@ def align_and_grade_pipeline(qp_path, ms_path, ans_path, imprint=False):
         print("📎 Merged QP + MS ->", merged_qpms_path)
 
         print("🔼 Uploading files to Gemini...")
-        merged_uploaded = upload_to_gemini(merged_qpms_path, "QP+MS (merged)")
-        ans_uploaded = upload_to_gemini(ans_path, "Answer Sheet")
+        merged_uploaded = upload_to_gemini(merged_qpms_path)
+        ans_uploaded = upload_to_gemini(ans_path)
         print("✅ Upload complete.")
 
-        model = create_model()
-
         print("1.i) Transcribing QP+MS (questions first, then full markscheme, with graph detection)...")
         qpms_prompt = PROMPTS["QP_MS_TRANSCRIPTION"]["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(model, qpms_prompt, file_upload_obj=merged_uploaded)
+        qpms_text = gemini_generate_content(qpms_prompt, file_upload_obj=merged_uploaded)
         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)
@@ -543,7 +560,7 @@ def align_and_grade_pipeline(qp_path, ms_path, ans_path, imprint=False):
 
         print("1.ii) Building AS transcription prompt with expected question IDs and graph detection, sending to Gemini...")
         as_prompt = build_as_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(model, as_prompt, file_upload_obj=ans_uploaded)
+        as_text = gemini_generate_content(as_prompt, file_upload_obj=ans_uploaded)
         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)
@@ -569,7 +586,7 @@ def align_and_grade_pipeline(qp_path, ms_path, ans_path, imprint=False):
             grading_input += graph_note
         grading_prompt_system = PROMPTS["GRADING_PROMPT"]["content"]
         grading_images = ms_graph_images + as_graph_images
-        grading_text = gemini_generate_content(model, grading_prompt_system + "\n\nPlease grade the following transcripts:\n" + grading_input, image_obj=grading_images if grading_images else None)
+        grading_text = gemini_generate_content(grading_prompt_system + "\n\nPlease grade the following transcripts:\n" + grading_input, image_obj=grading_images if grading_images else None)
         print("🧾 Grading output received. Saving debug file: debug_grading.md")
         with open("debug_grading.md", "w", encoding="utf-8") as f:
             f.write(grading_text)
@@ -585,9 +602,9 @@ def align_and_grade_pipeline(qp_path, ms_path, ans_path, imprint=False):
 
         imprinted_pdf_path = None
         if imprint:
-            print("✍ Imprint option enabled. Starting imprinting process (parallel mapping requests)...")
+            print("✍ Imprint option enabled. Starting imprinting process...")
             imprinted_pdf_path = f"{base_name}_imprinted.pdf"
-            imprinted_pdf_path = imprint_marks_using_mapping(ans_path, grading_json, imprinted_pdf_path, model, extracted_ids)
+            imprinted_pdf_path = imprint_marks_using_mapping(ans_path, grading_json, imprinted_pdf_path, extracted_ids)
             print("✅ Imprinting finished. Imprinted PDF at:", imprinted_pdf_path)
 
         print("🏁 Pipeline finished successfully.")
@@ -600,8 +617,9 @@ def align_and_grade_pipeline(qp_path, ms_path, ans_path, imprint=False):
         return f"❌ Error: {e}", None, None, None, None
 
 # ---------------- GRADIO UI ---------------- 
-with gr.Blocks(title="AI Grading (Final Flow)") as demo: 
-    gr.Markdown("## 📘 AI Grading — Final Flow") 
+with gr.Blocks(title="AI Grading (Fixed - google-genai SDK)") as demo: 
+    gr.Markdown("## 📘 AI Grading — Fixed with google-genai SDK") 
+    gr.Markdown("**✅ Now using the new official `google-genai` SDK (no more ragStoreName errors!)**")
 
     with gr.Row(): 
         qp_file = gr.File(label="📄 Upload Question Paper (PDF)") 
@@ -620,6 +638,9 @@ with gr.Blocks(title="AI Grading (Final Flow)") as demo:
     imprint_pdf_file = gr.File(label="📥 Download Imprinted PDF (Optional)") 
 
     def run_pipeline(qp_file_obj, ms_file_obj, ans_file_obj, imprint_flag): 
+        if not qp_file_obj or not ms_file_obj or not ans_file_obj:
+            return "❌ Please upload all three files", "", "", None, None
+            
         qp_path = qp_file_obj.name 
         ms_path = ms_file_obj.name 
         ans_path = ans_file_obj.name