Update app.py

app.py

@@ -5,8 +5,14 @@ import zipfile
 import asyncio
 import subprocess
 import tempfile
+import threading
+from datetime import datetime, timezone
 from pathlib import Path
 
+import openpyxl
+from openpyxl.styles import Font, PatternFill, Alignment
+from huggingface_hub import HfApi, hf_hub_download
+
 from fastapi import FastAPI, File, UploadFile, Form, HTTPException, Security
 from fastapi.responses import JSONResponse
 from fastapi.security.api_key import APIKeyHeader
@@ -25,13 +31,37 @@ app = FastAPI(
 # ---------------------------------------------------------------------------
 
 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_KEY = os.environ.get("SARVAM_API_KEY", "")
-API_SECRET_KEY = os.environ.get("API_SECRET_KEY", "")
+SARVAM_API_KEY  = os.environ.get("SARVAM_API_KEY", "")
+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
+# ---------------------------------------------------------------------------
+
+LOG_REPO_ID  = "PrathameshRaut/VisionModelLogs"   # dataset repo you create on HF
+LOG_FILENAME = "extraction_logs.xlsx"             # file inside that repo
+
+LOG_HEADERS = [
+    "timestamp_utc",
+    "filename",
+    "file_size_bytes",
+    "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"
+    "error_detail",
+]
+
+# A threading lock so concurrent requests don't corrupt the xlsx
+_log_lock = threading.Lock()
 
 # ---------------------------------------------------------------------------
 # Auth — X-API-Key header guard
@@ -48,6 +78,131 @@ async def verify_api_key(key: str = Security(api_key_header)):
             detail="Invalid or missing API key. Send it as request header: X-API-Key: <your-key>"
         )
 
+# ---------------------------------------------------------------------------
+# Excel log helpers
+# ---------------------------------------------------------------------------
+
+def _style_header_row(ws):
+    header_fill = PatternFill("solid", start_color="1F4E79")
+    header_font = Font(bold=True, color="FFFFFF", name="Arial", size=10)
+    for cell in ws[1]:
+        cell.fill = header_fill
+        cell.font = header_font
+        cell.alignment = Alignment(horizontal="center", vertical="center", wrap_text=True)
+
+
+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,
+            filename=LOG_FILENAME,
+            repo_type="dataset",
+            token=HF_TOKEN,
+            force_download=True,
+        )
+        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
+        }
+        for col, width in col_widths.items():
+            ws.column_dimensions[col].width = width
+        ws.row_dimensions[1].height = 28
+    return wb
+
+
+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,
+        path_in_repo=LOG_FILENAME,
+        repo_id=LOG_REPO_ID,
+        repo_type="dataset",
+        commit_message=f"Add log entry {datetime.now(timezone.utc).strftime('%Y-%m-%d %H:%M:%S')} UTC",
+    )
+
+
+def append_log(
+    filename: str,
+    file_size: int,
+    file_type: str,
+    document_type: str,
+    parameters: list[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 = [
+        datetime.now(timezone.utc).strftime("%Y-%m-%d %H:%M:%S"),
+        filename,
+        file_size,
+        file_type,
+        document_type,
+        ", ".join(parameters),
+        (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],
+    ]
+
+    with _log_lock:
+        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
+            for cell in ws[row_idx]:
+                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]:
+                    cell.fill = fill
+
+            _upload_workbook(wb)
+        except Exception as e:
+            # Logging must never crash the main API
+            print(f"[LOG WARNING] Failed to write log entry: {e}")
+
+
 # ---------------------------------------------------------------------------
 # File conversion — DOCX/PPTX → PDF via LibreOffice; images/PDFs pass through
 # ---------------------------------------------------------------------------
@@ -56,16 +211,10 @@ NATIVE_EXTENSIONS = {".pdf", ".png", ".jpg", ".jpeg"}
 NATIVE_MIME       = {"application/pdf", "image/png", "image/jpeg", "image/jpg"}
 
 def prepare_file(file_bytes: bytes, content_type: str, filename: str) -> tuple[bytes, str]:
-    """
-    Returns (file_bytes, filename) ready to upload to Sarvam.
-    Converts DOCX/PPTX/etc. to PDF first if needed.
-    """
     suffix = Path(filename).suffix.lower()
-
     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:
@@ -88,25 +237,15 @@ def prepare_file(file_bytes: bytes, content_type: str, filename: str) -> tuple[b
             return f.read(), Path(filename).stem + ".pdf"
 
 # ---------------------------------------------------------------------------
-# Step 1: Extract raw text from document via Sarvam Document Intelligence SDK
+# Step 1: Extract raw text via Sarvam Document Intelligence SDK
 # ---------------------------------------------------------------------------
 
 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)
 
-    # 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)
@@ -115,10 +254,7 @@ 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="md")
         job.upload_file(tmp_path)
         job.start()
         status = job.wait_until_complete()
