main.py

import contextlib
from fastapi import FastAPI, Request, BackgroundTasks
from fastapi.middleware.cors import CORSMiddleware
from pydantic import BaseModel
import json
from huggingface_hub import HfApi
import os
from dotenv import load_dotenv
import json

from datetime import datetime
from pathlib import Path

from huggingface_hub import CommitScheduler

load_dotenv()

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

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=["*"],
)

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


class Vote(BaseModel):
    dataset: str
    description: str
    vote: int
    userID: str


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.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)
    return {"message": "Vote submitted successfully"}