app.py

from fastapi import FastAPI, HTTPException
from pydantic import BaseModel
from typing import List
import uvicorn
import logging
import os

# Import the fixed predict function from your predict.py
from predict import predict as model_predict

logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)

app = FastAPI(
    title="Tox21 Toxicity Prediction API",
    version="2.1.0"
)

class PredictionRequest(BaseModel):
    smiles: List[str]

@app.get("/")
def read_root():
    return {"status": "running", "service": "Tox21 Classifier"}

@app.post("/predict")
def predict(request: PredictionRequest):
    try:
        # We call the wrapper function in predict.py 
        # which already formats everything into {"predictions": [[...]]}
        results = model_predict(request.smiles)
        return results
    except Exception as e:
        logger.error(f"Prediction error: {e}")
        raise HTTPException(status_code=500, detail=str(e))

@app.get("/metadata")
def get_metadata():
    # Importing TASKS here to show the leaderboard what we support
    from data import TASKS
    return {
        "model": "D-MPNN Ensemble",
        "tasks": TASKS,
        "version": "2.1.0"
    }

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