tracklight_server/api/artifacts.py

# tracklight_server/api/artifacts.py

from fastapi import APIRouter, Depends, UploadFile, File
from fastapi.responses import FileResponse
from .auth import verify_token
import os
import shutil

router = APIRouter()

# Get the absolute path to the artifacts directory
SERVER_DIR = os.path.dirname(os.path.dirname(__file__))
ARTIFACTS_DIR = os.path.join(SERVER_DIR, "artifacts")

# Create artifacts directory if it doesn't exist
os.makedirs(ARTIFACTS_DIR, exist_ok=True)

@router.post("/upload/{run_id}")
async def upload_artifact(run_id: str, file: UploadFile = File(...), token: str = Depends(verify_token)):
    """
    Uploads an artifact file for a given run.
    """
    run_artifact_dir = os.path.join(ARTIFACTS_DIR, run_id)
    os.makedirs(run_artifact_dir, exist_ok=True)
    
    file_path = os.path.join(run_artifact_dir, file.filename)
    
    with open(file_path, "wb") as buffer:
        shutil.copyfileobj(file.file, buffer)
        
    return {"filename": file.filename, "run_id": run_id}

@router.get("/download/{run_id}/{filename}")
async def download_artifact(run_id: str, filename: str, token: str = Depends(verify_token)):
    """
    Downloads an artifact file for a given run.
    """
    file_path = os.path.join(ARTIFACTS_DIR, run_id, filename)
    
    if not os.path.exists(file_path):
        return {"error": "File not found"}
        
    return FileResponse(file_path)