Refactor code and update dependencies

main.py

@@ -1,17 +1,17 @@
 import contextlib
-from fastapi import FastAPI, Request, BackgroundTasks
-from fastapi.middleware.cors import CORSMiddleware
-from pydantic import BaseModel
 import json
-from huggingface_hub import HfApi
 import os
-from dotenv import load_dotenv
-import json
 from contextlib import asynccontextmanager
 from datetime import datetime
 from pathlib import Path
+
+from dotenv import load_dotenv
+from fastapi import BackgroundTasks, FastAPI, Header, Request
+from fastapi.middleware.cors import CORSMiddleware
+from huggingface_hub import CommitScheduler, HfApi, whoami
+from huggingface_hub.utils import HTTPError
+from pydantic import BaseModel
 from starlette.responses import RedirectResponse
-from huggingface_hub import CommitScheduler
 
 load_dotenv()
 
@@ -52,13 +52,6 @@ app.add_middleware(
 )
 
 
-class Vote(BaseModel):
-    dataset: str
-    description: str
-    vote: int
-    userID: str
-
-
 def save_vote(vote_entry):
     with open(VOTES_FILE, "a") as file:
         # add time stamp to the vote entry
@@ -73,8 +66,30 @@ def root():
     return RedirectResponse(url="/docs")
 
 
+class Vote(BaseModel):
+    dataset: str
+    description: str
+    vote: int
+    userID: str
+
+
+def validate_token(token: str):
+    try:
+        user = whoami(token=token)
+        return True
+    except HTTPError as e:
+        # check for HTTPError: Invalid user token. If you didn't pass a user token, make sure you are properly logged in by executing `huggingface-cli login`, and if you did pass a user token, double-check it's correct. in the error message
+        if "Invalid user token" in str(e):
+            return False
+
+
 @app.post("/vote")
-async def receive_vote(vote: Vote, background_tasks: BackgroundTasks):
+async def receive_vote(
+    vote: Vote, background_tasks: BackgroundTasks, authorization: str = Header(...)
+):
+    token = authorization.split(" ")[1]
+    if not validate_token(token):
+        return {"message": "Invalid user token"}
     vote_entry = {
         "dataset": vote.dataset,
         "vote": vote.vote,

requirements.in

@@ -2,5 +2,4 @@ fastapi
 fastapi-cors
 huggingface_hub
 python-dotenv
-toolz
 uvicorn[standard]

requirements.txt

@@ -8,10 +8,16 @@ anyio==4.2.0
     #   watchfiles
 certifi==2023.11.17
     # via requests
+cffi==1.16.0
+    # via cryptography
 charset-normalizer==3.3.2
     # via requests
 click==8.1.7
     # via uvicorn
+cryptography==42.0.5
+    # via python-jose
+ecdsa==0.18.0
+    # via python-jose
 environs==11.0.0
     # via fastapi-cors
 fastapi==0.109.0
@@ -36,6 +42,12 @@ packaging==23.2
     # via
     #   huggingface-hub
     #   marshmallow
+pyasn1==0.6.0
+    # via
+    #   python-jose
+    #   rsa
+pycparser==2.22
+    # via cffi
 pydantic==2.5.3
     # via fastapi
 pydantic-core==2.14.6
@@ -44,17 +56,21 @@ python-dotenv==1.0.1
     # via
     #   environs
     #   uvicorn
+python-jose==3.3.0
 pyyaml==6.0.1
     # via
     #   huggingface-hub
     #   uvicorn
 requests==2.31.0
     # via huggingface-hub
+rsa==4.9
+    # via python-jose
+six==1.16.0
+    # via ecdsa
 sniffio==1.3.0
     # via anyio
 starlette==0.35.1
     # via fastapi
-toolz==0.12.0
 tqdm==4.66.1
     # via huggingface-hub
 typing-extensions==4.9.0