Update app.py

app.py

@@ -1,6 +1,240 @@
 import gradio as gr
+import torchimport gradio as gr
 import torch
 from transformers import AutoModelForCausalLM, AutoTokenizer
+from simple_salesforce import Salesforce
+from dotenv import load_dotenv
+import os
+import json
+from fastapi import FastAPI, HTTPException
+import uvicorn
+
+# Load environment variables from .env file
+load_dotenv()
+
+# Salesforce credentials
+SF_USERNAME = os.getenv('SF_USERNAME')
+SF_PASSWORD = os.getenv('SF_PASSWORD')
+SF_SECURITY_TOKEN = os.getenv('SF_SECURITY_TOKEN')
+SF_DOMAIN = os.getenv('SF_DOMAIN')
+HUGGINGFACE_API_KEY = os.getenv('HUGGINGFACE_API_KEY')
+
+# Initialize Salesforce connection
+try:
+    sf = Salesforce(
+        username=SF_USERNAME,
+        password=SF_PASSWORD,
+        security_token=SF_SECURITY_TOKEN,
+        domain=SF_DOMAIN
+    )
+except Exception as e:
+    print(f"Error connecting to Salesforce: {e}")
+
+# Initialize model and tokenizer
+model_name = "distilgpt2"
+try:
+    tokenizer = AutoTokenizer.from_pretrained(
+        model_name,
+        cache_dir="./model_cache",
+        token=HUGGINGFACE_API_KEY
+    )
+    model = AutoModelForCausalLM.from_pretrained(
+        model_name,
+        cache_dir="./model_cache",
+        token=HUGGINGFACE_API_KEY
+    )
+except Exception as e:
+    print(f"Error loading model: {e}")
+    tokenizer = AutoTokenizer.from_pretrained(model_name, token=HUGGINGFACE_API_KEY)
+    model = AutoModelForCausalLM.from_pretrained(model_name, token=HUGGINGFACE_API_KEY)
+
+# Set pad_token to eos_token if not already set
+if tokenizer.pad_token is None:
+    tokenizer.pad_token = tokenizer.eos_token
+    model.config.pad_token_id = tokenizer.eos_token_id
+
+# Simplified prompt template
+PROMPT_TEMPLATE = """Role: {role}
+Project: {project_id}
+Milestones: 
+- {milestones_list}
+Reflection: {reflection}
+Generate:
+Checklist:
+- {milestones_list}
+Suggestions:
+{suggestions_list}
+Quote:
+{your_motivational_quote}"""
+
+def generate_outputs(role, project_id, milestones, reflection):
+    # Input validation
+    if not all([role, project_id, milestones, reflection]):
+        return "Error: All fields are required.", "", ""
+    
+    # Process milestones
+    milestones_list = "\n- ".join([m.strip() for m in milestones.split(",") if m.strip()])
+    if not milestones_list:
+        return "Error: At least one valid milestone is required.", "", ""
+    
+    # Generate suggestions based on reflection
+    suggestions_list = []
+    reflection_lower = reflection.lower()
+    if "delays" in reflection_lower:
+        suggestions_list.extend(["Adjust timelines for delays.", "Communicate with stakeholders."])
+    if "weather" in reflection_lower:
+        suggestions_list.extend(["Ensure rain gear availability.", "Monitor weather updates."])
+    if "equipment" in reflection_lower:
+        suggestions_list.extend(["Inspect equipment.", "Schedule maintenance."])
+    suggestions_list = "\n- ".join(suggestions_list) if suggestions_list else "No specific suggestions."
+
+    # Format prompt
+    prompt = PROMPT_TEMPLATE.format(
+        role=role,
+        project_id=project_id,
+        milestones_list=milestones_list.replace("\n- ", "\n- "),
+        reflection=reflection,
+        suggestions_list=suggestions_list,
+        your_motivational_quote="Your motivational quote here"
+    )
+
+    # Tokenize input
+    try:
+        inputs = tokenizer(
+            prompt,
+            return_tensors="pt",
+            max_length=512,
+            truncation=True,
+            padding=True
+        )
+    except Exception as e:
+        return f"Error in tokenization: {e}", "", ""
+
+    # Generate output
+    try:
+        with torch.no_grad():
+            outputs = model.generate(
+                input_ids=inputs['input_ids'],
+                attention_mask=inputs['attention_mask'],
+                max_length=600,
+                num_return_sequences=1,
+                no_repeat_ngram_size=2,
+                top_k=50,
+                top_p=0.95,
+                temperature=0.7,
+                pad_token_id=tokenizer.eos_token_id,
+                do_sample=True
+            )
+        
+        generated_text = tokenizer.decode(outputs[0], skip_special_tokens=True)
+    except Exception as e:
+        return f"Error in generation: {e}", "", ""
+
+    # Extract sections
+    checklist = suggestions = quote = "Not generated."
+    try:
+        if "Checklist:" in generated_text:
+            checklist_start = generated_text.find("Checklist:") + 10
+            suggestions_start = generated_text.find("Suggestions:", checklist_start)
+            if suggestions_start == -1:
+                suggestions_start = len(generated_text)
+            checklist = generated_text[checklist_start:suggestions_start].strip()
+
+        if "Suggestions:" in generated_text:
+            suggestions_start = generated_text.find("Suggestions:") + 12
+            quote_start = generated_text.find("Quote:", suggestions_start)
+            if quote_start == -1:
+                quote_start = len(generated_text)
+            suggestions = generated_text[suggestions_start:quote_start].strip()
+
+        if "Quote:" in generated_text:
+            quote_start = generated_text.find("Quote:") + 7
+            quote = generated_text[quote_start:].strip()
+    except Exception as e:
+        return f"Error parsing output: {e}", "", ""
+
+    # Save to Salesforce
+    try:
+        project = sf.Project__c.get(project_id)
+        supervisor_id = project['Supervisor__c']
+        coaching_data = {
+            'Supervisor_ID__c': supervisor_id,
+            'Project_ID__c': project_id,
+            'Daily_Checklist__c': checklist,
+            'Suggested_Tips__c': suggestions,
+            'KPI_Flag__c': 'delay' in suggestions.lower() or 'issue' in suggestions.lower(),
+            'Last_Refresh_Date__c': datetime.now().isoformat()
+        }
+        sf.Supervisor_AI_Coaching__c.create(coaching_data)
+    except Exception as e:
+        print(f"Error saving to Salesforce: {e}")
+
+    return checklist, suggestions, quote
+
+def create_interface():
+    with gr.Blocks() as demo:
+        gr.Markdown("### Construction Supervisor AI Coach")
+        with gr.Row():
+            role = gr.Dropdown(choices=["Supervisor", "Foreman", "Project Manager"], label="Role", value="Supervisor")
+            project_id = gr.Textbox(label="Project ID", placeholder="e.g., PROJ-001")
+        milestones = gr.Textbox(
+            label="Milestones (comma-separated)",
+            placeholder="e.g., Foundation complete, Framing started, Roof installed"
+        )
+        reflection = gr.Textbox(
+            label="Reflection",
+            lines=3,
+            placeholder="e.g., Facing delays due to weather and equipment issues."
+        )
+        with gr.Row():
+            submit = gr.Button("Generate")
+            clear = gr.Button("Clear")
+        checklist_output = gr.Textbox(label="Checklist", lines=4)
+        suggestions_output = gr.Textbox(label="Suggestions", lines=4)
+        quote_output = gr.Textbox(label="Quote", lines=2)
+        
+        submit.click(
+            fn=generate_outputs,
+            inputs=[role, project_id, milestones, reflection],
+            outputs=[checklist_output, suggestions_output, quote_output]
+        )
+        clear.click(
+            fn=lambda: ("Supervisor", "", "", ""),
+            inputs=None,
+            outputs=[role, project_id, milestones, reflection]
+        )
+    
+    return demo
+
+# FastAPI for Salesforce integration
+app = FastAPI()
+
+@app.post("/api/predict")
+async def predict(data: dict):
+    try:
+        role = data.get('role')
+        project_id = data.get('project_id')
+        milestones = data.get('milestones')
+        reflection = data.get('reflection')
+        checklist, suggestions, quote = generate_outputs(role, project_id, milestones, reflection)
+        return {
+            'checklist': checklist,
+            'suggestions': suggestions,
+            'quote': quote
+        }
+    except Exception as e:
+        raise HTTPException(status_code=500, detail=f"Error: {str(e)}")
+
+if __name__ == "__main__":
+    try:
+        # For local testing with Gradio
+        demo = create_interface()
+        demo.launch(server_name="0.0.0.0", server_port=7860, share=False)
+    except Exception as e:
+        print(f"Error launching Gradio interface: {e}")
+        # Run FastAPI for production
+        uvicorn.run(app, host="0.0.0.0", port=7860)
+from transformers import AutoModelForCausalLM, AutoTokenizer
 
 # Initialize model and tokenizer
 model_name = "distilgpt2"