Update utils.py

utils.py

@@ -1,105 +1,109 @@
 import requests
 import base64
+import json
 import os
+from simple_salesforce import Salesforce
+from pdf2image import convert_from_path
+import pytesseract
+from transformers import LayoutLMv3Processor, LayoutLMv3ForTokenClassification
 from dotenv import load_dotenv
 
-# Load environment variables from .env file
+# Load environment variables
 load_dotenv()
 
-# Retrieve credentials from environment variables
-USERNAME = os.getenv("SF_USERNAME")
-PASSWORD = os.getenv("SF_PASSWORD")
-SECURITY_TOKEN = os.getenv("SF_SECURITY_TOKEN")
-CLIENT_ID = os.getenv("SF_CLIENT_ID")
-CLIENT_SECRET = os.getenv("SF_CLIENT_SECRET")
-DOMAIN = os.getenv("SF_DOMAIN", "test")  # Default to sandbox; use "login" for production
-
-# Validate environment variables
-if not all([USERNAME, PASSWORD, SECURITY_TOKEN, CLIENT_ID, CLIENT_SECRET]):
-    raise ValueError("❌ Missing required environment variables. Ensure SF_USERNAME, SF_PASSWORD, SF_SECURITY_TOKEN, SF_CLIENT_ID, and SF_CLIENT_SECRET are set.")
+# Salesforce Authentication
+def get_salesforce_client():
+    try:
+        sf = Salesforce(
+            username=os.getenv('SALESFORCE_USERNAME'),
+            password=os.getenv('SALESFORCE_PASSWORD'),
+            security_token=os.getenv('SALESFORCE_SECURITY_TOKEN'),
+            instance_url=os.getenv('SALESFORCE_INSTANCE_URL')
+        )
+        print("Salesforce client connected successfully")
+        return sf, None
+    except Exception as e:
+        print(f"Salesforce connection failed: {str(e)}")
+        return None, str(e)
 
-LOGIN_URL = f"https://{DOMAIN}.salesforce.com"
+# Fetch Salesforce Objects
+def get_salesforce_objects(sf):
+    try:
+        response = sf.restful('sobjects')
+        return [obj['name'] for obj in response['sobjects'] if obj['createable']], None
+    except Exception as e:
+        return [], str(e)
 
-def get_token():
-    url = f"{LOGIN_URL}/services/oauth2/token"
-    payload = {
-        'grant_type': 'password',
-        'client_id': CLIENT_ID,
-        'client_secret': CLIENT_SECRET,
-        'username': USERNAME,
-        'password': PASSWORD + SECURITY_TOKEN
-    }
+# Fetch Object Fields
+def get_object_fields(sf, object_name):
+    try:
+        desc = sf.__getattr__(object_name).describe()
+        return [field['name'] for field in desc['fields']], None
+    except Exception as e:
+        return [], str(e)
 
+# OCR for Text Extraction
+def extract_text_from_pdf(pdf_path):
     try:
-        response = requests.post(url, data=payload)
-        print(f"Response Status: {response.status_code}")
-        print(f"Response Body: {response.text}")
-        result = response.json()
+        images = convert_from_path(pdf_path)
+        text_data = [pytesseract.image_to_string(img) for img in images]
+        return {"pages": text_data}, None
+    except Exception as e:
+        return {}, str(e)
 
-        if response.status_code == 200 and 'access_token' in result:
-            return result['access_token'], result['instance_url']
-        else:
-            print(f"Login failed with response: {result}")
-            raise Exception(f"❌ Salesforce login failed: {result}")
-    except requests.RequestException as e:
-        raise Exception(f"🚨 Network error during Salesforce login: {str(e)}")
+# Key-Value Pair Extraction using LayoutLMv3
+def extract_key_value_pairs(pdf_path):
+    try:
+        processor = LayoutLMv3Processor.from_pretrained("microsoft/layoutlmv3-base")
+        model = LayoutLMv3ForTokenClassification.from_pretrained("microsoft/layoutlmv3-base-finetuned-funsd")
+        images = convert_from_path(pdf_path)
+        extracted_data = []
+        for img in images:
+            encoding = processor(img, truncation=True, return_tensors="pt")
+            outputs = model(**encoding)
+            # Simplified: Return dummy key-value pairs (real implementation needs post-processing)
+            extracted_data.append({"keys": ["Contract Number", "Date"], "values": ["12345", "2025-01-01"]})
+        return extracted_data, None
+    except Exception as e:
+        return [], str(e)
 
