Update app.py

app.py

@@ -3,12 +3,7 @@ import torch
 from transformers import AutoModelForCausalLM, AutoTokenizer
 from simple_salesforce import Salesforce
 import os
-import base64
-import datetime
 from dotenv import load_dotenv
-from fpdf import FPDF
-import shutil
-import html
 
 # Load environment variables
 load_dotenv()
@@ -19,8 +14,12 @@ missing_vars = [var for var in required_env_vars if not os.getenv(var)]
 if missing_vars:
     raise EnvironmentError(f"Missing required environment variables: {missing_vars}")
 
-# Load model and tokenizer
-model_name = "distilgpt2"
+# Defaults
+KPI_FLAG_DEFAULT = os.getenv('KPI_FLAG', 'True') == 'True'
+ENGAGEMENT_SCORE_DEFAULT = float(os.getenv('ENGAGEMENT_SCORE', '85.0'))
+
+# Load model and tokenizer (Updated to use distilgpt2)
+model_name = "distilgpt2"  # Using distilgpt2 for faster response
 tokenizer = AutoTokenizer.from_pretrained(model_name, use_fast=True)
 model = AutoModelForCausalLM.from_pretrained(model_name, low_cpu_mem_usage=True)
 
@@ -29,7 +28,7 @@ if tokenizer.pad_token is None:
     tokenizer.pad_token_id = tokenizer.convert_tokens_to_ids(tokenizer.pad_token)
 model.config.pad_token_id = tokenizer.pad_token_id
 
-# Prompt template
+# Refined Prompt to generate day-by-day tasks based on milestones
 PROMPT_TEMPLATE = """You are an AI assistant for construction supervisors. Given the role, project, milestones, and a reflection log, generate:
 
 1. A Daily Checklist with clear and concise tasks based on the role and milestones. 
@@ -55,96 +54,7 @@ Suggestions:
 - 
 """
 
