Update app.py

app.py

@@ -1,11 +1,11 @@
-
 import gradio as gr
 import logging
 from transformers import AutoModelForSeq2SeqLM, AutoTokenizer
 import os
-from fastapi import FastAPI, Request
+from fastapi import FastAPI
 from pydantic import BaseModel
-import uvicorn
+from simple_salesforce import Salesforce
+import json
 
 # Configure logging
 logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s')
@@ -13,19 +13,50 @@ logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(
 # Load Hugging Face token from environment variable
 HF_TOKEN = os.getenv("HF_TOKEN")
 if not HF_TOKEN:
-    logging.warning("HF_TOKEN not set. Private repositories may be inaccessible.")
+    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 model path
+MODEL_PATH = "your_username/fine_tuned_bart_construction"  # Replace with your actual model path
+FALLBACK_MODEL = "facebook/bart-large"  # Fallback public model
 
 try:
-    # Load model and tokenizer with authentication
+    # Try loading fine-tuned model
+    logging.info(f"Attempting to load model from {MODEL_PATH}")
     model = AutoModelForSeq2SeqLM.from_pretrained(MODEL_PATH, use_auth_token=HF_TOKEN)
     tokenizer = AutoTokenizer.from_pretrained(MODEL_PATH, use_auth_token=HF_TOKEN)
     logging.info(f"Model and tokenizer loaded successfully from {MODEL_PATH}")
 except Exception as e:
-    logging.error(f"Failed to load model or tokenizer: {e}")
-    raise e
+    logging.error(f"Failed to load model from {MODEL_PATH}: {e}")
+    logging.info(f"Falling back to public model: {FALLBACK_MODEL}")
+    try:
+        model = AutoModelForSeq2SeqLM.from_pretrained(FALLBACK_MODEL)
+        tokenizer = AutoTokenizer.from_pretrained(FALLBACK_MODEL)
+        logging.info(f"Fallback model {FALLBACK_MODEL} loaded successfully")
+    except Exception as fallback_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("SF_USERNAME")
+        sf_password = os.getenv("SF_PASSWORD")
+        sf_security_token = os.getenv("SF_SECURITY_TOKEN")
+        
+        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):
@@ -33,25 +64,52 @@ class ChecklistInput(BaseModel):
     project_id: str = "Unknown"
     project_name: str = "Unknown Project"
     milestones: str = "No milestones provided"
+    coaching_record_id: str = None  # Optional Salesforce record ID
 
-# Initialize FastAPI for JSON endpoint (Gradio mounts it)
+# Initialize FastAPI
 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: {data.milestones}"
-        logging.info(f"Generating checklist for inputs: {inputs[:100]}...")  # Truncate for logging
+        inputs = f"Role: {data.role} Project: {data.project_id} ({data.project_name}) Milestones: {milestones}"
+        logging.info(f"Generating checklist for inputs: {inputs[:100]}...")
 
         # Tokenize and generate
         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)
 
-        # Static tips (can be enhanced with model)
+        # 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:
+            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
+                })
+                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("Checklist and tips generated successfully")
         return {
             "checklist": checklist,
@@ -62,13 +120,29 @@ async def generate_checklist(data: ChecklistInput):
         return {"error": str(e)}
 
 # Gradio interface function
-def gradio_generate_checklist(role, project_id, project_name, milestones):
+def gradio_generate_checklist(role, project_id, project_name, milestones, coaching_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
+            })
+
         return checklist, tips
     except Exception as e:
         return f"Error: {str(e)}", ""
@@ -78,16 +152,17 @@ iface = gr.Interface(
     fn=gradio_generate_checklist,
     inputs=[
         gr.Textbox(label="Role", value="Supervisor"),
-        gr.Textbox(label="Project ID", value="P001"),
+        gr.Textbox(label="Project ID", value=""),
         gr.Textbox(label="Project Name", value="Building A"),
-        gr.Textbox(label="Milestones", value="Complete foundation by 5/15")
+        gr.Textbox(label="Milestones", value="Complete foundation by 5/15"),
+        gr.Textbox(label="Coaching Record ID", value="")
     ],
     outputs=[
         gr.Textbox(label="Checklist"),
         gr.Textbox(label="Tips")
     ],
     title="AI Coach for Site Supervisors",
-    description="Generate daily checklists and tips for site supervisors."
+    description="Generate daily checklists and tips, with Salesforce integration."
 )
 
 # Mount FastAPI to Gradio
@@ -95,9 +170,8 @@ iface.app = app
 
 if __name__ == "__main__":
     try:
-        # Launch Gradio with public URL for Hugging Face Spaces
         iface.launch(server_name="0.0.0.0", server_port=7860, share=False)
         logging.info("Gradio interface launched successfully")
     except Exception as e:
         logging.error(f"Failed to launch Gradio interface: {e}")
-        raise e
+        raise e