api/app.py

#!/usr/bin/env python3
"""Minimal FastAPI app for the two website endpoints."""
from __future__ import annotations

import sys
from pathlib import Path

from fastapi import FastAPI, HTTPException
from pydantic import BaseModel, Field

API_DIR = Path(__file__).resolve().parent
sys.path.insert(0, str(API_DIR))

from common import load_env  # noqa: E402
from generate_rna_and_motif import generate_rna_and_motif  # noqa: E402
from score_binding import score_binding  # noqa: E402

load_env()
app = FastAPI(title="ProRiboGen API", version="0.1.0")


class GenerateRequest(BaseModel):
    protein: str = Field(..., description="氨基酸序列")
    p_id: str = "QUERY"
    num_sequences: int = 256
    length_bp: int = 80
    run_motif: bool = True


class ScoreRequest(BaseModel):
    protein: str
    rna: str
    p_id: str = "QUERY"


@app.get("/health")
def health() -> dict:
    return {"status": "ok"}


@app.post("/v1/generate")
def api_generate(req: GenerateRequest) -> dict:
    try:
        return generate_rna_and_motif(
            req.protein,
            p_id=req.p_id,
            num_sequences=req.num_sequences,
            length_bp=req.length_bp,
            run_motif=req.run_motif,
        )
    except Exception as exc:  # noqa: BLE001
        raise HTTPException(status_code=500, detail=str(exc)) from exc


@app.post("/v1/score")
def api_score(req: ScoreRequest) -> dict:
    try:
        return score_binding(req.protein, req.rna, p_id=req.p_id)
    except Exception as exc:  # noqa: BLE001
        raise HTTPException(status_code=500, detail=str(exc)) from exc