Update main.py

main.py

@@ -201,7 +201,7 @@ def extract_fields_with_sarvam(raw_text: str, doc_type: str) -> dict:
     }
 
     if not SARVAM_API_KEY:
-        return {**empty_fields, "ocr_error": "SARVAM_API_KEY not configured"}
+        return {**empty_fields, "ocr_error": "SARVAM_API_KEY not configured in Space secrets"}
 
     system_prompt = (
         "You are a document field extractor for Indian identity documents. "
@@ -210,9 +210,9 @@ def extract_fields_with_sarvam(raw_text: str, doc_type: str) -> dict:
         "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\n"
+        "  - date_of_birth: In DD/MM/YYYY or YYYY format; use year_of_birth key if only year is present\n"
         "  - gender: Male / Female / Transgender\n"
-        "  - id_number: Aadhaar=12-digit, PAN=10-char alphanumeric, Voter ID=alphanumeric\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."
     )
@@ -242,38 +242,21 @@ def extract_fields_with_sarvam(raw_text: str, doc_type: str) -> dict:
             timeout=30
         )
 
-        # Log the raw response for debugging
-        raw_body = response.text[:500]  # truncate for safety
-
         if response.status_code != 200:
             return {
                 **empty_fields,
-                "sarvam_error": f"API returned {response.status_code}: {raw_body}"
-            }
-
-        resp_json = response.json()
-
-        # Defensive: check schema before indexing
-        choices = resp_json.get("choices")
-        if not choices or not isinstance(choices, list) or len(choices) == 0:
-            return {
-                **empty_fields,
-                "sarvam_error": f"Unexpected response schema: {raw_body}"
+                "sarvam_error": f"API returned {response.status_code}: {response.text[:300]}"
             }
 
-        message = choices[0].get("message", {})
-        content = message.get("content", "").strip()
+        content = response.json()["choices"][0]["message"]["content"].strip()
 
-        if not content:
-            return {
-                **empty_fields,
-                "sarvam_error": f"Empty content in response. Full body: {raw_body}"
-            }
-
-        # Strip markdown code fences if model wraps output
+        # 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()
+
         extracted = json.loads(content)
 
         fields = ["name", "address", "date_of_birth", "gender", "id_number", "mobile_number"]
@@ -285,11 +268,11 @@ def extract_fields_with_sarvam(raw_text: str, doc_type: str) -> dict:
         return normalized
 
     except json.JSONDecodeError as e:
-        return {**empty_fields, "sarvam_parse_error": f"JSON parse failed: {str(e)} | content was: '{content[:200]}'"}
+        return {**empty_fields, "sarvam_parse_error": f"JSON parse failed: {str(e)}"}
     except requests.exceptions.Timeout:
         return {**empty_fields, "sarvam_error": "Sarvam API request timed out"}
     except Exception as e:
-        return {**empty_fields, "sarvam_error": f"{type(e).__name__}: {str(e)}"}
+        return {**empty_fields, "sarvam_error": str(e)}
 
 # -----------------------------
 # FastAPI App