app.py

"""
This is the main entry point for the FastAPI application.
The app handles the request to predict toxicity for a list of SMILES strings.
"""

# ---------------------------------------------------------------------------------------
# Dependencies and global variable definition
import os
from typing import List, Dict, Optional
from fastapi import FastAPI, Header, HTTPException
from pydantic import BaseModel, Field

from predict import predict as predict_func

API_KEY = os.getenv("API_KEY")  # set via Space Secrets


# ---------------------------------------------------------------------------------------
class Request(BaseModel):
    smiles: List[str] = Field(min_items=1, max_items=1000)


class Response(BaseModel):
    predictions: dict
    model_info: Dict[str, str] = {}


app = FastAPI(title="toxicity-api")


@app.get("/")
def root():
    return {
        "message": "Toxicity Prediction API",
        "endpoints": {
            "/metadata": "GET - API metadata and capabilities",
            "/healthz": "GET - Health check",
            "/predict": "POST - Predict toxicity for SMILES",
        },
        "usage": "Send POST to /predict with {'smiles': ['your_smiles_here']} and Authorization header",
    }


@app.get("/metadata")
def metadata():
    return {
        "name": "Tox21_SNN",
        "version": "1.0.0",
        "max_batch_size": 256,
        "tox_endpoints": [
            "NR-AR",
            "NR-AR-LBD",
            "NR-AhR",
            "NR-Aromatase",
            "NR-ER",
            "NR-ER-LBD",
            "NR-PPAR-gamma",
            "SR-ARE",
            "SR-ATAD5",
            "SR-HSE",
            "SR-MMP",
            "SR-p53",
        ],
    }


@app.get("/healthz")
def healthz():
    return {"ok": True}


@app.post("/predict", response_model=Response)
def predict(request: Request):
    predictions = predict_func(request.smiles)
    return {
        "predictions": predictions,
        "model_info": {"name": "Tox21_SNN", "version": "1.0.0"},
    }