app.py

import os
import shutil
from fastapi import FastAPI
from typing import Dict, Any

app = FastAPI(title="Hugging Face Space Monitor", version="1.0.0")

@app.get("/")
def greet_json():
    return {"Hello": "World!"}

@app.get("/storage")
def get_storage_info() -> Dict[str, Any]:
    """取得目前 Space 的存儲使用情況"""
    try:
        # 取得根目錄的磁碟使用情況
        total, used, free = shutil.disk_usage("/")
        
        # 取得當前工作目錄大小
        app_size = get_directory_size("/app")
        
        # 取得 /tmp 目錄使用情況（臨時檔案）
        tmp_size = get_directory_size("/tmp")
        
        return {
            "disk_usage": {
                "total_bytes": total,
                "used_bytes": used,
                "free_bytes": free,
                "total_gb": round(total / (1024**3), 2),
                "used_gb": round(used / (1024**3), 2),
                "free_gb": round(free / (1024**3), 2),
                "usage_percentage": round((used / total) * 100, 2)
            },
            "application": {
                "app_size_bytes": app_size,
                "app_size_mb": round(app_size / (1024**2), 2),
                "tmp_size_bytes": tmp_size,
                "tmp_size_mb": round(tmp_size / (1024**2), 2)
            },
            "environment": {
                "current_dir": os.getcwd(),
                "space_id": os.getenv("SPACE_ID", "unknown"),
                "space_author": os.getenv("SPACE_AUTHOR_NAME", "unknown")
            }
        }
    except Exception as e:
        return {"error": str(e)}

def get_directory_size(path: str) -> int:
    """計算目錄大小"""
    total_size = 0
    try:
        for dirpath, dirnames, filenames in os.walk(path):
            for filename in filenames:
                filepath = os.path.join(dirpath, filename)
                try:
                    total_size += os.path.getsize(filepath)
                except (OSError, FileNotFoundError):
                    continue
    except (OSError, FileNotFoundError):
        pass
    return total_size

@app.get("/health")
def health_check():
    """健康檢查端點"""
    return {
        "status": "healthy",
        "message": "Space is running normally"
    }