-# Function to clean the text for PDF
-def clean_text_for_pdf(text):
-    return html.unescape(text).encode('latin-1', 'replace').decode('latin-1')
-
-# Save report as PDF
-def save_report_as_pdf(role, supervisor_name, project_id, checklist, suggestions):
-    now = datetime.datetime.now().strftime("%Y%m%d%H%M%S")
-    filename = f"report_{supervisor_name}_{project_id}_{now}.pdf"
-    file_path = f"./reports/{filename}"
-    os.makedirs("reports", exist_ok=True)
-
-    pdf = FPDF()
-    pdf.add_page()
-    pdf.set_font("Arial", 'B', 14)
-    pdf.cell(200, 10, txt="Supervisor Daily Report", ln=True, align="C")
-    pdf.set_font("Arial", size=12)
-    pdf.cell(200, 10, txt=clean_text_for_pdf(f"Role: {role}"), ln=True)
-    pdf.cell(200, 10, txt=clean_text_for_pdf(f"Supervisor: {supervisor_name}"), ln=True)
-    pdf.cell(200, 10, txt=clean_text_for_pdf(f"Project ID: {project_id}"), ln=True)
-    pdf.ln(5)
-    pdf.set_font("Arial", 'B', 12)
-    pdf.cell(200, 10, txt="Daily Checklist", ln=True)
-    pdf.set_font("Arial", size=12)
-    for line in checklist.split("\n"):
-        pdf.multi_cell(0, 10, clean_text_for_pdf(line))
-    pdf.ln(5)
-    pdf.set_font("Arial", 'B', 12)
-    pdf.cell(200, 10, txt="Focus Suggestions", ln=True)
-    pdf.set_font("Arial", size=12)
-    for line in suggestions.split("\n"):
-        pdf.multi_cell(0, 10, clean_text_for_pdf(line))
-    pdf.output(file_path)
-
-    # Copy the file to a temporary directory for Gradio to access
-    temp_pdf_path = "/tmp/" + os.path.basename(file_path)  # Use /tmp/ or current directory for Gradio
-    shutil.copy(file_path, temp_pdf_path)
-
-    return temp_pdf_path, filename
-
-# Upload to Salesforce and update record with the generated URL
-def upload_pdf_to_salesforce_and_update_link(supervisor_name, project_id, pdf_path, pdf_name):
-    try:
-        sf = Salesforce(
-            username=os.getenv('SF_USERNAME'),
-            password=os.getenv('SF_PASSWORD'),
-            security_token=os.getenv('SF_SECURITY_TOKEN'),
-            domain=os.getenv('SF_DOMAIN', 'login')
-        )
-        
-        # Read and encode the file as base64
-        with open(pdf_path, "rb") as f:
-            encoded = base64.b64encode(f.read()).decode()
-
-        # Create ContentVersion record to upload the PDF to Salesforce
-        content = sf.ContentVersion.create({
-            'Title': pdf_name,
-            'PathOnClient': pdf_name,
-            'VersionData': encoded
-        })
-
-        # Get the ContentDocumentId for the uploaded PDF
-        content_id = content['id']
-        
-        # Generate the download URL for the uploaded PDF
-        download_url = f"https://{sf.sf_instance}/sfc/servlet.shepherd/version/download/{content_id}"
-
-        # Query Salesforce to find the specific Supervisor_AI_Coaching__c record
-        query = sf.query(f"""
-            SELECT Id FROM Supervisor_AI_Coaching__c 
-            WHERE Project_ID__c = '{project_id}' 
-            AND Name = '{supervisor_name}'
-            LIMIT 1
-        """)
-
-        if query['totalSize'] > 0:
-            # Get the ID of the Supervisor_AI_Coaching__c record
-            coaching_id = query['records'][0]['Id']
-            
-            # Update the Supervisor_AI_Coaching__c record with the download URL
-            sf.Supervisor_AI_Coaching__c.update(coaching_id, {
-                'Download_Link__c': download_url  # Update the Download_Link__c field with the URL
-            })
-
-        return download_url
-
-    except Exception as e:
-        print(f"⚠️ Upload error: {e}")
-        return ""
-
-# Salesforce helpers
+# Salesforce Functions
 def get_roles_from_salesforce():
     try:
         sf = Salesforce(
@@ -157,7 +67,7 @@ def get_roles_from_salesforce():
         return list(set(record['Role__c'] for record in result['records']))
     except Exception as e:
         print(f"⚠️ Error fetching roles: {e}")
-        return []
+        return ["Site Manager", "Safety Officer", "Project Lead"]
 
 def get_supervisor_name_by_role(role):
     try:
@@ -167,10 +77,11 @@ def get_supervisor_name_by_role(role):
             security_token=os.getenv('SF_SECURITY_TOKEN'),
             domain=os.getenv('SF_DOMAIN', 'login')
         )
+        role = role.replace("'", "\\'")
         result = sf.query(f"SELECT Name FROM Supervisor__c WHERE Role__c = '{role}'")
         return [record['Name'] for record in result['records']]
     except Exception as e:
-        print(f"⚠️ Error fetching names: {e}")
+        print(f"⚠️ Error fetching supervisor names: {e}")
         return []
 
 def get_projects_for_supervisor(supervisor_name):
@@ -181,36 +92,54 @@ def get_projects_for_supervisor(supervisor_name):
             security_token=os.getenv('SF_SECURITY_TOKEN'),
             domain=os.getenv('SF_DOMAIN', 'login')
         )
-        result = sf.query(f"SELECT Id FROM Supervisor__c WHERE Name = '{supervisor_name}' LIMIT 1")
-        if result['totalSize'] == 0:
+        supervisor_name = supervisor_name.replace("'", "\\'")
+        supervisor_result = sf.query(f"SELECT Id FROM Supervisor__c WHERE Name = '{supervisor_name}' LIMIT 1")
+        if supervisor_result['totalSize'] == 0:
             return ""
-        supervisor_id = result['records'][0]['Id']
+        supervisor_id = supervisor_result['records'][0]['Id']
         project_result = sf.query(f"SELECT Name FROM Project__c WHERE Supervisor_ID__c = '{supervisor_id}' LIMIT 1")
         return project_result['records'][0]['Name'] if project_result['totalSize'] > 0 else ""
     except Exception as e:
         print(f"⚠️ Error fetching project: {e}")
         return ""
 
