Update app.py

app.py

@@ -1,83 +1,116 @@
-from flask import Flask, render_template, request, jsonify
-import requests
-import json
-from datetime import datetime
-import pandas as pd
+from fastapi import FastAPI, HTTPException, Depends, Security
+from fastapi.security import HTTPBearer, HTTPAuthorizationCredentials
+from pydantic import BaseModel
+import numpy as np
+from sklearn.linear_model import LinearRegression
+import uvicorn
+import os
+from typing import List, Dict
+import time
 
-app = Flask(__name__)
+app = FastAPI(title="AI Estimator API")
+security = HTTPBearer()
 
-# Hugging Face API configuration
-HF_API_URL = "https://api-inference.huggingface.co/models/your-model-endpoint"
-HF_TOKEN = "your_hugging_face_token"
+# Mock token for authentication (replace with Hugging Face Space secrets or env variable)
+API_TOKEN = os.getenv("API_TOKEN", "hf_mock_token_123")
 
-# Simulated project task logs (replace with actual data source, e.g., Salesforce or database)
-def get_task_logs():
-    return [
-        {
-            "project_id": "P001",
-            "task_type": "Concrete",
-            "weather_temp": 25.5,
-            "absenteeism_rate": 0.1,
-            "labour_count": 20
-        },
-        {
-            "project_id": "P002",
-            "task_type": "Masonry",
-            "weather_temp": 22.0,
-            "absenteeism_rate": 0.05,
-            "labour_count": 15
-        }
-    ]
+# Pydantic models for request/response
+class TaskInput(BaseModel):
+    project_id: str
+    date: str
+    task_type: str
+    historical_consumption: Dict[str, float]  # e.g., {"cement": 100, "steel": 50}
+    weather: Dict[str, float]  # e.g., {"temperature": 25, "rain_probability": 0.1}
+    absenteeism_rate: float
+    material_lead_time: Dict[str, float]  # e.g., {"cement": 2, "steel": 3}
 
-# Call Hugging Face API for forecasting
-def call_hf_api(task_log):
-    headers = {"Authorization": f"Bearer {HF_TOKEN}"}
-    payload = {
-        "project_id": task_log["project_id"],
-        "task_type": task_log["task_type"],
-        "weather_temp": task_log["weather_temp"],
-        "absenteeism_rate": task_log["absenteeism_rate"],
-        "labour_count": task_log["labour_count"],
-        "date": datetime.now().strftime("%Y-%m-%d")
-    }
-    try:
-        response = requests.post(HF_API_URL, headers=headers, json=payload, timeout=10)
-        response.raise_for_status()
-        return response.json()
-    except requests.RequestException as e:
-        return {"error": str(e), "project_id": task_log["project_id"]}
+class ForecastOutput(BaseModel):
+    project_id: str
+    date: str
+    task_type: str
+    material_needed: Dict[str, float]
+    labour_needed: Dict[str, float]
+    forecast_confidence: float
+    alert_flag: bool
+
+# Mock trained model (replace with actual trained model in production)
+class MockAIModel:
+    def __init__(self):
+        # Dummy coefficients for regression
+        self.material_model = LinearRegression()
+        self.labour_model = LinearRegression()
+        # Mock training (replace with actual training data)
+        X = np.array([[1, 25, 0.1, 0.05, 2], [2, 30, 0.2, 0.1, 1]])  # Features: task_type_id, temp, rain, absenteeism, lead_time
+        y_material = np.array([[100, 50], [120, 60]])  # cement, steel
+        y_labour = np.array([[10, 5], [12, 6]])  # trade1, trade2
+        self.material_model.fit(X, y_material)
+        self.labour_model.fit(X, y_labour)
+
+    def predict(self, features: np.ndarray):
+        material_pred = self.material_model.predict(features)[0]
+        labour_pred = self.labour_model.predict(features)[0]
+        return material_pred, labour_pred
 
-# Generate placeholder PDF link (replace with actual PDF generation logic)
-def generate_pdf_link(forecast):
-    return f"https://example.com/pdf/{forecast['project_id']}_{forecast['date']}.pdf"
+# Initialize model
+ai_model = MockAIModel()
 
-# Route for the dashboard
-@app.route("/")
-def dashboard():
-    task_logs = get_task_logs()
-    forecasts = []
+# Token authentication
+async def verify_token(credentials: HTTPAuthorizationCredentials = Security(security)):
+    if credentials.credentials != API_TOKEN:
+        raise HTTPException(status_code=401, detail="Invalid or missing API token")
+    return credentials.credentials
+
+@app.post("/forecast", response_model=ForecastOutput, dependencies=[Depends(verify_token)])
+async def generate_forecast(task: TaskInput):
+    start_time = time.time()
     
-    for log in task_logs:
-        forecast = call_hf_api(log)
-        if "error" not in forecast:
-            forecast["pdf_link"] = generate_pdf_link(forecast)
-            forecasts.append(forecast)
+    # Extract features for model
+    task_type_id = hash(task.task_type) % 100  # Mock task type encoding
+    features = np.array([[
+        task_type_id,
+        task.weather.get("temperature", 20),
+        task.weather.get("rain_probability", 0),
+        task.absenteeism_rate,
+        list(task.material_lead_time.values())[0]  # Use first material lead time
+    ]])
     
-    return render_template("index.html", forecasts=forecasts)
-
-# Route for API access (optional, for external integrations)
-@app.route("/forecast", methods=["POST"])
-def forecast():
-    data = request.json
-    if not data:
-        return jsonify({"error": "No input data provided"}), 400
+    # Generate predictions
+    material_pred, labour_pred = ai_model.predict(features)
+    
+    # Format predictions
+    material_needed = {
+        "cement": max(0, material_pred[0]),
+        "steel": max(0, material_pred[1]),
+        "brick": max(0, material_pred[0] * 2)  # Mock derived material
+    }
+    labour_needed = {
+        "carpenter": max(0, labour_pred[0]),
+        "mason": max(0, labour_pred[1]),
+        "laborer": max(0, labour_pred[0] + labour_pred[1])  # Mock derived trade
+    }
+    
+    # Calculate confidence (mock: based on input completeness)
+    confidence = 0.85 if all([task.historical_consumption, task.weather, task.material_lead_time]) else 0.75
+    
+    # Alert if resource gap exceeds 15% (mock comparison with historical)
+    historical_total = sum(task.historical_consumption.values())
+    predicted_total = sum(material_needed.values())
+    alert_flag = abs(predicted_total - historical_total) / historical_total > 0.15 if historical_total > 0 else False
     
-    forecast = call_hf_api(data)
-    if "error" in forecast:
-        return jsonify({"error": forecast["error"]}), 500
+    # Ensure latency < 10s
+    elapsed_time = time.time() - start_time
+    if elapsed_time > 10:
+        raise HTTPException(status_code=500, detail="Forecast generation exceeded 10s")
     
-    forecast["pdf_link"] = generate_pdf_link(forecast)
-    return jsonify(forecast)
+    return ForecastOutput(
+        project_id=task.project_id,
+        date=task.date,
+        task_type=task.task_type,
+        material_needed=material_needed,
+        labour_needed=labour_needed,
+        forecast_confidence=confidence,
+        alert_flag=alert_flag
+    )
 
 if __name__ == "__main__":
-    app.run(debug=True, host="0.0.0.0", port=5000)
+    uvicorn.run(app, host="0.0.0.0", port=8000)