Spaces:
Sleeping
Sleeping
| 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() | |
| 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)}") | |