app.py

from fastapi import FastAPI, HTTPException, Response
from pydantic import BaseModel
import gradio as gr
from rdkit.Chem import Draw, rdChemReactions

ERROR_STYLE = "color:#b91c1c;font-weight:600;text-align:center;"


class ReactionRequest(BaseModel):
    smiles: str


class ReactionResponse(BaseModel):
    svg: str
    html: str


def _error_message(message: str) -> str:
    """Format errors consistently for both UI and API users."""
    return f"<p style='{ERROR_STYLE}'>{message}</p>"


def _load_reaction(smiles_reaction: str):
    smiles = (smiles_reaction or "").strip()
    if not smiles:
        raise ValueError("Please provide a SMILES reaction string.")

    try:
        rxn = rdChemReactions.ReactionFromSmarts(smiles, useSmiles=True)
    except Exception as exc:
        raise ValueError(f"Unable to parse reaction string: {exc}") from exc

    if rxn is None or (
        rxn.GetNumReactantTemplates() == 0 and rxn.GetNumProductTemplates() == 0
    ):
        raise ValueError(
            "RDKit could not interpret the reaction. Verify the SMILES syntax."
        )
    return rxn


def _svg_string(rxn) -> str:
    try:
        svg = Draw.ReactionToImage(
            rxn,
            subImgSize=(260, 220),
            useSVG=True,
            highlightByReactant=False,
            continuousHighlight=False,
        )
    except TypeError:
        svg = Draw.ReactionToImage(
            rxn,
            subImgSize=(260, 220),
            useSVG=True,
        )
    return svg.decode("utf-8") if isinstance(svg, (bytes, bytearray)) else str(svg)


def _html_wrapper(svg_string: str) -> str:
    return f"<div style='display:flex;justify-content:center;'>{svg_string}</div>"


def smiles_to_reaction_image(smiles_reaction: str) -> str:
    """
    Convert a reaction SMILES/SMARTS string to an SVG reaction image.

    The returned HTML embeds the raw SVG so API callers receive SVG markup
    directly, while the UI can render it inline inside the browser.
    """
    try:
        rxn = _load_reaction(smiles_reaction)
        svg = _svg_string(rxn)
        return _html_wrapper(svg)
    except ValueError as exc:
        return _error_message(str(exc))


def _render_svg_payload(smiles: str) -> ReactionResponse:
    rxn = _load_reaction(smiles)
    svg = _svg_string(rxn)
    html = _html_wrapper(svg)
    return ReactionResponse(svg=svg, html=html)


fastapi_app = FastAPI(
    title="SMILES → Reaction SVG API",
    version="1.0.0",
    description="Single-step REST endpoint that turns reaction SMILES into SVG using RDKit.",
)


@fastapi_app.post("/api/render", response_model=ReactionResponse)
def render_reaction(payload: ReactionRequest):
    try:
        return _render_svg_payload(payload.smiles)
    except ValueError as exc:
        raise HTTPException(status_code=400, detail=str(exc)) from exc


@fastapi_app.post("/api/render/svg")
def render_reaction_svg(payload: ReactionRequest):
    """
    Return the raw SVG with an image/svg+xml content-type so consumers can
    treat the response as an image (e.g., display directly in browsers or
    write to disk without extra parsing).
    """
    try:
        svg = _svg_string(_load_reaction(payload.smiles))
    except ValueError as exc:
        raise HTTPException(status_code=400, detail=str(exc)) from exc

    return Response(content=svg, media_type="image/svg+xml")


demo = gr.Interface(
    fn=smiles_to_reaction_image,
    inputs=gr.Textbox(
        label="SMILES Reaction String",
        placeholder="Enter a reaction SMILES such as C=O.CC>>C=OCC",
        lines=3,
    ),
    outputs=gr.HTML(label="Reaction SVG"),
    title="SMILES ➜ Reaction SVG",
    description=(
        "Submit a reaction SMILES/SMARTS string and receive the rendered reaction as SVG. "
        "This interface can also be called programmatically; the returned payload is the SVG markup."
    ),
    examples=[
        ["CCO>>CC=O"],
        ["CCO.C=O>>CCOC=O"],
        ["O=C=O.O>>O=C(O)O"],
    ],
    flagging_mode="never",
    api_name="render",
)

app = gr.mount_gradio_app(fastapi_app, demo, path="/")

if __name__ == "__main__":
    import uvicorn

    uvicorn.run(app, host="0.0.0.0", port=7860)