utils/api_client.py

import os
import requests
import tempfile
import streamlit as st
from huggingface_hub import (
    HfApi,
    login,
    create_repo,
    delete_repo,
    upload_file,
    HfFolder,
)
from huggingface_hub.utils import RepositoryNotFoundError, RevisionNotFoundError


class HuggingFaceClient:
    def __init__(self, token=None):
        self.token = token
        self.api = HfApi(token=token)

    def authenticate(self, token):
        """Authenticate with Hugging Face API using token"""
        self.token = token
        self.api = HfApi(token=token)
        try:
            login(token=token)
            whoami = self.api.whoami()
            return True, whoami
        except Exception as e:
            return False, str(e)

    def get_user_models(self):
        """Get all models created by the logged-in user"""
        try:
            # First try to get username from whoami API call
            whoami = self.api.whoami()
            username = whoami.get("name")

            # Fallback to the HF_USERNAME secret if available
            if not username and os.environ.get("HF_USERNAME"):
                username = os.environ.get("HF_USERNAME")

            # Get all models for this user using the list_models API
            user_models = list(self.api.list_models(author=username))
            return user_models
        except Exception as e:
            st.error(f"Error fetching models: {str(e)}")
            return []

    def get_model_info(self, repo_id):
        """Get detailed information about a specific model"""
        try:
            model_info = self.api.model_info(repo_id)
            return model_info
        except RepositoryNotFoundError:
            st.error(f"Repository {repo_id} not found")
            return None
        except Exception as e:
            st.error(f"Error fetching model info: {str(e)}")
            return None

    def create_model_repository(
        self, repo_name, is_private=False, exist_ok=False, repo_type="model"
    ):
        """Create a new model repository on Hugging Face"""
        try:
            response = create_repo(
                repo_id=repo_name,
                token=self.token,
                private=is_private,
                exist_ok=exist_ok,
                repo_type=repo_type,
            )
            return True, response
        except Exception as e:
            return False, str(e)

    def delete_model_repository(self, repo_id):
        """Delete a model repository from Hugging Face"""
        try:
            response = delete_repo(repo_id=repo_id, token=self.token)
            return True, "Repository deleted successfully"
        except Exception as e:
            return False, str(e)

    def upload_model_files(self, repo_id, files, commit_message="Upload model files"):
        """Upload model files to a repository"""
        try:
            uploaded_files = []
            for file_path, file_content in files.items():
                with tempfile.NamedTemporaryFile(delete=False) as temp_file:
                    temp_file.write(file_content)
                    temp_file_path = temp_file.name

                upload_response = upload_file(
                    path_or_fileobj=temp_file_path,
                    path_in_repo=file_path,
                    repo_id=repo_id,
                    token=self.token,
                    commit_message=commit_message,
                )
                uploaded_files.append(upload_response)

                # Clean up temporary file
                os.unlink(temp_file_path)

            return True, uploaded_files
        except Exception as e:
            return False, str(e)

    def update_model_card(self, repo_id, model_card_content):
        """Update the README.md (model card) of a repository"""
        try:
            with tempfile.NamedTemporaryFile(delete=False, mode="w") as temp_file:
                temp_file.write(model_card_content)
                temp_file_path = temp_file.name

            upload_response = upload_file(
                path_or_fileobj=temp_file_path,
                path_in_repo="README.md",
                repo_id=repo_id,
                token=self.token,
                commit_message="Update model card",
            )

            # Clean up temporary file
            os.unlink(temp_file_path)

            return True, upload_response
        except Exception as e:
            return False, str(e)

    def get_model_tags(self):
        """Get available model tags from Hugging Face Hub"""
        try:
            # This is a simplified version; in a real app, you'd fetch actual tags from the HF API
            tags = [
                "text-classification",
                "token-classification",
                "question-answering",
                "translation",
                "summarization",
                "text-generation",
                "fill-mask",
                "conversational",
                "image-classification",
                "object-detection",
                "audio-classification",
                "automatic-speech-recognition",
            ]
            return tags
        except Exception as e:
            st.error(f"Error fetching tags: {str(e)}")
            return []