Update main.py

main.py

@@ -1,26 +1,47 @@
-from fastapi import FastAPI
+from fastapi import FastAPI, HTTPException
 from fastapi.responses import JSONResponse
 import firebase_admin
 from firebase_admin import credentials, firestore
-from transformers import pipeline
+from transformers import AutoTokenizer, AutoModelForSeq2SeqLM
 from pydantic import BaseModel
 import os
 from huggingface_hub import login
-from transformers import AutoTokenizer, AutoModelForSeq2SeqLM   # AutoModelForCausalLM,
+import traceback
+import logging
+
+# Configure logging
+logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s')
 
 # Load Firebase
-cred = credentials.Certificate("firebase_config.json")
-firebase_admin.initialize_app(cred)
-db = firestore.client()
+try:
+    cred = credentials.Certificate("firebase_config.json")
+    firebase_admin.initialize_app(cred)
+    db = firestore.client()
+    logging.info("Firebase initialized successfully.")
+except Exception as e:
+    logging.error(f"Error initializing Firebase: {e}")
+    db = None
 
 # Đăng nhập vào Hugging Face (nếu cần)
 HF_TOKEN = os.getenv("HF_TOKEN")
 if HF_TOKEN:
-    login(HF_TOKEN)
-
-# Load AI Model
-tokenizer = AutoTokenizer.from_pretrained("VietAI/vit5-large",token=HF_TOKEN)
-ai_model = AutoModelForSeq2SeqLM.from_pretrained("VietAI/vit5-large",token=HF_TOKEN)
+    try:
+        login(token=HF_TOKEN)  # Pass HF_TOKEN as keyword argument
+        logging.info("Hugging Face login successful.")
+    except Exception as e:
+        logging.error(f"Error logging into Hugging Face: {e}")
+
+# Initialize tokenizer and model outside the request handler
+tokenizer = None
+ai_model = None
+
+try:
+    # Load AI Model
+    tokenizer = AutoTokenizer.from_pretrained("VietAI/vit5-large")  # Removed token argument here
+    ai_model = AutoModelForSeq2SeqLM.from_pretrained("VietAI/vit5-large")  # Removed token argument here
+    logging.info("Tokenizer and model loaded successfully.")
+except Exception as e:
+    logging.error(f"Error loading tokenizer/model: {e}")
 
 app = FastAPI()
 
@@ -36,75 +57,91 @@ class FocusHistoryRequest(BaseModel):
 
 class AIRequest(BaseModel):
     user_id: str
+    bios: str
 
 # API cập nhật bios
 @app.post("/update_bios")
 async def update_bios(request: UpdateBiosRequest):
-    user_ref = db.collection("user_profiles").document(request.user_id)
-    user_ref.set({"bios": request.bios}, merge=True)
-    return {"message": "Cập nhật bios thành công"}
+    try:
+        user_ref = db.collection("user_profiles").document(request.user_id)
+        user_ref.set({"bios": request.bios}, merge=True)
+        return {"message": "Cập nhật bios thành công"}
+    except Exception as e:
+        logging.error(f"Error in /update_bios: {e}\n{traceback.format_exc()}")
+        raise HTTPException(status_code=500, detail=f"Error updating bios: {e}")
 
 # API thêm lịch sử focus
 @app.post("/add_focus_history")
 async def add_focus_history(request: FocusHistoryRequest):
-    user_ref = db.collection("user_profiles").document(request.user_id)
-    user_doc = user_ref.get()
-    data = user_doc.to_dict() or {}
-    focus_history = data.get("focus_history", [])
-    focus_history.append({"time_start": request.time_start, "total_time": request.total_time})
-    user_ref.set({"focus_history": focus_history}, merge=True)
-    return {"message": "Thêm lịch sử focus thành công"}
+    try:
+        user_ref = db.collection("user_profiles").document(request.user_id)
+        user_doc = user_ref.get()
+        data = user_doc.to_dict() or {}
+        focus_history = data.get("focus_history", [])
+        focus_history.append({"time_start": request.time_start, "total_time": request.total_time})
+        user_ref.set({"focus_history": focus_history}, merge=True)
+        return {"message": "Thêm lịch sử focus thành công"}
+    except Exception as e:
+        logging.error(f"Error in /add_focus_history: {e}\n{traceback.format_exc()}")
+        raise HTTPException(status_code=500, detail=f"Error adding focus history: {e}")
 
 # API lấy dữ liệu người dùng
 @app.get("/get_user_data")
 async def get_user_data(user_id: str):
