import gradio as gr from transformers import AutoModelForSeq2SeqLM, AutoTokenizer import os from fastapi import FastAPI from pydantic import BaseModel from simple_salesforce import Salesforce from dotenv import load_dotenv # Load environment variables load_dotenv() # Salesforce connection def get_salesforce_connection(): try: # Load credentials from environment variables username = os.getenv("SF_USERNAME") password = os.getenv("SF_PASSWORD") security_token = os.getenv("SF_SECURITY_TOKEN") domain = os.getenv("SF_DOMAIN", "login") # Default to production (login.salesforce.com) # Validate credentials if not all([username, password, security_token, domain]): missing = [] if not username: missing.append("SF_USERNAME") if not password: missing.append("SF_PASSWORD") if not security_token: missing.append("SF_SECURITY_TOKEN") if not domain: missing.append("SF_DOMAIN") raise ValueError(f"Missing environment variables: {', '.join(missing)}. Set them in .env or Space environment variables.") # Ensure all are strings if not all(isinstance(x, str) for x in [username, password, security_token, domain]): raise ValueError("All Salesforce credentials (SF_USERNAME, SF_PASSWORD, SF_SECURITY_TOKEN, SF_DOMAIN) must be strings.") sf = Salesforce( username=username, password=password, security_token=security_token, domain=domain ) # Test connection by fetching user info sf.User.get(sf.user_id) return sf except Exception as e: raise Exception(f"Failed to connect to Salesforce: {str(e)}") # Load Hugging Face token HF_TOKEN = os.getenv("HF_TOKEN") # Model configuration MODEL_PATH = "facebook/bart-large" # Public model # MODEL_PATH = "your_actual_username/fine_tuned_bart_construction" # Uncomment after uploading try: model = AutoModelForSeq2SeqLM.from_pretrained(MODEL_PATH, use_auth_token=HF_TOKEN if HF_TOKEN else None) tokenizer = AutoTokenizer.from_pretrained(MODEL_PATH, use_auth_token=HF_TOKEN if HF_TOKEN else None) except Exception as e: raise Exception(f"Failed to load model: {str(e)}") # Define input model for FastAPI class ChecklistInput(BaseModel): role: str = "Supervisor" project_id: str = "Unknown" project_name: str = "Unknown Project" milestones: str = "No milestones provided" record_id: str = None supervisor_id: str = None project_id_sf: str = None reflection_log: str = None download_link: str = None # Initialize FastAPI app = FastAPI() @app.post("/generate") async def generate_checklist(data: ChecklistInput): try: inputs = f"Role: {data.role} Project: {data.project_id} ({data.project_name}) Milestones: {data.milestones}" input_ids = tokenizer(inputs, return_tensors="pt", max_length=128, truncation=True).input_ids outputs = model.generate(input_ids, max_length=128, num_beams=4, early_stopping=True) checklist = tokenizer.decode(outputs[0], skip_special_tokens=True) tips = "1. Prioritize safety checks\n2. Review milestones\n3. Log progress" kpi_flag = "delay" in data.milestones.lower() or "behind" in data.milestones.lower() if data.record_id: sf = get_salesforce_connection() existing_record = sf.Supervisor_AI_Coaching__c.get(data.record_id, default={ 'Name': '', 'Supervisor_ID__c': None, 'Project_ID__c': None, 'Reflection_Log__c': '', 'Download_Link__c': '', 'Engagement_Score__c': 0, 'KPI_Flag__c': False, 'Daily_Checklist__c': '', 'Suggested_Tips__c': '' }) update_data = { 'Daily_Checklist__c': checklist, 'Suggested_Tips__c': tips, 'Engagement_Score__c': existing_record.get('Engagement_Score__c', 0) + 10, 'KPI_Flag__c': kpi_flag, 'Supervisor_ID__c': data.supervisor_id if data.supervisor_id else existing_record.get('Supervisor_ID__c'), 'Project_ID__c': data.project_id_sf if data.project_id_sf else existing_record.get('Project_ID__c'), 'Reflection_Log__c': data.reflection_log if data.reflection_log else existing_record.get('Reflection_Log__c', ''), 'Download_Link__c': data.download_link if data.download_link else existing_record.get('Download_Link__c', '') } sf.Supervisor_AI_Coaching__c.update(data.record_id, update_data) return { "checklist": checklist, "tips": tips, "kpi_flag": kpi_flag } except Exception as e: return {"error": str(e)} # Login and display records def login_and_display(project_id_sf): try: sf = get_salesforce_connection() query = f"SELECT Id, Name, Supervisor_ID__c, Project_ID__c, Daily_Checklist__c, Suggested_Tips__c, Reflection_Log__c, Engagement_Score__c, KPI_Flag__c, Download_Link__c FROM Supervisor_AI_Coaching__c WHERE Project_ID__c = '{project_id_sf}'" records = sf.query(query)["records"] if not records: return "No records found for Project ID.", "", False output = "Supervisor_AI_Coaching__c Records:\n" for record in records: output += ( f"Record ID: {record['Id']}\n" f"Name: {record['Name']}\n" f"Supervisor ID: {record['Supervisor_ID__c']}\n" f"Project ID: {record['Project_ID__c']}\n" f"Daily Checklist: {record['Daily_Checklist__c'] or 'N/A'}\n" f"Suggested Tips: {record['Suggested_Tips__c'] or 'N/A'}\n" f"Reflection Log: {record['Reflection_Log__c'] or 'N/A'}\n" f"Engagement Score: {record['Engagement_Score__c'] or 0}%\n" f"KPI Flag: {record['KPI_Flag__c']}\n" f"Download Link: {record['Download_Link__c'] or 'N/A'}\n" f"{'-'*50}\n" ) return output, "", False except Exception as e: return f"Error querying Salesforce: {str(e)}", "", False # Generate checklist from record def gradio_generate_checklist(record_id, role="Supervisor", project_id="Unknown", project_name="Unknown Project", milestones="No milestones provided", supervisor_id="", project_id_sf="", reflection_log="", download_link=""): try: sf = get_salesforce_connection() existing_record = sf.Supervisor_AI_Coaching__c.get(record_id, default={ 'Name': '', 'Supervisor_ID__c': None, 'Project_ID__c': None, 'Reflection_Log__c': '', 'Download_Link__c': '', 'Engagement_Score__c': 0, 'KPI_Flag__c': False, 'Daily_Checklist__c': '', 'Suggested_Tips__c': '' }) inputs = f"Role: {role} Project: {project_id} ({project_name}) Milestones: {milestones}" input_ids = tokenizer(inputs, return_tensors="pt", max_length=128, truncation=True).input_ids outputs = model.generate(input_ids, max_length=128, num_beams=4, early_stopping=True) checklist = tokenizer.decode(outputs[0], skip_special_tokens=True) tips = "1. Prioritize safety checks\n2. Review milestones\n3. Log progress" kpi_flag = "delay" in milestones.lower() or "behind" in milestones.lower() update_data = { 'Daily_Checklist__c': checklist, 'Suggested_Tips__c': tips, 'Engagement_Score__c': existing_record.get('Engagement_Score__c', 0) + 10, 'KPI_Flag__c': kpi_flag, 'Supervisor_ID__c': supervisor_id if supervisor_id else existing_record.get('Supervisor_ID__c'), 'Project_ID__c': project_id_sf if project_id_sf else existing_record.get('Project_ID__c'), 'Reflection_Log__c': reflection_log if reflection_log else existing_record.get('Reflection_Log__c', ''), 'Download_Link__c': download_link if download_link else existing_record.get('Download_Link__c', '') } sf.Supervisor_AI_Coaching__c.update(record_id, update_data) status = f"Updated Salesforce record {record_id}" return checklist, tips, kpi_flag, status except Exception as e: return f"Error: {str(e)}", "", False, "" # Define Gradio interface with gr.Blocks() as iface: gr.Markdown("# AI Coach for Site Supervisors") gr.Markdown("Enter a Project ID to view Supervisor_AI_Coaching__c records and generate checklists.") with gr.Tab("Login"): project_id_input = gr.Textbox(label="Project ID (Salesforce Project__c ID)", placeholder="Enter Project ID") login_button = gr.Button("Submit") records_output = gr.Textbox(label="Records", lines=10) login_button.click( fn=login_and_display, inputs=project_id_input, outputs=[records_output, gr.Textbox(visible=False), gr.Checkbox(visible=False)] ) with gr.Tab("Generate Checklist"): record_id = gr.Textbox(label="Record ID", placeholder="Enter Record ID from above") role = gr.Textbox(label="Role", value="Supervisor") project_id = gr.Textbox(label="Project ID", value="P001") project_name = gr.Textbox(label="Project Name", value="Building A") milestones = gr.Textbox(label="Milestones", value="Complete foundation by 5/15") supervisor_id = gr.Textbox(label="Supervisor ID (Salesforce User ID, optional)", value="") project_id_sf = gr.Textbox(label="Project ID (Salesforce Project__c ID, optional)", value="") reflection_log = gr.Textbox(label="Reflection Log (optional)", value="") download_link = gr.Textbox(label="Download Link (optional)", value="") generate_button = gr.Button("Generate and Update") checklist_output = gr.Textbox(label="Checklist") tips_output = gr.Textbox(label="Tips") kpi_flag_output = gr.Checkbox(label="KPI Flag") status_output = gr.Textbox(label="Salesforce Status") generate_button.click( fn=gradio_generate_checklist, inputs=[record_id, role, project_id, project_name, milestones, supervisor_id, project_id_sf, reflection_log, download_link], outputs=[checklist_output, tips_output, kpi_flag_output, status_output] ) # Mount FastAPI iface.app = app if __name__ == "__main__": try: iface.launch(server_name="0.0.0.0", server_port=7860, share=False) except Exception as e: print(f"Failed to launch Gradio: {str(e)}")