Update app.py

app.py

@@ -8,35 +8,35 @@ from dotenv import load_dotenv
 # Load environment variables
 load_dotenv()
 
-# Check if required environment variables are set
+# Check for required environment variables
 required_env_vars = ['SF_USERNAME', 'SF_PASSWORD', 'SF_SECURITY_TOKEN']
 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}")
 
-# Get configurable values for KPI_Flag__c and Engagement_Score__c
-KPI_FLAG_DEFAULT = os.getenv('KPI_FLAG', 'True') == 'True'  # Default to True if not set
-ENGAGEMENT_SCORE_DEFAULT = float(os.getenv('ENGAGEMENT_SCORE', '85.0'))  # Default to 85.0
+# Configurable defaults
+KPI_FLAG_DEFAULT = os.getenv('KPI_FLAG', 'True') == 'True'
+ENGAGEMENT_SCORE_DEFAULT = float(os.getenv('ENGAGEMENT_SCORE', '85.0'))
 
-# Initialize model and tokenizer
+# Initialize model and tokenizer with optimized settings
 model_name = "distilgpt2"
-tokenizer = AutoTokenizer.from_pretrained(model_name)
-model = AutoModelForCausalLM.from_pretrained(model_name)
+tokenizer = AutoTokenizer.from_pretrained(model_name, use_fast=True)  # Use fast tokenizer
+model = AutoModelForCausalLM.from_pretrained(model_name, low_cpu_mem_usage=True)  # Optimize memory usage
 
-# Avoid warnings by setting pad token
+# Set pad token
 if tokenizer.pad_token is None:
-    tokenizer.pad_token = tokenizer.eos_token if tokenizer.eos_token else "[PAD]"
+    tokenizer.pad_token = tokenizer.eos_token or "[PAD]"
     tokenizer.pad_token_id = tokenizer.convert_tokens_to_ids(tokenizer.pad_token)
 model.config.pad_token_id = tokenizer.pad_token_id
 
-# Prompt template for generating structured output
+# Prompt template
 PROMPT_TEMPLATE = """You are an AI coach for construction supervisors. Based on the following inputs, generate a daily checklist, focus suggestions, and a motivational quote.
 Inputs:
 Role: {role}
 Project: {project_id}
 Milestones: {milestones}
 Reflection: {reflection}
-Format your response clearly like this:
+Format:
 Checklist:
 - {milestones_list}
 Suggestions:
@@ -45,7 +45,7 @@ Quote:
 - Your motivational quote here
 """
 
-# Function to get all roles from Salesforce
+# Fetch roles from Salesforce
 def get_roles_from_salesforce():
     try:
         sf = Salesforce(
@@ -54,23 +54,14 @@ def get_roles_from_salesforce():
             security_token=os.getenv('SF_SECURITY_TOKEN'),
             domain=os.getenv('SF_DOMAIN', 'login')
         )
-
-        # Query distinct Role__c values
         result = sf.query("SELECT Role__c FROM Supervisor__c WHERE Role__c != NULL")
-
-        # Extract roles and remove duplicates
         roles = list(set(record['Role__c'] for record in result.get('records', [])))
-
-        print(f"✅ Fetched {len(roles)} unique roles from Salesforce")
         return roles
-
     except Exception as e:
-        print(f"⚠️ Error fetching roles from Salesforce: {e}")
-        print("Using fallback roles...")
-        return ["Site Manager", "Safety Officer", "Project Lead"]  # Match actual active roles
-
+        print(f"⚠️ Error fetching roles: {e}")
+        return ["Site Manager", "Safety Officer", "Project Lead"]
 
-# Function to get supervisor's Name (Auto Number) by role
+# Fetch supervisor names by role
 def get_supervisor_name_by_role(role):
     try:
         sf = Salesforce(
@@ -79,81 +70,51 @@ def get_supervisor_name_by_role(role):
             security_token=os.getenv('SF_SECURITY_TOKEN'),
             domain=os.getenv('SF_DOMAIN', 'login')
         )
