Update app.py

app.py

@@ -1,7 +1,4 @@
 import gradio as gr
-import torch
-from transformers import AutoModelForCausalLM, AutoTokenizer
-from simple_salesforce import Salesforce
 import os
 import base64
 import datetime
@@ -12,6 +9,7 @@ import html
 import io
 import matplotlib.pyplot as plt
 import numpy as np
+from simple_salesforce import Salesforce
 
 # Load environment variables
 load_dotenv()
@@ -22,15 +20,9 @@ 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"
-tokenizer = AutoTokenizer.from_pretrained(model_name, use_fast=True)
-model = AutoModelForCausalLM.from_pretrained(model_name, low_cpu_mem_usage=True)
-
-if tokenizer.pad_token is None:
-    tokenizer.pad_token = tokenizer.eos_token
-    tokenizer.pad_token_id = tokenizer.convert_tokens_to_ids(tokenizer.pad_token)
-model.config.pad_token_id = tokenizer.pad_token_id
+# Defaults
+KPI_FLAG_DEFAULT = os.getenv('KPI_FLAG', 'True') == 'True'
+ENGAGEMENT_SCORE_DEFAULT = float(os.getenv('ENGAGEMENT_SCORE', '85.0'))
 
 # Function to generate progress chart
 def show_dashboard_chart(start_date, end_date, tasks_completed):
@@ -38,7 +30,6 @@ def show_dashboard_chart(start_date, end_date, tasks_completed):
     labels = list(tasks_completed.keys())
     remaining_tasks = [5 - task for task in completed_tasks]  # Assuming 5 tasks per date
 
-    # Create the bar chart with completed and remaining tasks
     fig, ax = plt.subplots(figsize=(10, 6))
     ax.bar(labels, completed_tasks, color='green', label="Completed")
     ax.bar(labels, remaining_tasks, bottom=completed_tasks, color='gray', label="Remaining")
@@ -63,7 +54,6 @@ def get_dashboard_data_from_salesforce(supervisor_name, project_id):
             domain=os.getenv('SF_DOMAIN', 'login')
         )
 
-        # Get the start and end date from Salesforce
         query = sf.query(f"SELECT Start_Date__c, End_Date__c FROM Project__c WHERE Name = '{project_id}' LIMIT 1")
         if query['totalSize'] == 0:
             return "", "", None, "Project not found"
@@ -71,11 +61,9 @@ def get_dashboard_data_from_salesforce(supervisor_name, project_id):
         start_date_str = query['records'][0]['Start_Date__c']
         end_date_str = query['records'][0]['End_Date__c']
 
-        # Convert the string dates to datetime objects
         start_date = datetime.datetime.strptime(start_date_str, "%Y-%m-%d")
         end_date = datetime.datetime.strptime(end_date_str, "%Y-%m-%d")
 
-        # Generate task dates and simulated completion data
         task_dates = [start_date + datetime.timedelta(days=i) for i in range((end_date - start_date).days + 1)]
         tasks_completed = {str(task_dates[i].date()): np.random.randint(1, 6) for i in range(len(task_dates))}
 
@@ -169,21 +157,49 @@ def get_projects_for_supervisor(supervisor_name):
 
 # Function to generate daily checklist and focus suggestions
 def generate_checklist_and_suggestions(role, project_id, milestones, reflection):
