import json import logging import os from contextlib import asynccontextmanager from datetime import datetime from pathlib import Path from typing import Annotated from dotenv import load_dotenv from fastapi import BackgroundTasks, FastAPI, Header, HTTPException from fastapi.responses import JSONResponse from huggingface_hub import CommitScheduler, HfApi, hf_hub_download, whoami from huggingface_hub.utils._errors import HTTPError from pydantic import BaseModel, Field from starlette.responses import RedirectResponse load_dotenv() logger = logging.getLogger(__name__) app = FastAPI() VOTES_FILE = "votes/votes.jsonl" HF_TOKEN = os.getenv("HF_TOKEN") hf_api = HfApi(token=HF_TOKEN) scheduler = CommitScheduler( repo_id="davanstrien/summary-ratings", 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): if not Path("votes").exists(): Path("votes").mkdir() path = hf_hub_download( repo_id="davanstrien/summary-ratings", filename="data/votes.jsonl", repo_type="dataset", token=HF_TOKEN, local_dir="votes", local_dir_use_symlinks=False, ) logger.info(f"Downloaded votes.jsonl to {path}") yield # # 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: date_time = datetime.now().strftime("%Y-%m-%d %H:%M:%S") vote_entry["timestamp"] = date_time json.dump(vote_entry, file) file.write("\n") logger.info(f"Vote saved: {vote_entry}") @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)): try: whoami(token) return True except HTTPError: return False @app.post("/vote") async def receive_vote( vote: Vote, Authorization: Annotated[str, Header()], background_tasks: BackgroundTasks, ): if not validate_token(Authorization): logger.error("Invalid token") raise HTTPException(status_code=401, detail="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) return JSONResponse(content={"message": "Vote submitted successfully"})