main.py

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

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):
    logger.info("Running startup event")
    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


app = FastAPI(lifespan=lifespan)
# # 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"})