Update app.py

app.py

@@ -4,6 +4,7 @@ os.environ["OMP_NUM_THREADS"] = "1"
 import re
 import json
 import datetime as dt
+import csv
 from urllib.parse import urlparse
 import xml.etree.ElementTree as ET
 
@@ -15,17 +16,13 @@ from utils import extract_kyc_fields
 SF_USERNAME = "[email protected]"
 SF_PASSWORD = "Lic@2025"
 SF_SECURITY_TOKEN = "AmmfRcd6IiYaRtSGntBnzNMQU"
-SF_DOMAIN = "login"   # "login" for prod, "test" for sandbox (we will auto-try both)
-SF_API_VERSION = "60.0"  # REST & Partner SOAP API version
+SF_DOMAIN = "login"          # "login" (prod/dev) or "test" (sandbox)
+SF_API_VERSION = "60.0"      # Partner & REST API version
 # -------------------------------------------------------------------------------
 
-
 # ---------- helpers ----------
 def _parse_birthdate(dob_text: str):
-    """
-    Normalize DOB to YYYY-MM-DD (Salesforce Date).
-    Supports dd/mm/yyyy, dd-mm-yyyy, dd.mm.yyyy, yyyy-mm-dd, or just YYYY (mapped to mid-year).
-    """
+    """Normalize common DOB formats to YYYY-MM-DD (Salesforce Date)."""
     if not dob_text or dob_text == "Not found":
         return None
     s = dob_text.strip()
@@ -53,190 +50,146 @@ def _parse_birthdate(dob_text: str):
             return dt.date(y, 6, 15).isoformat()
         except ValueError:
             return None
-
     return None
 
 
-def _soap_login_once(username: str, password: str, domain: str, api_version: str):
-    """
-    One SOAP login attempt with given password (already concatenated or not).
-    Returns (instance_url, session_id) on success; raises RuntimeError on failure with fault details.
-    """
+def _extract_fault(xml_text: str):
+    """Pull faultcode/faultstring out of a SOAP envelope for clearer errors."""
+    try:
+        root = ET.fromstring(xml_text)
+    except Exception:
+        return None, None
+    fc = fs = None
+    for e in root.iter():
+        tag = e.tag.split('}', 1)[-1]
+        if tag == "faultcode": fc = e.text
+        if tag == "faultstring": fs = e.text
+    return fc, fs
+
+
+# ---------- AUTH (SOAP login) ----------
+def soap_login(username, password, token, domain="login", api_version=SF_API_VERSION):
+    """Log in via Partner SOAP. Returns dict with success flag + details."""
     endpoint = f"https://{domain}.salesforce.com/services/Soap/u/{api_version}"
-    envelope = f"""<?xml version="1.0" encoding="utf-8" ?>
+    payload = f"""<?xml version="1.0" encoding="utf-8" ?>
 <env:Envelope xmlns:xsd="http://www.w3.org/2001/XMLSchema"
               xmlns:xsi="http://www.w3.org/2001/XMLSchema-instance"
               xmlns:env="http://schemas.xmlsoap.org/soap/envelope/">
   <env:Body>
     <n1:login xmlns:n1="urn:partner.soap.sforce.com">
       <n1:username>{username}</n1:username>
