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"}