Update app.py

app.py

@@ -246,10 +246,8 @@ def clean_raw_text(text: str) -> str:
     """
     Remove base64-encoded image data that Sarvam Document Intelligence
     carries over from PDFs/images into the markdown output.
-
     These blobs are never useful for field extraction and can easily push
     the prompt far beyond sarvam-m's 7 192-token context window.
-
     Two patterns are handled:
       1. Markdown image syntax:  ![alt](data:image/png;base64,<data>)
       2. Bare data URIs that leaked outside markdown syntax
@@ -324,6 +322,19 @@ def extract_text_with_sdk(file_bytes: bytes, filename: str) -> str:
 # Step 2: Extract structured JSON via Sarvam Chat (sarvam-m)
 # ---------------------------------------------------------------------------
 
+def strip_think_tags(text: str) -> str:
+    """
+    Remove <think>…</think> reasoning blocks from sarvam-m output.
+    Handles both complete tags and truncated blocks where the closing
+    tag was cut off due to max_tokens being reached mid-response.
+    """
+    # Pass 1: remove complete <think>…</think> blocks
+    text = re.sub(r"<think>.*?</think>", "", text, flags=re.DOTALL)
+    # Pass 2: remove any unclosed/truncated <think> block that runs to end-of-string
+    text = re.sub(r"<think>.*$", "", text, flags=re.DOTALL)
+    return text.strip()
+
+
 async def extract_json_via_chat(doc_type: str, raw_text: str, parameters: list[str]) -> dict:
     description = SUPPORTED_TYPES[doc_type]
     param_list  = "\n".join(f"  - {p}" for p in parameters)
@@ -340,7 +351,6 @@ Rules:
 3. If a field is not found or not legible, set its value to null.
 4. Do not invent or infer values not explicitly present in the document.
 5. Resond everything in english language only. Also convert date into proper DD/MM/YYYY Foramt.
-
 Respond with the JSON object only."""
 
     async with httpx.AsyncClient(timeout=60) as http:
@@ -360,15 +370,19 @@ Respond with the JSON object only."""
 
     raw = resp.json()["choices"][0]["message"]["content"].strip()
 
-    raw = re.sub(r"<think>.*?</think>", "", raw, flags=re.DOTALL).strip()
+    # Remove <think> blocks — handles both complete and truncated closing tags
+    raw = strip_think_tags(raw)
 
+    # Strip markdown code fences if present
     if raw.startswith("```"):
         raw = raw.split("\n", 1)[-1]
         if raw.endswith("```"):
             raw = raw[: raw.rfind("```")]
 
+    raw = raw.strip()
+
     try:
-        return json.loads(raw.strip())
+        return json.loads(raw)
     except json.JSONDecodeError:
         return {"raw_response": raw, "parse_error": "Model did not return valid JSON"}