dschandra's picture
Update app.py
f8912a5 verified
import os
import pdfplumber
import gradio as gr
from transformers import pipeline
from simple_salesforce import Salesforce
from dotenv import load_dotenv
import base64
# Load environment variables from .env
load_dotenv()
# Salesforce credentials
SF_USERNAME = os.getenv("SF_USERNAME")
SF_PASSWORD = os.getenv("SF_PASSWORD")
SF_SECURITY_TOKEN = os.getenv("SF_SECURITY_TOKEN")
SF_LOGIN_URL = os.getenv("SF_LOGIN_URL", "https://login.salesforce.com")
SF_OBJECT_NAME = os.getenv("SF_OBJECT_NAME", "Agent_Prospect__c")
SF_SCORE_FIELD = os.getenv("SF_SCORE_FIELD", "Suitability_Score__c")
SF_LINK_FIELD = os.getenv("SF_RESUME_FIELD_LINK", "Resume_File_Link__c")
# Validate required credentials
required = ["SF_USERNAME", "SF_PASSWORD", "SF_SECURITY_TOKEN"]
missing = [var for var in required if not os.getenv(var)]
if missing:
raise ValueError(f"Missing required .env variables: {', '.join(missing)}")
# Determine domain
domain = "login" if "login" in SF_LOGIN_URL else "test"
# Connect to Salesforce
try:
sf = Salesforce(
username=SF_USERNAME,
password=SF_PASSWORD,
security_token=SF_SECURITY_TOKEN,
domain=domain,
version="59.0"
)
print("Salesforce connection successful.")
except Exception as e:
raise ValueError(f"Failed to connect to Salesforce: {str(e)}")
# Load Hugging Face model
classifier = pipeline("text-classification", model="nlptown/bert-base-multilingual-uncased-sentiment")
def process_resume(file):
try:
# Extract text from PDF
with pdfplumber.open(file.name) as pdf:
extracted_text = "\n".join([page.extract_text() or "" for page in pdf.pages])
print(f"Extracted text length: {len(extracted_text)} characters")
if not extracted_text.strip():
return "❌ No extractable text found in the PDF."
# Call Hugging Face model
result = classifier(extracted_text[:1000])
label = result[0]['label']
score = round(float(result[0]['score']) * 100, 2)
summary = f"Predicted Label: {label}\nSuitability Score: {score:.2f}"
print(f"Classifier result: {summary}")
# Encode PDF in base64
with open(file.name, "rb") as f:
encoded_pdf = base64.b64encode(f.read()).decode("utf-8")
print(f"Encoded PDF size: {len(encoded_pdf)}")
# Upload file as ContentVersion
content_result = sf.ContentVersion.create({
"Title": "Resume",
"PathOnClient": file.name,
"VersionData": encoded_pdf
})
version_id = content_result.get("id")
print(f"ContentVersion created: {version_id}")
# Get ContentDocumentId from ContentVersion
query_result = sf.query(
f"SELECT ContentDocumentId FROM ContentVersion WHERE Id = '{version_id}'"
)
print(f"Query result: {query_result}")
if not query_result["records"]:
return "❌ Failed to retrieve ContentDocumentId."
content_doc_id = query_result["records"][0]["ContentDocumentId"]
print(f"ContentDocumentId: {content_doc_id}")
# Create a new Agent_Prospect__c record
try:
record_result = sf.__getattr__(SF_OBJECT_NAME).create({
SF_SCORE_FIELD: score,
SF_LINK_FIELD: "" # Placeholder; will update with download link
})
record_id = record_result.get("id")
print(f"New Agent_Prospect__c record created: {record_id}")
except Exception as e:
print(f"Error creating Agent_Prospect__c record: {str(e)}")
raise
# Link file to the new Salesforce record
try:
sf.ContentDocumentLink.create({
"ContentDocumentId": content_doc_id,
"LinkedEntityId": record_id,
"ShareType": "V",
"Visibility": "AllUsers"
})
print("ContentDocumentLink created successfully.")
except Exception as e:
print(f"ContentDocumentLink Error: {str(e)}")
raise
# Create download link
download_link = f"https://{sf.sf_instance}/sfc/servlet.shepherd/document/download/{content_doc_id}"
print(f"Download link: {download_link}")
# Update the record with the download link
try:
sf.__getattr__(SF_OBJECT_NAME).update(record_id, {
SF_LINK_FIELD: download_link
})
print("Salesforce record updated with download link.")
except Exception as e:
print(f"Error updating record with download link: {str(e)}")
raise
return f"{summary}\n\nβœ… New record created and resume uploaded to Salesforce.\nπŸ“Ž [Download Resume]({download_link})\nRecord ID: {record_id}"
except Exception as e:
return f"❌ Error: {str(e)}"
# Gradio Interface
gr.Interface(
fn=process_resume,
inputs=gr.File(label="Upload Resume (PDF)", file_types=[".pdf"]),
outputs="text",
title="LIC Resume AI Scorer",
description="Upload a resume PDF. A new record will be created in Salesforce with the score and resume."
).launch(share=False)