Update app.py

app.py

@@ -1,31 +1,35 @@
 import os
-import base64
+import io
 import json
-import httpx
-from fastapi import FastAPI, File, UploadFile, Form, HTTPException, Security
-from fastapi.responses import JSONResponse
-from fastapi.security.api_key import APIKeyHeader
+import zipfile
+import asyncio
 import subprocess
 import tempfile
 from pathlib import Path
 
+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="1.0.0"
+    version="3.0.0"
 )
 
 # ---------------------------------------------------------------------------
-# Supported document types — no default parameters, fully user-defined
+# Supported document types
 # ---------------------------------------------------------------------------
 
 SUPPORTED_TYPES = {
-    "income_certificate": "Income certificate issued by a government authority",
-    "caste_certificate": "Caste certificate issued by a government authority",
-    "domicile_certificate": "Domicile / residence certificate issued by a government authority",
+    "income_certificate":  "Income certificate issued by a government authority",
+    "caste_certificate":   "Caste certificate issued by a government authority",
+    "domicile_certificate":"Domicile / residence certificate issued by a government authority",
 }
 
-SARVAM_API_URL = "https://api.sarvam.ai/v1/chat/completions"
 SARVAM_API_KEY = os.environ.get("SARVAM_API_KEY", "")
 API_SECRET_KEY = os.environ.get("API_SECRET_KEY", "")
 
@@ -35,13 +39,9 @@ API_SECRET_KEY = os.environ.get("API_SECRET_KEY", "")
 
 api_key_header = APIKeyHeader(name="X-API-Key", auto_error=False)
 
-
 async def verify_api_key(key: str = Security(api_key_header)):
     if not API_SECRET_KEY:
-        raise HTTPException(
-            status_code=500,
-            detail="API_SECRET_KEY is not configured on the server"
-        )
+        raise HTTPException(status_code=500, detail="API_SECRET_KEY is not configured on the server")
     if not key or key != API_SECRET_KEY:
         raise HTTPException(
             status_code=401,
@@ -49,120 +49,146 @@ async def verify_api_key(key: str = Security(api_key_header)):
         )
 
 # ---------------------------------------------------------------------------
-# Helpers
+# File conversion — DOCX/PPTX → PDF via LibreOffice; images/PDFs pass through
 # ---------------------------------------------------------------------------
 
-SUPPORTED_IMAGE_EXTENSIONS = {".jpg", ".jpeg", ".png", ".gif", ".webp", ".tiff", ".bmp"}
-SUPPORTED_IMAGE_MIME = {
-    "image/jpeg", "image/jpg", "image/png", "image/gif",
-    "image/webp", "image/tiff", "image/bmp"
-}
-
+NATIVE_EXTENSIONS = {".pdf", ".png", ".jpg", ".jpeg"}
+NATIVE_MIME       = {"application/pdf", "image/png", "image/jpeg", "image/jpg"}
 
-def convert_to_images(file_bytes: bytes, content_type: str, filename: str) -> list[str]:
+def prepare_file(file_bytes: bytes, content_type: str, filename: str) -> tuple[bytes, str]:
     """
-    Convert any uploaded file into a list of base64-encoded PNG images.
-    Native images -> single-item list.
-    PDF / DOCX / PPTX / etc. -> LibreOffice -> PDF -> pdftoppm -> pages.
+    Returns (file_bytes, filename) ready to upload to Sarvam.
+    Converts DOCX/PPTX/etc. to PDF first if needed.
     """
     suffix = Path(filename).suffix.lower()
 
-    if content_type in SUPPORTED_IMAGE_MIME or suffix in SUPPORTED_IMAGE_EXTENSIONS:
-        return [base64.b64encode(file_bytes).decode()]
+    if suffix in NATIVE_EXTENSIONS or content_type in NATIVE_MIME:
+        return file_bytes, filename
 
+    # Convert to PDF via LibreOffice
     with tempfile.TemporaryDirectory() as tmpdir:
         src = os.path.join(tmpdir, "input" + suffix)
         with open(src, "wb") as f:
             f.write(file_bytes)
 
