bioai/web/api.py

"""bioai.web.api -- FastAPI backend for the Biopesticide-AI web UI.

Replaces the Gradio UI with a clean REST API + static frontend architecture.

Endpoints:
    GET  /                -> serve static index.html
    GET  /api/health      -> {"status": "ok"}
    GET  /api/status      -> backend status (model, LLM, device, species counts)
    POST /api/design      -> run the design pipeline, return JSON results
    GET  /static/*        -> static assets (CSS, JS, images)

The frontend (static/index.html) calls these endpoints via fetch() and
renders the UI with vanilla JS. No build step required.
"""

from __future__ import annotations

import argparse
import json
import sys
import time
from pathlib import Path

# Make sure we can import bioai from anywhere
_PROJECT_ROOT = Path(__file__).resolve().parents[2]
if str(_PROJECT_ROOT) not in sys.path:
    sys.path.insert(0, str(_PROJECT_ROOT))

from fastapi import FastAPI, HTTPException
from fastapi.responses import HTMLResponse, JSONResponse, FileResponse
from fastapi.staticfiles import StaticFiles
from pydantic import BaseModel

# NOTE: BiopesticideOrchestrator is imported lazily inside get_orchestrator()
# to avoid loading torch + models at module import time. This keeps the
# FastAPI app startup fast (<2 seconds) so HuggingFace doesn't kill the
# container before uvicorn binds to port 7860.
from bioai.sequence_utils import SAFETY_SPECIES, PEST_SPECIES

# ─────────────────────────────────────────────────────────────────────────────
# Singleton orchestrator
# ─────────────────────────────────────────────────────────────────────────────
_ORCHESTRATOR = None  # lazily initialized in get_orchestrator()


def get_orchestrator():
    """Lazily initialize the orchestrator on first API call.

    This avoids loading torch + models at import time, keeping FastAPI
    startup fast enough for HuggingFace's 60-second container timeout.
    """
    global _ORCHESTRATOR
    if _ORCHESTRATOR is None:
        print("[api] initializing orchestrator (first request)...", flush=True)
        from bioai.orchestrator import BiopesticideOrchestrator
        _ORCHESTRATOR = BiopesticideOrchestrator()
        print(f"[api] backend = {type(_ORCHESTRATOR.ranker.sirna_model).__name__}, degraded_mode = {_ORCHESTRATOR.degraded_mode}", flush=True)
    return _ORCHESTRATOR


# ─────────────────────────────────────────────────────────────────────────────
# Pydantic models
# ─────────────────────────────────────────────────────────────────────────────
class DesignRequest(BaseModel):
    user_text: str
    top_k: int = 10
    pest_species: str | None = None  # if provided, skip LLM parsing


class SimulateRequest(BaseModel):
    candidates: list
    pest_species: str = "unknown"


# ─────────────────────────────────────────────────────────────────────────────
# FastAPI app
# ─────────────────────────────────────────────────────────────────────────────
STATIC_DIR = Path(__file__).parent / "static"

app = FastAPI(
    title="Biopesticide-AI",
    description="dsRNA biopesticide design pipeline API",
    version="2.0.0",
)

# Mount static files (CSS, JS, images)
app.mount("/static", StaticFiles(directory=str(STATIC_DIR)), name="static")


@app.get("/api/health")
async def health():
    return {"status": "ok"}


@app.get("/api/status")
async def status():
    """Return backend status for the frontend status strip."""
    orch = get_orchestrator()
    model_class = type(orch.ranker.sirna_model).__name__
    # Check if checkpoints actually loaded (not random init)
    from bioai.paths import SIRNA_CHECKPOINT, PINN_CHECKPOINT
    sirna_loaded = SIRNA_CHECKPOINT.exists()
    pinn_loaded = PINN_CHECKPOINT.exists()
    return {
        "model_backend": "Caduceus-Ph-1" if model_class == "CaduceusAdapter" else "Dilated CNN (HyenaDNA-inspired)",
        "llm_status": "Ollama Llama 3.2 3B (local)" if not orch.degraded_mode else "Degraded mode (Ollama not running)",
        "llm_active": not orch.degraded_mode,
        "device": str(orch.ranker.device),
        "safety_species_count": len(SAFETY_SPECIES),
        "pest_species_count": len(PEST_SPECIES),
        "safety_species": SAFETY_SPECIES,
        "pest_species": PEST_SPECIES,
        "checkpoints": {
            "sirna_loaded": sirna_loaded,
            "pinn_loaded": pinn_loaded,
            "sirna_path": str(SIRNA_CHECKPOINT),
            "pinn_path": str(PINN_CHECKPOINT),
        },
    }


