Update app.py

app.py

@@ -18,13 +18,12 @@ 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 Vision",
-    version="3.0.0"
+    description="Extract structured data from documents using Sarvam Document Intelligence",
+    version="4.0.0"
 )
 
 # ---------------------------------------------------------------------------
@@ -42,11 +41,11 @@ API_SECRET_KEY  = os.environ.get("API_SECRET_KEY", "")
 HF_TOKEN        = os.environ.get("HF_TOKEN", "")
 
 # ---------------------------------------------------------------------------
-# Logging config — edit these two constants to match your HF dataset repo
+# Logging config
 # ---------------------------------------------------------------------------
 
-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",
@@ -55,17 +54,16 @@ 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"
+    "raw_vision_output",   # raw output from Document Intelligence
+    "extracted_json",      # final JSON 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 — X-API-Key header guard
+# Auth
 # ---------------------------------------------------------------------------
 
 api_key_header = APIKeyHeader(name="X-API-Key", auto_error=False)
@@ -93,10 +91,8 @@ 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,
@@ -107,24 +103,15 @@ 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,  # 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
@@ -133,11 +120,9 @@ 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,
@@ -154,14 +139,12 @@ def append_log(
     file_type: str,
     document_type: str,
     parameters: list[str],
-    raw_thinking: str,
+    raw_vision_output: 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 = [
@@ -171,7 +154,7 @@ def append_log(
         file_type,
         document_type,
         ", ".join(parameters),
-        (raw_thinking or "")[:5000],            # cap at 5 000 chars to keep xlsx sane
+        (raw_vision_output or "")[:5000],
         json.dumps(extracted_json, ensure_ascii=False) if extracted_json else "",
         status,
         error_detail[:1000],
@@ -181,10 +164,8 @@ 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
@@ -192,7 +173,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]:
@@ -200,7 +180,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}")
 
 
@@ -239,49 +218,44 @@ def prepare_file(file_bytes: bytes, content_type: str, filename: str) -> tuple[b
 
 
 # ---------------------------------------------------------------------------
-# Text cleanup — strip base64-encoded image blobs from OCR output
+# THE ONLY MODEL CALL: Sarvam Document Intelligence with extraction instructions
 # ---------------------------------------------------------------------------
 
-def clean_raw_text(text: str) -> str:
+def extract_with_document_intelligence(
+    file_bytes: bytes,
+    filename: str,
+    doc_type: str,
+    parameters: list[str],
+) -> tuple[dict, 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.
+    Single call to Sarvam Document Intelligence.
+    The vision model reads the document AND returns structured JSON directly.
+    No second model call needed.
 
-    Two patterns are handled:
-      1. Markdown image syntax:  ![alt](data:image/png;base64,<data>)
-      2. Bare data URIs that leaked outside markdown syntax
+    Returns: (extracted_json_dict, raw_output_string)
     """
-    # 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)
@@ -290,7 +264,11 @@ def extract_text_with_sdk(file_bytes: bytes, filename: str) -> str:
     zip_path = tmp_path + "_output.zip"
 
     try:
-        job = client.document_intelligence.create_job(language="en-IN", output_format="md")
+        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.upload_file(tmp_path)
         job.start()
         status = job.wait_until_complete()
@@ -303,15 +281,43 @@ def extract_text_with_sdk(file_bytes: bytes, filename: str) -> str:
 
         job.download_output(zip_path)
 
+        raw_output = ""
         with zipfile.ZipFile(zip_path, "r") as z:
-            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
-            )
+            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()
 
-        return full_text
+        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"
+            }
+
+        return extracted, raw_output
 
     finally:
         for path in (tmp_path, zip_path):
@@ -320,55 +326,6 @@ def extract_text_with_sdk(file_bytes: bytes, filename: str) -> str:
             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
@@ -377,11 +334,12 @@ Respond with the JSON object only."""
 @app.get("/")
 def root():
     return {
-        "message": "Document Extraction API powered by Sarvam Vision",
+        "message": "Document Extraction API powered by Sarvam Document Intelligence",
         "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()}}
@@ -395,6 +353,7 @@ 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**
     """
 
@@ -422,9 +381,7 @@ async def extract_document(
     filename     = file.filename or "document"
     file_size    = len(file_bytes)
 
-    raw_text   = ""
-    extracted  = {}
-
+    # Step 1 — convert DOCX/PPTX to PDF if needed
     try:
         processed_bytes, processed_name = prepare_file(file_bytes, content_type, filename)
     except HTTPException as exc:
@@ -435,34 +392,25 @@ async def extract_document(
         )
         raise
 
+    # Step 2 — single vision model call: read + extract in one shot
     try:
-        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,
-            filename, file_size, content_type, document_type, param_list,
-            "", None, "error", str(exc.detail)
+        extracted, raw_output = await asyncio.to_thread(
+            extract_with_document_intelligence,
+            processed_bytes, processed_name, document_type, param_list
         )
-        raise
-
-    # 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)
+            "", None, "error", str(exc.detail)
         )
         raise
 
-    # Fire-and-forget logging (don't block the response)
+    # Fire-and-forget logging
     asyncio.get_event_loop().run_in_executor(
         None, append_log,
         filename, file_size, content_type, document_type, param_list,
-        raw_text, extracted, "success", ""
+        raw_output, extracted, "success", ""
     )
 
     return JSONResponse(content={