Update app.py

app.py

@@ -18,62 +18,26 @@ sf = Salesforce(
     security_token=os.getenv("SF_SECURITY_TOKEN")
 )
 
-def fetch_salesforce_data(supervisor_id, project_id):
+def generate_ai_data(supervisor_id, project_id, supervisor_data, 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.
+    Generate AI coaching data and reports based on supervisor and project data.
     
     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']}. 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']}."
+            f"{project_data['Name']} with milestones {project_data['Milestones__c']} and schedule "
+            f"{project_data['Project_Schedule__c']}."
         )
 
         # Generate AI output
@@ -92,7 +56,7 @@ def generate_ai_data(supervisor_id, project_id):
             "3. Schedule a team review."
         )
         risk_alerts = "Risk of delay: Rain expected on May 22, 2025."
-        upcoming_milestones = project_data['Milestones__c'].split(';')[0] if project_data['Milestones__c'] else "No upcoming milestones"
+        upcoming_milestones = project_data['Milestones__c'].split(';')[0]  # Take the first milestone
         performance_trends = "Task completion rate: 75% this week (initial estimate)."
 
         # Save AI data to AI_Coaching_Data__c
@@ -122,8 +86,6 @@ def generate_ai_data(supervisor_id, project_id):
         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
         }
@@ -138,15 +100,14 @@ def generate_ai_data(supervisor_id, project_id):
 iface = gr.Interface(
     fn=generate_ai_data,
     inputs=[
-        gr.Textbox(label="Supervisor ID (from Supervisor_Profile__c)"),
-        gr.Textbox(label="Project ID (from Project_Details__c)")
+        gr.Textbox(label="Supervisor ID"),
+        gr.Textbox(label="Project ID"),
+        gr.JSON(label="Supervisor Data"),
+        gr.JSON(label="Project Data")
     ],
     outputs=gr.JSON(label="Result"),
     title="AI Coach Data Generator",
-    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."
-    )
+    description="Generate AI coaching data and reports based on supervisor and project details."
 )
 
 # Launch the Gradio app