-def get_salesforce_objects(token, instance_url):
-    headers = {'Authorization': f'Bearer {token}'}
-    response = requests.get(f"{instance_url}/services/data/v56.0/sobjects", headers=headers)
-    return response.json().get('sobjects', [])
+# Map Extracted Data to Salesforce Fields
+def map_fields(extracted_data, salesforce_fields):
+    mappings = {}
+    confidence_scores = {}
+    for key in extracted_data[0]["keys"]:  # Simplified: Using first page
+        for field in salesforce_fields:
+            if key.lower() in field.lower():
+                mappings[key] = field
+                confidence_scores[key] = 0.9  # Dummy confidence score
+    return mappings, confidence_scores, None
 
-def create_record(object_name, data, token, instance_url):
-    headers = {
-        'Authorization': f'Bearer {token}',
-        'Content-Type': 'application/json'
-    }
-    url = f"{instance_url}/services/data/v56.0/sobjects/{object_name}/"
-    response = requests.post(url, json=data, headers=headers)
-    return response.json()
+# Create Salesforce Record
+def create_record(sf, object_api_name, data):
+    try:
+        result = sf.__getattr__(object_api_name).create(data)
+        return result['id'], None
+    except Exception as e:
+        return None, str(e)
 
-def attach_pdf(record_id, file_path, token, instance_url):
+# Attach PDF to Salesforce Record
+def attach_pdf(sf, record_id, file_path):
     try:
         with open(file_path, "rb") as f:
-            body = base64.b64encode(f.read()).decode()
-
-        data = {
+            encoded_file = base64.b64encode(f.read()).decode()
+        attachment = {
             "ParentId": record_id,
             "Name": os.path.basename(file_path),
-            "Body": body
+            "Body": encoded_file
         }
-        headers = {
-            'Authorization': f'Bearer {token}',
-            'Content-Type': 'application/json'
-        }
-        url = f"{instance_url}/services/data/v56.0/sobjects/Attachment"
-        response = requests.post(url, json=data, headers=headers)
-        return response.json()
-    except FileNotFoundError:
-        raise Exception(f"❌ File not found: {file_path}")
+        sf.Attachment.create(attachment)
+        return "PDF Attached", None
     except Exception as e:
-        raise Exception(f"🚨 Error attaching PDF: {str(e)}")
+        return None, str(e)
 
-# Example usage
-if __name__ == "__main__":
-    try:
-        # Get access token and instance URL
-        token, instance_url = get_token()
-        print(f"✅ Successfully authenticated. Token: {token[:10]}..., Instance URL: {instance_url}")
-
-        # Example: List Salesforce objects
-        objects = get_salesforce_objects(token, instance_url)
-        print(f"Found {len(objects)} Salesforce objects")
-
-        # Example: Create a record
-        sample_data = {"Name": "Test Account"}
-        result = create_record("Account", sample_data, token, instance_url)
-        print(f"Created record: {result}")
-
-        # Example: Attach a PDF
-        record_id = result.get("id")
-        if record_id:
-            attach_pdf(record_id, "sample.pdf", token, instance_url)
-            print("✅ PDF attached successfully")
-    except Exception as e:
-        print(f"Error: {str(e)}")
+# Log Failed Migration
+def log_failure(pdf_path, object_name, error):
+    log_entry = {"pdf": pdf_path, "object": object_name, "error": error}
+    with open("failures.json", "a") as f:
+        json.dump(log_entry, f)
+        f.write("\n")