Update app.py

app.py

@@ -20,7 +20,6 @@ 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)
@@ -30,7 +29,6 @@ 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
 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. 
@@ -99,31 +97,23 @@ def upload_pdf_to_salesforce_and_update_link(supervisor_name, project_id, pdf_pa
             security_token=os.getenv('SF_SECURITY_TOKEN'),
             domain=os.getenv('SF_DOMAIN', 'login')
         )
-
         with open(pdf_path, "rb") as f:
             encoded = base64.b64encode(f.read()).decode()
-
         content = sf.ContentVersion.create({
             'Title': pdf_name,
             'PathOnClient': pdf_name,
             'VersionData': encoded
         })
-
         content_id = content['id']
         download_url = f"https://{sf.sf_instance}/sfc/servlet.shepherd/version/download/{content_id}"
-
         query = sf.query(f"SELECT Id FROM Supervisor__c WHERE Name = '{supervisor_name}' LIMIT 1")
         if query['totalSize'] == 0:
             return ""
-
         supervisor_id = query['records'][0]['Id']
-
         project_query = sf.query(f"SELECT Id FROM Project__c WHERE Name = '{project_id}' LIMIT 1")
         if project_query['totalSize'] == 0:
             return ""
-
         project_id_sf = project_query['records'][0]['Id']
-
         sf.Supervisor_AI_Coaching__c.create({
             'Project_ID__c': project_id_sf,
             'Supervisor_ID__c': supervisor_id,
@@ -131,10 +121,8 @@ def upload_pdf_to_salesforce_and_update_link(supervisor_name, project_id, pdf_pa
             'Suggested_Tips__c': suggestions,
             'Download_Link__c': download_url
         })
-
         return download_url
-
-    except Exception as e:
+    except:
         return ""
 
 def get_roles_from_salesforce():
@@ -147,7 +135,7 @@ def get_roles_from_salesforce():
         )
         result = sf.query("SELECT Role__c FROM Supervisor__c WHERE Role__c != NULL")
         return list(set(record['Role__c'] for record in result['records']))
-    except Exception as e:
+    except:
         return []
 
 def get_supervisor_name_by_role(role):
@@ -160,7 +148,7 @@ def get_supervisor_name_by_role(role):
         )
         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:
+    except:
         return []
 
 def get_projects_for_supervisor(supervisor_name):
@@ -177,7 +165,7 @@ def get_projects_for_supervisor(supervisor_name):
         supervisor_id = 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:
+    except:
         return ""
 
 def get_dashboard_data_from_salesforce(supervisor_name, project_id):
@@ -208,17 +196,6 @@ def get_dashboard_data_from_salesforce(supervisor_name, project_id):
     except Exception as e:
         return f"Error loading dashboard: {e}", "", "", ""
 
-# Function to compare start and end date
-def compare_dates(start_date, end_date):
-    current_date = datetime.datetime.now()
-    if current_date >= end_date:
-        return "Completed"
-    elif current_date < start_date:
-        return "Not Started"
-    else:
-        return "In Progress"
-
-# Render bar chart (updated)
 def render_bar_chart(completed, total):
     fig, ax = plt.subplots(figsize=(4, 3))
     ax.bar(["Completed", "Remaining"], [completed, max(0, total - completed)], color=["green", "gray"])
@@ -231,45 +208,47 @@ def render_bar_chart(completed, total):
     encoded = base64.b64encode(buf.read()).decode("utf-8")
     return f"<img src='data:image/png;base64,{encoded}'/>"
 
-# Updated show_dashboard_html function
-def show_dashboard_html(sup_name, proj_id):
-    checklist, suggestions, link, date = get_dashboard_data_from_salesforce(sup_name, proj_id)
-    
-    # Parse date from Salesforce (assuming 'CreatedDate' is the project end date)
-    end_date = datetime.datetime.strptime(date, "%Y-%m-%dT%H:%M:%S.%f+0000")  # Adjust format if needed
-    start_date = datetime.datetime.now()  # Example: assume project starts today
-    
-    # Compare the start and end dates
-    status = compare_dates(start_date, end_date)
-
-    # Progress calculation (dummy logic here)
-    completed_tasks = checklist.count("- ")  # Count bullet points
-    total_tasks = 5  # You can replace this with actual task count from the checklist
-    chart_html = render_bar_chart(completed_tasks, total_tasks)
-    
-    # Create the final HTML content
-    link_html = f'<a href="{link}" target="_blank">Download PDF</a>' if link else "No report available"
-    
-    html_card = f"""
-    <div style='border:1px solid #ccc; border-radius:10px; padding:20px; background-color:#f9f9f9'>
-        <h3>📋 <u>Dashboard Summary</u></h3>
-        <p><b>Supervisor:</b> {sup_name}</p>
-        <p><b>Project ID:</b> {proj_id}</p>
-        <p><b>🗓️ Status:</b> {status}</p>
-        <hr>
-        <h4>✅ <u>Daily Checklist</u></h4>
-        <pre style='background:#eef; padding:10px; border-radius:5px;'>{checklist}</pre>
-        <h4>💡 <u>Focus Suggestions</u></h4>
-        <pre style='background:#ffe; padding:10px; border-radius:5px;'>{suggestions}</pre>
-        <p>{link_html}</p>
-        <hr>
-        <h4>📊 <u>Progress Chart</u></h4>
-        {chart_html}
-    </div>
-    """
-    return html_card
-
-# Main Gradio Interface
+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, ""
+    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.7,
+                pad_token_id=tokenizer.pad_token_id
+            )
+        result = tokenizer.decode(outputs[0], skip_special_tokens=True)
+    except:
+        return "", "", None, ""
+
+    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 ""
+
+    checklist = extract_between(result, "Checklist:\n", "Suggestions:")
+    suggestions = extract_between(result, "Suggestions:\n", None)
+
+    if not checklist.strip():
+        checklist = "- Perform daily safety inspection"
+    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, checklist, suggestions)
+
+    if pdf_url:
+        suggestions += f"\n\n🔗 [Download PDF Report]({pdf_url})"
+
+    return checklist, suggestions, pdf_path, pdf_name
+
 def create_interface():
     roles = get_roles_from_salesforce()
     with gr.Blocks(theme="soft", css=".footer { display: none; }") as demo: