Update app.py

app.py

@@ -8,44 +8,50 @@ from dotenv import load_dotenv
 # Load environment variables
 load_dotenv()
 
-# Check for required environment variables
+# Required env vars check
 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}")
 
-# Configurable defaults
+# Defaults
 KPI_FLAG_DEFAULT = os.getenv('KPI_FLAG', 'True') == 'True'
 ENGAGEMENT_SCORE_DEFAULT = float(os.getenv('ENGAGEMENT_SCORE', '85.0'))
 
-# Initialize model and tokenizer with optimized settings
+# Load model and tokenizer
 model_name = "distilgpt2"
-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
+tokenizer = AutoTokenizer.from_pretrained(model_name, use_fast=True)
+model = AutoModelForCausalLM.from_pretrained(model_name, low_cpu_mem_usage=True)
 
-# Set pad token
 if tokenizer.pad_token is None:
-    tokenizer.pad_token = tokenizer.eos_token or "[PAD]"
+    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
 
-# 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.
+# Better Prompt
+PROMPT_TEMPLATE = """You are an AI coach for construction supervisors. Given the role, project, milestones, and a reflection log, generate:
+
+1. A Daily Checklist based on the milestones and role.
+2. Focus Suggestions based on concerns or keywords in the reflection log.
+3. A Motivational Quote for the day.
+
 Inputs:
 Role: {role}
-Project: {project_id}
+Project ID: {project_id}
 Milestones: {milestones}
-Reflection: {reflection}
-Format:
+Reflection Log: {reflection}
+
+Output Format:
 Checklist:
-- {milestones_list}
+- 
+
 Suggestions:
-- {suggestions_list}
+- 
+
 Quote:
-- Your motivational quote here
 """
 
-# Fetch roles from Salesforce
+# Salesforce Functions
 def get_roles_from_salesforce():
     try:
         sf = Salesforce(
@@ -55,13 +61,11 @@ def get_roles_from_salesforce():
             domain=os.getenv('SF_DOMAIN', 'login')
         )
         result = sf.query("SELECT Role__c FROM Supervisor__c WHERE Role__c != NULL")
-        roles = list(set(record['Role__c'] for record in result.get('records', [])))
-        return roles
+        return list(set(record['Role__c'] for record in result['records']))
     except Exception as e:
         print(f"⚠️ Error fetching roles: {e}")
         return ["Site Manager", "Safety Officer", "Project Lead"]
 
-# Fetch supervisor names by role
 def get_supervisor_name_by_role(role):
     try:
         sf = Salesforce(
@@ -72,14 +76,11 @@ def get_supervisor_name_by_role(role):
         )
         role = role.replace("'", "\\'")
         result = sf.query(f"SELECT Name FROM Supervisor__c WHERE Role__c = '{role}'")
-        if result['totalSize'] == 0:
-            return []
         return [record['Name'] for record in result['records']]
     except Exception as e:
         print(f"⚠️ Error fetching supervisor names: {e}")
         return []
 
