main.py

import contextlib
import json
import os
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

load_dotenv()

HF_TOKEN = os.getenv("HF_TOKEN")
hf_api = HfApi(token=HF_TOKEN)

scheduler = CommitScheduler(
    repo_id="davanstrien/votes",
    repo_type="dataset",
    folder_path="votes",
    path_in_repo="data",
    every=1,
    hf_api=hf_api,
)


@asynccontextmanager
async def lifespan(app: FastAPI):
    Path("votes").mkdir(exist_ok=True)

    yield


app = FastAPI()
VOTES_FILE = "votes/votes.jsonl"
# Configure CORS
origins = [
    "https://huggingface.co",
]


app.add_middleware(
    CORSMiddleware,
    allow_origins=origins,
    allow_credentials=True,
    allow_methods=["POST"],
    allow_headers=["*"],
)


def save_vote(vote_entry):
    with open(VOTES_FILE, "a") as file:
        # add time stamp to the vote entry
        date_time = datetime.now().strftime("%Y-%m-%d %H:%M:%S")
        vote_entry["timestamp"] = date_time
        json.dump(vote_entry, file)
        file.write("\n")


@app.get("/", include_in_schema=False)
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, 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,
        "description": vote.description,
        "userID": vote.userID,
    }
    # Append the vote entry to the JSONL file
    background_tasks.add_task(save_vote, vote_entry)
    return {"message": "Vote submitted successfully"}