Update main.py

main.py

@@ -25,7 +25,6 @@ type_mapping = {
     "voter_id": "Voter Id"
 }
 
-# Document types that support OCR extraction
 EXTRACTABLE_TYPES = {"Aadhaar Card", "Pan Card", "Voter Id"}
 
 # -----------------------------
@@ -163,7 +162,6 @@ def predict_file(path):
 
 # -----------------------------
 # OCR: Extract raw text from file
-# Uses eng+hin Tesseract langs to handle both English and Hindi text
 # -----------------------------
 def extract_text_from_file(path: str) -> str:
     ext = os.path.splitext(path)[1].lower()
@@ -176,7 +174,6 @@ def extract_text_from_file(path: str) -> str:
             all_text.append(text)
 
         elif ext == ".pdf":
-            # Use higher DPI for better OCR accuracy on PDFs
             pages = convert_from_path(path, dpi=250)
             for page in pages:
                 text = pytesseract.image_to_string(page, lang="eng+hin")
@@ -205,22 +202,31 @@ def extract_fields_with_sarvam(raw_text: str, doc_type: str) -> dict:
 
     system_prompt = (
         "You are a document field extractor for Indian identity documents. "
-        "Given raw OCR text, extract specific fields and return ONLY a valid JSON object. "
-        "No markdown, no explanation, no extra text — just the JSON.\n\n"
+        "Given raw OCR text, extract specific fields and return the result wrapped in <output> tags like this:\n\n"
+        "<output>\n"
+        "{\n"
+        '  "name": "...",\n'
+        '  "address": "...",\n'
+        '  "date_of_birth": "...",\n'
+        '  "gender": "...",\n'
+        '  "id_number": "...",\n'
+        '  "mobile_number": "..."\n'
+        "}\n"
+        "</output>\n\n"
         "Fields to extract:\n"
         "  - name: Full name of the document holder\n"
         "  - address: Full address (combine multi-line if needed)\n"
-        "  - date_of_birth: In DD/MM/YYYY or YYYY format; use year_of_birth key if only year is present\n"
+        "  - date_of_birth: In DD/MM/YYYY or YYYY format\n"
         "  - gender: Male / Female / Transgender\n"
         "  - id_number: Aadhaar=12-digit number, PAN=10-char alphanumeric, Voter ID=alphanumeric card number\n"
         "  - mobile_number: 10-digit mobile number if present\n\n"
-        "Use null for any field not found. Return a flat JSON object only."
+        "Use null for any field not found. Only output the <output> block, nothing else outside it."
     )
 
     user_prompt = (
         f"Document Type: {doc_type}\n\n"
         f"Raw OCR Text:\n{raw_text}\n\n"
-        "Extract the fields and return JSON."
+        "Extract the fields and return inside <output> tags."
     )
 
     try:
@@ -250,18 +256,22 @@ def extract_fields_with_sarvam(raw_text: str, doc_type: str) -> dict:
 
         content = response.json()["choices"][0]["message"]["content"].strip()
 
-        # Strip markdown code fences
-        content = re.sub(r"^```(?:json)?\s*", "", content)
-        content = re.sub(r"\s*```$", "", content).strip()
-
-        # Strip <think>...</think> reasoning block (sarvam-m is a reasoning model)
-        content = re.sub(r"<think>.*?</think>\s*", "", content, flags=re.DOTALL).strip()
+        # Extract JSON from <output>...</output> tag
+        output_match = re.search(r"<output>(.*?)</output>", content, flags=re.DOTALL)
+        if not output_match:
+            return {
+                **empty_fields,
+                "sarvam_error": f"No <output> tag found in response: {content[:300]}"
+            }
 
-        extracted = json.loads(content)
+        json_str = output_match.group(1).strip()
+        extracted = json.loads(json_str)
 
+        # Normalize keys
         fields = ["name", "address", "date_of_birth", "gender", "id_number", "mobile_number"]
         normalized = {f: extracted.get(f) for f in fields}
 
+        # Handle year_of_birth fallback
         if normalized["date_of_birth"] is None and "year_of_birth" in extracted:
             normalized["date_of_birth"] = str(extracted["year_of_birth"])
 
@@ -308,7 +318,7 @@ async def predict(
             else:
                 result["Authentication"] = "Not valid"
 
-        # Step 3: OCR + Sarvam field extraction (only for recognized ID types)
+        # Step 3: OCR + Sarvam field extraction
         doc_type = result.get("type")
         if doc_type in EXTRACTABLE_TYPES:
             raw_text = extract_text_from_file(temp_path)