snap2sim/backend.py

"""Backend client for Modal inference with explicit local sample mode."""

from __future__ import annotations

import base64
import os
from dataclasses import dataclass
from io import BytesIO
from typing import Any

import requests
from PIL import Image

from snap2sim.schema import EXAMPLE_ANALYSIS, select_render_mode, validate_analysis


@dataclass(frozen=True)
class Settings:
    backend: str = os.getenv("INFERENCE_BACKEND", "modal")
    analyze_url: str = os.getenv("MODAL_ANALYZE_URL", "")
    generate_url: str = os.getenv("MODAL_GENERATE_URL", "")
    api_token: str = os.getenv("SNAP2SIM_API_TOKEN", "")
    timeout_seconds: int = int(os.getenv("INFERENCE_TIMEOUT_SECONDS", "300"))


def encode_image(image: Image.Image) -> str:
    buffer = BytesIO()
    image.convert("RGB").save(buffer, format="JPEG", quality=92)
    return base64.b64encode(buffer.getvalue()).decode("ascii")


class InferenceClient:
    def __init__(self, settings: Settings | None = None) -> None:
        self.settings = settings or Settings()

    def analyze_image(self, image: Image.Image | None) -> dict[str, Any]:
        if self.settings.backend == "modal":
            if image is None:
                raise ValueError("Upload an image before analysis.")
            return self._post_json(
                self.settings.analyze_url,
                {"image_base64": encode_image(image)},
            )

        return validate_analysis(dict(EXAMPLE_ANALYSIS))

    def generate_scene(self, analysis: dict[str, Any], threshold: Any = None) -> dict[str, Any]:
        valid_analysis = validate_analysis(analysis)
        render_mode = select_render_mode(valid_analysis, threshold)
        renderer = "three" if render_mode == "three" else "photo"
        return {"renderer": renderer, "render_mode": render_mode, "analysis": valid_analysis}

    def _post_json(self, url: str, payload: dict[str, Any]) -> dict[str, Any]:
        if not url:
            raise RuntimeError("Modal backend selected but endpoint URL is not configured.")
        if not self.settings.api_token:
            raise RuntimeError("Modal backend selected but SNAP2SIM_API_TOKEN is not configured.")
        response = requests.post(
            url,
            json=payload,
            headers={"Authorization": f"Bearer {self.settings.api_token}"},
            timeout=self.settings.timeout_seconds,
        )
        response.raise_for_status()
        data = response.json()
        if not isinstance(data, dict):
            raise RuntimeError("Inference backend returned a non-object JSON payload.")
        return data