Update app.py

app.py

@@ -34,7 +34,6 @@ Error Avoidance:
 - **No misclassification of errors:** Distinguish correctly between "Conceptual Errors" and "Silly Mistakes."
 - **Follow markscheme logic exactly:** Especially regarding when to withhold accuracy marks if method marks are not earned."""
 
-
 # ---------- HELPER: Save to PDF ----------
 def save_as_pdf(text, filename="output.pdf"):
     styles = getSampleStyleSheet()
@@ -43,44 +42,67 @@ 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
+
+    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}"
+
+    return resp.text, None
 
 # ---------- STEP 1: TRANSCRIPTION ----------
 def transcribe(ans_file):
     try:
         ans_uploaded = genai.upload_file(path=ans_file, display_name="Answer Sheet")
-        model = genai.GenerativeModel("gemini-2.5-pro", generation_config={"temperature": 0})
-
-        resp = model.generate_content([TRANSCRIPTION_PROMPT, ans_uploaded])
-
-        transcription = getattr(resp, "text", None)
-        if not transcription and resp.candidates:
-            transcription = resp.candidates[0].content.parts[0].text
+        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,
+            }
+        )
+
+        transcription, error = safe_generate(model, [TRANSCRIPTION_PROMPT, ans_uploaded])
+        if error:
+            return error, None
 
         pdf_path = save_as_pdf(transcription, "transcription.pdf")
         return transcription, pdf_path
     except Exception as e:
         return f"❌ Error during transcription: {e}", None
 
-
 # ---------- STEP 2: GRADING ----------
 def grade(qp_file, ms_file, transcription):
     try:
         qp_uploaded = genai.upload_file(path=qp_file, display_name="Question Paper")
         ms_uploaded = genai.upload_file(path=ms_file, display_name="Marking Scheme")
-        model = genai.GenerativeModel("gemini-2.5-pro", generation_config={"temperature": 0})
-
-        resp = model.generate_content([GRADING_PROMPT, qp_uploaded, ms_uploaded, transcription])
-
-        grading = getattr(resp, "text", None)
-        if not grading and resp.candidates:
-            grading = resp.candidates[0].content.parts[0].text
+        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,
+            }
+        )
+
+        grading, error = safe_generate(model, [GRADING_PROMPT, qp_uploaded, ms_uploaded, transcription])
+        if error:
+            return error, None
 
         pdf_path = save_as_pdf(grading, "grading.pdf")
         return grading, pdf_path
     except Exception as e:
         return f"❌ Error during grading: {e}", None
 
-
 # ---------- GRADIO APP ----------
 with gr.Blocks(title="LeadIB AI Grading") as demo:
     gr.Markdown("## LeadIB AI Grading\nUpload exam documents to transcribe and grade student answers step by step.")