-
-        # Escape single quotes in the role to prevent SOQL injection
         role = role.replace("'", "\\'")
-
-        # Query all supervisors for the selected role
         result = sf.query(f"SELECT Name FROM Supervisor__c WHERE Role__c = '{role}'")
         if result['totalSize'] == 0:
-            print("❌ No matching supervisors found.")
             return []
-
-        # Extract all supervisor names
-        supervisor_names = [record['Name'] for record in result['records']]
-        print(f"✅ Found supervisors: {supervisor_names} for role: {role}")
-        return supervisor_names
-
+        return [record['Name'] for record in result['records']]
     except Exception as e:
         print(f"⚠️ Error fetching supervisor names: {e}")
         return []
 
-
-# Function to get project IDs and names assigned to selected supervisor
+# Fetch project for supervisor
 def get_projects_for_supervisor(supervisor_name):
     try:
-        # Use the selected supervisor name to fetch the associated project
         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')
         )
-
-        # Escape single quotes in the supervisor_name
         supervisor_name = supervisor_name.replace("'", "\\'")
-
-        # Step 1: Get the Salesforce record ID of the supervisor based on the Name
         supervisor_result = sf.query(f"SELECT Id FROM Supervisor__c WHERE Name = '{supervisor_name}' LIMIT 1")
         if supervisor_result['totalSize'] == 0:
-            print("❌ No supervisor found with the given name.")
             return ""
-
         supervisor_id = supervisor_result['records'][0]['Id']
-
-        # Step 2: Query Project__c records where Supervisor_ID__c matches the supervisor's record ID
         project_result = sf.query(f"SELECT Name FROM Project__c WHERE Supervisor_ID__c = '{supervisor_id}' LIMIT 1")
-
         if project_result['totalSize'] == 0:
-            print("❌ No project found for supervisor.")
             return ""
-
-        project_name = project_result['records'][0]['Name']
-        print(f"✅ Found project: {project_name} for supervisor: {supervisor_name}")
-        return project_name
-
+        return project_result['records'][0]['Name']
     except Exception as e:
-        print(f"⚠️ Error fetching project for supervisor: {e}")
+        print(f"⚠️ Error fetching project: {e}")
         return ""
 
-
-# Function to generate AI-based coaching output
+# Generate AI outputs
 def generate_outputs(role, supervisor_name, project_id, milestones, reflection):
     if not all([role, supervisor_name, project_id, milestones, reflection]):
         return "Error: All fields are required.", "", ""
 
-    # Format the prompt
     milestones_list = "\n- ".join([m.strip() for m in milestones.split(",")])
-
     suggestions_list = ""
     if "delays" in reflection.lower():
-        suggestions_list = "- Consider adjusting timelines to accommodate delays.\n- Communicate delays to all relevant stakeholders."
+        suggestions_list = "- Adjust timelines for delays.\n- Inform stakeholders."
     elif "weather" in reflection.lower():
-        suggestions_list = "- Ensure team has rain gear.\n- Monitor weather updates for possible further delays."
+        suggestions_list = "- Ensure rain gear availability.\n- Monitor weather updates."
     elif "equipment" in reflection.lower():
-        suggestions_list = "- Inspect all equipment to ensure no malfunctions.\n- Schedule maintenance if necessary."
+        suggestions_list = "- Inspect equipment.\n- Schedule maintenance."
 
-    # Fill in the prompt template
     prompt = PROMPT_TEMPLATE.format(
         role=role,
         project_id=project_id,
@@ -163,30 +124,28 @@ def generate_outputs(role, supervisor_name, project_id, milestones, reflection):
         suggestions_list=suggestions_list
     )
 
-    # Tokenize input
-    inputs = tokenizer(prompt, return_tensors="pt", max_length=512, truncation=True, padding=True)
+    # Tokenize with optimized settings
+    inputs = tokenizer(prompt, return_tensors="pt", max_length=256, truncation=True, padding=True)
 
