Create app.py

app.py

@@ -0,0 +1,233 @@
+
+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)}")