-        if suffix == ".pdf":
-            pdf_path = src
-        else:
-            result = subprocess.run(
-                ["libreoffice", "--headless", "--convert-to", "pdf", "--outdir", tmpdir, src],
-                capture_output=True, timeout=120
-            )
-            if result.returncode != 0:
-                raise HTTPException(
-                    status_code=500,
-                    detail=f"File conversion failed: {result.stderr.decode()}"
-                )
-            pdfs = [f for f in os.listdir(tmpdir) if f.endswith(".pdf")]
-            if not pdfs:
-                raise HTTPException(status_code=500, detail="PDF conversion produced no output")
-            pdf_path = os.path.join(tmpdir, pdfs[0])
-
-        out_prefix = os.path.join(tmpdir, "page")
         result = subprocess.run(
-            ["pdftoppm", "-r", "150", "-png", pdf_path, out_prefix],
+            ["libreoffice", "--headless", "--convert-to", "pdf", "--outdir", tmpdir, src],
             capture_output=True, timeout=120
         )
         if result.returncode != 0:
             raise HTTPException(
                 status_code=500,
-                detail=f"PDF rasterisation failed: {result.stderr.decode()}"
+                detail=f"File conversion to PDF failed: {result.stderr.decode()}"
             )
+        pdfs = [f for f in os.listdir(tmpdir) if f.endswith(".pdf")]
+        if not pdfs:
+            raise HTTPException(status_code=500, detail="PDF conversion produced no output")
+        out_path = os.path.join(tmpdir, pdfs[0])
+        with open(out_path, "rb") as f:
+            return f.read(), Path(filename).stem + ".pdf"
 
-        pages = sorted([
-            os.path.join(tmpdir, f)
-            for f in os.listdir(tmpdir)
-            if f.startswith("page") and f.endswith(".png")
-        ])
-        if not pages:
-            raise HTTPException(status_code=500, detail="No pages extracted from document")
+# ---------------------------------------------------------------------------
+# Step 1: Extract raw text from document via Sarvam Document Intelligence SDK
+# ---------------------------------------------------------------------------
 
-        return [base64.b64encode(open(p, "rb").read()).decode() for p in pages]
+def extract_text_with_sdk(file_bytes: bytes, filename: str) -> str:
+    """
+    Uses the official sarvamai SDK to:
+      1. Create a job
+      2. Upload the file
+      3. Start processing
+      4. Wait for completion
+      5. Download ZIP output
+      6. Return concatenated markdown text
+    """
+    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)
 