-    # Generate response
+    # Faster generation with adjusted parameters
     try:
         with torch.no_grad():
             outputs = model.generate(
                 inputs['input_ids'],
-                max_length=1024,  # Increased to allow for longer outputs
+                max_length=512,  # Reduced for speed
                 num_return_sequences=1,
                 no_repeat_ngram_size=2,
                 do_sample=True,
-                top_p=0.9,
-                temperature=0.8,
+                top_p=0.85,  # Slightly tighter for faster convergence
+                temperature=0.7,  # Lower for more deterministic output
                 pad_token_id=tokenizer.pad_token_id
             )
-
         generated_text = tokenizer.decode(outputs[0], skip_special_tokens=True)
-
     except Exception as e:
-        print(f"⚠️ Error during model generation: {e}")
+        print(f"⚠️ Error during generation: {e}")
         return "Error: Failed to generate outputs.", "", ""
 
-    # Parse sections
+    # Extract sections
     def extract_section(text, start_marker, end_marker):
         start = text.find(start_marker)
         if start == -1:
@@ -204,20 +163,16 @@ def generate_outputs(role, supervisor_name, project_id, milestones, reflection):
 
     return checklist, suggestions, quote
 
-
-# Function to check if a field exists in a Salesforce object
+# Check if Salesforce field exists
 def field_exists(sf, object_name, field_name):
     try:
-        # Describe the object to get its fields
         obj_desc = getattr(sf, object_name).describe()
-        fields = [field['name'] for field in obj_desc['fields']]
-        return field_name in fields
+        return field_name in [field['name'] for field in obj_desc['fields']]
     except Exception as e:
-        print(f"⚠️ Error checking if field {field_name} exists in {object_name}: {e}")
+        print(f"⚠️ Error checking field {field_name}: {e}")
         return False
 
