Spaces:
Sleeping
Sleeping
| import gradio as gr | |
| import torch | |
| from transformers import AutoModelForCausalLM, AutoTokenizer | |
| from simple_salesforce import Salesforce | |
| import os | |
| import base64 | |
| import datetime | |
| from dotenv import load_dotenv | |
| from fpdf import FPDF | |
| import shutil | |
| import html | |
| # Load environment variables | |
| load_dotenv() | |
| # 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}") | |
| # 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 | |
| # 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. | |
| Split the checklist into day-by-day tasks for a specified time period (e.g., one week). | |
| 2. Focus Suggestions based on concerns or keywords in the reflection log. Provide at least 2 suggestions. | |
| Inputs: | |
| Role: {role} | |
| Project ID: {project_id} | |
| Milestones: {milestones} | |
| Reflection Log: {reflection} | |
| Output Format: | |
| Checklist (Day-by-Day): | |
| - Day 1: | |
| - Task 1 | |
| - Task 2 | |
| - Day 2: | |
| - Task 1 | |
| - Task 2 | |
| ... | |
| Suggestions: | |
| - | |
| """ | |
| # Function to clean the text for PDF | |
| def clean_text_for_pdf(text): | |
| return html.unescape(text).encode('latin-1', 'replace').decode('latin-1') | |
| # Save report as PDF | |
| def save_report_as_pdf(role, supervisor_name, project_id, checklist, suggestions): | |
| now = datetime.datetime.now().strftime("%Y%m%d%H%M%S") | |
| filename = f"report_{supervisor_name}_{project_id}_{now}.pdf" | |
| file_path = f"./reports/{filename}" | |
| os.makedirs("reports", exist_ok=True) | |
| pdf = FPDF() | |
| pdf.add_page() | |
| pdf.set_font("Arial", 'B', 14) | |
| pdf.cell(200, 10, txt="Supervisor Daily Report", ln=True, align="C") | |
| pdf.set_font("Arial", size=12) | |
| pdf.cell(200, 10, txt=clean_text_for_pdf(f"Role: {role}"), ln=True) | |
| pdf.cell(200, 10, txt=clean_text_for_pdf(f"Supervisor: {supervisor_name}"), ln=True) | |
| pdf.cell(200, 10, txt=clean_text_for_pdf(f"Project ID: {project_id}"), ln=True) | |
| pdf.ln(5) | |
| pdf.set_font("Arial", 'B', 12) | |
| pdf.cell(200, 10, txt="Daily Checklist", ln=True) | |
| pdf.set_font("Arial", size=12) | |
| for line in checklist.split("\n"): | |
| pdf.multi_cell(0, 10, clean_text_for_pdf(line)) | |
| pdf.ln(5) | |
| pdf.set_font("Arial", 'B', 12) | |
| pdf.cell(200, 10, txt="Focus Suggestions", ln=True) | |
| pdf.set_font("Arial", size=12) | |
| for line in suggestions.split("\n"): | |
| pdf.multi_cell(0, 10, clean_text_for_pdf(line)) | |
| pdf.output(file_path) | |
| # Copy the file to a temporary directory for Gradio to access | |
| temp_pdf_path = "/tmp/" + os.path.basename(file_path) # Use /tmp/ or current directory for Gradio | |
| shutil.copy(file_path, temp_pdf_path) | |
| return temp_pdf_path, filename | |
| # Upload to Salesforce and update record with the generated URL | |
| def upload_pdf_to_salesforce_and_update_link(supervisor_name, project_id, pdf_path, pdf_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') | |
| ) | |
| # Read and encode the file as base64 | |
| with open(pdf_path, "rb") as f: | |
| encoded = base64.b64encode(f.read()).decode() | |
| # Create ContentVersion record to upload the PDF to Salesforce | |
| content = sf.ContentVersion.create({ | |
| 'Title': pdf_name, | |
| 'PathOnClient': pdf_name, | |
| 'VersionData': encoded | |
| }) | |
| # Get the ContentDocumentId for the uploaded PDF | |
| content_id = content['id'] | |
| # Generate the download URL for the uploaded PDF | |
| download_url = f"https://{sf.sf_instance}/sfc/servlet.shepherd/version/download/{content_id}" | |
| # Query Salesforce to find the specific Supervisor_AI_Coaching__c record | |
| query = sf.query(f""" | |
| SELECT Id FROM Supervisor_AI_Coaching__c | |
| WHERE Project_ID__c = '{project_id}' | |
| AND Name = '{supervisor_name}' | |
| LIMIT 1 | |
| """) | |
| if query['totalSize'] > 0: | |
| # Get the ID of the Supervisor_AI_Coaching__c record | |
| coaching_id = query['records'][0]['Id'] | |
| # Update the Supervisor_AI_Coaching__c record with the download URL | |
| sf.Supervisor_AI_Coaching__c.update(coaching_id, { | |
| 'Download_Link__c': download_url # Update the Download_Link__c field with the URL | |
| }) | |
| return download_url | |
| except Exception as e: | |
| print(f"⚠️ Upload error: {e}") | |
| return "" | |
| # Salesforce helpers | |
| 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') | |
| ) | |
| 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: | |
| print(f"⚠️ Error fetching roles: {e}") | |
| return [] | |
| 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') | |
| ) | |
| 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: | |
| print(f"⚠️ Error fetching names: {e}") | |
| return [] | |
| def get_projects_for_supervisor(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') | |
| ) | |
| result = sf.query(f"SELECT Id FROM Supervisor__c WHERE Name = '{supervisor_name}' LIMIT 1") | |
| if result['totalSize'] == 0: | |
| return "" | |
| 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: | |
| print(f"⚠️ Error fetching project: {e}") | |
| return "" | |
| # Generate Salesforce dashboard URL | |
| def generate_salesforce_dashboard_url(supervisor_name, project_id): | |
| return f"https://aicoachforsitesupervisors-dev-ed--c.develop.vf.force.com/apex/DashboardPage?supervisorName={supervisor_name}&projectId={project_id}" | |
| def open_dashboard(role, supervisor_name, project_id): | |
| url = generate_salesforce_dashboard_url(supervisor_name, project_id) | |
| return f'<a href="{url}" target="_blank">Open Salesforce Dashboard</a>' | |
| # Generate AI output | |
| 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 Exception as e: | |
| print(f"⚠️ Generation error: {e}") | |
| 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) | |
| if pdf_url: | |
| suggestions += f"\n\n🔗 [Download PDF Report]({pdf_url})" | |
| return checklist, suggestions, pdf_path, pdf_path | |
| # Gradio Interface | |
| def create_interface(): | |
| roles = get_roles_from_salesforce() | |
| with gr.Blocks(theme="soft") 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) | |
| milestones = gr.Textbox(label="Milestones (comma-separated KPIs)") | |
| reflection = gr.Textbox(label="Reflection Log", lines=4) | |
| with gr.Row(): | |
| generate = gr.Button("Generate") | |
| clear = gr.Button("Clear") | |
| refresh = gr.Button("🔄 Refresh Roles") | |
| dashboard_btn = gr.Button("Dashboard") | |
| checklist_output = gr.Textbox(label="✅ Daily Checklist") | |
| suggestions_output = gr.Textbox(label="💡 Focus Suggestions") | |
| download_button = gr.File(label="⬇ Download Report") | |
| pdf_link = gr.HTML() | |
| dashboard_link = gr.HTML() | |
| 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) | |
| def handle_generate(role, supervisor_name, project_id, milestones, reflection): | |
| checklist, suggestions, pdf_path, _ = generate_outputs(role, supervisor_name, project_id, milestones, reflection) | |
| return checklist, suggestions, pdf_path, pdf_path | |
| generate.click(fn=handle_generate, | |
| inputs=[role, supervisor_name, project_id, milestones, reflection], | |
| outputs=[checklist_output, suggestions_output, download_button, pdf_link]) | |
| 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) | |
| dashboard_btn.click(fn=open_dashboard, | |
| inputs=[role, supervisor_name, project_id], | |
| outputs=dashboard_link) | |
| return demo | |
| if __name__ == "__main__": | |
| app = create_interface() | |
| app.launch() | |