Update app.py

app.py

@@ -1,3 +1,4 @@
+```python
 import gradio as gr
 import logging
 from transformers import AutoModelForSeq2SeqLM, AutoTokenizer
@@ -5,18 +6,35 @@ import os
 from fastapi import FastAPI
 from pydantic import BaseModel
 from simple_salesforce import Salesforce
-import json
+from dotenv import load_dotenv
+
+# Load environment variables from .env
+load_dotenv()
 
 # Configure logging
 logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s')
 
-# Load Hugging Face token from environment variable
+# Salesforce connection
+def get_salesforce_connection():
+    try:
+        sf = Salesforce(
+            username=os.getenv("SF_USERNAME"),
+            password=os.getenv("SF_PASSWORD"),
+            security_token=os.getenv("SF_SECURITY_TOKEN")
+        )
+        logging.info("Salesforce connection established successfully")
+        return sf
+    except Exception as e:
+        logging.error(f"Failed to connect to Salesforce: {e}")
+        raise e
+
+# Load Hugging Face token
 HF_TOKEN = os.getenv("HF_TOKEN")
 if not HF_TOKEN:
     logging.warning("HF_TOKEN not set. Attempting to load public models only.")
 
 # Model configuration
-MODEL_PATH = "your_username/fine_tuned_bart_construction"  # Replace with your actual model path
+MODEL_PATH = "your_actual_username/fine_tuned_bart_construction"  # Replace with your model path
 FALLBACK_MODEL = "facebook/bart-large"  # Fallback public model
 
 try:
@@ -36,35 +54,13 @@ except Exception as e:
         logging.error(f"Failed to load fallback model {FALLBACK_MODEL}: {fallback_e}")
         raise fallback_e
 
-# Salesforce connection
-def get_salesforce_connection():
-    try:
-        sf_username = os.getenv("onteddugeetha104@gmail.com")
-        sf_password = os.getenv("9032590825g@")
-        sf_security_token = os.getenv("tTPLQduw8wDpxdKOMJ3d9dM3o")
-        
-        if not all([sf_username, sf_password, sf_security_token]):
-            logging.error("Salesforce credentials missing in environment variables")
-            raise ValueError("Salesforce credentials not configured")
-        
-        sf = Salesforce(
-            username=sf_username,
-            password=sf_password,
-            security_token=sf_security_token
-        )
-        logging.info("Connected to Salesforce successfully")
-        return sf
-    except Exception as e:
-        logging.error(f"Failed to connect to Salesforce: {e}")
-        raise e
-
 # Define input model for FastAPI
 class ChecklistInput(BaseModel):
     role: str = "Supervisor"
     project_id: str = "Unknown"
     project_name: str = "Unknown Project"
     milestones: str = "No milestones provided"
-    coaching_record_id: str = None  # Optional Salesforce record ID
+    record_id: str = None  # Salesforce record ID for updates
 
 # Initialize FastAPI
 app = FastAPI()
@@ -72,22 +68,8 @@ app = FastAPI()
 @app.post("/generate")
 async def generate_checklist(data: ChecklistInput):
     try:
-        # Connect to Salesforce
-        sf = get_salesforce_connection()
-        
-        # Fetch milestones from Project__c if project_id is provided
-        milestones = data.milestones
-        if data.project_id != "Unknown":
-            try:
-                project = sf.query(f"SELECT Milestones__c FROM Project__c WHERE Id = '{data.project_id}'")
-                if project["totalSize"] > 0:
-                    milestones = project["records"][0]["Milestones__c"] or milestones
-                    logging.info(f"Fetched milestones for project {data.project_id}: {milestones}")
-            except Exception as e:
-                logging.warning(f"Failed to fetch project milestones: {e}")
-
         # Prepare input for the model
-        inputs = f"Role: {data.role} Project: {data.project_id} ({data.project_name}) Milestones: {milestones}"
+        inputs = f"Role: {data.role} Project: {data.project_id} ({data.project_name}) Milestones: {data.milestones}"
         logging.info(f"Generating checklist for inputs: {inputs[:100]}...")
 
         # Tokenize and generate
@@ -98,17 +80,18 @@ async def generate_checklist(data: ChecklistInput):
         # Static tips (replace with model-generated tips after fine-tuning)
         tips = "1. Prioritize safety checks\n2. Review milestones\n3. Log progress"
 
-        # Update Salesforce record if coaching_record_id is provided
-        if data.coaching_record_id:
+        # Update Salesforce record if record_id is provided
+        if data.record_id:
             try:
-                sf.Supervisor_AI_Coaching__c.update(data.coaching_record_id, {
-                    "Daily_Checklist__c": checklist,
-                    "Suggested_Tips__c": tips,
-                    "Engagement_Score__c": 10  # Increment or set as needed
+                sf = get_salesforce_connection()
+                sf.Supervisor_AI_Coaching__c.update(data.record_id, {
+                    'Daily_Checklist__c': checklist,
+                    'Suggested_Tips__c': tips,
+                    'Engagement_Score__c': 10 if not sf.Supervisor_AI_Coaching__c.get(data.record_id).get('Engagement_Score__c') else sf.Supervisor_AI_Coaching__c.get(data.record_id).get('Engagement_Score__c') + 10
                 })
-                logging.info(f"Updated Salesforce record {data.coaching_record_id}")
-            except Exception as e:
-                logging.error(f"Failed to update Salesforce record: {e}")
+                logging.info(f"Updated Salesforce record {data.record_id}")
+            except Exception as sf_e:
+                logging.error(f"Failed to update Salesforce: {sf_e}")
 
         logging.info("Checklist and tips generated successfully")
         return {
@@ -120,49 +103,44 @@ async def generate_checklist(data: ChecklistInput):
         return {"error": str(e)}
 
 # Gradio interface function
-def gradio_generate_checklist(role, project_id, project_name, milestones, coaching_record_id=""):
+def gradio_generate_checklist(role, project_id, project_name, milestones, record_id=""):
     try:
-        sf = get_salesforce_connection()
-        
-        # Fetch milestones if project_id is valid
-        if project_id:
-            project = sf.query(f"SELECT Milestones__c FROM Project__c WHERE Id = '{project_id}'")
-            if project["totalSize"] > 0:
-                milestones = project["records"][0]["Milestones__c"] or milestones
-
         inputs = f"Role: {role} Project: {project_id} ({project_name}) Milestones: {milestones}"
         input_ids = tokenizer(inputs, return_tensors="pt", max_length=128, truncation=True).input_ids
         outputs = model.generate(input_ids, max_length=128, num_beams=4, early_stopping=True)
         checklist = tokenizer.decode(outputs[0], skip_special_tokens=True)
         tips = "1. Prioritize safety checks\n2. Review milestones\n3. Log progress"
 
-        if coaching_record_id:
-            sf.Supervisor_AI_Coaching__c.update(coaching_record_id, {
-                "Daily_Checklist__c": checklist,
-                "Suggested_Tips__c": tips,
-                "Engagement_Score__c": 10
+        # Update Salesforce if record_id is provided
+        if record_id:
+            sf = get_salesforce_connection()
+            sf.Supervisor_AI_Coaching__c.update(record_id, {
+                'Daily_Checklist__c': checklist,
+                'Suggested_Tips__c': tips,
+                'Engagement_Score__c': 10 if not sf.Supervisor_AI_Coaching__c.get(record_id).get('Engagement_Score__c') else sf.Supervisor_AI_Coaching__c.get(record_id).get('Engagement_Score__c') + 10
             })
-
-        return checklist, tips
+            return checklist, tips, f"Updated Salesforce record {record_id}"
+        return checklist, tips, "No Salesforce update (record_id not provided)"
     except Exception as e:
-        return f"Error: {str(e)}", ""
+        return f"Error: {str(e)}", "", ""
 
 # Define Gradio interface
 iface = gr.Interface(
     fn=gradio_generate_checklist,
     inputs=[
         gr.Textbox(label="Role", value="Supervisor"),
-        gr.Textbox(label="Project ID", value=""),
+        gr.Textbox(label="Project ID", value="P001"),
         gr.Textbox(label="Project Name", value="Building A"),
         gr.Textbox(label="Milestones", value="Complete foundation by 5/15"),
-        gr.Textbox(label="Coaching Record ID", value="")
+        gr.Textbox(label="Salesforce Record ID (optional)", value="")
     ],
     outputs=[
         gr.Textbox(label="Checklist"),
-        gr.Textbox(label="Tips")
+        gr.Textbox(label="Tips"),
+        gr.Textbox(label="Salesforce Status")
     ],
     title="AI Coach for Site Supervisors",
-    description="Generate daily checklists and tips, with Salesforce integration."
+    description="Generate daily checklists and tips, and update Salesforce records."
 )
 
 # Mount FastAPI to Gradio
@@ -174,4 +152,5 @@ if __name__ == "__main__":
         logging.info("Gradio interface launched successfully")
     except Exception as e:
         logging.error(f"Failed to launch Gradio interface: {e}")
-        raise e
+        raise e
+```