Update app.py

app.py

@@ -15,7 +15,6 @@ from fastapi import FastAPI, Form, File, UploadFile
 from fastapi.middleware.cors import CORSMiddleware
 from fastapi.responses import JSONResponse
 from fastapi.staticfiles import StaticFiles
-import time  # to simulate processing delay
 
 # Configure logging to show detailed messages
 logging.basicConfig(level=logging.DEBUG)
@@ -95,7 +94,105 @@ def generate_pdf(record_data):
         logger.error(f"Error generating PDF: {str(e)}", exc_info=True)
         return None
 
-# Function to generate project timeline and risk tags
+# Function to upload PDF to Salesforce and get its URL
+def upload_pdf_to_salesforce(pdf_file, project_title, record_id=None):
+    try:
+        sf = get_salesforce_connection()
+        if not sf:
+            logger.error("Salesforce connection failed. Cannot upload PDF.")
+            return None, None
+
+        encoded_pdf_data = base64.b64encode(pdf_file.getvalue()).decode('utf-8')
+        logger.debug(f"Uploading PDF for project: {project_title}, record ID: {record_id}")
+        content_version_data = {
+            "Title": f"{project_title} - Gantt Chart PDF",
+            "PathOnClient": f"{project_title}_Gantt_Chart.pdf",
+            "VersionData": encoded_pdf_data,
+        }
+        
+        if record_id:
+            content_version_data["FirstPublishLocationId"] = record_id
+
+        content_version = sf.ContentVersion.create(content_version_data)
+        content_version_id = content_version["id"]
+        logger.info(f"PDF uploaded to Salesforce with ContentVersion ID: {content_version_id}")
+
+        result = sf.query(f"SELECT Id, ContentDocumentId FROM ContentVersion WHERE Id = '{content_version_id}'")
+        if not result['records']:
+            logger.error("No records returned for ContentVersion query")
+            return content_version_id, None
+            
+        content_document_id = result['records'][0]['ContentDocumentId']
+        file_url = f"https://{sf.sf_instance}/sfc/servlet.shepherd/version/download/{content_version_id}"
+        logger.debug(f"Generated PDF URL: {file_url}")
+        return content_version_id, file_url
+    except Exception as e:
+        logger.error(f"Error uploading PDF to Salesforce: {str(e)}", exc_info=True)
+        return None, None
+
+# Function to create or update project timeline in Salesforce
+def send_to_salesforce(project_title, gantt_chart_url, ai_plan_score, estimated_duration, status="Draft", record_id=None, location="", weather_type="", work_items=None, work_items_id=None):
+    try:
+        logger.debug("Starting send_to_salesforce function...")
+        sf = get_salesforce_connection()
+        if not sf:
+            logger.error("Salesforce connection failed. Cannot proceed with record creation/update.")
+            return None
+
+        try:
+            obj_description = sf.AI_Project_Timeline__c.describe()
+            logger.debug("AI_Project_Timeline__c object exists and is accessible.")
+            available_fields = [field['name'] for field in obj_description['fields']]
+            logger.debug(f"Available fields on AI_Project_Timeline__c: {available_fields}")
+        except Exception as e:
+            logger.error(f"Error: AI_Project_Timeline__c object not found or inaccessible: {str(e)}")
+            return None
+
+        sf_data = {
+            "Name": project_title[:80],
+            "Project_Title__c": project_title,
+            "Estimated_Duration__c": estimated_duration,
+            "AI_Plan_Score__c": ai_plan_score,
+            "Status__c": status,
+            "Location__c": location,
+            "Weather_Type__c": weather_type,
+        }
+        
+        if gantt_chart_url:
+            sf_data["Gantt_Chart_PDF__c"] = gantt_chart_url
+            
+        if work_items_id:
+            sf_data["Work_Items__c"] = work_items_id
+
+        logger.debug(f"Prepared Salesforce data: {sf_data}")
+
+        if record_id:
+            try:
+                logger.info(f"Attempting to update Salesforce record with ID: {record_id}")
+                sf.AI_Project_Timeline__c.update(record_id, sf_data)
+                logger.info(f"Successfully updated Salesforce record with ID: {record_id}")
+                return record_id
+            except Exception as e:
+                logger.error(f"Error updating record {record_id}: {str(e)}")
+                record_id = None
+
+        logger.info("Creating new Salesforce record...")
+        project_record = sf.AI_Project_Timeline__c.create(sf_data)
+        if not project_record.get('id'):
+            logger.error("Failed to create record, no ID returned")
+            return None
+            
+        new_record_id = project_record['id']
+        logger.info(f"Created new Salesforce record with ID: {new_record_id}")
+        return new_record_id
+    
+    except Exception as e:
+        logger.error(f"Error sending data to Salesforce: {str(e)}", exc_info=True)
+        if hasattr(e, 'content') and e.content:
+            logger.error(f"Salesforce API response: {e.content}")
+        return None
+
+# Function to generate Gantt chart
 def generate_project_timeline(boq_file, weather, workforce, location, project_title):
     temp_dir = None
     try:
@@ -143,57 +240,234 @@ def generate_project_timeline(boq_file, weather, workforce, location, project_ti
 
 # Gradio interface function (for Gradio UI)
 def gradio_interface(boq_file, weather, workforce, location, project_title):
-    # Simulate a long-running task (remove this line in your actual code)
-    time.sleep(2)
-
+    temp_dir = None
     try:
         logger.info("Starting gradio_interface...")
 
-        # Your existing logic here...
+        if not boq_file:
+            return None, "Error: No BOQ file uploaded"
 
-        # Simulate an image URL for output
-        image_url = "/path/to/your/generated_gantt_chart.png"  # Replace with actual URL or path
+        file_path, risk_tags, temp_dir = generate_project_timeline(boq_file, weather, workforce, location, project_title)
+        if not file_path:
+            return None, f"Error: Failed to generate timeline: {risk_tags}"
 
-        risk_tags = "Risk tags based on weather and task durations"  # Replace with actual risk tags generation
+        if hasattr(boq_file, 'name'):
+            df = pd.read_csv(boq_file.name)
+        else:
+            df = pd.read_csv(boq_file)
+        estimated_duration = sum(df["Duration"])
+        ai_plan_score = min(100, max(0, 100 - (estimated_duration / 100)))
+        logger.debug(f"Estimated duration: {estimated_duration}, AI plan score: {ai_plan_score}")
+
+        record_id = send_to_salesforce(
+            project_title=project_title,
+            gantt_chart_url="",
+            ai_plan_score=ai_plan_score,
+            estimated_duration=estimated_duration,
+            status="Draft",
+            record_id=None,
+            location=location,
+            weather_type=weather
+        )
+        
+        if not record_id:
+            return None, f"Error: Failed to create Salesforce record - check logs for details\n\nRisk Tags:\n{risk_tags}"
+
+        boq_file_path = boq_file.name if hasattr(boq_file, 'name') else boq_file
+        work_items_id = upload_file_to_salesforce(boq_file_path, "Boq_data.csv", record_id)
+        if not work_items_id:
+            logger.warning("Failed to upload BOQ file, but proceeding with record creation")
+
+        record_data = {
+            "project_title": project_title,
+            "estimated_duration": estimated_duration,
+            "ai_plan_score": ai_plan_score,
+            "status": "Draft",
+            "risk_tags": risk_tags,
+        }
+        pdf_file = generate_pdf(record_data)
+        if not pdf_file:
+            logger.warning("Failed to generate PDF, but proceeding with record creation")
+
+        pdf_content_id, pdf_url = None, None
+        if pdf_file:
+            pdf_content_id, pdf_url = upload_pdf_to_salesforce(pdf_file, project_title, record_id)
+            if not pdf_content_id:
+                logger.warning("Failed to upload PDF, but proceeding with record creation")
+
+        update_result = send_to_salesforce(
+            project_title=project_title,
+            gantt_chart_url=pdf_url if pdf_url else "",
+            ai_plan_score=ai_plan_score,
+            estimated_duration=estimated_duration,
+            status="Draft",
+            record_id=record_id,
+            location=location,
+            weather_type=weather,
+            work_items_id=work_items_id if work_items_id else ""
+        )
+        if not update_result:
+            logger.warning("Failed to update record with PDF URL, but record was created")
+
+        # Upload Gantt chart image to Salesforce for a public URL
+        image_content_id = upload_file_to_salesforce(file_path, f"{project_title}_Gantt_Chart.png", record_id)
+        image_url = None
+        if image_content_id:
+            sf = get_salesforce_connection()
+            if sf:
+                image_url = f"https://{sf.sf_instance}/sfc/servlet.shepherd/version/download/{image_content_id}"
+                logger.debug(f"Generated image URL: {image_url}")
 
-        # Returning results
-        return image_url, risk_tags
+        logger.info("Gradio interface completed successfully.")
+        return image_url if image_url else file_path, f"Successfully created Salesforce record ID: {record_id}\n\nRisk Tags:\n{risk_tags}"
     except Exception as e:
         logger.error(f"Error in Gradio interface: {str(e)}", exc_info=True)
         return None, f"Error in Gradio interface: {str(e)}"
+    finally:
+        if temp_dir and os.path.exists(temp_dir):
+            shutil.rmtree(temp_dir)
+            logger.debug(f"Cleaned up temporary directory: {temp_dir}")
+
+# Create a FastAPI app and mount the Gradio app with CORS support
+app = FastAPI()
+app.add_middleware(
+    CORSMiddleware,
+    allow_origins=["https://aiplannerforcivilworktimel2-dev-ed.develop.lightning.force.com"],
+    allow_credentials=True,
+    allow_methods=["*"],
+    allow_headers=["*"],
+)
+app.mount("/static", StaticFiles(directory=tempfile.gettempdir()), name="static")
+
+# FastAPI endpoint for Salesforce controller (defined before mounting Gradio)
+@app.post("/api/gradio_interface")
+async def api_gradio_interface(
+    boq_file: UploadFile = File(...),
+    weather: str = Form(...),
+    workforce: int = Form(...),
+    location: str = Form(...),
+    project_title: str = Form(...)
+):
+    temp_dir = None
+    try:
+        logger.info("Starting api_gradio_interface...")
+
+        # Save uploaded file to temporary path
+        temp_dir = tempfile.mkdtemp()
+        boq_file_path = os.path.join(temp_dir, boq_file.filename)
+        with open(boq_file_path, "wb") as f:
+            f.write(boq_file.file.read())
+
+        file_path, risk_tags, temp_dir = generate_project_timeline(boq_file_path, weather, workforce, location, project_title)
+        if not file_path:
+            return JSONResponse({"error": f"Failed to generate timeline: {risk_tags}"}, status_code=400)
+
+        df = pd.read_csv(boq_file_path)
+        estimated_duration = sum(df["Duration"])
+        ai_plan_score = min(100, max(0, 100 - (estimated_duration / 100)))
+        logger.debug(f"Estimated duration: {estimated_duration}, AI plan score: {ai_plan_score}")
+
+        record_id = send_to_salesforce(
+            project_title=project_title,
+            gantt_chart_url="",
+            ai_plan_score=ai_plan_score,
+            estimated_duration=estimated_duration,
+            status="Draft",
+            record_id=None,
+            location=location,
+            weather_type=weather
+        )
+        
+        if not record_id:
+            return JSONResponse({
+                "error": f"Failed to create Salesforce record - check logs for details",
+                "text": f"Risk Tags:\n{risk_tags}"
+            }, status_code=500)
+
+        work_items_id = upload_file_to_salesforce(boq_file_path, "Boq_data.csv", record_id)
+        if not work_items_id:
+            logger.warning("Failed to upload BOQ file, but proceeding with record creation")
+
+        record_data = {
+            "project_title": project_title,
+            "estimated_duration": estimated_duration,
+            "ai_plan_score": ai_plan_score,
+            "status": "Draft",
+            "risk_tags": risk_tags,
+        }
+        pdf_file = generate_pdf(record_data)
+        if not pdf_file:
+            logger.warning("Failed to generate PDF, but proceeding with record creation")
 
-# Modified Gradio Interface:
-def create_gradio_interface():
-    with gr.Blocks() as demo:
-        gr.Markdown("## AI Civil Work Planner")
-        gr.Markdown("Generate a project timeline (Gantt chart) and risk tags based on BOQ data and site parameters.")
-
-        with gr.Row():
-            with gr.Column():
-                boq_file = gr.File(label="Upload BOQ Data (CSV format)")
-                weather = gr.Dropdown(label="Weather", choices=["sunny", "rainy", "cloudy"], value="sunny")
-                workforce = gr.Number(label="Workforce Size", value=10, precision=0)
-                location = gr.Textbox(label="Location", placeholder="Enter project location")
-                project_title = gr.Textbox(label="Project Title", placeholder="Enter project title")
-                submit_btn = gr.Button("Generate Timeline")
-
-            with gr.Column():
-                output_image = gr.Image(label="Gantt Chart")
-                risk_tags = gr.Textbox(label="Risk Tags and Salesforce Status")
-
-        submit_btn.click(
-            fn=gradio_interface,
-            inputs=[boq_file, weather, workforce, location, project_title],
-            outputs=[output_image, risk_tags],
-            show_progress=True  # Show progress while the function is running
+        pdf_content_id, pdf_url = None, None
+        if pdf_file:
+            pdf_content_id, pdf_url = upload_pdf_to_salesforce(pdf_file, project_title, record_id)
+            if not pdf_content_id:
+                logger.warning("Failed to upload PDF, but proceeding with record creation")
+
+        update_result = send_to_salesforce(
+            project_title=project_title,
+            gantt_chart_url=pdf_url if pdf_url else "",
+            ai_plan_score=ai_plan_score,
+            estimated_duration=estimated_duration,
+            status="Draft",
+            record_id=record_id,
+            location=location,
+            weather_type=weather,
+            work_items_id=work_items_id if work_items_id else ""
         )
+        if not update_result:
+            logger.warning("Failed to update record with PDF URL, but record was created")
+
+        image_content_id = upload_file_to_salesforce(file_path, f"{project_title}_Gantt_Chart.png", record_id)
+        image_url = None
+        if image_content_id:
+            sf = get_salesforce_connection()
+            if sf:
+                image_url = f"https://{sf.sf_instance}/sfc/servlet.shepherd/version/download/{image_content_id}"
+                logger.debug(f"Generated image URL: {image_url}")
+
+        logger.info("API gradio interface completed successfully.")
+        return JSONResponse({
+            "image": image_url if image_url else f"/static/{os.path.basename(file_path)}",
+            "text": f"Successfully created Salesforce record ID: {record_id}\n\nRisk Tags:\n{risk_tags}"
+        })
+    except Exception as e:
+        logger.error(f"Error in API gradio interface: {str(e)}", exc_info=True)
+        return JSONResponse({"error": f"Error in API gradio interface: {str(e)}"}, status_code=500)
+    finally:
+        if temp_dir and os.path.exists(temp_dir):
+            shutil.rmtree(temp_dir)
+            logger.debug(f"Cleaned up temporary directory: {temp_dir}")
+
+# Create Gradio interface
+demo = gr.Blocks(theme="default")
+with demo:
+    gr.Markdown("## AI Civil Work Planner")
+    gr.Markdown("Generate a project timeline (Gantt chart) and risk tags based on BOQ data and site parameters.")
+
+    with gr.Row():
+        with gr.Column():
+            boq_file = gr.File(label="Upload BOQ Data (CSV format)")
+            weather = gr.Dropdown(label="Weather", choices=["sunny", "rainy", "cloudy"], value="sunny")
+            workforce = gr.Number(label="Workforce Size", value=10, precision=0)
+            location = gr.Textbox(label="Location", placeholder="Enter project location")
+            project_title = gr.Textbox(label="Project Title", placeholder="Enter project title")
+            submit_btn = gr.Button("Generate Timeline")
+
+        with gr.Column():
+            output_image = gr.Image(label="Gantt Chart")
+            risk_tags = gr.Textbox(label="Risk Tags and Salesforce Status")
 
-    return demo
+    submit_btn.click(
+        fn=gradio_interface,
+        inputs=[boq_file, weather, workforce, location, project_title],
+        outputs=[output_image, risk_tags],
+    )
 
-# Call the function to create the Gradio app
-demo = create_gradio_interface()
+# Mount Gradio app (use a different variable to avoid overwriting 'app')
+gradio_app = gr.mount_gradio_app(app, demo, path="/")
 
-# Start the app
 if __name__ == "__main__":
     import uvicorn
-    uvicorn.run(demo, host="0.0.0.0", port=7860)
+    uvicorn.run(gradio_app, host="0.0.0.0", port=7860)