Update app.py

app.py

@@ -1,3 +1,4 @@
+```python
 import gradio as gr
 import logging
 from transformers import AutoModelForSeq2SeqLM, AutoTokenizer
@@ -7,7 +8,7 @@ from pydantic import BaseModel
 from simple_salesforce import Salesforce
 from dotenv import load_dotenv
 
-# Load environment variables from .env
+# Load environment variables
 load_dotenv()
 
 # Configure logging
@@ -33,8 +34,8 @@ if not HF_TOKEN:
     logging.warning("HF_TOKEN not set. Using public models only.")
 
 # Model configuration
-MODEL_PATH = "facebook/bart-large"  # Public model for testing
-# MODEL_PATH = "your_actual_username/fine_tuned_bart_construction"  # Uncomment and replace after uploading your model
+MODEL_PATH = "facebook/bart-large"  # Public model
+# MODEL_PATH = "your_actual_username/fine_tuned_bart_construction"  # Uncomment after uploading
 
 try:
     logging.info(f"Attempting to load model from {MODEL_PATH}")
@@ -51,11 +52,11 @@ class ChecklistInput(BaseModel):
     project_id: str = "Unknown"
     project_name: str = "Unknown Project"
     milestones: str = "No milestones provided"
-    record_id: str = None  # Supervisor_AI_Coaching__c record ID
-    supervisor_id: str = None  # Supervisor_ID__c
-    project_id_sf: str = None  # Project_ID__c
-    reflection_log: str = None  # Reflection_Log__c
-    download_link: str = None  # Download_Link__c
+    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()
@@ -63,7 +64,6 @@ app = FastAPI()
 @app.post("/generate")
 async def generate_checklist(data: ChecklistInput):
     try:
-        # Generate checklist and tips
         inputs = f"Role: {data.role} Project: {data.project_id} ({data.project_name}) Milestones: {data.milestones}"
         logging.info(f"Generating checklist for inputs: {inputs[:100]}...")
         input_ids = tokenizer(inputs, return_tensors="pt", max_length=128, truncation=True).input_ids
@@ -72,7 +72,6 @@ async def generate_checklist(data: ChecklistInput):
         tips = "1. Prioritize safety checks\n2. Review milestones\n3. Log progress"
         kpi_flag = "delay" in data.milestones.lower() or "behind" in data.milestones.lower()
 
-        # Update Salesforce if record_id is provided
         if data.record_id:
             try:
                 sf = get_salesforce_connection()
@@ -96,14 +95,12 @@ async def generate_checklist(data: ChecklistInput):
                     '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', '')
-                    # Name is read-only, not updated
                 }
                 sf.Supervisor_AI_Coaching__c.update(data.record_id, update_data)
                 logging.info(f"Updated Salesforce record {data.record_id} with fields: {update_data}")
             except Exception as sf_e:
                 logging.error(f"Failed to update Salesforce: {sf_e}")
 
-        logging.info("Checklist and tips generated successfully")
         return {
             "checklist": checklist,
             "tips": tips,
@@ -113,9 +110,49 @@ async def generate_checklist(data: ChecklistInput):
         logging.error(f"Error generating checklist: {e}")
         return {"error": str(e)}
 
-# Gradio interface function
-def gradio_generate_checklist(role, project_id, project_name, milestones, record_id="", supervisor_id="", project_id_sf="", reflection_log="", download_link=""):
+# 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)
@@ -123,63 +160,60 @@ def gradio_generate_checklist(role, project_id, project_name, milestones, record
         tips = "1. Prioritize safety checks\n2. Review milestones\n3. Log progress"
         kpi_flag = "delay" in milestones.lower() or "behind" in milestones.lower()
 
-        status = "No Salesforce update (record_id not provided)"
-        if record_id:
-            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': ''
-            })
-            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', '')
-                # Name is read-only
-            }
-            sf.Supervisor_AI_Coaching__c.update(record_id, update_data)
-            status = f"Updated Salesforce record {record_id} with fields: {update_data}"
+        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} with fields: {update_data}"
 
         return checklist, tips, kpi_flag, status
     except Exception as e:
         return f"Error: {str(e)}", "", False, ""
 
 # Define Gradio interface
-iface = gr.Interface(
-    fn=gradio_generate_checklist,
-    inputs=[
-        gr.Textbox(label="Role", value="Supervisor"),
-        gr.Textbox(label="Project ID", value="P001"),
-        gr.Textbox(label="Project Name", value="Building A"),
-        gr.Textbox(label="Milestones", value="Complete foundation by 5/15"),
-        gr.Textbox(label="Salesforce Record ID (optional)", value=""),
-        gr.Textbox(label="Supervisor ID (Salesforce User ID, optional)", value=""),
-        gr.Textbox(label="Project ID (Salesforce Project__c ID, optional)", value=""),
-        gr.Textbox(label="Reflection Log (optional)", value=""),
-        gr.Textbox(label="Download Link (optional)", value="")
-    ],
-    outputs=[
-        gr.Textbox(label="Checklist"),
-        gr.Textbox(label="Tips"),
-        gr.Checkbox(label="KPI Flag"),
-        gr.Textbox(label="Salesforce Status")
-    ],
-    title="AI Coach for Site Supervisors",
-    description="Generate daily checklists and tips, and update Salesforce records."
-)
-
-# Mount FastAPI to Gradio
+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), gr.Textbox(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__":
@@ -189,3 +223,4 @@ if __name__ == "__main__":
     except Exception as e:
         logging.error(f"Failed to launch Gradio interface: {e}")
         raise e
+```