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

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


@app.post("/vote")
async def receive_vote(
    vote: Vote,
    background_tasks: BackgroundTasks,
):
    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-Control-Allow-Headers"] = "*"
    return response