Update app.py

app.py

@@ -1,10 +1,12 @@
-from flask import Flask, request, jsonify
-from transformers import AutoModelForCausalLM, AutoTokenizer
+
+
 import json
 import logging
 import os
 import threading
 import time
+from flask import Flask, request, jsonify
+from transformers import AutoModelForCausalLM, AutoTokenizer
 
 # Set up logging to stdout only
 logging.basicConfig(
@@ -54,6 +56,46 @@ def start_background_tasks():
     thread.daemon = True
     thread.start()
 
+# Utility function to wait for model loading with a timeout
+def wait_for_model(timeout=60):
+    start_time = time.time()
+    while time.time() - start_time < timeout:
+        if model_load_status in ["local_model_loaded", "fallback_model_loaded"]:
+            return True
+        elif "failed" in model_load_status:
+            return False
+        time.sleep(1)
+    return False
+
+# Utility function to parse raw text into structured JSON response
+def parse_raw_text_to_json(raw_text):
+    lines = raw_text.strip().split("\n")
+    checklist = []
+    tips = []
+    quote = "Every step forward counts!"
+
+    # Simple parsing logic: look for keywords or assume structure
+    for line in lines:
+        line = line.strip()
+        if line.startswith("- "):  # Assume checklist items start with "- "
+            checklist.append(line[2:].strip())
+        elif line.startswith("* "):  # Assume tips start with "* "
+            tips.append(line[2:].strip())
+        elif line and not checklist and not tips:  # If no checklist or tips, treat as a quote
+            quote = line
+
+    # Fallback if parsing fails
+    if not checklist:
+        checklist = ["Inspect safety equipment", "Review milestone progress"]
+    if not tips:
+        tips = ["Prioritize team communication", "Check weather updates"]
+
+    return {
+        "checklist": checklist,
+        "tips": tips,
+        "quote": quote
+    }
+
 @app.route("/")
 def root():
     logger.debug("Root endpoint accessed")
@@ -70,7 +112,7 @@ def health_check():
 @app.route("/generate_coaching", methods=["POST"])
 def generate_coaching():
     logger.debug("Generate coaching endpoint accessed")
-    # Manual validation of request data (replacing Pydantic)
+    # Manual validation of request data
     data = request.get_json()
     if not data:
         logger.error("Invalid request: No JSON data provided")
@@ -82,15 +124,14 @@ def generate_coaching():
         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
-        response_json = {
+    # Wait for the model to load (up to 60 seconds)
+    if not wait_for_model(timeout=60):
+        logger.warning("Model failed to load within timeout")
+        return jsonify({
             "checklist": ["Inspect safety equipment", "Review milestone progress"],
             "tips": ["Prioritize team communication", "Check weather updates"],
             "quote": "Every step forward counts!"
-        }
-        return jsonify(response_json)
+        })
 
     try:
         # Prepare input text
@@ -98,10 +139,10 @@ def generate_coaching():
             f"Role: {data['role']}, Project: {data['project_id']}, "
             f"Milestones: {data['milestones']}, Reflection: {data['reflection_log']}"
         )
-        
+       
         # Tokenize input
         inputs = tokenizer(input_text, return_tensors="pt", max_length=512, truncation=True)
-        
+       
         # Generate output
         outputs = model.generate(
             inputs["input_ids"],
@@ -111,30 +152,23 @@ def generate_coaching():
             do_sample=True,
             temperature=0.7
         )
-        
-        # Decode and parse response
+       
+        # Decode response
         response_text = tokenizer.decode(outputs[0], skip_special_tokens=True)
-        
-        # Since distilgpt2 may not output JSON, parse the response manually or use fallback
-        if not response_text.startswith("{"):
-            checklist = ["Inspect safety equipment", "Review milestone progress"]
-            tips = ["Prioritize team communication", "Check weather updates"]
-            quote = "Every step forward counts!"
-            response_json = {"checklist": checklist, "tips": tips, "quote": quote}
-            logger.warning("Model output is not JSON, using default response")
-        else:
-            try:
-                response_json = json.loads(response_text)
-            except json.JSONDecodeError:
-                response_json = {
-                    "checklist": ["Inspect safety equipment", "Review milestone progress"],
-                    "tips": ["Prioritize team communication", "Check weather updates"],
-                    "quote": "Every step forward counts!"
-                }
-                logger.warning("Failed to parse model output as JSON, using default response")
-        
+       
+        # Try parsing as JSON first
+        try:
+            response_json = json.loads(response_text)
+            # Validate required fields in response
+            if not all(key in response_json for key in ["checklist", "tips", "quote"]):
+                raise ValueError("Missing required fields in model output")
+        except (json.JSONDecodeError, ValueError):
+            # If not JSON or invalid JSON, parse raw text
+            logger.warning("Model output is not valid JSON, parsing raw text")
+            response_json = parse_raw_text_to_json(response_text)
+       
         return jsonify(response_json)
-    
+   
     except Exception as e:
         logger.error(f"Error generating coaching response: {str(e)}")
         return jsonify({"error": f"Internal server error: {str(e)}"}), 500
@@ -145,4 +179,5 @@ if __name__ == "__main__":
     # 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)
+    serve(app, host="0.0.0.0", port=7860)
+