Add authentication to endpoints

auth/authentication.py

@@ -0,0 +1,35 @@
+from fastapi.security import OAuth2PasswordBearer
+from fastapi import HTTPException, Depends
+from jose import JWTError, jwt
+from datetime import datetime, timedelta
+
+from svc.schemas import User
+
+
+SECRET_KEY = "llmbenchmark_tr" # your secret key
+ALGORITHM = "HS256"
+ACCESS_TOKEN_EXPIRE_MINUTES = 30
+
+oauth2_scheme = OAuth2PasswordBearer(tokenUrl="token")
+
+def create_access_token(data: dict):
+    to_encode = data.copy()
+    expire = datetime.now() + timedelta(minutes=ACCESS_TOKEN_EXPIRE_MINUTES)
+    to_encode.update({"exp": expire})
+    encoded_jwt = jwt.encode(to_encode, SECRET_KEY, algorithm=ALGORITHM)
+    return encoded_jwt
+
+def get_current_user(token: str = Depends(oauth2_scheme)):
+    credentials_exception = HTTPException(
+        status_code=401,
+        detail="Could not validate credentials",
+        headers={"WWW-Authenticate": "Bearer"},
+    )
+    try:
+        payload = jwt.decode(token, SECRET_KEY, algorithms=[ALGORITHM])
+        username: str = payload.get("sub")
+        if username is None:
+            raise credentials_exception
+        return username
+    except JWTError:
+        raise credentials_exception

requirements.txt

@@ -1,3 +1,4 @@
 fastapi
 uvicorn[standard]
-lm_eval
+lm_eval
+jose

svc/router.py

@@ -2,17 +2,35 @@ from fastapi import APIRouter, HTTPException
 import logging
 
 from lm_eval import evaluator, utils
-from svc.schemas import LMHarnessTaskRequest, LMHarnessTaskResponse
+from svc.schemas import LMHarnessTaskRequest, LMHarnessTaskResponse, OAuth2PasswordRequestForm, User
+from fastapi import FastAPI, Depends, HTTPException
+from auth.authentication import oauth2_scheme, get_current_user, create_access_token
+from data.users import users_db
 
 router = APIRouter()
 
 logging.basicConfig(level=logging.INFO, format="%(asctime)s - %(levelname)s - %(message)s")
 logger = logging.getLogger(__name__)
 
+from dotenv import load_dotenv
+import os
+
+load_dotenv()
+@router.post("/token")
+async def login_for_access_token(form_data: OAuth2PasswordRequestForm = Depends()):
+    if os.getenv("HF_TOKEN") != form_data.hf_token:
+        raise HTTPException(status_code=400, detail="Incorrect username or password")
+    access_token = create_access_token(data={"sub": form_data.username})
+    return {"access_token": access_token, "token_type": "bearer"}
+
+
+@router.get("/protected")
+async def protected_route(username: str = Depends(get_current_user)):
+    return {"message": f"Hello, {username}! This is a protected resource."}
 
 
 @router.post("/chat",  response_model=LMHarnessTaskResponse)
-def inference_model(request: LMHarnessTaskRequest):
+def inference_model(request: LMHarnessTaskRequest = Depends(get_current_user)):
     try:
         results = evaluator.simple_evaluate(
             model=request.model,

svc/schemas.py

@@ -1,6 +1,19 @@
 from pydantic import BaseModel
 from typing import List, Optional, Union, Any
 
+
+from pydantic import BaseModel
+
+
+class OAuth2PasswordRequestForm(BaseModel):
+    username: str
+    hf_token: str
+
+class User(BaseModel):
+    username: str
+    hf_token: str
+
+
 class LMHarnessTaskRequest(BaseModel):
     model: str
     model_args: Optional[Union[str, dict]] = None