-    prompt = f"""
-    You are a supervisor assistant. Given the role {role}, project {project_id}, milestones {milestones}, and reflection log {reflection}, generate:
-    
-    1. A Daily Checklist with clear and concise tasks.
-    2. Focus Suggestions based on concerns or keywords from the reflection log.
-    """
-    inputs = tokenizer(prompt, return_tensors="pt")
-    outputs = model.generate(inputs['input_ids'], max_length=200, num_return_sequences=1)
-    generated_text = tokenizer.decode(outputs[0], skip_special_tokens=True)
-    
-    # Split generated text into checklist and suggestions
-    parts = generated_text.split("\n")
-    checklist = "\n".join(parts[:len(parts)//2])
-    suggestions = "\n".join(parts[len(parts)//2:])
-    
+    if not all([role, project_id, milestones, reflection]):
+        return "Please fill all fields.", ""
+
+    # Generate Checklist
+    checklist_items = []
+    if role.lower() == "site supervisor":
+        checklist_items.append("- Conduct daily site safety briefing")
+        checklist_items.append("- Monitor progress against project timeline")
+
+    milestones_lower = milestones.lower()
+    if "excavation" in milestones_lower:
+        checklist_items.append("- Verify excavation progress and safety measures")
+    if "foundation work" in milestones_lower:
+        checklist_items.append("- Inspect foundation work for quality")
+    if "pour the foundation concrete" in milestones_lower:
+        checklist_items.append("- Ensure proper concrete pouring and curing")
+    if "set up reinforcement bars" in milestones_lower:
+        checklist_items.append("- Check reinforcement bar installation")
+    if "test soil compaction" in milestones_lower:
+        checklist_items.append("- Conduct soil compaction tests")
+
+    if not checklist_items:
+        checklist_items.append("- Perform general site inspection")
+
+    checklist = "\n".join(checklist_items)
+
+    # Generate Suggestions
+    suggestions_items = []
+    reflection_lower = reflection.lower()
+    if "weather conditions" in reflection_lower or "rainfall" in reflection_lower:
+        suggestions_items.append("- Adjust schedule for weather delays")
+        suggestions_items.append("- Implement weather protection measures")
+    if "delays" in reflection_lower:
+        suggestions_items.append("- Identify causes of delays and mitigate")
+    if "leveling" in reflection_lower or "base" in reflection_lower:
+        suggestions_items.append("- Verify leveling accuracy and resolve issues")
+
+    if not suggestions_items:
+        suggestions_items.append("- Monitor team performance")
+        suggestions_items.append("- Review safety protocols")
+
+    suggestions = "\n".join(suggestions_items)
+
     return checklist, suggestions
 
 # Function to upload the report and create the Supervisor AI Coaching record in Salesforce
@@ -196,7 +212,6 @@ def upload_report_and_create_supervisor_ai_coaching(supervisor_name, project_id,
             domain=os.getenv('SF_DOMAIN', 'login')
         )
 
-        # Upload the PDF file to Salesforce as Content Version
         with open(pdf_path, "rb") as f:
             encoded = base64.b64encode(f.read()).decode()
 
@@ -209,7 +224,6 @@ def upload_report_and_create_supervisor_ai_coaching(supervisor_name, project_id,
         content_id = content['id']
         download_url = f"https://{sf.sf_instance}/sfc/servlet.shepherd/version/download/{content_id}"
 
-        # Create a Supervisor AI Coaching record
         query = sf.query(f"SELECT Id FROM Supervisor__c WHERE Name = '{supervisor_name}' LIMIT 1")
         supervisor_id = query['records'][0]['Id'] if query['totalSize'] > 0 else None
         if not supervisor_id:
@@ -220,7 +234,6 @@ def upload_report_and_create_supervisor_ai_coaching(supervisor_name, project_id,
         if not project_id_sf:
             return "Project not found."
 
-        # Create Supervisor AI Coaching record with all necessary fields
         sf.Supervisor_AI_Coaching__c.create({
             'Project_ID__c': project_id_sf,
             'Supervisor_ID__c': supervisor_id,
@@ -230,23 +243,22 @@ def upload_report_and_create_supervisor_ai_coaching(supervisor_name, project_id,
         })
 
         return "Supervisor AI Coaching record created and report uploaded successfully."
-
     except Exception as e:
         return f"Error: {str(e)}"
 
 # Gradio interface
 def create_interface():
-    roles = get_roles_from_salesforce()  # Get roles from Salesforce dynamically
+    roles = get_roles_from_salesforce()
     with gr.Blocks(theme="soft", css=".footer { display: none; }") as demo:
         gr.Markdown("## 🧠 AI-Powered Supervisor Assistant")
 
         with gr.Row():
-            role = gr.Dropdown(choices=roles, label="Role")
-            supervisor_name = gr.Dropdown(choices=[], label="Supervisor Name")
-            project_id = gr.Textbox(label="Project ID", interactive=False)
+            role = gr.Dropdown(choices=roles, label="Role", value="Site Supervisor")
+            supervisor_name = gr.Dropdown(choices=["SUP-056"], label="Supervisor Name", value="SUP-056")
+            project_id = gr.Textbox(label="Project ID", value="PROJ-078", 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)", value="Excavation and Foundation Work, Pour the foundation concrete, Set up reinforcement bars for the foundation, Test soil compaction")
+        reflection = gr.Textbox(label="Reflection Log", lines=4, value="The excavation and foundation work was slower than expected due to unexpected weather conditions. We faced some delays because of heavy rainfall. The foundation team had some minor issues with leveling the base but they managed to resolve it in the end.")
 
         with gr.Row():
             generate = gr.Button("Generate")
@@ -263,9 +275,12 @@ def create_interface():
 
         def handle_generate(role, supervisor_name, project_id, milestones, reflection):
             checklist, suggestions = generate_checklist_and_suggestions(role, project_id, milestones, reflection)
+            if not checklist or not suggestions:
+                return "Error generating checklist or suggestions.", "Error generating checklist or suggestions.", None, None, "Error: Generation failed."
             pdf_path, pdf_name = save_report_as_pdf(role, supervisor_name, project_id, checklist, suggestions)
             supervisor_ai_coaching_response = upload_report_and_create_supervisor_ai_coaching(supervisor_name, project_id, checklist, suggestions, pdf_path, pdf_name)
-            return checklist, suggestions, pdf_path, pdf_name, supervisor_ai_coaching_response
+            pdf_link_content = f'<a href="file={pdf_path}" download="{pdf_name}">Download PDF</a>'
+            return checklist, suggestions, pdf_path, pdf_link_content, supervisor_ai_coaching_response
 
         generate.click(fn=handle_generate,
                        inputs=[role, supervisor_name, project_id, milestones, reflection],
@@ -277,7 +292,6 @@ def create_interface():
 
         refresh.click(fn=lambda: gr.update(choices=get_roles_from_salesforce()), outputs=role)
 
-        # Supervisor Dashboard Tab
         with gr.Tab("📊 Supervisor Dashboard"):
             dash_supervisor = gr.Textbox(label="Supervisor Name", placeholder="e.g., SUP-056")
             dash_project = gr.Textbox(label="Project ID", placeholder="e.g., PROJ-078")
@@ -298,6 +312,4 @@ def create_interface():
 
 if __name__ == "__main__":
     app = create_interface()
-    app.launch()
-
-
+    app.launch()