-      <n1:password>{password}</n1:password>
+      <n1:password>{password}{token}</n1:password>
     </n1:login>
   </env:Body>
 </env:Envelope>"""
+    headers = {"Content-Type": "text/xml; charset=UTF-8", "SOAPAction": "login"}
+    r = requests.post(endpoint, data=payload.encode("utf-8"), headers=headers, timeout=30)
 
-    headers = {
-        "Content-Type": "text/xml; charset=UTF-8",
-        "SOAPAction": "login",
-    }
-
-    resp = requests.post(endpoint, data=envelope.encode("utf-8"), headers=headers, timeout=30)
-    if resp.status_code != 200:
-        # Try to parse a SOAP fault if present
-        faultcode, faultstring = _extract_fault(resp.text)
-        msg = f"HTTP {resp.status_code}"
-        if faultcode or faultstring:
-            msg += f" | faultcode={faultcode} | faultstring={faultstring}"
-        raise RuntimeError(msg)
+    if r.status_code != 200:
+        fc, fs = _extract_fault(r.text)
+        return {
+            "success": False,
+            "status_code": r.status_code,
+            "url": endpoint,
+            "faultcode": fc, "faultstring": fs,
+            "response_text": r.text[:800]
+        }
 
-    # Parse XML to find sessionId and serverUrl
     try:
-        root = ET.fromstring(resp.text)
+        root = ET.fromstring(r.text)
     except ET.ParseError as e:
-        raise RuntimeError(f"SOAP login parse error: {e}")
-
-    session_id = None
-    server_url = None
-    for elem in root.iter():
-        tag = elem.tag.split('}', 1)[-1]  # strip namespace
-        if tag == "sessionId":
-            session_id = elem.text
-        elif tag == "serverUrl":
-            server_url = elem.text
-
-    if not session_id or not server_url:
-        faultcode, faultstring = _extract_fault(resp.text)
-        raise RuntimeError(f"SOAP login failed: sessionId/serverUrl not found | faultcode={faultcode} | faultstring={faultstring}")
-
-    parsed = urlparse(server_url)
+        return {"success": False, "error": f"SOAP parse error: {e}"}
+
+    sid = srv = None
+    for e in root.iter():
+        tag = e.tag.split('}', 1)[-1]
+        if tag == "sessionId": sid = e.text
+        if tag == "serverUrl": srv = e.text
+
+    if not sid or not srv:
+        fc, fs = _extract_fault(r.text)
+        return {"success": False, "error": "SOAP login missing sessionId/serverUrl", "faultcode": fc, "faultstring": fs}
+
+    parsed = urlparse(srv)
     instance_url = f"{parsed.scheme}://{parsed.netloc}"
-    return instance_url, session_id
+    return {"success": True, "instance_url": instance_url, "access_token": sid}
 
 
-def _extract_fault(xml_text: str):
-    """
-    Extract SOAP faultcode and faultstring for clearer errors.
-    """
-    try:
-        root = ET.fromstring(xml_text)
-    except Exception:
-        return None, None
-    faultcode = None
-    faultstring = None
-    for elem in root.iter():
-        tag = elem.tag.split('}', 1)[-1]
-        if tag == "faultcode":
-            faultcode = elem.text
-        elif tag == "faultstring":
-            faultstring = elem.text
-    return faultcode, faultstring
-
-
-def soap_login_all_paths(username: str, password: str, token: str, preferred_domain: str = "login", api_version: str = SF_API_VERSION):
-    """
-    Try multiple safe login permutations:
-      1) preferred_domain with password+token
-      2) preferred_domain with password only
-      3) alternate domain with password+token
-      4) alternate domain with password only
-
-    Returns (instance_url, session_id, diagnostics) on success.
-    On failure, raises RuntimeError with aggregated diagnostics.
-    """
-    domains = [preferred_domain] + [d for d in ["login", "test"] if d != preferred_domain]
-    attempts = []
-    for domain in domains:
-        for mode in ["pw_token", "pw_only"]:
-            pw = f"{password}{token}" if mode == "pw_token" else password
-            try:
-                instance_url, session_id = _soap_login_once(username, pw, domain, api_version)
-                diag = {"domain": domain, "mode": mode, "result": "success"}
-                attempts.append(diag)
-                return instance_url, session_id, attempts
-            except Exception as e:
-                # Collect reason but keep trying
-                attempts.append({"domain": domain, "mode": mode, "result": "fail", "reason": str(e)})
-
-    # If we’re here, all attempts failed
-    raise RuntimeError(json.dumps({
-        "message": "All SOAP login attempts failed",
-        "attempts": attempts
-    }))
-
-
-def rest_create(instance_url: str, access_token: str, object_api: str, payload: dict, api_version: str = SF_API_VERSION):
-    """
-    Create a record via REST API. Returns dict with success flag and full details.
-    """
+def rest_create(instance_url: str, access_token: str, object_api: str, payload: dict, api_version=SF_API_VERSION):
+    """Create a record via REST. Returns dict with success flag + details."""
     url = f"{instance_url}/services/data/v{api_version}/sobjects/{object_api}"
-    headers = {
-        "Authorization": f"Bearer {access_token}",
-        "Content-Type": "application/json",
-    }
-    resp = requests.post(url, headers=headers, data=json.dumps(payload), timeout=30)
+    headers = {"Authorization": f"Bearer {access_token}", "Content-Type": "application/json"}
+    r = requests.post(url, headers=headers, data=json.dumps(payload), timeout=30)
 
     try:
-        body = resp.json()
+        body = r.json()
     except Exception:
         body = None
 
-    if 200 <= resp.status_code < 300:
+    if 200 <= r.status_code < 300:
         rec_id = body.get("id") if isinstance(body, dict) else None
-        return {"success": True, "id": rec_id, "status_code": resp.status_code, "url": url, "response_json": body}
-    else:
-        return {
-            "success": False,
-            "status_code": resp.status_code,
-            "url": url,
-            "response_json": body,
-            "response_text": resp.text,
-        }
+        return {"success": True, "id": rec_id, "status_code": r.status_code, "url": url, "response_json": body}
+    return {"success": False, "status_code": r.status_code, "url": url, "response_json": body, "response_text": r.text}
 
 
-def sf_push_kyc_record(ocr_results: dict):
-    """
-    Combine Aadhaar + PAN into one KYC_Record__c and create via REST.
-    Fields expected (API names):
-      Aadhaar_Number__c, Aadhaar_Name__c, Aadhaar_DOB__c (Date)
-      PAN_Number__c,    Pan_Name__c,     Pan_DOB__c     (Date)
-    """
+# ---------- Build payload + push (and always export CSV) ----------
+CSV_HEADERS = [
+    "Aadhaar_Number__c", "Aadhaar_Name__c", "Aadhaar_DOB__c",
+    "PAN_Number__c", "Pan_Name__c", "Pan_DOB__c"
+]
+
+def build_kyc_row(ocr_results: dict):
+    """Map OCR results to your KYC_Record__c fields."""
     a = ocr_results.get("aadhaar") or {}
     p = ocr_results.get("pan") or {}
 
-    aadhaar_number = a.get("aadhaar_number") if (a.get("card_type") == "AADHAAR") else None
-    aadhaar_name   = a.get("name")           if (a.get("card_type") == "AADHAAR") else None
-    aadhaar_dob    = _parse_birthdate(a.get("dob")) if (a.get("card_type") == "AADHAAR") else None
-
-    pan_number     = p.get("pan_number") if (p.get("card_type") == "PAN") else None
-    pan_name       = p.get("name")       if (p.get("card_type") == "PAN") else None
-    pan_dob        = _parse_birthdate(p.get("dob")) if (p.get("card_type") == "PAN") else None
-
-    payload = {
-        "Aadhaar_Number__c": aadhaar_number,
-        "Aadhaar_Name__c":   aadhaar_name,
-        "Aadhaar_DOB__c":    aadhaar_dob,
-        "PAN_Number__c":     pan_number,
-        "Pan_Name__c":       pan_name,
-        "Pan_DOB__c":        pan_dob,
+    return {
+        "Aadhaar_Number__c": a.get("aadhaar_number") if (a.get("card_type") == "AADHAAR") else None,
+        "Aadhaar_Name__c":   a.get("name") if (a.get("card_type") == "AADHAAR") else None,
+        "Aadhaar_DOB__c":    _parse_birthdate(a.get("dob")) if (a.get("card_type") == "AADHAAR") else None,
+        "PAN_Number__c":     p.get("pan_number") if (p.get("card_type") == "PAN") else None,
+        "Pan_Name__c":       p.get("name") if (p.get("card_type") == "PAN") else None,
+        "Pan_DOB__c":        _parse_birthdate(p.get("dob")) if (p.get("card_type") == "PAN") else None,
     }
-    payload = {k: v for k, v in payload.items() if v is not None}
 
-    # 1) SOAP login (auto-tries domains & password modes)
-    instance_url, access_token, diagnostics = soap_login_all_paths(
-        SF_USERNAME, SF_PASSWORD, SF_SECURITY_TOKEN, preferred_domain=SF_DOMAIN, api_version=SF_API_VERSION
-    )
-    # 2) REST create
-    rest_res = rest_create(instance_url, access_token, "KYC_Record__c", payload, SF_API_VERSION)
-    rest_res["login_diagnostics"] = diagnostics
-    return rest_res
+def write_csv(row: dict, path: str):
+    """Write a single-row CSV with your exact field API names."""
+    os.makedirs(os.path.dirname(path), exist_ok=True)
+    with open(path, "w", newline="", encoding="utf-8") as f:
+        w = csv.DictWriter(f, fieldnames=CSV_HEADERS)
+        w.writeheader()
+        # Keep only the headers in order; missing keys become empty.
+        cleaned = {k: (row.get(k) or "") for k in CSV_HEADERS}
+        w.writerow(cleaned)
+    return path
 
 
-# ---------- gradio callback ----------
-def process_documents(aadhaar_file, pan_file, push_to_sf):
+def push_kyc_record(ocr_results: dict):
     """
