Update app.py

app.py

@@ -44,8 +44,8 @@ HF_TOKEN        = os.environ.get("HF_TOKEN", "")
 # Logging config — edit these two constants to match your HF dataset repo
 # ---------------------------------------------------------------------------
 
-LOG_REPO_ID  = "PrathameshRaut/VisionModelLogs"   # dataset repo you create on HF
-LOG_FILENAME = "extraction_logs.xlsx"             # file inside that repo
+LOG_REPO_ID  = "PrathameshRaut/VisionModelLogs"
+LOG_FILENAME = "extraction_logs.xlsx"
 
 LOG_HEADERS = [
     "timestamp_utc",
@@ -54,9 +54,9 @@ LOG_HEADERS = [
     "file_type",
     "document_type",
     "parameters_requested",
-    "model_raw_thinking",       # raw markdown text from Sarvam Document Intelligence
-    "extracted_json",           # final JSON response sent to caller
-    "status",                   # "success" or "error"
+    "model_raw_thinking",
+    "extracted_json",
+    "status",
     "error_detail",
 ]
 
@@ -112,18 +112,17 @@ def _download_or_create_workbook() -> openpyxl.Workbook:
         ws.title = "Logs"
         ws.append(LOG_HEADERS)
         _style_header_row(ws)
-        # Set column widths
         col_widths = {
-            "A": 22,  # timestamp
-            "B": 30,  # filename
-            "C": 16,  # file_size
-            "D": 14,  # file_type
-            "E": 22,  # document_type
-            "F": 35,  # parameters
-            "G": 60,  # model_raw_thinking
-            "H": 60,  # extracted_json
-            "I": 12,  # status
-            "J": 40,  # error_detail
+            "A": 22,
+            "B": 30,
+            "C": 16,
+            "D": 14,
+            "E": 22,
+            "F": 35,
+            "G": 60,
+            "H": 60,
+            "I": 12,
+            "J": 40,
         }
         for col, width in col_widths.items():
             ws.column_dimensions[col].width = width
@@ -160,7 +159,6 @@ def append_log(
 ):
     """Thread-safe: download → append row → upload."""
     if not HF_TOKEN:
-        # Logging silently skipped if token not configured
         return
 
     row = [
@@ -170,7 +168,7 @@ def append_log(
         file_type,
         document_type,
         ", ".join(parameters),
-        (raw_thinking or "")[:5000],            # cap at 5 000 chars to keep xlsx sane
+        (raw_thinking or "")[:5000],
         json.dumps(extracted_json, ensure_ascii=False) if extracted_json else "",
         status,
         error_detail[:1000],
@@ -183,7 +181,6 @@ def append_log(
 
             ws.append(row)
 
-            # Style the new data row
             data_font = Font(name="Arial", size=10)
             wrap_align = Alignment(vertical="top", wrap_text=True)
             row_idx = ws.max_row
@@ -191,7 +188,6 @@ def append_log(
                 cell.font = data_font
                 cell.alignment = wrap_align
 
-            # Zebra striping: light blue on even rows
             if row_idx % 2 == 0:
                 fill = PatternFill("solid", start_color="DCE6F1")
                 for cell in ws[row_idx]:
@@ -199,7 +195,6 @@ def append_log(
 
             _upload_workbook(wb)
         except Exception as e:
-            # Logging must never crash the main API
             print(f"[LOG WARNING] Failed to write log entry: {e}")
 
 
@@ -242,15 +237,6 @@ def prepare_file(file_bytes: bytes, content_type: str, filename: str) -> tuple[b
 # ---------------------------------------------------------------------------
 
 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
-    """
     # Pattern 1 — markdown image embed
     text = re.sub(
         r'!\[[^\]]*\]\(data:[a-zA-Z]+/[a-zA-Z+\-]+;base64,[A-Za-z0-9+/=\s]+\)',
@@ -258,8 +244,7 @@ def clean_raw_text(text: str) -> str:
         text,
         flags=re.DOTALL,
     )
-    # Pattern 2 — bare data URI (must be at least 100 chars of base64 to avoid
-    #              accidentally removing short legitimate strings)
+    # Pattern 2 — bare data URI
     text = re.sub(
         r'data:[a-zA-Z]+/[a-zA-Z+\-]+;base64,[A-Za-z0-9+/=\s]{100,}',
         '',
@@ -277,7 +262,11 @@ def extract_text_with_sdk(file_bytes: bytes, filename: str) -> str:
     if not SARVAM_API_KEY:
         raise HTTPException(status_code=500, detail="SARVAM_API_KEY is not configured on the server")
 
-    client = SarvamAI(api_subscription_key=SARVAM_API_KEY)
+    # FIX: increased timeout so the polling loop doesn't give up too early
+    client = SarvamAI(
+        api_subscription_key=SARVAM_API_KEY,
+        timeout_in_seconds=180,
+    )
 
     suffix = Path(filename).suffix.lower()
     with tempfile.NamedTemporaryFile(suffix=suffix, delete=False) as tmp:
@@ -317,19 +306,13 @@ def extract_text_with_sdk(file_bytes: bytes, filename: str) -> str:
             except FileNotFoundError:
                 pass
 
+
 # ---------------------------------------------------------------------------
 # Step 2: Extract structured JSON via SarvamAI SDK chat completions
 # ---------------------------------------------------------------------------
 
 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()
 
@@ -341,6 +324,9 @@ async def extract_json_via_chat(doc_type: str, raw_text: str, parameters: list[s
     description = SUPPORTED_TYPES[doc_type]
     param_list  = "\n".join(f"  - {p}" for p in parameters)
 
+    # Truncate to avoid blowing context window
+    raw_text = raw_text[:6000] if len(raw_text) > 6000 else raw_text
+
     prompt = f"""You are a document data extraction assistant. Below is the full text extracted from a {description}.
 --- DOCUMENT TEXT START ---
 {raw_text}
@@ -355,24 +341,43 @@ Rules:
 5. Respond everything in english language only. Also convert date into proper DD/MM/YYYY format.
 Respond with the JSON object only."""
 
-    # Run the synchronous SDK call in a thread so we don't block the event loop
     def _call_sdk() -> str:
-        client = SarvamAI(api_subscription_key=SARVAM_API_KEY)
-        response = client.chat.completions(
-            model="sarvam-105b",
-            messages=[{"role": "user", "content": prompt}],
-            temperature=0,
-            top_p=1,
-            max_tokens=2048,
+        # FIX: increased timeout — sarvam-105b reasoning can take >60s
+        client = SarvamAI(
+            api_subscription_key=SARVAM_API_KEY,
+            timeout_in_seconds=180,
         )
-        return response.choices[0].message.content.strip()
+
+        # FIX: retry up to 3 times — 105b occasionally returns null content
+        for attempt in range(3):
+            response = client.chat.completions(
+                model="sarvam-105b",
+                messages=[{"role": "user", "content": prompt}],
+                temperature=0,
+                top_p=1,
+                max_tokens=2048,
+            )
+
+            # FIX: guard against null content before calling .strip()
+            content = response.choices[0].message.content
+            if content:
+                return content.strip()
+
+            print(
+                f"[WARN] Attempt {attempt + 1}/3: Sarvam returned null content "
+                f"(finish_reason={response.choices[0].finish_reason!r}), retrying..."
+            )
+
+        raise ValueError("Sarvam returned null content after 3 attempts")
 
     try:
         raw = await asyncio.to_thread(_call_sdk)
+    except ValueError as exc:
+        raise HTTPException(status_code=502, detail=str(exc)) from exc
     except Exception as exc:
         raise HTTPException(status_code=500, detail=f"Sarvam chat error: {exc}") from exc
 
-    # Remove <think> blocks — handles both complete and truncated closing tags
+    # Remove <think> blocks
     raw = strip_think_tags(raw)
 
     # Strip markdown code fences if present
@@ -388,6 +393,7 @@ Respond with the JSON object only."""
     except json.JSONDecodeError:
         return {"raw_response": raw, "parse_error": "Model did not return valid JSON"}
 
+
 # ---------------------------------------------------------------------------
 # Routes
 # ---------------------------------------------------------------------------
@@ -440,8 +446,8 @@ async def extract_document(
     filename     = file.filename or "document"
     file_size    = len(file_bytes)
 
-    raw_text   = ""
-    extracted  = {}
+    raw_text  = ""
+    extracted = {}
 
     try:
         processed_bytes, processed_name = prepare_file(file_bytes, content_type, filename)
@@ -463,7 +469,7 @@ async def extract_document(
         )
         raise
 
-    # Strip base64 image blobs before sending to sarvam-m
+    # Strip base64 image blobs before sending to chat model
     raw_text = clean_raw_text(raw_text)
 
     try: