Update app.py

app.py

@@ -9,6 +9,7 @@ from dotenv import load_dotenv
 from fpdf import FPDF
 import shutil
 import html
+from io import BytesIO
 
 # Load environment variables
 load_dotenv()
@@ -61,9 +62,9 @@ def clean_text_for_pdf(text):
 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)
-
+    
+    # In-memory PDF creation
+    pdf_buffer = BytesIO()
     pdf = FPDF()
     pdf.add_page()
     pdf.set_font("Arial", 'B', 14)
@@ -84,13 +85,13 @@ def save_report_as_pdf(role, supervisor_name, project_id, checklist, suggestions
     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)
+    
+    pdf.output(pdf_buffer)
+    pdf_buffer.seek(0)
 
-    temp_pdf_path = "/tmp/" + os.path.basename(file_path)
-    shutil.copy(file_path, temp_pdf_path)
-    return temp_pdf_path,filename
+    return pdf_buffer, filename
 
-def upload_pdf_to_salesforce_and_update_link(supervisor_name, project_id, pdf_path, pdf_name, checklist, suggestions):
+def upload_pdf_to_salesforce_and_update_link(supervisor_name, project_id, pdf_buffer, pdf_name, checklist, suggestions):
     try:
         sf = Salesforce(
             username=os.getenv('SF_USERNAME'),
@@ -99,9 +100,8 @@ def upload_pdf_to_salesforce_and_update_link(supervisor_name, project_id, pdf_pa
             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()
+        # Encode PDF to base64
+        encoded = base64.b64encode(pdf_buffer.read()).decode()
 
         # Create ContentVersion record to upload the PDF to Salesforce
         content = sf.ContentVersion.create({
@@ -110,40 +110,24 @@ def upload_pdf_to_salesforce_and_update_link(supervisor_name, project_id, pdf_pa
             'VersionData': encoded
         })
 
-        # Get the ContentDocumentId for the uploaded PDF
+        # Get ContentDocumentId
         content_id = content['id']
         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__c 
-            WHERE Name = '{supervisor_name}'
-            LIMIT 1
+        # Query for supervisor and project
+        supervisor_query = sf.query(f"""
+            SELECT Id FROM Supervisor__c WHERE Name = '{supervisor_name}' LIMIT 1
         """)
+        supervisor_id = supervisor_query['records'][0]['Id']
 
-        # Ensure that the Supervisor is found
-        if query['totalSize'] == 0:
-            print("⚠️ Supervisor not found!")
-            return ""
-
-        supervisor_id = query['records'][0]['Id']
-
-        # Query to get the Project ID in Salesforce
         project_query = sf.query(f"""
-            SELECT Id FROM Project__c WHERE Name = '{project_id}'
-            LIMIT 1
+            SELECT Id FROM Project__c WHERE Name = '{project_id}' LIMIT 1
         """)
-
-        # Ensure that the Project ID is found
-        if project_query['totalSize'] == 0:
-            print(f"⚠️ Project '{project_id}' not found in Salesforce!")
-            return ""
-
         project_id_sf = project_query['records'][0]['Id']
 
-        # Create a new Supervisor_AI_Coaching__c record
+        # Create record in Supervisor_AI_Coaching__c object
         sf.Supervisor_AI_Coaching__c.create({
-            'Project_ID__c': project_id_sf,  # Use the Salesforce Project ID
+            'Project_ID__c': project_id_sf,
             'Supervisor_ID__c': supervisor_id,
             'Daily_Checklist__c': checklist,
             'Suggested_Tips__c': suggestions,
@@ -151,7 +135,6 @@ def upload_pdf_to_salesforce_and_update_link(supervisor_name, project_id, pdf_pa
         })
 
         return download_url
-
     except Exception as e:
         print(f"⚠️ Upload error: {e}")
         return ""
@@ -170,58 +153,21 @@ def get_roles_from_salesforce():
         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 ""
-
-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>'
-
 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)
+    # Shortened input for faster processing
+    prompt = PROMPT_TEMPLATE.format(role=role, project_id=project_id, milestones=milestones[:256], reflection=reflection[:256])
     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,
+                max_new_tokens=100,
                 no_repeat_ngram_size=2,
-                do_sample=True,
+                do_sample=False,
                 top_p=0.9,
                 temperature=0.7,
                 pad_token_id=tokenizer.pad_token_id
@@ -244,13 +190,13 @@ def generate_outputs(role, supervisor_name, project_id, milestones, reflection):
     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)
+    pdf_buffer, 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_buffer, pdf_name, checklist, suggestions)
 
     if pdf_url:
         suggestions += f"\n\n🔗 [Download PDF Report]({pdf_url})"
 
-    return checklist, suggestions, pdf_path, pdf_name
+    return checklist, suggestions, pdf_buffer, pdf_name
 
 def create_interface():
     roles = get_roles_from_salesforce()
@@ -281,8 +227,8 @@ def create_interface():
         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, pdf_name = generate_outputs(role, supervisor_name, project_id, milestones, reflection)
-            return checklist, suggestions, pdf_path, pdf_name
+            checklist, suggestions, pdf_buffer, pdf_name = generate_outputs(role, supervisor_name, project_id, milestones, reflection)
+            return checklist, suggestions, pdf_buffer, pdf_name
 
         generate.click(fn=handle_generate,
                        inputs=[role, supervisor_name, project_id, milestones, reflection],