Update app.py

app.py

@@ -1,11 +1,10 @@
-from fastapi import FastAPI, HTTPException
-from pydantic import BaseModel
+from flask import Flask, request, jsonify
 from transformers import AutoModelForCausalLM, AutoTokenizer
-from contextlib import asynccontextmanager
 import json
 import logging
 import os
-import asyncio
+import threading
+import time
 
 # Set up logging to stdout only
 logging.basicConfig(
@@ -17,6 +16,9 @@ logging.basicConfig(
 )
 logger = logging.getLogger(__name__)
 
+# Initialize Flask app
+app = Flask(__name__)
+
 # Global variables for model and tokenizer
 model = None
 tokenizer = None
@@ -26,8 +28,8 @@ model_load_status = "not_loaded"
 model_path = "/app/fine-tuned-construction-llm"
 fallback_model = "distilgpt2"
 
-# Asynchronous function to load model in the background
-async def load_model_background():
+# Function to load model in the background
+def load_model_background():
     global model, tokenizer, model_load_status
     try:
         if os.path.isdir(model_path):
@@ -45,41 +47,41 @@ async def load_model_background():
         logger.error(f"Failed to load model or tokenizer: {str(e)}")
         model_load_status = f"failed: {str(e)}"
 
-# Lifespan event handler to manage startup and shutdown
-@asynccontextmanager
-async def lifespan(app: FastAPI):
-    logger.debug("FastAPI application starting")
-    # Start the background task for model loading
-    asyncio.create_task(load_model_background())
-    yield
-    logger.debug("FastAPI application shutting down")
-
-# Initialize FastAPI app with lifespan handler
-app = FastAPI(lifespan=lifespan)
-
-# Define input model for validation
-class CoachingInput(BaseModel):
-    role: str
-    project_id: str
-    milestones: str
-    reflection_log: str
+# Start model loading in a background thread
+def start_background_tasks():
+    logger.debug("Starting background tasks")
+    thread = threading.Thread(target=load_model_background)
+    thread.daemon = True
+    thread.start()
 
-@app.get("/")
-async def root():
+@app.route("/")
+def root():
     logger.debug("Root endpoint accessed")
-    return {"message": "Supervisor AI Coach is running"}
+    return jsonify({"message": "Supervisor AI Coach is running"})
 
-@app.get("/health")
-async def health_check():
+@app.route("/health")
+def health_check():
     logger.debug("Health endpoint accessed")
-    return {
+    return jsonify({
         "status": "healthy" if model_load_status in ["local_model_loaded", "fallback_model_loaded"] else "starting",
         "model_load_status": model_load_status
-    }
+    })
 
-@app.post("/generate_coaching")
-async def generate_coaching(data: CoachingInput):
+@app.route("/generate_coaching", methods=["POST"])
+def generate_coaching():
     logger.debug("Generate coaching endpoint accessed")
+    # Manual validation of request data (replacing Pydantic)
+    data = request.get_json()
+    if not data:
+        logger.error("Invalid request: No JSON data provided")
+        return jsonify({"error": "Invalid request: JSON data required"}), 400
+
+    required_fields = ["role", "project_id", "milestones", "reflection_log"]
+    missing_fields = [field for field in required_fields if field not in data]
+    if missing_fields:
+        logger.error(f"Missing required fields: {missing_fields}")
+        return jsonify({"error": f"Missing required fields: {missing_fields}"}), 400
+
     if model is None or tokenizer is None:
         logger.warning("Model or tokenizer not loaded")
         # Return a static response if the model isn't loaded yet
@@ -88,13 +90,13 @@ async def generate_coaching(data: CoachingInput):
             "tips": ["Prioritize team communication", "Check weather updates"],
             "quote": "Every step forward counts!"
         }
-        return response_json
+        return jsonify(response_json)
 
     try:
         # Prepare input text
         input_text = (
-            f"Role: {data.role}, Project: {data.project_id}, "
-            f"Milestones: {data.milestones}, Reflection: {data.reflection_log}"
+            f"Role: {data['role']}, Project: {data['project_id']}, "
+            f"Milestones: {data['milestones']}, Reflection: {data['reflection_log']}"
         )
         
         # Tokenize input
@@ -131,8 +133,16 @@ async def generate_coaching(data: CoachingInput):
                 }
                 logger.warning("Failed to parse model output as JSON, using default response")
         
-        return response_json
+        return jsonify(response_json)
     
     except Exception as e:
         logger.error(f"Error generating coaching response: {str(e)}")
-        raise HTTPException(status_code=500, detail=f"Internal server error: {str(e)}")
+        return jsonify({"error": f"Internal server error: {str(e)}"}), 500
+
+if __name__ == "__main__":
+    # Start background tasks before the app runs
+    start_background_tasks()
+    # Run Flask app with waitress for production-ready WSGI server
+    from waitress import serve
+    logger.debug("Starting Flask app with Waitress")
+    serve(app, host="0.0.0.0", port=7860)