Update app.py

app.py

@@ -5,13 +5,12 @@ from google.generativeai.types import HarmCategory, HarmBlockThreshold
 from markdown_pdf import MarkdownPdf, Section
 
 # -------------------- CONFIG --------------------
-API_KEY = os.getenv("GOOGLE_AI_STUDIO_API_KEY")
-genai.configure(api_key=API_KEY)
+genai.configure(api_key=os.getenv("GEMINI_API_KEY"))
 
-# ---------- PROMPTS IN JSON ----------
+# ---------- PROMPTS (JSON Style) ----------
 PROMPTS = {
     "TRANSCRIPTION_PROMPT": {
-        "role": "system",
+        "role": "Student Transcription",
         "content": """Your Role: You are an expert technical transcriber specializing in mathematical and scientific documents. 
 Your mission is to convert handwritten solutions from a provided image or PDF into a clean, accurate, and logically structured Markdown format.
 Instructions:
@@ -22,9 +21,8 @@ Instructions:
 - Do not recreate graphs, only describe them.
 """
     },
-
     "MARKSCHEME_TRANSCRIPTION_PROMPT": {
-        "role": "system",
+        "role": "Markscheme Transcription",
         "content": """Your Role: You are an expert transcriber. 
 Convert the official marking scheme from the provided PDF into clean, structured Markdown.
 Instructions:
@@ -36,9 +34,8 @@ Instructions:
 - Use code blocks for equations.
 """
     },
-
     "GRADING_PROMPT": {
-        "role": "system",
+        "role": "Examiner Grading",
         "content": """You are an official examiner. Use the following grading rules strictly.
 Abbreviations:  
 - M: Marks awarded for attempting to use a correct Method.  
@@ -79,9 +76,9 @@ Produce a GitHub-flavored Markdown table with 3 columns:
 |---------------|---------------|--------|  
 Special Formatting Rule:  
 - Whenever a mark is lost (M0, A0, R0 etc.), wrap it in red using: `<span style=\"color:red\">M0</span>`.  
-- Also make the corresponding **Reason column text red** when a mark is lost.  
+- Also wrap the corresponding Reason in red using `<span style=\"color:red\">reason text</span>`.  
 - Keep awarded marks (M1, A1, etc.) in plain text.  
-- If mixed (e.g., M1A0A1), only highlight the lost marks (`A0`).  
+- If mixed (e.g., M1A0A1), only highlight the lost marks (`A0`) and its reason.  
 After the table, provide:  
 ### Summary & Final Mark  
 - Total marks obtained vs total available  
@@ -98,24 +95,27 @@ def save_as_pdf(text, filename="output.pdf"):
     pdf.save(filename)
     return filename
 
+# ---------- COMMON MODEL CALL ----------
+def call_model(inputs):
+    model = genai.GenerativeModel("gemini-2.5-pro", generation_config={"temperature": 0})
+    response = model.generate_content(
+        inputs,
+        safety_settings={
+            HarmCategory.HARM_CATEGORY_HATE_SPEECH: HarmBlockThreshold.BLOCK_LOW_AND_ABOVE,
+            HarmCategory.HARM_CATEGORY_HARASSMENT: HarmBlockThreshold.BLOCK_NONE,
+        }
+    )
+    if getattr(response, "text", None):
+        return response.text
+    elif response.candidates:
+        return response.candidates[0].content.parts[0].text
+    return None
+
 # ---------- STEP 1: TRANSCRIBE STUDENT ----------
 def transcribe_student(ans_file):
     try:
         ans_uploaded = genai.upload_file(path=ans_file, display_name="Answer Sheet")
-        model = genai.GenerativeModel("gemini-2.5-pro")
-
-        resp = model.generate_content(
-            [PROMPTS["TRANSCRIPTION_PROMPT"]["content"], ans_uploaded],
-            safety_settings={
-                HarmCategory.HARM_CATEGORY_HATE_SPEECH: HarmBlockThreshold.BLOCK_LOW_AND_ABOVE,
-                HarmCategory.HARM_CATEGORY_HARASSMENT: HarmBlockThreshold.BLOCK_NONE,
-            }
-        )
-
-        transcription = getattr(resp, "text", None)
-        if not transcription and resp.candidates:
-            transcription = resp.candidates[0].content.parts[0].text
-
+        transcription = call_model([PROMPTS["TRANSCRIPTION_PROMPT"]["content"], ans_uploaded])
         pdf_path = save_as_pdf(transcription, "student_transcription.pdf")
         return transcription, pdf_path
     except Exception as e:
@@ -125,20 +125,7 @@ def transcribe_student(ans_file):
 def transcribe_ms(ms_file):
     try:
         ms_uploaded = genai.upload_file(path=ms_file, display_name="Markscheme")
-        model = genai.GenerativeModel("gemini-2.5-pro")
-
-        resp = model.generate_content(
-            [PROMPTS["MARKSCHEME_TRANSCRIPTION_PROMPT"]["content"], ms_uploaded],
-            safety_settings={
-                HarmCategory.HARM_CATEGORY_HATE_SPEECH: HarmBlockThreshold.BLOCK_LOW_AND_ABOVE,
-                HarmCategory.HARM_CATEGORY_HARASSMENT: HarmBlockThreshold.BLOCK_NONE,
-            }
-        )
-
-        ms_transcription = getattr(resp, "text", None)
-        if not ms_transcription and resp.candidates:
-            ms_transcription = resp.candidates[0].content.parts[0].text
-
+        ms_transcription = call_model([PROMPTS["MARKSCHEME_TRANSCRIPTION_PROMPT"]["content"], ms_uploaded])
         pdf_path = save_as_pdf(ms_transcription, "ms_transcription.pdf")
         return ms_transcription, pdf_path
     except Exception as e:
@@ -148,22 +135,12 @@ def transcribe_ms(ms_file):
 def grade(qp_file, ms_transcription, student_transcription):
     try:
         qp_uploaded = genai.upload_file(path=qp_file, display_name="Question Paper")
-        model = genai.GenerativeModel("gemini-2.5-pro")
-
-        response = model.generate_content(
-            [PROMPTS["GRADING_PROMPT"]["content"], qp_uploaded,
-             "### Markscheme Transcription:\n" + ms_transcription,
-             "### Student Transcription:\n" + student_transcription],
-            safety_settings={
-                HarmCategory.HARM_CATEGORY_HATE_SPEECH: HarmBlockThreshold.BLOCK_LOW_AND_ABOVE,
-                HarmCategory.HARM_CATEGORY_HARASSMENT: HarmBlockThreshold.BLOCK_NONE,
-            }
-        )
-
-        grading = getattr(response, "text", None)
-        if not grading and response.candidates:
-            grading = response.candidates[0].content.parts[0].text
-
+        grading = call_model([
+            PROMPTS["GRADING_PROMPT"]["content"],
+            qp_uploaded,
+            "### Markscheme Transcription:\n" + ms_transcription,
+            "### Student Transcription:\n" + student_transcription
+        ])
         pdf_path = save_as_pdf(grading, "grading.pdf")
         return grading, pdf_path
     except Exception as e: