Update app.py

app.py

@@ -18,12 +18,13 @@ from fastapi import FastAPI, File, UploadFile, Form, HTTPException, Security
 from fastapi.responses import JSONResponse
 from fastapi.security.api_key import APIKeyHeader
 
+import httpx
 from sarvamai import SarvamAI
 
 app = FastAPI(
     title="Document Extraction API",
-    description="Extract structured data from documents using Sarvam Document Intelligence",
-    version="4.0.0"
+    description="Extract structured data from documents using Sarvam Vision",
+    version="3.0.0"
 )
 
 # ---------------------------------------------------------------------------
@@ -41,11 +42,11 @@ API_SECRET_KEY  = os.environ.get("API_SECRET_KEY", "")
 HF_TOKEN        = os.environ.get("HF_TOKEN", "")
 
 # ---------------------------------------------------------------------------
-# Logging config
+# Logging config — edit these two constants to match your HF dataset repo
 # ---------------------------------------------------------------------------
 
-LOG_REPO_ID  = "PrathameshRaut/VisionModelLogs"
-LOG_FILENAME = "extraction_logs.xlsx"
+LOG_REPO_ID  = "PrathameshRaut/VisionModelLogs"   # dataset repo you create on HF
+LOG_FILENAME = "extraction_logs.xlsx"             # file inside that repo
 
 LOG_HEADERS = [
     "timestamp_utc",
@@ -54,16 +55,17 @@ LOG_HEADERS = [
     "file_type",
     "document_type",
     "parameters_requested",
-    "raw_vision_output",   # raw output from Document Intelligence
-    "extracted_json",      # final JSON sent to caller
-    "status",              # "success" or "error"
+    "model_raw_thinking",       # raw markdown text from Sarvam Document Intelligence
+    "extracted_json",           # final JSON response sent to caller
+    "status",                   # "success" or "error"
     "error_detail",
 ]
 
+# A threading lock so concurrent requests don't corrupt the xlsx
 _log_lock = threading.Lock()
 
 # ---------------------------------------------------------------------------
-# Auth
+# Auth — X-API-Key header guard
 # ---------------------------------------------------------------------------
 
 api_key_header = APIKeyHeader(name="X-API-Key", auto_error=False)
@@ -91,8 +93,10 @@ def _style_header_row(ws):
 
 
 def _download_or_create_workbook() -> openpyxl.Workbook:
+    """Download the existing log workbook from HF, or create a fresh one."""
     if not HF_TOKEN:
         raise RuntimeError("HF_TOKEN secret is not set — cannot write logs")
+
     try:
         local_path = hf_hub_download(
             repo_id=LOG_REPO_ID,
@@ -103,15 +107,24 @@ def _download_or_create_workbook() -> openpyxl.Workbook:
         )
         wb = openpyxl.load_workbook(local_path)
     except Exception:
+        # File doesn't exist yet — start fresh
         wb = openpyxl.Workbook()
         ws = wb.active
         ws.title = "Logs"
         ws.append(LOG_HEADERS)
         _style_header_row(ws)
+        # Set column widths
         col_widths = {
-            "A": 22, "B": 30, "C": 16, "D": 14,
-            "E": 22, "F": 35, "G": 60, "H": 60,
-            "I": 12, "J": 40,
+            "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
         }
         for col, width in col_widths.items():
             ws.column_dimensions[col].width = width
@@ -120,9 +133,11 @@ def _download_or_create_workbook() -> openpyxl.Workbook:
 
 
 def _upload_workbook(wb: openpyxl.Workbook):
+    """Save workbook to a buffer and push to HF dataset repo."""
     buf = io.BytesIO()
     wb.save(buf)
     buf.seek(0)
+
     api = HfApi(token=HF_TOKEN)
     api.upload_file(
         path_or_fileobj=buf,
@@ -139,12 +154,14 @@ def append_log(
     file_type: str,
     document_type: str,
     parameters: list[str],
-    raw_vision_output: str,
+    raw_thinking: str,
     extracted_json: dict | None,
     status: str,
     error_detail: str = "",
 ):
+    """Thread-safe: download → append row → upload."""
     if not HF_TOKEN:
+        # Logging silently skipped if token not configured
         return
 
     row = [
@@ -154,7 +171,7 @@ def append_log(
         file_type,
         document_type,
         ", ".join(parameters),
-        (raw_vision_output or "")[:5000],
+        (raw_thinking or "")[:5000],            # cap at 5 000 chars to keep xlsx sane
         json.dumps(extracted_json, ensure_ascii=False) if extracted_json else "",
         status,
         error_detail[:1000],
@@ -164,8 +181,10 @@ def append_log(
         try:
             wb = _download_or_create_workbook()
             ws = wb.active
+
             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
@@ -173,6 +192,7 @@ 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]:
@@ -180,6 +200,7 @@ 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}")
 
 
@@ -218,44 +239,49 @@ def prepare_file(file_bytes: bytes, content_type: str, filename: str) -> tuple[b
 
 
 # ---------------------------------------------------------------------------
-# THE ONLY MODEL CALL: Sarvam Document Intelligence with extraction instructions
+# Text cleanup — strip base64-encoded image blobs from OCR output
 # ---------------------------------------------------------------------------
 
-def extract_with_document_intelligence(
-    file_bytes: bytes,
-    filename: str,
-    doc_type: str,
-    parameters: list[str],
-) -> tuple[dict, str]:
+def clean_raw_text(text: str) -> str:
     """
-    Single call to Sarvam Document Intelligence.
-    The vision model reads the document AND returns structured JSON directly.
-    No second model call needed.
+    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.
 
-    Returns: (extracted_json_dict, raw_output_string)
+    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]+\)',
+        '',
+        text,
+        flags=re.DOTALL,
+    )
+    # Pattern 2 — bare data URI (must be at least 100 chars of base64 to avoid
+    #              accidentally removing short legitimate strings)
+    text = re.sub(
+        r'data:[a-zA-Z]+/[a-zA-Z+\-]+;base64,[A-Za-z0-9+/=\s]{100,}',
+        '',
+        text,
+        flags=re.DOTALL,
+    )
+    return text.strip()
+
+
+# ---------------------------------------------------------------------------
+# Step 1: Extract raw text via Sarvam Document Intelligence SDK
+# ---------------------------------------------------------------------------
+
+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)
 
-    description   = SUPPORTED_TYPES[doc_type]
-    param_list    = ", ".join(parameters)
-    param_bullets = "\n".join(f"- {p}" for p in parameters)
-
-    # This instruction goes directly to the vision model upfront.
-    # It reads the document AND extracts — no second model needed.
-    instructions = (
-        f"This document is a {description}.\n"
-        f"Your task: extract ONLY the following fields and return a single valid JSON object.\n"
-        f"Fields to extract:\n{param_bullets}\n\n"
-        f"Rules:\n"
-        f"1. Return ONLY a valid JSON object. No explanation, no markdown, no code fences.\n"
-        f"2. Use exactly these field names as JSON keys: {param_list}\n"
-        f"3. If a field is not found or not legible, set its value to null.\n"
-        f"4. Do not add any fields not listed above."
-    )
-
     suffix = Path(filename).suffix.lower()
     with tempfile.NamedTemporaryFile(suffix=suffix, delete=False) as tmp:
         tmp.write(file_bytes)
@@ -264,11 +290,7 @@ def extract_with_document_intelligence(
     zip_path = tmp_path + "_output.zip"
 
     try:
-        job = client.document_intelligence.create_job(
-            language="en-IN",
-            output_format="json",       # ask for JSON output
-            description=instructions,   # tell it exactly what fields to extract
-        )
+        job = client.document_intelligence.create_job(language="en-IN", output_format="md")
         job.upload_file(tmp_path)
         job.start()
         status = job.wait_until_complete()
@@ -281,43 +303,15 @@ def extract_with_document_intelligence(
 
         job.download_output(zip_path)
 
-        raw_output = ""
         with zipfile.ZipFile(zip_path, "r") as z:
-            all_files  = sorted(z.namelist())
-            json_files = [n for n in all_files if n.endswith(".json")]
-            md_files   = [n for n in all_files if n.endswith(".md")]
-
-            # Prefer .json output files first
-            for name in json_files:
-                content = z.read(name).decode("utf-8", errors="replace").strip()
-                if content:
-                    raw_output = content
-                    break
-
-            # Fallback to .md if no JSON file found
-            if not raw_output and md_files:
-                raw_output = "\n\n".join(
-                    z.read(n).decode("utf-8", errors="replace") for n in md_files
-                )
-
-        if not raw_output:
-            raise HTTPException(status_code=500, detail="No output found in Sarvam result")
-
-        # Clean up any <think>...</think> blocks or markdown fences the model may add
-        clean = re.sub(r"<think>.*?</think>", "", raw_output, flags=re.DOTALL).strip()
-        clean = re.sub(r"^```[a-z]*\n?", "", clean, flags=re.MULTILINE)
-        clean = clean.rstrip("`").strip()
-
-        try:
-            extracted = json.loads(clean)
-        except json.JSONDecodeError:
-            # Graceful degradation — return raw so caller still gets something
-            extracted = {
-                "raw_response": clean,
-                "parse_error": "Model did not return valid JSON"
-            }
+            md_files = sorted(n for n in z.namelist() if n.endswith(".md"))
+            if not md_files:
+                raise HTTPException(status_code=500, detail="No markdown output found in Sarvam result")
+            full_text = "\n\n".join(
+                z.read(name).decode("utf-8", errors="replace") for name in md_files
+            )
 
-        return extracted, raw_output
+        return full_text
 
     finally:
         for path in (tmp_path, zip_path):
@@ -326,6 +320,55 @@ def extract_with_document_intelligence(
             except FileNotFoundError:
                 pass
 
+# ---------------------------------------------------------------------------
+# Step 2: Extract structured JSON via Sarvam Chat (sarvam-m)
+# ---------------------------------------------------------------------------
+
+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)
+
+    prompt = f"""You are a document data extraction assistant. Below is the full text extracted from a {description}.
+--- DOCUMENT TEXT START ---
+{raw_text}
+--- DOCUMENT TEXT END ---
+Extract the following fields from the document text above:
+{param_list}
+Rules:
+1. Return ONLY a valid JSON object. No explanations, no markdown, no code fences.
+2. Use exactly the field names listed above as JSON keys.
+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.
+Respond with the JSON object only."""
+
+    async with httpx.AsyncClient(timeout=60) as http:
+        resp = await http.post(
+            "https://api.sarvam.ai/v1/chat/completions",
+            headers={"api-subscription-key": SARVAM_API_KEY, "Content-Type": "application/json"},
+            json={
+                "model": "sarvam-m",
+                "messages": [{"role": "user", "content": prompt}],
+                "max_tokens": 2048,
+                "temperature": 0
+            }
+        )
+
+    if resp.status_code != 200:
+        raise HTTPException(status_code=resp.status_code, detail=f"Sarvam chat error: {resp.text}")
+
+    raw = resp.json()["choices"][0]["message"]["content"].strip()
+
+    raw = re.sub(r"<think>.*?</think>", "", raw, flags=re.DOTALL).strip()
+
+    if raw.startswith("```"):
+        raw = raw.split("\n", 1)[-1]
+        if raw.endswith("```"):
+            raw = raw[: raw.rfind("```")]
+
+    try:
+        return json.loads(raw.strip())
+    except json.JSONDecodeError:
+        return {"raw_response": raw, "parse_error": "Model did not return valid JSON"}
 
 # ---------------------------------------------------------------------------
 # Routes
@@ -334,12 +377,11 @@ def extract_with_document_intelligence(
 @app.get("/")
 def root():
     return {
-        "message": "Document Extraction API powered by Sarvam Document Intelligence",
+        "message": "Document Extraction API powered by Sarvam Vision",
         "supported_document_types": list(SUPPORTED_TYPES.keys()),
         "docs": "/docs"
     }
 
-
 @app.get("/types")
 def get_supported_types():
     return {"supported_types": {k: {"description": v} for k, v in SUPPORTED_TYPES.items()}}
@@ -353,7 +395,6 @@ async def extract_document(
 ):
     """
     Extract structured JSON data from an uploaded document.
-    Uses a single Sarvam Document Intelligence call — no second model.
     Requires header: **X-API-Key: your-secret-key**
     """
 
@@ -381,7 +422,9 @@ async def extract_document(
     filename     = file.filename or "document"
     file_size    = len(file_bytes)
 
-    # Step 1 — convert DOCX/PPTX to PDF if needed
+    raw_text   = ""
+    extracted  = {}
+
     try:
         processed_bytes, processed_name = prepare_file(file_bytes, content_type, filename)
     except HTTPException as exc:
@@ -392,12 +435,8 @@ async def extract_document(
         )
         raise
 
-    # Step 2 — single vision model call: read + extract in one shot
     try:
-        extracted, raw_output = await asyncio.to_thread(
-            extract_with_document_intelligence,
-            processed_bytes, processed_name, document_type, param_list
-        )
+        raw_text = await asyncio.to_thread(extract_text_with_sdk, processed_bytes, processed_name)
     except HTTPException as exc:
         asyncio.get_event_loop().run_in_executor(
             None, append_log,
@@ -406,11 +445,24 @@ async def extract_document(
         )
         raise
 
-    # Fire-and-forget logging
+    # Strip base64 image blobs before sending to sarvam-m
+    raw_text = clean_raw_text(raw_text)
+
+    try:
+        extracted = await extract_json_via_chat(document_type, raw_text, param_list)
+    except HTTPException as exc:
+        asyncio.get_event_loop().run_in_executor(
+            None, append_log,
+            filename, file_size, content_type, document_type, param_list,
+            raw_text, None, "error", str(exc.detail)
+        )
+        raise
+
+    # Fire-and-forget logging (don't block the response)
     asyncio.get_event_loop().run_in_executor(
         None, append_log,
         filename, file_size, content_type, document_type, param_list,
-        raw_output, extracted, "success", ""
+        raw_text, extracted, "success", ""
     )
 
     return JSONResponse(content={