-    user_doc = db.collection("user_profiles").document(user_id).get()
-    data = user_doc.to_dict() or {}
-    return {
-        "bios": data.get("bios", "Chưa có bios."),
-        "focus_history": data.get("focus_history", [])
-    }
+    try:
+        user_doc = db.collection("user_profiles").document(user_id).get()
+        data = user_doc.to_dict() or {}
+        return {
+            "bios": data.get("bios", "Chưa có bios."),
+            "focus_history": data.get("focus_history", [])
+        }
+    except Exception as e:
+        logging.error(f"Error in /get_user_data: {e}\n{traceback.format_exc()}")
+        raise HTTPException(status_code=500, detail=f"Error getting user data: {e}")
 
 # API AI tư vấn
 @app.post("/ai_personal_advice")
 async def ai_personal_advice(request: AIRequest):
-    user_doc = db.collection("user_profiles").document(request.user_id).get()
-    data = user_doc.to_dict() or {}
-    bios = data.get("bios", "Chưa có bios.")
-    focus_history = data.get("focus_history", [])
-    focus_text = "\n".join([f"- {f['time_start']}: {f['total_time']} phút" for f in focus_history])
-    
-    prompt = f"""
-    Thông tin người dùng:
-    - Bios: {bios}
-    - Lịch sử focus:
-    {focus_text}
-    
-    Hãy tư vấn cách cải thiện hiệu suất làm việc dựa trên thông tin trên.
-    """
-    input_ids = tokenizer(prompt, return_tensors="pt")
-    response = ai_model.generate(**input_ids, max_new_tokens=500)
-    # response = ai_model(prompt, max_length=200)
-    return {"advice": tokenizer.decode(response[0], skip_special_tokens=True)}
+    try:
+        if tokenizer is None or ai_model is None:
+            logging.error("Tokenizer or AI model not loaded.")
+            raise HTTPException(status_code=500, detail="Tokenizer or AI model not loaded.")
+
+        if db is None:
+            logging.error("Firebase not initialized.")
+            raise HTTPException(status_code=500, detail="Firebase not initialized.")
+
+        user_doc = db.collection("user_profiles").document(request.user_id).get()
+        if not user_doc.exists:
+            logging.warning(f"User profile not found for user_id: {request.user_id}")
+            return {"advice": "Không tìm thấy thông tin người dùng."}
+
+        data = user_doc.to_dict() or {}
+        bios = request.bios if request.bios else data.get("bios", "Chưa có bios.")
+        focus_history = data.get("focus_history", [])
+        focus_text = "\n".join([f"- {f['time_start']}: {f['total_time']} phút" for f in focus_history])
+
+        prompt = f"""
+        Thông tin người dùng:
+        - Bios: {bios}
+        - Lịch sử focus:
+        {focus_text}
+
+        Hãy tư vấn cách cải thiện hiệu suất làm việc dựa trên thông tin trên.
+        """
+        input_ids = tokenizer(prompt, return_tensors="pt", truncation=True, max_length=1024)
+        response = ai_model.generate(**input_ids, max_new_tokens=500)
+        advice = tokenizer.decode(response[0], skip_special_tokens=True)
+
+        return {"advice": advice}
+
+    except Exception as e:
+        error_message = f"Error in /ai_personal_advice: {e}"
+        logging.error(error_message)
+        logging.error(traceback.format_exc())
+        raise HTTPException(status_code=500, detail=error_message)
 
 # Trang chủ
 @app.get("/")
 async def home():
     return JSONResponse(content={"message": "Welcome to the Recommendation API!"})
-
-
-
-# from transformers import AutoTokenizer, AutoModelForSeq2SeqLM
-
-# # Kiểm tra biến môi trường HF_TOKEN
-# import os
-# HF_TOKEN = os.getenv("HF_TOKEN")
-# print("HF_TOKEN:", HF_TOKEN is not None)
-
-# # Load model
-# model_name = "google/flan-t5-small"
-# tokenizer = AutoTokenizer.from_pretrained(model_name, token=HF_TOKEN)
-# ai_model = AutoModelForSeq2SeqLM.from_pretrained(model_name, token=HF_TOKEN)
-
-# print("Model loaded successfully!")
-