-
-# Function to create a record in Salesforce
+# Save to Salesforce
 def save_to_salesforce(role, project_id, milestones, reflection, checklist, suggestions, quote, supervisor_name):
     try:
         sf = Salesforce(
@@ -226,125 +181,89 @@ def save_to_salesforce(role, project_id, milestones, reflection, checklist, sugg
             security_token=os.getenv('SF_SECURITY_TOKEN'),
             domain=os.getenv('SF_DOMAIN', 'login')
         )
-
-        # Escape single quotes in supervisor_name and project_id
         supervisor_name = supervisor_name.replace("'", "\\'")
         project_id = project_id.replace("'", "\\'")
-
-        # Step 1: Get the Salesforce record ID for the supervisor
         supervisor_result = sf.query(f"SELECT Id FROM Supervisor__c WHERE Name = '{supervisor_name}' LIMIT 1")
         if supervisor_result['totalSize'] == 0:
-            print(f"❌ No supervisor found with Name: {supervisor_name}")
+            print(f"❌ No supervisor found: {supervisor_name}")
             return
-
         supervisor_id = supervisor_result['records'][0]['Id']
-
-        # Step 2: Get the Salesforce record ID for the project
         project_result = sf.query(f"SELECT Id FROM Project__c WHERE Name = '{project_id}' LIMIT 1")
         if project_result['totalSize'] == 0:
-            print(f"❌ No project found with Name: {project_id}")
+            print(f"❌ No project found: {project_id}")
             return
-
         project_record_id = project_result['records'][0]['Id']
 
-        # Truncate text fields to avoid exceeding Salesforce field length limits (assuming 255 characters for simplicity)
         MAX_TEXT_LENGTH = 255
         checklist = checklist[:MAX_TEXT_LENGTH] if checklist else ""
         suggestions = suggestions[:MAX_TEXT_LENGTH] if suggestions else ""
         reflection = reflection[:MAX_TEXT_LENGTH] if reflection else ""
 
-        # Prepare data for Salesforce with explicit mapping
         data = {
-            'Supervisor_ID__c': supervisor_id,       # Lookup field expects the record ID of Supervisor__c
-            'Project_ID__c': project_record_id,      # Lookup field expects the record ID of Project__c
-            'Daily_Checklist__c': checklist,         # Maps to the generated Daily Checklist
-            'Suggested_Tips__c': suggestions,        # Maps to the generated Focus Suggestions
-            'Reflection_Log__c': reflection,         # Maps to the Reflection Log input
-            'KPI_Flag__c': KPI_FLAG_DEFAULT,         # Configurable via .env
-            'Engagement_Score__c': ENGAGEMENT_SCORE_DEFAULT  # Configurable via .env
+            'Supervisor_ID__c': supervisor_id,
+            'Project_ID__c': project_record_id,
+            'Daily_Checklist__c': checklist,
+            'Suggested_Tips__c': suggestions,
+            'Reflection_Log__c': reflection,
+            'KPI_Flag__c': KPI_FLAG_DEFAULT,
+            'Engagement_Score__c': ENGAGEMENT_SCORE_DEFAULT
         }
 
-        # Check if Milestones_KPIs__c field exists before mapping
         if field_exists(sf, 'Supervisor_AI_Coaching__c', 'Milestones_KPIs__c'):
-            # Truncate milestones as well if the field exists
             milestones = milestones[:MAX_TEXT_LENGTH] if milestones else ""
             data['Milestones_KPIs__c'] = milestones
-        else:
-            print("⚠️ Milestones_KPIs__c field does not exist in Supervisor_AI_Coaching__c. Skipping mapping.")
-
-        # Create record
-        response = sf.Supervisor_AI_Coaching__c.create(data)
-        print("✅ Record created successfully in Salesforce.")
-        print("Record ID:", response['id'])
 
+        sf.Supervisor_AI_Coaching__c.create(data)
+        print("✅ Record created in Salesforce.")
     except Exception as e:
         print(f"❌ Error saving to Salesforce: {e}")
-        print("Data being sent:", data)
-        if hasattr(e, 'content'):
-            print("Salesforce API response:", e.content)
 
-
-# Gradio Interface
+# Create Gradio interface
 def create_interface():
-    # Fetch roles from Salesforce
     roles = get_roles_from_salesforce()
-    print(f"Fetched Roles: {roles}")
-
     with gr.Blocks(theme="soft") as demo:
         gr.Markdown("# 🏗️ Construction Supervisor AI Coach")
-        gr.Markdown("Enter details to generate a daily checklist, focus suggestions, and a motivational quote.")
-
+        gr.Markdown("Enter details for daily checklist, suggestions, and quote.")
         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)
-
         milestones = gr.Textbox(label="Milestones (comma-separated KPIs)")
         reflection = gr.Textbox(label="Reflection Log", lines=4)
-
         with gr.Row():
             submit = gr.Button("Generate", variant="primary")
             clear = gr.Button("Clear")
             refresh_btn = gr.Button("🔄 Refresh Roles")
-
         checklist_output = gr.Textbox(label="✅ Daily Checklist")
         suggestions_output = gr.Textbox(label="💡 Focus Suggestions")
         quote_output = gr.Textbox(label="✨ Motivational Quote")
 
-        # Event: When role changes, update supervisor name dropdown
         role.change(
             fn=lambda r: gr.update(choices=get_supervisor_name_by_role(r)),
             inputs=[role],
             outputs=[supervisor_name]
         )
-
-        # Event: When supervisor name changes, update project ID
         supervisor_name.change(
             fn=get_projects_for_supervisor,
             inputs=[supervisor_name],
             outputs=[project_id]
         )
-
         submit.click(
             fn=generate_outputs,
             inputs=[role, supervisor_name, project_id, milestones, reflection],
             outputs=[checklist_output, suggestions_output, quote_output]
         )
-
         clear.click(
             fn=lambda: ("", "", "", "", ""),
             inputs=None,
             outputs=[role, supervisor_name, project_id, milestones, reflection]
         )
-
         refresh_btn.click(
             fn=lambda: gr.update(choices=get_roles_from_salesforce()),
             outputs=role
         )
-
     return demo
 
-
 if __name__ == "__main__":
     app = create_interface()
     app.launch()