gated_model_controller.py

"""
Manage Gated Access Requests for bbkdevops/Fiber-MoE-Symplectic-Gating-Research
Uses official huggingface_hub Gated Model API and HTTP Endpoints.
"""

from huggingface_hub import HfApi, get_token
import requests
import json

REPO_ID = "bbkdevops/Fiber-MoE-Symplectic-Gating-Research"
BASE_URL = f"https://huggingface.co/api/models/{REPO_ID}/user-access-request"

class GatedModelController:
    def __init__(self, repo_id: str = REPO_ID):
        self.repo_id = repo_id
        self.token = get_token()
        self.api = HfApi(token=self.token)

    def _headers(self):
        return {"Authorization": f"Bearer {self.token}"}

    def list_pending(self):
        """Retrieve list of pending access requests."""
        url = f"{BASE_URL}/pending"
        res = requests.get(url, headers=self._headers())
        if res.status_code == 200:
            return res.json()
        return {"status_code": res.status_code, "msg": res.text}

    def list_accepted(self):
        """Retrieve list of accepted access requests."""
        url = f"{BASE_URL}/accepted"
        res = requests.get(url, headers=self._headers())
        if res.status_code == 200:
            return res.json()
        return {"status_code": res.status_code, "msg": res.text}

    def grant_user(self, username: str):
        """Directly grant a user access to the model."""
        url = f"{BASE_URL}/grant"
        payload = {"user": username}
        res = requests.post(url, headers=self._headers(), json=payload)
        return {"status_code": res.status_code, "response": res.text}

    def handle_request(self, username: str, status: str = "accepted", reason: str = ""):
        """Handle request: status in ['accepted', 'rejected', 'pending', 'reset']."""
        url = f"{BASE_URL}/handle"
        payload = {"status": status, "user": username}
        if reason:
            if status == "rejected":
                payload["rejectionReason"] = reason[:200]
            elif status == "reset":
                payload["resetReason"] = reason[:200]
        res = requests.post(url, headers=self._headers(), json=payload)
        return {"status_code": res.status_code, "response": res.text}

if __name__ == "__main__":
    controller = GatedModelController()
    print("Checking Gated Model Requests for:", REPO_ID)
    pending = controller.list_pending()
    print("Pending Requests:", json.dumps(pending, indent=2))
    accepted = controller.list_accepted()
    print("Accepted Requests:", json.dumps(accepted, indent=2))