@@ -131,14 +267,12 @@ def extract_text_with_sdk(file_bytes: bytes, filename: str) -> str:
 
         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
+                z.read(name).decode("utf-8", errors="replace") for name in md_files
             )
 
         return full_text
@@ -151,7 +285,7 @@ def extract_text_with_sdk(file_bytes: bytes, filename: str) -> str:
                 pass
 
 # ---------------------------------------------------------------------------
-# Step 2: Extract structured JSON from raw text via Sarvam Chat (sarvam-m)
+# 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:
@@ -159,20 +293,16 @@ async def extract_json_via_chat(doc_type: str, raw_text: str, parameters: list[s
     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:
@@ -192,11 +322,9 @@ Respond with the JSON object only."""
 
     raw = resp.json()["choices"][0]["message"]["content"].strip()
 
-    # Strip <think>...</think> reasoning block if present
     import re
     raw = re.sub(r"<think>.*?</think>", "", raw, flags=re.DOTALL).strip()
 
-    # Strip markdown fences if the model wraps its response
     if raw.startswith("```"):
         raw = raw.split("\n", 1)[-1]
         if raw.endswith("```"):
@@ -221,10 +349,8 @@ def root():
 
 @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()}
-    }
+    return {"supported_types": {k: {"description": v} for k, v in SUPPORTED_TYPES.items()}}
+
 
 @app.post("/extract", dependencies=[Security(verify_api_key)])
 async def extract_document(
@@ -234,15 +360,9 @@ async def extract_document(
 ):
     """
     Extract structured JSON data from an uploaded document.
-
     Requires header: **X-API-Key: your-secret-key**
-
-    - **file**: The document (image, PDF, DOCX, PPTX, etc.)
-    - **document_type**: Must be a supported type
-    - **parameters**: You define what to extract — any field names, comma-separated
     """
 
-    # Validate document type
     if document_type not in SUPPORTED_TYPES:
         raise HTTPException(
             status_code=400,
@@ -252,7 +372,6 @@ async def extract_document(
             }
         )
 
-    # Parse parameters
     param_list = [p.strip() for p in parameters.split(",") if p.strip()]
     if not param_list:
         raise HTTPException(
@@ -260,32 +379,53 @@ async def extract_document(
             detail="'parameters' must contain at least one field name (e.g. 'name,income,date_of_issue')"
         )
 
-    # Read uploaded file
-    file_bytes = await file.read()
+    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"
+    file_size    = len(file_bytes)
+
+    raw_text   = ""
+    extracted  = {}
 
-    # Convert to PDF/image if needed
     try:
         processed_bytes, processed_name = prepare_file(file_bytes, content_type, filename)
-    except HTTPException:
+    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)
+        )
         raise
-    except Exception as e:
-        raise HTTPException(status_code=500, detail=f"File preparation error: {str(e)}")
 
-    # 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:
+    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)
+        )
         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)
+    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_text, extracted, "success", ""
+    )
 
     return JSONResponse(content={
         "document_type": document_type,
@@ -294,6 +434,7 @@ async def extract_document(
         "extracted_data": extracted
     })
 
+
 @app.get("/health")
 def health():
     return {"status": "ok"}