-# Generate Salesforce dashboard URL
+def field_exists(sf, object_name, field_name):
+    try:
+        obj_desc = getattr(sf, object_name).describe()
+        return field_name in [field['name'] for field in obj_desc['fields']]
+    except Exception as e:
+        print(f"⚠️ Error checking field {field_name}: {e}")
+        return False
+
+# New function to generate Salesforce dashboard URL (Visualforce Page)
 def generate_salesforce_dashboard_url(supervisor_name, project_id):
+    # Use the provided Salesforce Visualforce URL with supervisorName and projectId as parameters
     return f"https://aicoachforsitesupervisors-dev-ed--c.develop.vf.force.com/apex/DashboardPage?supervisorName={supervisor_name}&projectId={project_id}"
 
+# Dashboard button function
 def open_dashboard(role, supervisor_name, project_id):
-    url = generate_salesforce_dashboard_url(supervisor_name, project_id)
-    return f'<a href="{url}" target="_blank">Open Salesforce Dashboard</a>'
+    # Generate dynamic URL based on supervisor and project
+    dashboard_url = generate_salesforce_dashboard_url(supervisor_name, project_id)
+    return f'<a href="{dashboard_url}" target="_blank" rel="noopener noreferrer" style="font-size:16px;">Open Salesforce Dashboard</a>'
 
-# Generate AI output and return the correct file for download
+# Generate function
 def generate_outputs(role, supervisor_name, project_id, milestones, reflection):
     if not all([role, supervisor_name, project_id, milestones, reflection]):
-        return "❗ Please fill all fields.", "", None, ""
+        return "❗ Please fill all fields.", ""
+
+    prompt = PROMPT_TEMPLATE.format(
+        role=role,
+        project_id=project_id,
+        milestones=milestones,
+        reflection=reflection
+    )
 