-    - Runs OCR on Aadhaar and PAN separately.
-    - Optionally pushes a single KYC_Record__c to Salesforce via SOAP+REST.
+    Try API push (SOAP login + REST create). Return detailed result.
+    If API login is blocked (INVALID_LOGIN), you'll still have the CSV file from write_csv().
     """
+    # 1) Build payload for REST
+    payload = build_kyc_row(ocr_results)
+    payload = {k: v for k, v in payload.items() if v not in (None, "")}
+
+    # 2) Try SOAP login on preferred domain, then alternate
+    primary = soap_login(SF_USERNAME, SF_PASSWORD, SF_SECURITY_TOKEN, SF_DOMAIN)
+    if not primary.get("success"):
+        alt_domain = "test" if SF_DOMAIN == "login" else "login"
+        alternate = soap_login(SF_USERNAME, SF_PASSWORD, SF_SECURITY_TOKEN, alt_domain)
+        return {"success": False, "auth": {"primary": primary, "alternate": alternate}}
+
+    # 3) Create record via REST
+    instance_url = primary["instance_url"]
+    access_token = primary["access_token"]
+    created = rest_create(instance_url, access_token, "KYC_Record__c", payload)
+    created["auth_used"] = {"domain": SF_DOMAIN}
+    return created
+
+
+# ---------- Gradio callback ----------
+def process_documents(aadhaar_file, pan_file, push_to_sf):
     results = {"aadhaar": None, "pan": None}
 
     if not aadhaar_file and not pan_file:
@@ -260,19 +213,17 @@ def process_documents(aadhaar_file, pan_file, push_to_sf):
         except Exception as e:
             results["pan"] = {"error": f"PAN OCR failed: {str(e)}", "card_type": "UNKNOWN"}
 
-    output = {"ocr": results}
+    # Always export CSV (works even if API login is blocked)
+    csv_row = build_kyc_row(results)
+    csv_path = write_csv(csv_row, "/mnt/data/KYC_Record_upload.csv")
+
+    output = {"ocr": results, "csv_file": csv_path}
 
     if push_to_sf:
-        try:
-            created = sf_push_kyc_record(results)
-            output["salesforce"] = {"pushed": created.get("success", False), **created}
-        except Exception as e:
-            output["salesforce"] = {
-                "pushed": False,
-                "error": {"type": e.__class__.__name__, "message": str(e)}
-            }
+        push_res = push_kyc_record(results)
+        output["salesforce"] = {"pushed": push_res.get("success", False), **push_res}
 
-    return output
+    return output, csv_path
 
 
 # ---------- UI ----------
@@ -280,41 +231,29 @@ with gr.Blocks(title="Smart KYC OCR → Salesforce (KYC_Record__c)") as demo:
     gr.Markdown(
         """
         # 🧾 Smart KYC OCR → Salesforce
