app.py

from fastapi import FastAPI, UploadFile, File, HTTPException, Depends
from fastapi.responses import FileResponse
from fastapi.security import OAuth2PasswordBearer
from pathlib import Path
import psutil
import shutil
import os

app = FastAPI()

# Directory where files will be uploaded
UPLOAD_DIRECTORY = Path("uploads")
UPLOAD_DIRECTORY.mkdir(parents=True, exist_ok=True)

# OAuth2 scheme for token-based authorization
oauth2_scheme = OAuth2PasswordBearer(tokenUrl="token")

# Hugging Face secret token from environment variables
VALID_TOKEN = os.environ.get("SECRET_TOKEN")

# Constants to represent the limits for the container resources
MAX_CPU_CORES = 2  # vCPUs
MAX_MEMORY_GB = 16  # GB
MAX_DISK_GB = 50  # GB

# Helper function to calculate the size of a directory
def get_directory_size(directory: Path) -> int:
    total_size = 0
    for dirpath, dirnames, filenames in os.walk(directory):
        for f in filenames:
            fp = os.path.join(dirpath, f)
            total_size += os.path.getsize(fp)
    return total_size

# Dependency to verify the token
def get_current_token(token: str = Depends(oauth2_scheme)):
    if token != VALID_TOKEN:
        raise HTTPException(status_code=401, detail="Invalid token")
    return token

# Health check endpoint
@app.get("/health")
def health_check(token: str = Depends(get_current_token)):
    return {
        "status": "healthy"
    }

# System metrics endpoint: CPU, RAM, and disk usage for uploads folder relative to 2 vCPUs, 16GB RAM, and 50GB disk
@app.get("/metrics")
def get_metrics(token: str = Depends(get_current_token)):
    # CPU percentage relative to 2 vCPUs
    cpu_percent = round((psutil.cpu_percent(interval=1) / 100) * MAX_CPU_CORES, 2)

    # Memory usage relative to 16GB of RAM
    memory = psutil.virtual_memory()
    memory_used_gb = (memory.total - memory.available) / (1024 ** 3)

    # Make sure to cap the used memory at MAX_MEMORY_GB
    memory_used_gb_relative = min(memory_used_gb, MAX_MEMORY_GB)
    memory_usage_percent = (memory_used_gb_relative / MAX_MEMORY_GB) * 100

    # Disk usage for uploads directory relative to 50GB storage
    uploads_size_bytes = get_directory_size(UPLOAD_DIRECTORY)
    uploads_size_gb = uploads_size_bytes / (1024 ** 3)
    uploads_usage_percent = (uploads_size_gb / MAX_DISK_GB) * 100

    return {
        "cpu_usage_cores": cpu_percent,  # Relative CPU usage as cores (out of 2 vCPUs)
        "memory": {
            "used_gb": round(memory_used_gb_relative, 2),  # Capped at 16GB max
            "used_percent_of_16gb": round(memory_usage_percent, 2)
        },
        "disk": {
            "uploads_folder_size_gb": round(uploads_size_gb, 2),
            "uploads_usage_percent_of_50gb": round(uploads_usage_percent, 2)
        }
    }

# File upload endpoint
@app.post("/uploadfile/")
async def upload_file(file: UploadFile = File(...), token: str = Depends(get_current_token)):
    file_location = UPLOAD_DIRECTORY / file.filename

    with file_location.open("wb") as buffer:
        shutil.copyfileobj(file.file, buffer)

    return {"info": f"file '{file.filename}' saved at '{file_location}'"}

# File download endpoint
@app.get("/downloadfile/{filename}")
def download_file(filename: str, token: str = Depends(get_current_token)):
    file_location = UPLOAD_DIRECTORY / filename

    if not file_location.exists():
        raise HTTPException(status_code=404, detail="File not found")

    return FileResponse(path=file_location, filename=filename)

# Show current files endpoint
@app.get("/files/")
def list_files(token: str = Depends(get_current_token)):
    files = [f for f in os.listdir(UPLOAD_DIRECTORY) if os.path.isfile(UPLOAD_DIRECTORY / f)]
    return {"files": files}