Update app.py

app.py

@@ -1,7 +1,9 @@
 import os
 import gradio as gr
+from datetime import datetime
 from utils import extract_kyc_fields
 from simple_salesforce import Salesforce
+from paddleocr import PaddleOCR
 
 # ------------------------
 # CONFIGURATION
@@ -13,11 +15,13 @@ SF_TOKEN = os.getenv("SF_TOKEN", "AmmfRcd6IiYaRtSGntBnzNMQU")
 # Limit Paddle threads to avoid container crashes
 os.environ["OMP_NUM_THREADS"] = "1"
 
+# Initialize OCR
+ocr = PaddleOCR(use_angle_cls=True, lang='en')
+
 # ------------------------
 # CONNECT TO SALESFORCE
 # ------------------------
 def connect_salesforce():
-    """Connects to Salesforce (Production first, then Sandbox fallback)"""
     try:
         sf = Salesforce(
             username=SF_USERNAME,
@@ -26,7 +30,7 @@ def connect_salesforce():
             domain="login"  # Production
         )
         print("✅ Connected to Salesforce (Production)")
-        return {"status": "success", "env": "Production", "sf": sf}
+        return sf
     except Exception as e:
         print("❌ Production login failed:", e)
         try:
@@ -37,50 +41,76 @@ def connect_salesforce():
                 domain="test"  # Sandbox
             )
             print("✅ Connected to Salesforce (Sandbox)")
-            return {"status": "success", "env": "Sandbox", "sf": sf}
+            return sf
         except Exception as e2:
             print("❌ Sandbox login failed:", e2)
-            return {"status": "error", "message": "Salesforce login failed"}
+            return None
+
+sf = connect_salesforce()
+
+# ------------------------
+# CREATE KYC RECORD IN SALESFORCE
+# ------------------------
+def create_kyc_record(sf, agent_id, kyc_data, file_name=None):
+    try:
+        if not sf:
+            return {"status": "error", "message": "Salesforce connection failed"}
+
+        record = {
+            "Agent__c": agent_id,
+            "Name": kyc_data.get("name", "Unknown KYC"),
+            "Document_Number__c": kyc_data.get("pan_number") or kyc_data.get("aadhaar_number"),
+            "Document_Type__c": kyc_data.get("card_type", "UNKNOWN"),
+            "Upload_Date__c": datetime.now().isoformat()
+        }
+
+        # Optional: store file name / URL if needed
+        if file_name:
+            record["KYC_Document_URL__c"] = f"Uploaded File: {file_name}"
 
-sf_status = connect_salesforce()
+        response = sf.KYC_Document__c.create(record)
+        return {"status": "success", "record_id": response.get("id")}
+    except Exception as e:
+        return {"status": "error", "message": str(e)}
 
 # ------------------------
 # PROCESS IMAGE FUNCTION
 # ------------------------
-def process_image(image_file):
-    """Process uploaded KYC image and extract PAN/Aadhaar details"""
-    if sf_status["status"] != "success":
-        return {"salesforce_status": "error", "message": sf_status.get("message", "Login failed")}
+def process_image(image_file, agent_id):
+    if not sf:
+        return {"status": "error", "message": "Salesforce connection failed"}
+
+    if not image_file:
+        return {"status": "error", "message": "No image uploaded"}
 
-    if image_file is None:
-        return {"ocr_status": "error", "message": "No image uploaded"}
+    if not agent_id:
+        return {"status": "error", "message": "Agent ID is required"}
 
     try:
-        # Pass the filepath string directly, no .name
-        ocr_result = extract_kyc_fields(image_file)
+        # Extract KYC data using utils.py
+        kyc_data = extract_kyc_fields(image_file)
+        # Create Salesforce record
+        sf_response = create_kyc_record(sf, agent_id, kyc_data, os.path.basename(image_file))
         return {
-            "salesforce_status": "connected",
-            "sf_env": sf_status.get("env"),
-            "ocr_result": ocr_result
+            "ocr_result": kyc_data,
+            "salesforce_result": sf_response
         }
     except Exception as e:
-        return {"ocr_status": "error", "message": f"OCR extraction failed: {str(e)}"}
-
+        return {"status": "error", "message": f"OCR/Salesforce error: {str(e)}"}
 
 # ------------------------
 # GRADIO INTERFACE
 # ------------------------
 iface = gr.Interface(
     fn=process_image,
-    inputs=gr.Image(type="filepath"),  # Use filepath for PaddleOCR
+    inputs=[
+        gr.Image(type="filepath", label="Upload PAN / Aadhaar"),
+        gr.Textbox(label="Agent Salesforce ID", placeholder="Enter Agent record ID")
+    ],
     outputs=gr.JSON(),
     title="AI KYC OCR & Salesforce Integration",
-    description="Upload PAN or Aadhaar card images to extract text and verify Salesforce connection."
+    description="Upload PAN or Aadhaar card images to extract text and automatically create KYC records in Salesforce."
 )
 
-# ------------------------
-# LAUNCH APP
-# ------------------------
 if __name__ == "__main__":
-    # Launch Gradio server
-    iface.launch(server_name="0.0.0.0", server_port=7860, debug=True)
+    iface.launch(server_name="0.0.0.0", server_port=7860)