Update utils.py

utils.py

@@ -1,16 +1,17 @@
 import os
 import re
 from datetime import datetime
+from typing import Optional, Dict, Any
 from simple_salesforce import Salesforce
 from paddleocr import PaddleOCR
 
-# -----------------------------------
-# OCR SETUP
-# -----------------------------------
-os.environ.setdefault("OMP_NUM_THREADS", "1")  # limit threads for stability
-ocr = PaddleOCR(use_angle_cls=True, lang='en')
+# ------------------------------
+# OCR setup (stable on small CPU)
+# ------------------------------
+os.environ.setdefault("OMP_NUM_THREADS", "1")
+_ocr = PaddleOCR(use_angle_cls=True, lang='en')
 
-# Regex patterns
+# Regexes
 PAN_REGEX = r'\b[A-Z]{5}[0-9]{4}[A-Z]\b'
 AADHAAR_REGEX = r'\b\d{4}[\s-]?\d{4}[\s-]?\d{4}\b'
 DOB_REGEXES = [
@@ -19,23 +20,20 @@ DOB_REGEXES = [
     r'\b\d{2}[./-](JAN|FEB|MAR|APR|MAY|JUN|JUL|AUG|SEP|OCT|NOV|DEC)[./-]\d{4}\b',
     r'\b(19|20)\d{2}\b'
 ]
-GENDERS = ["MALE", "FEMALE", "TRANSGENDER"]  # kept for completeness (not stored)
 
-# -----------------------------------
-# OCR HELPERS
-# -----------------------------------
-def extract_kyc_fields(file_path, force_type=None):
+# ------------------------------
+# OCR helpers
+# ------------------------------
+def extract_kyc_fields(file_path: str, force_type: Optional[str] = None) -> Dict[str, Any]:
     """
-    Returns a dict with:
-      card_type: PAN | AADHAAR | UNKNOWN
-      pan_number / aadhaar_number
-      name (best-guess)
-      dob  (best-guess for the detected card)
+    Extracts KYC fields from an Aadhaar or PAN image.
+    Returns a dict with keys:
+      card_type, name, dob, aadhaar_number, pan_number
+    (number keys present only when relevant)
     """
     try:
-        result = ocr.ocr(file_path, cls=True)
+        result = _ocr.ocr(file_path, cls=True)
         lines = []
-
         for block in result:
             for line in block:
                 text = re.sub(r'\s+', ' ', line[1][0].strip())
@@ -53,40 +51,35 @@ def extract_kyc_fields(file_path, force_type=None):
             elif re.search(AADHAAR_REGEX, full_text):
                 card_type = "AADHAAR"
 
-        response = {"card_type": card_type}
+        out = {"card_type": card_type}
 
         if card_type == "PAN":
-            response["pan_number"] = _first_match(PAN_REGEX, full_text) or "Not found"
-            response["dob"] = _extract_dob(lines)
-            response["name"] = _extract_pan_name(lines)
-
+            out["pan_number"] = _first(PAN_REGEX, full_text) or "Not found"
+            out["dob"] = _extract_dob(lines)
+            out["name"] = _extract_pan_name(lines)
         elif card_type == "AADHAAR":
-            response["aadhaar_number"] = _first_match(AADHAAR_REGEX, full_text) or "Not found"
-            response["dob"] = _extract_dob(lines)
-            response["name"] = _extract_aadhaar_name(lines)
-
+            out["aadhaar_number"] = _first(AADHAAR_REGEX, full_text) or "Not found"
+            out["dob"] = _extract_dob(lines)
+            out["name"] = _extract_aadhaar_name(lines)
         else:
-            response["error"] = "Could not identify document as PAN or Aadhaar."
-            # best-effort generic fields
-            response["dob"] = _extract_dob(lines)
-            response["name"] = _extract_generic_name(lines)
+            out["name"] = _extract_generic_name(lines)
+            out["dob"] = _extract_dob(lines)
+            out["error"] = "Could not identify document as PAN or Aadhaar."
+        return out
 
-        return response
     except Exception as e:
-        return {"error": f"OCR processing failed: {str(e)}"}
+        return {"error": f"OCR processing failed: {str(e)}", "card_type": "UNKNOWN"}
 
-def _first_match(pattern, text, flags=0):
+def _first(pattern: str, text: str, flags: int = 0) -> Optional[str]:
     m = re.search(pattern, text, flags)
     return m.group(0) if m else None
 
 def _extract_dob(lines):
-    # Try common formats
     for line in lines:
         for pattern in DOB_REGEXES[:-1]:
             m = re.search(pattern, line, re.IGNORECASE)
             if m:
                 return m.group(0)
-    # Year-only with labels
     for line in lines:
         m = re.search(DOB_REGEXES[-1], line)
         if m and any(lbl in line.upper() for lbl in ["YOB", "YEAR", "BIRTH", "DOB"]):
@@ -104,14 +97,12 @@ def _extract_pan_name(lines):
     return "Not found"
 
 def _extract_aadhaar_name(lines):
-    # Heuristic: Name usually above DOB
     for i, line in enumerate(lines):
         if any(re.search(p, line, re.IGNORECASE) for p in DOB_REGEXES):
             if i > 0:
                 candidate = lines[i - 1].strip()
                 if _looks_like_name(candidate):
                     return candidate
-    # Fallback
     for line in lines:
         if _looks_like_name(line.strip()):
             return line.strip()
@@ -123,7 +114,7 @@ def _extract_generic_name(lines):
             return line.strip()
     return "Not found"
 
-def _looks_like_name(text):
+def _looks_like_name(text: str) -> bool:
     if re.search(r'\d', text):
         return False
     if len(text.split()) < 2:
@@ -131,73 +122,60 @@ def _looks_like_name(text):
     banned = ["DOB", "INDIA", "MALE", "FEMALE", "GOVERNMENT"]
     return not any(b in text.upper() for b in banned)
 
-# -----------------------------------
-# SALESFORCE HELPERS
-# -----------------------------------
-SF_USERNAME = os.getenv("SF_USERNAME", "")
-SF_PASSWORD = os.getenv("SF_PASSWORD", "")
-SF_TOKEN    = os.getenv("SF_TOKEN", "")
-SF_DOMAIN   = os.getenv("SF_DOMAIN", "login")  # "login"=prod, "test"=sandbox
-
-def connect_salesforce():
-    try:
-        sf = Salesforce(
-            username=SF_USERNAME,
-            password=SF_PASSWORD,
-            security_token=SF_TOKEN,
-            domain=SF_DOMAIN
+# ------------------------------
+# Salesforce helpers
+# ------------------------------
+def connect_salesforce(
+    username: Optional[str] = None,
+    password: Optional[str] = None,
+    token: Optional[str] = None,
+    domain: Optional[str] = None,
+):
+    """
+    Connect using provided args; fallback to environment variables.
+    Returns a Salesforce client or raises an Exception with the root cause.
+    """
+    sf_username = username or os.getenv("SF_USERNAME", "")
+    sf_password = password or os.getenv("SF_PASSWORD", "")
+    sf_token    = token    or os.getenv("SF_TOKEN", "")
+    sf_domain   = (domain or os.getenv("SF_DOMAIN", "login")).strip() or "login"
+
+    if not (sf_username and sf_password and sf_token):
+        raise RuntimeError(
+            "Missing Salesforce credentials. Provide username/password/token (via UI or env vars)."
         )
-        print(f"✅ Connected to Salesforce ({SF_DOMAIN})")
-        return sf
-    except Exception as e:
-        print("❌ Salesforce login failed:", e)
-        return None
 
-def create_kyc_record(sf, kyc_data, file_name=None, agent_id=None):
+    return Salesforce(
+        username=sf_username,
+        password=sf_password,
+        security_token=sf_token,
+        domain=sf_domain  # "login"=prod, "test"=sandbox
+    )
+
+def create_kyc_record(
+    sf,
+    aadhaar: Optional[Dict[str, Any]],
+    pan: Optional[Dict[str, Any]],
+    agent_id: Optional[str] = None
+) -> Dict[str, Any]:
     """
-    Creates a record in KYC_Record__c with the fields:
-      Aadhaar_Name__c, Aadhaar_DOB__c, Aadhaar_Number__c
-      Pan_Name__c,    Pan_DOB__c,    PAN_Number__c
-    Optionally includes Agent__c if you pass agent_id and that field exists.
+    Create a single KYC_Record__c, filling the Aadhaar_* and Pan_* fields
+    from the respective OCR outputs. Fields (exact API names):
+
+      Aadhaar_Name__c
+      Aadhaar_DOB__c
+      Aadhaar_Number__c
+      Pan_Name__c
+      Pan_DOB__c
+      PAN_Number__c
     """
     try:
-        if not sf:
-            return {"status": "error", "message": "Salesforce not connected"}
-
-        # Normalize values
-        def val_or_blank(key): return (kyc_data.get(key) or "").replace("Not found", "")
+        def safe(v): 
+            if not v or v == "Not found":
+                return ""
+            return v
 
         record = {
-            "Aadhaar_Name__c":  "",
-            "Aadhaar_DOB__c":   "",
-            "Aadhaar_Number__c":"",
-            "Pan_Name__c":      "",
-            "Pan_DOB__c":       "",
-            "PAN_Number__c":    "",
-        }
-
-        ct = (kyc_data.get("card_type") or "").upper()
-        if ct == "AADHAAR":
-            record["Aadhaar_Name__c"]   = val_or_blank("name")
-            record["Aadhaar_DOB__c"]    = val_or_blank("dob")
-            record["Aadhaar_Number__c"] = val_or_blank("aadhaar_number")
-        elif ct == "PAN":
-            record["Pan_Name__c"]    = val_or_blank("name")
-            record["Pan_DOB__c"]     = val_or_blank("dob")
-            record["PAN_Number__c"]  = val_or_blank("pan_number")
-        else:
-            # Unknown: best effort — fill name/dob into Aadhaar side to avoid losing data
-            record["Aadhaar_Name__c"] = val_or_blank("name")
-            record["Aadhaar_DOB__c"]  = val_or_blank("dob")
-
-        # Optionally include Agent__c if provided (and exists in your org)
-        if agent_id:
-            record["Agent__c"] = agent_id
-
-        # Optionally store file name in a text field if you have one (not required by you):
-        # record["KYC_File_Name__c"] = file_name or ""
-
-        resp = sf.KYC_Record__c.create(record)
-        return {"status": "success", "record_id": resp.get("id")}
-    except Exception as e:
-        return {"status": "error", "message": str(e)}
+            "Aadhaar_Name__c":   safe(aadhaar.get("name") if aadhaar else ""),
+            "Aadhaar_DOB__c":    safe(aadhaar.get("dob") if aadhaar else ""),
+            "Aadhaar_Number__c": safe(aadhaar.get("aa_