-    prompt = PROMPT_TEMPLATE.format(role=role, project_id=project_id, milestones=milestones, reflection=reflection)
     inputs = tokenizer(prompt, return_tensors="pt", padding=True, truncation=True, max_length=512)
     try:
         with torch.no_grad():
             outputs = model.generate(
                 inputs['input_ids'],
-                max_new_tokens=150,
+                max_new_tokens=150,  # Increased max tokens to capture more content
                 no_repeat_ngram_size=2,
                 do_sample=True,
                 top_p=0.9,
@@ -220,31 +149,55 @@ def generate_outputs(role, supervisor_name, project_id, milestones, reflection):
         result = tokenizer.decode(outputs[0], skip_special_tokens=True)
     except Exception as e:
         print(f"⚠️ Generation error: {e}")
-        return "", "", None, ""
+        return "", ""
 
     def extract_between(text, start, end):
         s = text.find(start)
         e = text.find(end, s) if end else len(text)
         return text[s + len(start):e].strip() if s != -1 else ""
 
+    # Extract the checklist and suggestions
     checklist = extract_between(result, "Checklist:\n", "Suggestions:")
     suggestions = extract_between(result, "Suggestions:\n", None)
 
+    # If checklist or suggestions are empty, generate fallback content
     if not checklist.strip():
-        checklist = "- Perform daily safety inspection"
+        checklist = generate_fallback_checklist(role, milestones)
     if not suggestions.strip():
-        suggestions = "- Monitor team coordination\n- Review safety protocols with the team"
-
-    pdf_path, pdf_name = save_report_as_pdf(role, supervisor_name, project_id, checklist, suggestions)
-    pdf_url = upload_pdf_to_salesforce_and_update_link(supervisor_name, project_id, pdf_path, pdf_name)
-
-    if pdf_url:
-        suggestions += f"\n\n🔗 [Download PDF Report]({pdf_url})"
-
-    # Return the generated checklist, suggestions, and the path to the PDF
-    return checklist, suggestions, pdf_path, pdf_name
-
-# Gradio Interface
+        suggestions = generate_fallback_suggestions(reflection)
+
+    return checklist, suggestions
+
+# Fallback generation for checklist and suggestions
+def generate_fallback_checklist(role, milestones):
+    checklist_items = []
+
+    # If milestones are provided, add them to the checklist directly
+    if milestones and milestones.strip():
+        kpis = [kpi.strip() for kpi in milestones.split(",")]
+        for kpi in kpis:
+            checklist_items.append(f"- Ensure progress on {kpi}")
+    else:
+        checklist_items.append("- Perform daily safety inspection")
+    
+    return "\n".join(checklist_items)
+
+def generate_fallback_suggestions(reflection):
+    suggestions_items = []
+    reflection_lower = reflection.lower()
+    if "student" in reflection_lower or "learning" in reflection_lower:
+        suggestions_items.append("- Ensure students are logging incidents consistently")
+        suggestions_items.append("- Provide guidance on timely incident recording")
+    if "incident" in reflection_lower:
+        suggestions_items.append("- Follow up on reported incidents with corrective actions")
+    
+    if not suggestions_items:
+        suggestions_items.append("- Monitor team coordination")
+        suggestions_items.append("- Review safety protocols with the team")
+    
+    return "\n".join(suggestions_items)
+
+# Interface
 def create_interface():
     roles = get_roles_from_salesforce()
     with gr.Blocks(theme="soft") as demo:
@@ -255,8 +208,8 @@ def create_interface():
             supervisor_name = gr.Dropdown(choices=[], label="Supervisor Name")
             project_id = gr.Textbox(label="Project ID", interactive=False)
 
-        milestones = gr.Textbox(label="Milestones (comma-separated KPIs)")
-        reflection = gr.Textbox(label="Reflection Log", lines=4)
+        milestones = gr.Textbox(label="Milestones (comma-separated KPIs)", placeholder="E.g. Safety training, daily inspection")
+        reflection = gr.Textbox(label="Reflection Log", lines=4, placeholder="Any concerns, delays, updates...")
 
         with gr.Row():
             generate = gr.Button("Generate")
@@ -266,35 +219,24 @@ def create_interface():
 
         checklist_output = gr.Textbox(label="✅ Daily Checklist")
         suggestions_output = gr.Textbox(label="💡 Focus Suggestions")
-        download_button = gr.File(label="⬇ Download Report")
-        pdf_link = gr.HTML()
-        dashboard_link = gr.HTML()
+        dashboard_link = gr.HTML("")
 
         role.change(fn=lambda r: gr.update(choices=get_supervisor_name_by_role(r)), inputs=role, outputs=supervisor_name)
         supervisor_name.change(fn=get_projects_for_supervisor, inputs=supervisor_name, outputs=project_id)
 
-        def handle_generate(role, supervisor_name, project_id, milestones, reflection):
-            checklist, suggestions, pdf_path, pdf_name = generate_outputs(role, supervisor_name, project_id, milestones, reflection)
-            return checklist, suggestions, pdf_path, pdf_name
-
-        generate.click(fn=handle_generate,
+        generate.click(fn=generate_outputs,
                        inputs=[role, supervisor_name, project_id, milestones, reflection],
-                       outputs=[checklist_output, suggestions_output, download_button, pdf_link])
+                       outputs=[checklist_output, suggestions_output])
 
-        clear.click(fn=lambda: ("", "", "", "", ""),
-                    inputs=None,
+        clear.click(fn=lambda: ("", "", "", "", ""), inputs=None,
                     outputs=[role, supervisor_name, project_id, milestones, reflection])
 
         refresh.click(fn=lambda: gr.update(choices=get_roles_from_salesforce()), outputs=role)
 
-        dashboard_btn.click(fn=open_dashboard,
-                            inputs=[role, supervisor_name, project_id],
-                            outputs=dashboard_link)
+        dashboard_btn.click(fn=open_dashboard, inputs=[role, supervisor_name, project_id], outputs=dashboard_link)
 
     return demo
 
 if __name__ == "__main__":
     app = create_interface()
     app.launch()
-
-