Update auth.py

auth.py

@@ -14,6 +14,7 @@ from jose import JWTError, jwt
 from passlib.context import CryptContext
 from pymongo import MongoClient
 import gridfs
+from bson import ObjectId  # Ensure ObjectId is imported
 
 from models import User, UserUpdate, Token, LoginResponse
 from config import CONNECTION_STRING, SECRET_KEY, ACCESS_TOKEN_EXPIRE_MINUTES, REFRESH_TOKEN_EXPIRE_DAYS
@@ -101,6 +102,7 @@ async def signup(
     name: str = Form(...),
     email: str = Form(...),
     password: str = Form(...),
+    role: str = Form(...),  # <-- MODIFICATION: Added role
     avatar: Optional[UploadFile] = File(None)
 ):
     try:
@@ -111,18 +113,24 @@ async def signup(
     if get_user(email):
         logger.warning(f"Attempt to register already existing email: {email}")
         raise HTTPException(status_code=400, detail="Email already registered")
+    
     hashed_password = get_password_hash(password)
+    
     user_data = {
         "name": name,
         "email": email,
         "hashed_password": hashed_password,
+        "role": role,  # <-- MODIFICATION: Added role to user data
         "chat_histories": []
     }
+    
     if avatar:
         file_id = await save_avatar_file_to_gridfs(avatar)
         user_data["avatar"] = file_id
+        
     users_collection.insert_one(user_data)
-    logger.info(f"New user registered: {email}")
+    logger.info(f"New user registered: {email} with role: {role}")
+    
     return {
         "access_token": create_access_token(email),
         "refresh_token": create_refresh_token(email),
@@ -135,16 +143,20 @@ async def login(request: Request, form_data: OAuth2PasswordRequestForm = Depends
     if not user:
         logger.warning(f"Failed login attempt for: {form_data.username}")
         raise HTTPException(status_code=401, detail="Incorrect username or password")
+    
     logger.info(f"User logged in: {user['email']}")
+    
     avatar_url = None
     if "avatar" in user and user["avatar"]:
         avatar_url = f"/auth/avatar/{user['avatar']}"
+        
     return {
         "access_token": create_access_token(user["email"]),
         "refresh_token": create_refresh_token(user["email"]),
         "token_type": "bearer",
         "name": user["name"],
-        "avatar": avatar_url
+        "avatar": avatar_url,
+        "role": user.get("role", "user")  # <-- MODIFICATION: Return role (default to "user")
     }
 
 @router.get("/user/data")
@@ -152,10 +164,12 @@ async def get_user_data(request: Request, current_user: dict = Depends(get_curre
     avatar_url = None
     if "avatar" in current_user and current_user["avatar"]:
         avatar_url = f"/auth/avatar/{current_user['avatar']}"
+        
     return {
         "name": current_user["name"],
         "email": current_user["email"],
         "avatar": avatar_url,
+        "role": current_user.get("role", "user"),  # <-- MODIFICATION: Return role (default to "user")
         "chat_histories": current_user.get("chat_histories", [])
     }
 
@@ -180,14 +194,18 @@ async def update_user(
             logger.error(f"Password validation error during update: {e}")
             raise HTTPException(status_code=400, detail=str(e))
         update_data["hashed_password"] = get_password_hash(password)
+        
     if avatar:
         file_id = await save_avatar_file_to_gridfs(avatar)
         update_data["avatar"] = file_id
+        
     if not update_data:
         logger.info("No update parameters provided")
         raise HTTPException(status_code=400, detail="No update parameters provided")
+        
     users_collection.update_one({"email": current_user["email"]}, {"$set": update_data})
     logger.info(f"User updated: {current_user['email']}")
+    
     return {"message": "User updated successfully"}
 
 @router.post("/logout")
@@ -195,8 +213,6 @@ async def logout(request: Request, current_user: dict = Depends(get_current_user
     logger.info(f"User logged out: {current_user['email']}")
     return {"message": "User logged out successfully"}
 
-from bson import ObjectId
-
 @router.get("/avatar/{file_id}")
 async def get_avatar(file_id: str):
     try:
@@ -205,4 +221,4 @@ async def get_avatar(file_id: str):
         return StreamingResponse(file, media_type=file.content_type)
     except Exception as e:
         logger.error(f"Avatar not found for file_id {file_id}: {e}")
-        raise HTTPException(status_code=404, detail="Avatar not found")
+        raise HTTPException(status_code=404, detail="Avatar not found")