-# Fetch project for supervisor
 def get_projects_for_supervisor(supervisor_name):
     try:
         sf = Salesforce(
@@ -94,76 +95,11 @@ def get_projects_for_supervisor(supervisor_name):
             return ""
         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")
-        if project_result['totalSize'] == 0:
-            return ""
-        return project_result['records'][0]['Name']
+        return project_result['records'][0]['Name'] if project_result['totalSize'] > 0 else ""
     except Exception as e:
         print(f"⚠️ Error fetching project: {e}")
         return ""
 
-# 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.", "", ""
-
-    milestones_list = "\n- ".join([m.strip() for m in milestones.split(",")])
-    suggestions_list = ""
-    if "delays" in reflection.lower():
-        suggestions_list = "- Adjust timelines for delays.\n- Inform stakeholders."
-    elif "weather" in reflection.lower():
-        suggestions_list = "- Ensure rain gear availability.\n- Monitor weather updates."
-    elif "equipment" in reflection.lower():
-        suggestions_list = "- Inspect equipment.\n- Schedule maintenance."
-
-    prompt = PROMPT_TEMPLATE.format(
-        role=role,
-        project_id=project_id,
-        milestones=milestones,
-        reflection=reflection,
-        milestones_list=milestones_list,
-        suggestions_list=suggestions_list
-    )
-
-    # Tokenize with optimized settings
-    inputs = tokenizer(prompt, return_tensors="pt", max_length=256, truncation=True, padding=True)
-
-    # Faster generation with adjusted parameters
-    try:
-        with torch.no_grad():
-            outputs = model.generate(
-                inputs['input_ids'],
-                max_length=512,  # Reduced for speed
-                num_return_sequences=1,
-                no_repeat_ngram_size=2,
-                do_sample=True,
-                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 generation: {e}")
-        return "Error: Failed to generate outputs.", "", ""
-
-    # Extract 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
-
-# Check if Salesforce field exists
 def field_exists(sf, object_name, field_name):
     try:
         obj_desc = getattr(sf, object_name).describe()
@@ -172,7 +108,6 @@ def field_exists(sf, object_name, field_name):
         print(f"⚠️ Error checking field {field_name}: {e}")
         return False
 
-# Save to Salesforce
 def save_to_salesforce(role, project_id, milestones, reflection, checklist, suggestions, quote, supervisor_name):
     try:
         sf = Salesforce(
@@ -181,89 +116,113 @@ def save_to_salesforce(role, project_id, milestones, reflection, checklist, sugg
             security_token=os.getenv('SF_SECURITY_TOKEN'),
             domain=os.getenv('SF_DOMAIN', 'login')
         )
-        supervisor_name = supervisor_name.replace("'", "\\'")
-        project_id = project_id.replace("'", "\\'")
         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: {supervisor_name}")
+            print("❌ No supervisor found.")
             return
         supervisor_id = supervisor_result['records'][0]['Id']
         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: {project_id}")
+            print("❌ No project found.")
             return
         project_record_id = project_result['records'][0]['Id']
 
-        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 ""
-
         data = {
             'Supervisor_ID__c': supervisor_id,
             'Project_ID__c': project_record_id,
-            'Daily_Checklist__c': checklist,
-            'Suggested_Tips__c': suggestions,
-            'Reflection_Log__c': reflection,
+            'Daily_Checklist__c': checklist[:255],
+            'Suggested_Tips__c': suggestions[:255],
+            'Reflection_Log__c': reflection[:255],
             'KPI_Flag__c': KPI_FLAG_DEFAULT,
             'Engagement_Score__c': ENGAGEMENT_SCORE_DEFAULT
         }
 
         if field_exists(sf, 'Supervisor_AI_Coaching__c', 'Milestones_KPIs__c'):
-            milestones = milestones[:MAX_TEXT_LENGTH] if milestones else ""
-            data['Milestones_KPIs__c'] = milestones
+            data['Milestones_KPIs__c'] = milestones[:255]
 
         sf.Supervisor_AI_Coaching__c.create(data)
         print("✅ Record created in Salesforce.")
     except Exception as e:
         print(f"❌ Error saving to Salesforce: {e}")
 
-# Create Gradio interface
+# 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.", "", ""
+
+    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,
+                no_repeat_ngram_size=2,
+                do_sample=True,
+                top_p=0.9,
+                temperature=0.8,
+                pad_token_id=tokenizer.pad_token_id
+            )
+        result = tokenizer.decode(outputs[0], skip_special_tokens=True)
+    except Exception as e:
+        print(f"⚠️ Generation error: {e}")
+        return "Error during text generation", "", ""
+
+    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 "Not found"
+
+    checklist = extract_between(result, "Checklist:\n", "Suggestions:")
+    suggestions = extract_between(result, "Suggestions:\n", "Quote:")
+    quote = extract_between(result, "Quote:\n", None)
+
+    save_to_salesforce(role, project_id, milestones, reflection, checklist, suggestions, quote, supervisor_name)
+    return checklist, suggestions, quote
+
+# Interface
 def create_interface():
     roles = get_roles_from_salesforce()
     with gr.Blocks(theme="soft") as demo:
-        gr.Markdown("# 🏗️ Construction Supervisor AI Coach")
-        gr.Markdown("Enter details for daily checklist, suggestions, and quote.")
+        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)
-        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():
-            submit = gr.Button("Generate", variant="primary")
+            generate = gr.Button("Generate")
             clear = gr.Button("Clear")
-            refresh_btn = gr.Button("🔄 Refresh Roles")
+            refresh = 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")
 
-        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]
-        )
-        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
-        )
+        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)
+
+        generate.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.click(fn=lambda: gr.update(choices=get_roles_from_salesforce()), outputs=role)
+
     return demo
 
 if __name__ == "__main__":
     app = create_interface()
-    app.launch()
+    app.launch()