utils/hub_storage.py

from pathlib import Path
from huggingface_hub import HfApi
from .token_handler import TokenHandler

class HubStorage:
    def __init__(self, repo_id):
        self.token_handler = TokenHandler()
        self.token = self.token_handler.get_verified_token()
        self.api = HfApi()
        self.repo_id = repo_id

    def save_to_hub(self, file_content, path_in_repo, commit_message):
        """Save file to HuggingFace Hub"""
        return self.api.upload_file(
            path_or_fileobj=file_content,
            path_in_repo=path_in_repo,
            repo_id=self.repo_id,
            repo_type="space",
            commit_message=commit_message,
            token=self.token
        )

    def load_from_hub(self, path_in_repo):
        """Load file from HuggingFace Hub"""
        return self.api.hf_hub_download(
            repo_id=self.repo_id,
            filename=path_in_repo,
            repo_type="space",
            token=self.token
        )
    
    def get_file_content(self, file_path):
        """
        Get content of a file from the repository
        """
        try:
            api = HfApi()
            content = api.file_download(
                repo_id=self.repo_id,
                repo_type="space",
                filename=file_path
            )
            return content.read().decode('utf-8')
        except Exception as e:
            print(f"Error reading file {file_path}: {str(e)}")
            return None