-        Upload **Aadhaar** and **PAN** in separate boxes, then (optional) push one **KYC_Record__c**.
+        Upload **Aadhaar** and **PAN** separately.  
+        Click **Extract KYC Info** to get OCR + CSV (for Data Import Wizard).  
+        Optionally, toggle **Push to Salesforce** (API) — may fail on free/Essentials orgs.
         """
     )
 
     with gr.Row():
         with gr.Column(scale=1):
-            aadhaar_uploader = gr.File(
-                label="📤 Aadhaar Upload",
-                file_types=[".jpg", ".jpeg", ".png"]
-            )
+            aadhaar_uploader = gr.File(label="📤 Aadhaar Upload", file_types=[".jpg", ".jpeg", ".png"])
         with gr.Column(scale=1):
-            pan_uploader = gr.File(
-                label="📤 PAN Upload",
-                file_types=[".jpg", ".jpeg", ".png"]
-            )
-
-    push_to_sf = gr.Checkbox(label="Push to Salesforce (create KYC_Record__c)", value=False)
+            pan_uploader = gr.File(label="📤 PAN Upload", file_types=[".jpg", ".jpeg", ".png"])
 
+    push_to_sf = gr.Checkbox(label="Push to Salesforce via API (create KYC_Record__c)", value=False)
     submit_btn = gr.Button("🔍 Extract KYC Info", variant="primary")
-    output_json = gr.JSON(label="📋 Output (OCR + Salesforce)")
+
+    output_json = gr.JSON(label="📋 Output (OCR + Salesforce/Diagnostics)")
+    csv_download = gr.File(label="⬇️ CSV for Data Import Wizard")
 
     submit_btn.click(
         fn=process_documents,
         inputs=[aadhaar_uploader, pan_uploader, push_to_sf],
-        outputs=output_json,
-    )
-
-    gr.Markdown("---")
-    gr.Markdown(
-        """
-        On failure you’ll see a `login_diagnostics` list (each attempt with domain+mode+reason)
-        and REST `status_code/url/response_json`. This eliminates the SDK TypeError and pinpoints issues fast.
-        """
+        outputs=[output_json, csv_download],
     )
 
-# Keep `demo` available for local run / Spaces
 if __name__ == "__main__":
     demo.launch()