first auth support

Author: aliasaria

api.py

@@ -50,6 +50,7 @@
     batched_prompts,
     recipes,
     remote,
+    auth2,
 )
 import torch
 
@@ -81,15 +82,7 @@
 
 from dotenv import load_dotenv
 
-from transformerlab.models.users import (
-    fastapi_users,
-    auth_backend,
-    current_active_user,
-    UserRead,
-    UserCreate,
-    UserUpdate,
-)
-from transformerlab.routers.test_users import router as users_router
+
 from transformerlab.shared.models.user_model import create_db_and_tables, User
 
 load_dotenv()
@@ -241,25 +234,7 @@ async def validation_exception_handler(request, exc):
 app.include_router(batched_prompts.router)
 app.include_router(remote.router)
 app.include_router(fastchat_openai_api.router)
-
-# Include Auth and Registration Routers
-app.include_router(
-    fastapi_users.get_auth_router(auth_backend),
-    prefix="/auth/jwt",
-    tags=["auth"],
-)
-app.include_router(
-    fastapi_users.get_register_router(UserRead, UserCreate),
-    prefix="/auth",
-    tags=["auth"],
-)
-# Include User Management Router (allows authenticated users to view/update their profile)
-app.include_router(
-    fastapi_users.get_users_router(UserRead, UserUpdate),
-    prefix="/users",
-    tags=["users"],
-)
-app.include_router(users_router)
+app.include_router(auth2.router)
 
 # Authentication and session management routes
 if os.getenv("TFL_MULTITENANT") == "true":

transformerlab/models/users.py

@@ -1,12 +1,14 @@
 # users.py
 import uuid
-from typing import Optional, AsyncGenerator
+from typing import Optional
 from fastapi import Depends, Request
 from fastapi_users import BaseUserManager, FastAPIUsers, UUIDIDMixin, schemas
 from fastapi_users.authentication import AuthenticationBackend, BearerTransport, JWTStrategy
 from fastapi_users.db import SQLAlchemyUserDatabase
 from transformerlab.shared.models.user_model import User, get_async_session
 from sqlalchemy.ext.asyncio import AsyncSession
+from jose import jwt as _jose_jwt
+from datetime import datetime, timedelta
 
 
 # --- Pydantic Schemas for API interactions ---
@@ -25,6 +27,8 @@ class UserUpdate(schemas.BaseUserUpdate):
 
 # --- User Manager (Handles registration, password reset, etc.) ---
 SECRET = "YOUR_STRONG_SECRET"  # !! CHANGE THIS IN PRODUCTION !!
+REFRESH_SECRET = "YOUR_REFRESH_TOKEN_SECRET"  # !! USE A DIFFERENT SECRET !!
+REFRESH_LIFETIME = 60 * 60 * 24 * 7  # 7 days
 
 
 class UserManager(UUIDIDMixin, BaseUserManager[User, uuid.UUID]):
@@ -68,3 +72,67 @@ def get_jwt_strategy() -> JWTStrategy:
 # --- Dependency for Protected Routes ---
 # This is what you'll use in your route decorators
 current_active_user = fastapi_users.current_user(active=True)
+
+
+def get_refresh_strategy() -> JWTStrategy:
+    return JWTStrategy(secret=REFRESH_SECRET, lifetime_seconds=REFRESH_LIFETIME)
+
+
+# --- Small helper to create access + refresh tokens for manual flows (e.g. refresh endpoint) ---
+
+
+class _JWTAuthenticationHelper:
+    """Minimal helper that mirrors a login response (access + refresh token).
+
+    We keep this small and explicit so callers (like the `refresh` endpoint in
+    `routers/auth.py`) can create new access tokens when given a valid
+    refresh token.
+    """
+
+    def __init__(
+        self,
+        access_secret: str,
+        refresh_secret: str,
+        access_lifetime: int = 3600,
+        refresh_lifetime: int = REFRESH_LIFETIME,
+    ):
+        self.access_secret = access_secret
+        self.refresh_secret = refresh_secret
+        self.access_lifetime = access_lifetime
+        self.refresh_lifetime = refresh_lifetime
+
+    def _create_token(self, user, secret: str, lifetime_seconds: int) -> str:
+        now = datetime.utcnow()
+        exp = now + timedelta(seconds=lifetime_seconds)
+        payload = {
+            "sub": str(user.id),
+            "email": getattr(user, "email", None),
+            "exp": int(exp.timestamp()),
+        }
+        return _jose_jwt.encode(payload, secret, algorithm="HS256")
+
+    def get_login_response(self, user) -> dict:
+        """Return a dict similar to what FastAPI-Users returns on login.
+
+        Keys:
+        - access_token: short-lived JWT
+        - refresh_token: long-lived JWT (can be validated with refresh strategy)
+        - token_type: 'bearer'
+        - expires_in: seconds until access token expiry
+        """
+        access = self._create_token(user, self.access_secret, self.access_lifetime)
+        refresh = self._create_token(user, self.refresh_secret, self.refresh_lifetime)
+        return {
+            "access_token": access,
+            "refresh_token": refresh,
+            "token_type": "bearer",
+            "expires_in": self.access_lifetime,
+        }
+
+
+# Module-level helpers for imports elsewhere
+jwt_authentication = _JWTAuthenticationHelper(
+    SECRET, REFRESH_SECRET, access_lifetime=3600, refresh_lifetime=REFRESH_LIFETIME
+)
+access_strategy = get_jwt_strategy()
+refresh_strategy = get_refresh_strategy()

