Update app.py

app.py

@@ -1,45 +1,150 @@
 import gradio as gr
 import torch
 from transformers import AutoModelForCausalLM, AutoTokenizer
-import datetime
+from simple_salesforce import Salesforce
+import os
+from dotenv import load_dotenv
 
-# Initialize model and tokenizer (preloading them for quicker response)
+# Load environment variables
+load_dotenv()
+
+# Check if required environment variables are set
+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
+
+# Initialize model and tokenizer
 model_name = "distilgpt2"
 tokenizer = AutoTokenizer.from_pretrained(model_name)
 model = AutoModelForCausalLM.from_pretrained(model_name)
 
-# Set pad_token_id to eos_token_id to avoid warnings
-tokenizer.pad_token = tokenizer.eos_token
-model.config.pad_token_id = tokenizer.eos_token_id
-
-# Define a more contextual 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. Format your response with clear labels as follows:
+# Avoid warnings by setting pad token
+if tokenizer.pad_token is None:
+    tokenizer.pad_token = tokenizer.eos_token if tokenizer.eos_token else "[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 = """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:
 Checklist:
 - {milestones_list}
-
 Suggestions:
 - {suggestions_list}
-
 Quote:
 - Your motivational quote here
-
-Inputs:
-Role: {role}
-Project: {project_id}
-Milestones: {milestones}
-Reflection: {reflection}
 """
 
-# Function to generate outputs based on inputs
-def generate_outputs(role, project_id, milestones, reflection):
-    # Validate inputs to ensure no missing fields
-    if not all([role, project_id, milestones, reflection]):
+# Function to get all roles from Salesforce
+def get_roles_from_salesforce():
+    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')
+        )
+
+        # 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
+
+
+# Function to get supervisor's Name (Auto Number) by role
+def get_supervisor_name_by_role(role):
+    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')
+        )
+
+        # 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
+
+    except Exception as e:
+        print(f"⚠️ Error fetching supervisor names: {e}")
+        return []
+
+
+# Function to get project IDs and names assigned to selected 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
+
+    except Exception as e:
+        print(f"⚠️ Error fetching project for supervisor: {e}")
+        return ""
+
+
+# Function to generate AI-based coaching output
+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.", "", ""
-    
-    # Create prompt from template
+
+    # 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."
@@ -48,7 +153,7 @@ def generate_outputs(role, project_id, milestones, reflection):
     elif "equipment" in reflection.lower():
         suggestions_list = "- Inspect all equipment to ensure no malfunctions.\n- Schedule maintenance if necessary."
 
-    # Create final prompt
+    # Fill in the prompt template
     prompt = PROMPT_TEMPLATE.format(
         role=role,
         project_id=project_id,
@@ -58,81 +163,188 @@ def generate_outputs(role, project_id, milestones, reflection):
         suggestions_list=suggestions_list
     )
 
-    # Tokenize inputs for model processing
+    # Tokenize input
     inputs = tokenizer(prompt, return_tensors="pt", max_length=512, truncation=True, padding=True)
-    
-    # Generate response from the model
-    with torch.no_grad():
-        outputs = model.generate(
-            inputs['input_ids'],
-            max_length=512,
-            num_return_sequences=1,
-            no_repeat_ngram_size=2,
-            do_sample=True,
-            top_p=0.9,
-            temperature=0.8,
-            pad_token_id=tokenizer.eos_token_id
-        )
-    
-    # Decode the response
-    generated_text = tokenizer.decode(outputs[0], skip_special_tokens=True)
-
-    # Parse the output and ensure it is structured
-    checklist = "No checklist generated."
-    suggestions = "No suggestions generated."
-    quote = "No quote generated."
-
-    if "Checklist:" in generated_text:
-        checklist_start = generated_text.find("Checklist:") + len("Checklist:")
-        suggestions_start = generated_text.find("Suggestions:")
-        checklist = generated_text[checklist_start:suggestions_start].strip()
-    
-    if "Suggestions:" in generated_text:
-        suggestions_start = generated_text.find("Suggestions:") + len("Suggestions:")
-        quote_start = generated_text.find("Quote:")
-        suggestions = generated_text[suggestions_start:quote_start].strip()
-    
-    if "Quote:" in generated_text:
-        quote_start = generated_text.find("Quote:") + len("Quote:")
-        quote = generated_text[quote_start:].strip()
-
-    # Return structured outputs
+
+    # Generate response
+    try:
+        with torch.no_grad():
+            outputs = model.generate(
+                inputs['input_ids'],
+                max_length=1024,  # Increased to allow for longer outputs
+                num_return_sequences=1,
+                no_repeat_ngram_size=2,
+                do_sample=True,
+                top_p=0.9,
+                temperature=0.8,
+                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}")
+        return "Error: Failed to generate outputs.", "", ""
+
+    # Parse sections
+    def extract_section(text, start_marker, end_marker):
+        start = text.find(start_marker)
+        if start == -1:
+            return "Not found"
+        start += len(start_marker)
+        end = text.find(end_marker, start) if end_marker else len(text)
+        return text[start:end].strip()
+
+    checklist = extract_section(generated_text, "Checklist:\n", "Suggestions:")
+    suggestions = extract_section(generated_text, "Suggestions:\n", "Quote:")
+    quote = extract_section(generated_text, "Quote:\n", None)
+
+    # Save to Salesforce
+    save_to_salesforce(role, project_id, milestones, reflection, checklist, suggestions, quote, supervisor_name)
+
     return checklist, suggestions, quote
 
-# Gradio interface for fast user interaction
+
+# Function to check if a field exists in a Salesforce object
+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
+    except Exception as e:
+        print(f"⚠️ Error checking if field {field_name} exists in {object_name}: {e}")
+        return False
+
+
+# Function to create a record in Salesforce
+def save_to_salesforce(role, project_id, milestones, reflection, checklist, suggestions, quote, supervisor_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')
+        )
+
+        # 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}")
+            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}")
+            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
+        }
+
+        # 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'])
+
+    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
 def create_interface():
-    with gr.Blocks() as demo:
-        gr.Markdown("# Construction Supervisor AI Coach")
+    # 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.")
-        
+
         with gr.Row():
-            role = gr.Dropdown(choices=["Supervisor", "Foreman", "Project Manager"], label="Role")
-            project_id = gr.Textbox(label="Project ID")
-        
+            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=5)
-        
+        reflection = gr.Textbox(label="Reflection Log", lines=4)
+
         with gr.Row():
-            submit = gr.Button("Generate")
+            submit = gr.Button("Generate", variant="primary")
             clear = gr.Button("Clear")
-        
-        checklist_output = gr.Textbox(label="Daily Checklist")
-        suggestions_output = gr.Textbox(label="Focus Suggestions")
-        quote_output = gr.Textbox(label="Motivational Quote")
-        
+            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, project_id, milestones, reflection],
+            inputs=[role, supervisor_name, project_id, milestones, reflection],
             outputs=[checklist_output, suggestions_output, quote_output]
         )
+
         clear.click(
-            fn=lambda: ("", "", "", ""),
+            fn=lambda: ("", "", "", "", ""),
             inputs=None,
-            outputs=[role, project_id, milestones, reflection]
+            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__":
-    demo = create_interface()
-    demo.launch()
+    app = create_interface()
+    app.launch()