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 from fastapi.middleware.cors import CORSMiddleware from fastapi.responses import JSONResponse from huggingface_hub import CommitScheduler, HfApi, whoami from huggingface_hub.utils._errors import HTTPError from pydantic import BaseModel, Field from starlette.responses import RedirectResponse from typing import Annotated from fastapi import FastAPI, Header 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, token=HF_TOKEN, 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", "chrome-extension://ogbhjlfpmjgjbjoiffagjogbhgaipopf", # Replace with your Chrome plugin ID ] 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 = Field(..., ge=-1, le=1) userID: str def validate_token(token: str = Header(None)): return token == "ABCD" @app.post("/vote") async def receive_vote( vote: Vote, Authorization: Annotated[str, Header()], background_tasks: BackgroundTasks, ): print(f"Received vote: {vote}") print(Authorization) if not validate_token(Authorization): return JSONResponse(status_code=401, content={"message": "Invalid 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) response = JSONResponse(content={"message": "Vote submitted successfully"}) response.headers[ "Access-Control-Allow-Origin" ] = "chrome-extension://ogbhjlfpmjgjbjoiffagjogbhgaipopf" # TODO Replace with your Chrome plugin ID response.headers["Access-Control-Allow-Methods"] = "POST" response.headers["Access-Contr`ol-Allow-Headers"] = "*" return response