@app.post("/api/design")
async def design(req: DesignRequest):
    """Run the design pipeline and return results as JSON."""
    if not req.user_text or not req.user_text.strip():
        raise HTTPException(status_code=400, detail="user_text is required")

    orch = get_orchestrator()
    t0 = time.time()
    try:
        result = orch.design(req.user_text, top_k=int(req.top_k), pest_species_override=req.pest_species)
    except Exception as e:
        import traceback
        raise HTTPException(
            status_code=500,
            detail=f"Pipeline error: {type(e).__name__}: {e}\n{traceback.format_exc()}"
        )
    elapsed = time.time() - t0

    # Build response with metadata
    response = {
        "elapsed_seconds": round(elapsed, 2),
        "pest_report": result.get("pest_report", {}),
        "candidates": result.get("candidates", []),
        "safety_cards": result.get("safety_cards", ""),
        "regulatory_memo": result.get("regulatory_memo", ""),
        "n_transcripts": result.get("n_transcripts", 0),
        "n_precursors": result.get("n_precursors", 0),
        "n_sirnas": result.get("n_sirnas", 0),
        "total_cost_estimate": result.get("total_cost_estimate", 0.0),
        "degraded_mode": result.get("degraded_mode", True),
    }
    return JSONResponse(content=response)


@app.get("/")
async def index():
    """Serve the main HTML page."""
    index_path = STATIC_DIR / "index.html"
    if not index_path.exists():
        raise HTTPException(status_code=404, detail="index.html not found")
    return HTMLResponse(content=index_path.read_text(encoding="utf-8"), status_code=200)


@app.post("/api/simulate")
async def simulate(req: SimulateRequest):
    """Run the virtual wet-lab simulation on the top candidates.

    Each candidate should have: sirna_seq, efficacy, half_life_hours.
    Returns 1000-trial Monte Carlo results per candidate.
    """
    from bioai.simulation.wet_lab import WetLabSimulator, result_to_dict

    if not req.candidates:
        raise HTTPException(status_code=400, detail="candidates is required")

    sim = WetLabSimulator(n_trials=1000, rng_seed=42)
    results = sim.simulate_batch(req.candidates, pest_species=req.pest_species)
    return JSONResponse(content={
        "trials": [result_to_dict(r) for r in results],
        "pest_species": req.pest_species,
        "n_trials_per_candidate": 1000,
        "stages": [
            "delivery", "uptake", "dicer", "risc_loading", "cleavage", "phenotype"
        ],
    })


# ─────────────────────────────────────────────────────────────────────────────
# Entry point
# ─────────────────────────────────────────────────────────────────────────────
def main():
    parser = argparse.ArgumentParser(description="Biopesticide-AI Web UI (FastAPI)")
    parser.add_argument("--host", default="0.0.0.0", help="bind host (default 0.0.0.0)")
    parser.add_argument("--port", type=int, default=7860, help="bind port (default 7860)")
    parser.add_argument("--reload", action="store_true", help="auto-reload on file changes (dev mode)")
    args = parser.parse_args()

    import uvicorn

    # DON'T pre-initialize the orchestrator here. On HuggingFace Spaces, the
    # 14-species k-mer indexing + Ollama probe takes >60 seconds on the slow
    # CPU, which causes HF to kill the container before uvicorn binds to port
    # 7860. Instead, we start uvicorn immediately and initialize the
    # orchestrator lazily on the first API request. The first request will
    # be slow (~60s) but the server stays alive.
    print("[api] launching uvicorn immediately (orchestrator initializes lazily on first request)")
    print(f"[api] launching on http://{args.host}:{args.port}")
    uvicorn.run(
        "bioai.web.api:app",
        host=args.host,
        port=args.port,
        reload=args.reload,
        log_level="info",
    )


if __name__ == "__main__":
    main()