api.py

from fastapi import FastAPI, HTTPException, BackgroundTasks
from fastapi.middleware.cors import CORSMiddleware
from pydantic import BaseModel
from typing import Dict, Any, Optional
import json
import os
from datetime import datetime
from huggingface_hub import HfApi, Repository
import hashlib
import uuid

# FastAPI app
app = FastAPI(title="ML Tracker API", version="1.0.0")

# CORS middleware
app.add_middleware(
    CORSMiddleware,
    allow_origins=["*"],
    allow_credentials=True,
    allow_methods=["*"],
    allow_headers=["*"],
)

# Pydantic models
class LogRequest(BaseModel):
    api_key: str
    experiment: str
    step: int
    metrics: Dict[str, Any]
    timestamp: Optional[float] = None
    config: Optional[Dict[str, Any]] = None

class ExperimentResponse(BaseModel):
    experiment: str
    total_steps: int
    created_at: str
    last_updated: str
    metrics: list

# In-memory storage (in production, use HF Hub or database)
experiments_db = {}
api_keys_db = {}

def verify_api_key(api_key: str) -> bool:
    """Verify if the API key is valid"""
    # In production, verify against HF Hub or database
    return True  # Simplified for demo

def get_user_from_api_key(api_key: str) -> str:
    """Get username from API key"""
    # In production, lookup from database
    return hashlib.sha256(api_key.encode()).hexdigest()[:8]

@app.get("/")
async def root():
    return {"message": "ML Tracker API - Free W&B Alternative", "version": "1.0.0"}

@app.post("/api/log")
async def log_metrics(request: LogRequest, background_tasks: BackgroundTasks):
    """Log metrics for an experiment"""
    try:
        # Verify API key
        if not verify_api_key(request.api_key):
            raise HTTPException(status_code=401, detail="Invalid API key")
        
        user_id = get_user_from_api_key(request.api_key)
        experiment_key = f"{user_id}_{request.experiment}"
        
        # Initialize experiment if not exists
        if experiment_key not in experiments_db:
            experiments_db[experiment_key] = {
                "experiment": request.experiment,
                "user_id": user_id,
                "created_at": datetime.now().isoformat(),
                "metrics": [],
                "config": request.config or {}
            }
        
        # Add metrics
        metric_entry = {
            "step": request.step,
            "timestamp": request.timestamp or datetime.now().timestamp(),
            "metrics": request.metrics
        }
        
        experiments_db[experiment_key]["metrics"].append(metric_entry)
        experiments_db[experiment_key]["last_updated"] = datetime.now().isoformat()
        
        # Update config if provided
        if request.config:
            experiments_db[experiment_key]["config"].update(request.config)
        
        # Background task to save to HF Hub (simplified)
        background_tasks.add_task(save_to_hub, experiment_key)
        
        return {
            "status": "success",
            "experiment": request.experiment,
            "step": request.step,
            "metrics_logged": len(request.metrics)
        }
    
    except Exception as e:
        raise HTTPException(status_code=500, detail=f"Error logging metrics: {str(e)}")

@app.get("/api/experiments")
async def get_experiments(api_key: str):
    """Get all experiments for a user"""
    try:
        if not verify_api_key(api_key):
            raise HTTPException(status_code=401, detail="Invalid API key")
        
        user_id = get_user_from_api_key(api_key)
        user_experiments = []
        
        for exp_key, exp_data in experiments_db.items():
            if exp_data["user_id"] == user_id:
                user_experiments.append({
                    "experiment": exp_data["experiment"],
                    "total_steps": len(exp_data["metrics"]),
                    "created_at": exp_data["created_at"],
                    "last_updated": exp_data.get("last_updated", exp_data["created_at"])
                })
        
        return {"experiments": user_experiments}
    
    except Exception as e:
        raise HTTPException(status_code=500, detail=f"Error fetching experiments: {str(e)}")

@app.get("/api/experiment/{experiment_name}")
async def get_experiment(experiment_name: str, api_key: str):
    """Get detailed data for a specific experiment"""
    try:
        if not verify_api_key(api_key):
            raise HTTPException(status_code=401, detail="Invalid API key")
        
        user_id = get_user_from_api_key(api_key)
        experiment_key = f"{user_id}_{experiment_name}"
        
        if experiment_key not in experiments_db:
            raise HTTPException(status_code=404, detail="Experiment not found")
        
        exp_data = experiments_db[experiment_key]
        
        return {
            "experiment": exp_data["experiment"],
            "created_at": exp_data["created_at"],
            "last_updated": exp_data.get("last_updated", exp_data["created_at"]),
            "total_steps": len(exp_data["metrics"]),
            "config": exp_data["config"],
            "metrics": exp_data["metrics"]
        }
    
    except Exception as e:
        raise HTTPException(status_code=500, detail=f"Error fetching experiment: {str(e)}")

@app.delete("/api/experiment/{experiment_name}")
async def delete_experiment(experiment_name: str, api_key: str):
    """Delete an experiment"""
    try:
        if not verify_api_key(api_key):
            raise HTTPException(status_code=401, detail="Invalid API key")
        
        user_id = get_user_from_api_key(api_key)
        experiment_key = f"{user_id}_{experiment_name}"
        
        if experiment_key not in experiments_db:
            raise HTTPException(status_code=404, detail="Experiment not found")
        
        del experiments_db[experiment_key]
        
        return {"status": "success", "message": f"Experiment '{experiment_name}' deleted"}
    
    except Exception as e:
        raise HTTPException(status_code=500, detail=f"Error deleting experiment: {str(e)}")

async def save_to_hub(experiment_key: str):
    """Save experiment data to HuggingFace Hub (background task)"""
    try:
        # In production, save to HF Hub datasets
        # For demo, we'll just log
        print(f"Saving experiment {experiment_key} to HF Hub...")
        
        # Example HF Hub integration:
        # api = HfApi(token=user_token)
        # repo_id = f"{username}/ml-tracker-data"
        # api.upload_file(
        #     path_or_fileobj=json.dumps(experiments_db[experiment_key]),
        #     path_in_repo=f"experiments/{experiment_key}.json",
        #     repo_id=repo_id,
        #     repo_type="dataset"
        # )
        
    except Exception as e:
        print(f"Error saving to hub: {str(e)}")

@app.get("/api/health")
async def health_check():
    """Health check endpoint"""
    return {
        "status": "healthy",
        "timestamp": datetime.now().isoformat(),
        "total_experiments": len(experiments_db)
    }

if __name__ == "__main__":
    import uvicorn
    uvicorn.run(app, host="0.0.0.0", port=8000)