transformerlab/routers/auth2.py

@@ -0,0 +1,88 @@
+from fastapi import APIRouter, Depends, HTTPException
+from transformerlab.shared.models.user_model import User
+from transformerlab.models.users import (
+    fastapi_users,
+    auth_backend,
+    current_active_user,
+    UserRead,
+    UserCreate,
+    UserUpdate,
+    get_user_manager,
+    get_refresh_strategy,
+    jwt_authentication,
+)
+
+from jose import jwt, JWTError
+
+router = APIRouter(tags=["users"])
+
+
+# Include Auth and Registration Routers
+router.include_router(
+    fastapi_users.get_auth_router(auth_backend),
+    prefix="/auth/jwt",
+    tags=["auth"],
+)
+router.include_router(
+    fastapi_users.get_register_router(UserRead, UserCreate),
+    prefix="/auth",
+    tags=["auth"],
+)
+# Include User Management Router (allows authenticated users to view/update their profile)
+router.include_router(
+    fastapi_users.get_users_router(UserRead, UserUpdate),
+    prefix="/users",
+    tags=["users"],
+)
+
+
+@router.get("/test-users/authenticated-route")
+async def authenticated_route(user: User = Depends(current_active_user)):
+    return {"message": f"Hello, {user.email}! You are authenticated."}
+
+
+# To test this, register a new user via /auth/register
+# curl -X POST 'http://127.0.0.1:8338/auth/register' \
+#  -H 'Content-Type: application/json' \
+#  -d '{
+#    "email": "test@example.com",
+#    "password": "password123"
+# }'
+
+# Then
+# curl -X POST 'http://127.0.0.1:8338/auth/jwt/login' \
+#  -H 'Content-Type: application/x-www-form-urlencoded' \
+#  -d 'username=test@example.com&password=password123'
+
+# This will return a token, which you can use to access the authenticated route:
+# curl -X GET 'http://127.0.0.1:8338/authenticated-route' \
+#  -H 'Authorization: Bearer <YOUR_ACCESS_TOKEN>'
+
+
+@router.post("/auth/refresh")
+async def refresh_access_token(
+    refresh_token: str,  # Sent by the client in the request body
+    user_manager=Depends(get_user_manager),
+):
+    try:
+        # 1. Decode and Validate the Refresh Token
+        # Get a fresh refresh strategy instance and use its secret to decode
+        refresh_strategy = get_refresh_strategy()
+        payload = jwt.decode(refresh_token, str(refresh_strategy.secret), algorithms=["HS256"])
+        user_id = payload.get("sub")
+
+        if user_id is None:
+            raise HTTPException(status_code=401, detail="Invalid refresh token payload")
+
+        # 2. Get the user object from the database
+        user = await user_manager.get(user_id)
+        if user is None or not user.is_active:
+            raise HTTPException(status_code=401, detail="User inactive or not found")
+
+        # 3. Create a NEW Access Token (using the short-lived strategy from the main JWT)
+        new_access_token = jwt_authentication.get_login_response(user)  # Needs custom helper
+
+        return {"access_token": new_access_token["access_token"], "token_type": "bearer"}
+
+    except JWTError:
+        raise HTTPException(status_code=401, detail="Expired or invalid refresh token")