backend/hf_logging.py

# backend/hf_logging.py

import json
import io
import os
from datetime import datetime
from typing import Dict, Any

from huggingface_hub import HfApi, CommitOperationAdd

from .config import RATINGS_DATASET_ID

# Hugging Face API for logging - get token from environment or use default auth
hf_token = os.getenv("HF_TOKEN") or os.getenv("HUGGING_FACE_HUB_TOKEN")
hf_api = HfApi(token=hf_token)


def push_session_to_hub(session_id: str, export_data: Dict[str, Any]):
    """
    Save one session's ratings to the HF dataset as sessions/<session_id>.json
    """
    # Add human-readable timestamps and summary to session metadata
    if "session_metadata" in export_data:
        metadata = export_data["session_metadata"]
        if "created_at" in metadata and metadata["created_at"]:
            metadata["created_at_readable"] = datetime.fromtimestamp(metadata["created_at"]).isoformat()
        if "exported_at" in metadata and metadata["exported_at"]:
            metadata["exported_at_readable"] = datetime.fromtimestamp(metadata["exported_at"]).isoformat()
    
    # Add a summary section for easy viewing
    export_data["summary"] = {
        "session_id": session_id,
        "evaluation_completed": export_data.get("session_metadata", {}).get("completed", False),
        "total_mos_ratings": len(export_data.get("mos_ratings", [])),
        "total_ab_comparisons": len(export_data.get("ab_comparisons", [])),
        "models_evaluated": list(set(
            [r.get("model") for r in export_data.get("mos_ratings", [])] +
            [r.get("clip_a_model") for r in export_data.get("ab_comparisons", [])] +
            [r.get("clip_b_model") for r in export_data.get("ab_comparisons", [])]
        )),
        "evaluation_date": export_data.get("session_metadata", {}).get("created_at_readable", ""),
    }
    
    # Add readable timestamps to individual responses
    for response_list in [export_data.get("mos_ratings", []), export_data.get("ab_comparisons", [])]:
        for response in response_list:
            if "response_timestamp" in response and response["response_timestamp"]:
                response["response_timestamp_readable"] = datetime.fromtimestamp(response["response_timestamp"]).isoformat()
    
    # Serialize to bytes with nice formatting
    json_bytes = json.dumps(export_data, indent=2, default=str, ensure_ascii=False).encode("utf-8")
    fileobj = io.BytesIO(json_bytes)

    # Path inside the dataset repo
    path_in_repo = f"sessions/{session_id}.json"

    # Create a git commit adding/updating that file
    try:
        hf_api.create_commit(
            repo_id=RATINGS_DATASET_ID,
            repo_type="dataset",
            operations=[
                CommitOperationAdd(
                    path_in_repo=path_in_repo,
                    path_or_fileobj=fileobj,
                )
            ],
            commit_message=f"Add MOS ratings for session {session_id}",
            token=hf_token,  # Explicitly pass token
        )
    except Exception as e:
        print(f"[ERROR] Failed to push to Hub: {e}")
        # If no token or authentication fails, save locally as fallback
        fallback_path = f"./session_exports/{session_id}.json"
        os.makedirs("./session_exports", exist_ok=True)
        with open(fallback_path, "w") as f:
            json.dump(export_data, f, indent=2, default=str, ensure_ascii=False)
        print(f"[INFO] Saved session data locally to {fallback_path}")
        return

    print(f"[LOG] Pushed session {session_id} to {RATINGS_DATASET_ID}/{path_in_repo}")