Update app.py

app.py

@@ -18,26 +18,62 @@ sf = Salesforce(
     security_token=os.getenv("SF_SECURITY_TOKEN")
 )
 
-def generate_ai_data(supervisor_id, project_id, supervisor_data, project_data):
+def fetch_salesforce_data(supervisor_id, project_id):
     """
-    Generate AI coaching data and reports based on supervisor and project data.
+    Fetch supervisor and project data from Salesforce based on provided IDs.
+    
+    Args:
+        supervisor_id (str): ID of the supervisor from Supervisor_Profile__c
+        project_id (str): ID of the project from Project_Details__c
+    
+    Returns:
+        tuple: (supervisor_data, project_data) as dictionaries
+    """
+    try:
+        # Fetch supervisor data
+        supervisor_query = f"SELECT Id, Role__c, Location__c FROM Supervisor_Profile__c WHERE Id = '{supervisor_id}'"
+        supervisor_result = sf.query(supervisor_query)
+        if not supervisor_result['records']:
+            raise ValueError(f"No supervisor found with ID: {supervisor_id}")
+        supervisor_data = supervisor_result['records'][0]
+        
+        # Fetch project data
+        project_query = f"SELECT Id, Name, Start_Date__c, End_Date__c, Milestones__c, Project_Schedule__c FROM Project_Details__c WHERE Id = '{project_id}'"
+        project_result = sf.query(project_query)
+        if not project_result['records']:
+            raise ValueError(f"No project found with ID: {project_id}")
+        project_data = project_result['records'][0]
+        
+        # Remove Salesforce attributes we don't need
+        supervisor_data = {k: v for k, v in supervisor_data.items() if not k.startswith('attributes')}
+        project_data = {k: v for k, v in project_data.items() if not k.startswith('attributes')}
+        
+        return supervisor_data, project_data
+    
+    except Exception as e:
+        raise Exception(f"Error fetching data from Salesforce: {str(e)}")
+
+def generate_ai_data(supervisor_id, project_id):
+    """
+    Generate AI coaching data and reports based on supervisor and project IDs.
     
     Args:
         supervisor_id (str): ID of the supervisor from Supervisor_Profile__c
         project_id (str): ID of the project from Project_Details__c
-        supervisor_data (dict): Contains Role__c, Location__c
-        project_data (dict): Contains Name, Start_Date__c, End_Date__c, Milestones__c, Project_Schedule__c
     
     Returns:
         dict: Status and generated data
     """
     try:
+        # Fetch data from Salesforce
+        supervisor_data, project_data = fetch_salesforce_data(supervisor_id, project_id)
+        
         # Construct prompt for AI generation
         prompt = (
             f"Generate daily checklist, tips, risk alerts, upcoming milestones, and performance trends for a "
             f"{supervisor_data['Role__c']} at {supervisor_data['Location__c']} working on project "
-            f"{project_data['Name']} with milestones {project_data['Milestones__c']} and schedule "
-            f"{project_data['Project_Schedule__c']}."
+            f"{project_data['Name']}. Project runs from {project_data['Start_Date__c']} to {project_data['End_Date__c']} "
+            f"with milestones {project_data['Milestones__c']} and schedule {project_data['Project_Schedule__c']}."
         )
 
         # Generate AI output
@@ -56,7 +92,7 @@ def generate_ai_data(supervisor_id, project_id, supervisor_data, project_data):
             "3. Schedule a team review."
         )
         risk_alerts = "Risk of delay: Rain expected on May 22, 2025."
-        upcoming_milestones = project_data['Milestones__c'].split(';')[0]  # Take the first milestone
+        upcoming_milestones = project_data['Milestones__c'].split(';')[0] if project_data['Milestones__c'] else "No upcoming milestones"
         performance_trends = "Task completion rate: 75% this week (initial estimate)."
 
         # Save AI data to AI_Coaching_Data__c
@@ -86,6 +122,8 @@ def generate_ai_data(supervisor_id, project_id, supervisor_data, project_data):
         return {
             "status": "success",
             "message": "AI data and report generated successfully",
+            "supervisor_data": supervisor_data,
+            "project_data": project_data,
             "ai_data": ai_data,
             "report_data": report_data
         }
@@ -100,14 +138,15 @@ def generate_ai_data(supervisor_id, project_id, supervisor_data, project_data):
 iface = gr.Interface(
     fn=generate_ai_data,
     inputs=[
-        gr.Textbox(label="Supervisor ID"),
-        gr.Textbox(label="Project ID"),
-        gr.JSON(label="Supervisor Data"),
-        gr.JSON(label="Project Data")
+        gr.Textbox(label="Supervisor ID (from Supervisor_Profile__c)"),
+        gr.Textbox(label="Project ID (from Project_Details__c)")
     ],
     outputs=gr.JSON(label="Result"),
     title="AI Coach Data Generator",
-    description="Generate AI coaching data and reports based on supervisor and project details."
+    description=(
+        "Generate AI coaching data and reports based on supervisor and project details. "
+        "Enter the Supervisor ID from Supervisor_Profile__c and Project ID from Project_Details__c."
+    )
 )
 
 # Launch the Gradio app