-def build_prompt(doc_type: str, parameters: list[str]) -> str:
+    # Write file to a temp path — SDK's upload_file expects a file path
+    suffix = Path(filename).suffix.lower()
+    with tempfile.NamedTemporaryFile(suffix=suffix, delete=False) as tmp:
+        tmp.write(file_bytes)
+        tmp_path = tmp.name
+
+    zip_path = tmp_path + "_output.zip"
+
+    try:
+        job = client.document_intelligence.create_job(
+            language="en-IN",
+            output_format="md"
+        )
+        job.upload_file(tmp_path)
+        job.start()
+        status = job.wait_until_complete()
+
+        if status.job_state not in ("Completed", "PartiallyCompleted"):
+            raise HTTPException(
+                status_code=500,
+                detail=f"Sarvam document processing ended with state: {status.job_state}"
+            )
+
+        job.download_output(zip_path)
+
+        # Unpack ZIP and join all markdown pages in order
+        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
+            )
+
+        return full_text
+
+    finally:
+        for path in (tmp_path, zip_path):
+            try:
+                os.unlink(path)
+            except FileNotFoundError:
+                pass
+
+# ---------------------------------------------------------------------------
+# Step 2: Extract structured JSON from raw text 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)
-    return f"""You are a document extraction assistant. The document provided is a {description}.
+    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 ---
 
-Your task is to carefully read the document and extract the following fields:
+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 in the document, set its value to null.
-4. Do not invent or infer values that are not explicitly present in the document.
+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 def call_sarvam_vision(images_b64: list[str], prompt: str) -> dict:
-    if not SARVAM_API_KEY:
-        raise HTTPException(status_code=500, detail="SARVAM_API_KEY is not configured on the server")
-
-    content = [
-        {"type": "image_url", "image_url": {"url": f"data:image/png;base64,{img}"}}
-        for img in images_b64
-    ]
-    content.append({"type": "text", "text": prompt})
-
-    payload = {
-        "model": "sarvam-m",
-        "messages": [{"role": "user", "content": content}],
-        "max_tokens": 2048,
-        "temperature": 0
-    }
-
-    async with httpx.AsyncClient(timeout=120) as client:
-        resp = await client.post(
-            SARVAM_API_URL,
-            json=payload,
-            headers={
-                "Authorization": f"Bearer {SARVAM_API_KEY}",
-                "Content-Type": "application/json"
+    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 API error: {resp.text}"
-        )
+        raise HTTPException(status_code=resp.status_code, detail=f"Sarvam chat error: {resp.text}")
 
     raw = resp.json()["choices"][0]["message"]["content"].strip()
 
@@ -177,7 +203,6 @@ async def call_sarvam_vision(images_b64: list[str], prompt: str) -> dict:
     except json.JSONDecodeError:
         return {"raw_response": raw, "parse_error": "Model did not return valid JSON"}
 
-
 # ---------------------------------------------------------------------------
 # Routes
 # ---------------------------------------------------------------------------
@@ -190,18 +215,13 @@ def root():
         "docs": "/docs"
     }
 
-
 @app.get("/types")
 def get_supported_types():
     """List all supported document types."""
     return {
-        "supported_types": {
-            k: {"description": v}
-            for k, v in SUPPORTED_TYPES.items()
-        }
+        "supported_types": {k: {"description": v} for k, v in SUPPORTED_TYPES.items()}
     }
 
-
 @app.post("/extract", dependencies=[Security(verify_api_key)])
 async def extract_document(
     file: UploadFile = File(..., description="Document file (image, PDF, DOCX, PPTX, etc.)"),
@@ -228,7 +248,7 @@ async def extract_document(
             }
         )
 
-    # Parse and validate parameters
+    # Parse parameters
     param_list = [p.strip() for p in parameters.split(",") if p.strip()]
     if not param_list:
         raise HTTPException(
@@ -236,35 +256,40 @@ async def extract_document(
             detail="'parameters' must contain at least one field name (e.g. 'name,income,date_of_issue')"
         )
 
-    # Read file
+    # Read uploaded file
     file_bytes = await file.read()
     if not file_bytes:
         raise HTTPException(status_code=400, detail="Uploaded file is empty")
 
     content_type = file.content_type or ""
-    filename = file.filename or "document"
+    filename     = file.filename or "document"
 
-    # Convert to images
+    # Convert to PDF/image if needed
     try:
-        images_b64 = convert_to_images(file_bytes, content_type, filename)
+        processed_bytes, processed_name = prepare_file(file_bytes, content_type, filename)
     except HTTPException:
         raise
     except Exception as e:
-        raise HTTPException(status_code=500, detail=f"File processing error: {str(e)}")
+        raise HTTPException(status_code=500, detail=f"File preparation error: {str(e)}")
 
-    # Call model
-    prompt = build_prompt(document_type, param_list)
-    extracted = await call_sarvam_vision(images_b64, prompt)
+    # Step 1: OCR + text extraction via Sarvam Document Intelligence (blocking SDK call in thread)
+    try:
+        raw_text = await asyncio.to_thread(extract_text_with_sdk, processed_bytes, processed_name)
+    except HTTPException:
+        raise
+    except Exception as e:
+        raise HTTPException(status_code=500, detail=f"Document intelligence error: {str(e)}")
+
+    # Step 2: JSON extraction via Sarvam Chat
+    extracted = await extract_json_via_chat(document_type, raw_text, param_list)
 
     return JSONResponse(content={
         "document_type": document_type,
         "filename": filename,
-        "pages_processed": len(images_b64),
         "parameters_requested": param_list,
         "extracted_data": extracted
     })
 
-
 @app.get("/health")
 def health():
     return {"status": "ok"}