Update app.py

app.py

@@ -4,16 +4,14 @@ import google.generativeai as genai
 from reportlab.platypus import SimpleDocTemplate, Paragraph
 from reportlab.lib.styles import getSampleStyleSheet
 from reportlab.lib.pagesizes import A4
+from google.generativeai.types import HarmCategory, HarmBlockThreshold, SafetySetting
 
 # -------------------- CONFIG --------------------
 genai.configure(api_key=os.getenv("GEMINI_API_KEY"))
 
 # ---------- PROMPTS ----------
-TRANSCRIPTION_PROMPT = """ Transcribe the PDF into a clean, question-wise format. While doing so, ensure the following:
-Exclude any text that has been cut, striked, or crossed out (since these may represent student mistakes).
-Clearly distinguish between a "step cut" and when variables cancel out during a step.
-Present the transcription neatly without including unnecessary markings. """
-GRADING_PROMPT = """ Instructions to Examiners:
+TRANSCRIPTION_PROMPT = """Convert the student’s PDF into structured text with each question clearly separated. Ignore crossed-out steps or markings."""
+GRADING_PROMPT = """Instructions to Examiners:
 Abbreviations:
 - M: Marks for correct Method.
 - A: Marks for Answer or Accuracy (often depends on preceding M mark).
@@ -42,17 +40,41 @@ def save_as_pdf(text, filename="output.pdf"):
     doc.build(story)
     return filename
 
-# ---------- HELPER: Safe Generate ----------
-def safe_generate(model, inputs):
-    resp = model.generate_content(inputs)
-    cand = resp.candidates[0] if resp.candidates else None
+# ---------- HELPER: Safe Generate with Retry ----------
+def safe_generate(model, inputs, fallback_prompt=None):
+    try:
+        # Try normal generate
+        resp = model.generate_content(inputs)
+        cand = resp.candidates[0] if resp.candidates else None
+
+        if cand and cand.content and cand.content.parts:
+            return resp.text, None
 
-    if not cand or not cand.content or not cand.content.parts:
         reason = getattr(cand, "finish_reason", "None")
-        ratings = getattr(cand, "safety_ratings", "None")
-        return None, f"❌ Empty/blocked response. finish_reason={reason}, safety_ratings={ratings}"
+        if reason == "1":  # SAFETY block
+            # Retry with streaming
+            chunks = []
+            stream_resp = model.generate_content(inputs, stream=True)
+            for chunk in stream_resp:
+                if chunk.candidates and chunk.candidates[0].content.parts:
+                    chunks.append(chunk.text)
+            if chunks:
+                return "".join(chunks), None
+
+            # Retry with simplified prompt if provided
+            if fallback_prompt:
+                retry_resp = model.generate_content([fallback_prompt] + inputs[1:], stream=True)
+                chunks = []
+                for chunk in retry_resp:
+                    if chunk.candidates and chunk.candidates[0].content.parts:
+                        chunks.append(chunk.text)
+                if chunks:
+                    return "".join(chunks), None
+
+        return None, f"❌ Empty/blocked response. finish_reason={reason}, safety_ratings={getattr(cand, 'safety_ratings', None)}"
 
-    return resp.text, None
+    except Exception as e:
+        return None, f"❌ Exception: {e}"
 
 # ---------- STEP 1: TRANSCRIPTION ----------
 def transcribe(ans_file):
@@ -61,15 +83,15 @@ def transcribe(ans_file):
         model = genai.GenerativeModel(
             "gemini-2.5-pro",
             generation_config={"temperature": 0},
-            safety_settings={
-                genai.types.HarmCategory.HARM_CATEGORY_DANGEROUS_CONTENT: genai.types.HarmBlockThreshold.BLOCK_NONE,
-                genai.types.HarmCategory.HARM_CATEGORY_SEXUALLY_EXPLICIT: genai.types.HarmBlockThreshold.BLOCK_NONE,
-                genai.types.HarmCategory.HARM_CATEGORY_HATE_SPEECH: genai.types.HarmBlockThreshold.BLOCK_NONE,
-                genai.types.HarmCategory.HARM_CATEGORY_HARASSMENT: genai.types.HarmBlockThreshold.BLOCK_NONE,
-            }
+            safety_settings=[
+                SafetySetting(category=HarmCategory.HARM_CATEGORY_DANGEROUS_CONTENT, threshold=HarmBlockThreshold.BLOCK_NONE),
+                SafetySetting(category=HarmCategory.HARM_CATEGORY_SEXUALLY_EXPLICIT, threshold=HarmBlockThreshold.BLOCK_NONE),
+                SafetySetting(category=HarmCategory.HARM_CATEGORY_HATE_SPEECH, threshold=HarmBlockThreshold.BLOCK_NONE),
+                SafetySetting(category=HarmCategory.HARM_CATEGORY_HARASSMENT, threshold=HarmBlockThreshold.BLOCK_NONE),
+            ]
         )
 
-        transcription, error = safe_generate(model, [TRANSCRIPTION_PROMPT, ans_uploaded])
+        transcription, error = safe_generate(model, [TRANSCRIPTION_PROMPT, ans_uploaded], fallback_prompt="Convert the PDF into structured plain text with questions separated.")
         if error:
             return error, None
 
@@ -86,15 +108,15 @@ def grade(qp_file, ms_file, transcription):
         model = genai.GenerativeModel(
             "gemini-2.5-pro",
             generation_config={"temperature": 0},
-            safety_settings={
-                genai.types.HarmCategory.HARM_CATEGORY_DANGEROUS_CONTENT: genai.types.HarmBlockThreshold.BLOCK_NONE,
-                genai.types.HarmCategory.HARM_CATEGORY_SEXUALLY_EXPLICIT: genai.types.HarmBlockThreshold.BLOCK_NONE,
-                genai.types.HarmCategory.HARM_CATEGORY_HATE_SPEECH: genai.types.HarmBlockThreshold.BLOCK_NONE,
-                genai.types.HarmCategory.HARM_CATEGORY_HARASSMENT: genai.types.HarmBlockThreshold.BLOCK_NONE,
-            }
+            safety_settings=[
+                SafetySetting(category=HarmCategory.HARM_CATEGORY_DANGEROUS_CONTENT, threshold=HarmBlockThreshold.BLOCK_NONE),
+                SafetySetting(category=HarmCategory.HARM_CATEGORY_SEXUALLY_EXPLICIT, threshold=HarmBlockThreshold.BLOCK_NONE),
+                SafetySetting(category=HarmCategory.HARM_CATEGORY_HATE_SPEECH, threshold=HarmBlockThreshold.BLOCK_NONE),
+                SafetySetting(category=HarmCategory.HARM_CATEGORY_HARASSMENT, threshold=HarmBlockThreshold.BLOCK_NONE),
+            ]
         )
 
-        grading, error = safe_generate(model, [GRADING_PROMPT, qp_uploaded, ms_uploaded, transcription])
+        grading, error = safe_generate(model, [GRADING_PROMPT, qp_uploaded, ms_uploaded, transcription], fallback_prompt="Grade the answers according to the marking scheme. Show marks step by step.")
         if error:
             return error, None
 
@@ -139,4 +161,4 @@ with gr.Blocks(title="LeadIB AI Grading") as demo:
     )
 
 if __name__ == "__main__":
-    demo.launch()
+    demo.launch()