main.py

import asyncio
import json
import logging
import os
import re
from contextlib import asynccontextmanager
from datetime import datetime
from pathlib import Path
from typing import Annotated, List

from cashews import NOT_NONE, cache
from dotenv import load_dotenv
from fastapi import BackgroundTasks, FastAPI, Header, HTTPException
from fastapi.middleware.cors import CORSMiddleware
from fastapi.responses import JSONResponse
from httpx import AsyncClient
from huggingface_hub import CommitScheduler, DatasetCard, HfApi, hf_hub_download, whoami
from huggingface_hub.utils import disable_progress_bars, logging
from huggingface_hub.utils._errors import HTTPError
from langfuse.openai import AsyncOpenAI  # OpenAI integration
from pydantic import BaseModel, Field
from starlette.responses import RedirectResponse
from card_processing import parse_markdown, try_load_text, is_empty_template

disable_progress_bars()
load_dotenv()
logger = logging.get_logger(__name__)

Gb = 1073741824
cache.setup("disk://", size_limit=16 * Gb)  # configure as in-memory cache

VOTES_FILE = "data/votes.jsonl"

HF_TOKEN = os.getenv("HF_TOKEN")

hf_api = HfApi(token=HF_TOKEN)
async_httpx_client = AsyncClient()
scheduler = CommitScheduler(
    repo_id="davanstrien/summary-ratings",
    repo_type="dataset",
    folder_path="data",
    path_in_repo="data",
    every=5,
    token=HF_TOKEN,
    hf_api=hf_api,
)


@asynccontextmanager
async def lifespan(app: FastAPI):
    logger.info("Running startup event")
    if not Path(VOTES_FILE).exists():
        path = hf_hub_download(
            repo_id="davanstrien/summary-ratings",
            filename="data/votes.jsonl",
            repo_type="dataset",
            token=HF_TOKEN,
            local_dir=".",
            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://deckahggoiaphiebdipfbiinmaihfpbk",  # Replace with your Chrome plugin ID
]


# Configure CORS settings
app.add_middleware(
    CORSMiddleware,
    allow_origins=[
        "https://huggingface.co/datasets/*"
    ],  # Update with your frontend URL
    allow_credentials=True,
    allow_methods=["*"],
    allow_headers=["*"],
)


def save_vote(vote_entry):
    with scheduler.lock:
        with open(VOTES_FILE, "a") as file:
            date_time = datetime.now().strftime("%Y-%m-%d %H:%M:%S")
            vote_entry["timestamp"] = date_time
            file.write(json.dumps(vote_entry) + "\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)) -> bool:
    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"})


def format_prompt(card: str) -> str:
    return f"""
Write a tl;dr summary of a dataset based on the dataset card. Focus on the most critical aspects of the dataset. 
The summary should aim to concisely describe the dataset.

CARD: \n\n{card[:6000]}
---

\n\nInstructions:
If the card provides the necessary information, say what the dataset can be used for.
You do not need to mention that the dataset is hosted or available on the Hugging Face Hub.
Do not mention the license of the dataset.
Do not mention the number of examples in the training or test split.
Only mention size if there is extensive discussion of the scale of the dataset in the dataset card.
Do not speculate on anything not explicitly mentioned in the dataset card.
In general avoid references to the quality of the dataset i.e. don't use phrases like 'a high-quality dataset' in the summary.

\n\nOne sentence summary:"""


async def check_when_dataset_last_modified(dataset_id: str) -> datetime | None:
    try:
        response = await async_httpx_client.get(
            f"https://huggingface.co/api/datasets/{dataset_id}"
        )
        if last_modified := response.json().get("lastModified"):
            return datetime.fromisoformat(last_modified)
        return None
    except Exception as e:
        logger.error(e)
        return None


@cache(ttl="48h", condition=NOT_NONE, key="predict:{dataset_id}")
async def predict(card: str, dataset_id: str) -> str | None:
    try:
        prompt = format_prompt(card)
        client = AsyncOpenAI(
            base_url="https://api-inference.huggingface.co/models/mistralai/Mixtral-8x7B-Instruct-v0.1/v1",
            api_key=HF_TOKEN,
        )

        chat_completion = await client.chat.completions.create(
            model="tgi",
            messages=[
                {"role": "user", "content": prompt},
            ],
            stream=False,
            tags=["tldr-summaries"],
        )
        return chat_completion.choices[0].message.content.strip()
    except Exception as e:
        logger.error(e)
        return None


@app.get("/summary")
async def get_summary(dataset_id: str) -> str | None:
    """
    Get a summary for a dataset based on the provided dataset ID.

    Args:
        dataset_id (str): The ID of the dataset to retrieve the summary for.

    Returns:
        str | None: The generated summary for the dataset, or None if no summary is available or an error occurs."""

    try:
        # dataset_id = request.dataset_id
        card_text = await async_httpx_client.get(
            f"https://huggingface.co/datasets/{dataset_id}/raw/main/README.md"
        )
        card_text = card_text.text
        card = DatasetCard(card_text)
        text = card.text
        parsed_text = parse_markdown(text)
        if is_empty_template(parsed_text):
            return None
        cache_key = f"predict:{dataset_id}"
        cached_data = await cache.get(cache_key)

        if cached_data is not None:
            cached_summary, cached_last_modified_time = cached_data
            # Get the current last modified time of the dataset
            current_last_modified_time = await check_when_dataset_last_modified(
                dataset_id
            )

            if (
                current_last_modified_time is None
                or cached_last_modified_time >= current_last_modified_time
            ):
                # Use the cached summary if the cached last modified time is greater than or equal to the current last modified time
                logger.info("Using cached summary")
                return cached_summary
        summary = await predict(parsed_text, dataset_id)
        current_last_modified_time = await check_when_dataset_last_modified(dataset_id)
        await cache.set(cache_key, (summary, current_last_modified_time))
        return summary
    except Exception as e:
        logger.error(e)
        return None


class SummariesRequest(BaseModel):
    dataset_ids: List[str]


@cache(ttl="1h", condition=NOT_NONE)
@app.post("/summaries")
async def get_summaries(request: SummariesRequest) -> dict:
    """
    Get summaries for a list of datasets based on the provided dataset IDs.

    Args:
        dataset_ids (List[str]): A list of dataset IDs to retrieve the summaries for.

    Returns:
        dict: A dictionary mapping dataset IDs to their corresponding summaries.
    """
    dataset_ids = request.dataset_ids

    async def get_summary_wrapper(dataset_id):
        return dataset_id, await get_summary(dataset_id)

    summary_tasks = [get_summary_wrapper(dataset_id) for dataset_id in dataset_ids]
    summaries = dict(await asyncio.gather(*summary_tasks))
    for dataset_id in dataset_ids:
        if summaries[dataset_id] is None:
            summaries[dataset_id] = "No summary available"
    return summaries