hf_client.py

from __future__ import annotations

from dataclasses import dataclass
from datetime import datetime
from pathlib import Path
from typing import Any

import pandas as pd
from huggingface_hub import HfApi, hf_hub_download
from huggingface_hub.utils import EntryNotFoundError, HfHubHTTPError

from config import AppConfig


README_CANDIDATES = ("README.md", "readme.md")


@dataclass(frozen=True)
class RepoRevision:
    repo_id: str
    revision: str
    committed_at: datetime | None


class HubClient:
    def __init__(self, config: AppConfig):
        self.config = config
        self.api = HfApi(token=config.hf_token)

    def discover_dataset_repos(self) -> list[str]:
        repos = set(self.config.repo_allowlist)

        if self.config.discovery_enabled:
            try:
                datasets = self.api.list_datasets(
                    author=self.config.hf_owner,
                    search=self.config.repo_id_contains or None,
                    limit=max(self.config.max_repos * 3, self.config.max_repos),
                    full=False,
                )
                for dataset in datasets:
                    repo_id = getattr(dataset, "id", None)
                    if not repo_id:
                        continue
                    if self._repo_matches_filters(repo_id):
                        repos.add(repo_id)
                    if len(repos) >= self.config.max_repos:
                        break
            except HfHubHTTPError:
                pass

        filtered = [repo for repo in repos if self._repo_matches_filters(repo)]
        return sorted(filtered)[: self.config.max_repos]

    def list_dataset_commits(self, repo_id: str, max_commits: int) -> list[RepoRevision]:
        revisions: list[RepoRevision] = []
        try:
            commits = self.api.list_repo_commits(repo_id=repo_id, repo_type="dataset")
            for commit in commits[:max_commits]:
                revision = getattr(commit, "commit_id", None) or getattr(commit, "oid", None)
                if not revision:
                    continue
                committed_at = _coerce_datetime(getattr(commit, "created_at", None) or getattr(commit, "date", None))
                revisions.append(
                    RepoRevision(
                        repo_id=repo_id,
                        revision=revision,
                        committed_at=committed_at,
                    )
                )
        except HfHubHTTPError:
            revisions = []

        if revisions:
            return revisions

        try:
            info = self.api.dataset_info(repo_id=repo_id)
            fallback_revision = getattr(info, "sha", None) or "main"
            fallback_timestamp = _coerce_datetime(getattr(info, "lastModified", None))
        except HfHubHTTPError:
            fallback_revision = "main"
            fallback_timestamp = None

        return [RepoRevision(repo_id=repo_id, revision=fallback_revision, committed_at=fallback_timestamp)]

    def fetch_dataset_readme(self, repo_id: str, revision: str | None = None) -> str | None:
        for filename in README_CANDIDATES:
            try:
                path = hf_hub_download(
                    repo_id=repo_id,
                    filename=filename,
                    repo_type="dataset",
                    revision=revision,
                    token=self.config.hf_token,
                )
                return Path(path).read_text(encoding="utf-8", errors="replace")
            except EntryNotFoundError:
                continue
            except HfHubHTTPError:
                continue
        return None

    def _repo_matches_filters(self, repo_id: str) -> bool:
        lowered_repo_id = repo_id.lower()

        if self.config.repo_id_contains and self.config.repo_id_contains.lower() not in lowered_repo_id:
            return False

        if self.config.repo_prefixes and not any(repo_id.startswith(prefix) for prefix in self.config.repo_prefixes):
            return False

        return True


def _coerce_datetime(value: Any) -> datetime | None:
    if value is None:
        return None

    timestamp = pd.to_datetime(value, errors="coerce", utc=True)
    if pd.isna(timestamp):
        return None
    return timestamp.to_pydatetime()