satquery_space_bridge.py

"""Deployment-only adapter between Gradio file inputs and SatQuery's agent API.

This module intentionally delegates every scientific decision to the existing
``satquery_agent.api.agent_image_query`` implementation.  It only translates
transport types and exposes generated preview files to Gradio.
"""

from __future__ import annotations

import asyncio
import mimetypes
import os
from pathlib import Path
from typing import Any, Iterable

from fastapi import HTTPException
from starlette.datastructures import Headers, UploadFile


APP_ROOT = Path(__file__).resolve().parent
CACHE_ROOT = APP_ROOT / ".cache"


def configure_environment() -> None:
    """Set safe Space defaults before importing SatQuery model singletons."""

    cache_root = CACHE_ROOT.resolve()
    os.environ.setdefault("HF_HOME", str(cache_root / "huggingface"))
    os.environ.setdefault("HUGGINGFACE_HUB_CACHE", str(cache_root / "huggingface" / "hub"))
    os.environ.setdefault("TRANSFORMERS_CACHE", str(cache_root / "huggingface" / "hub"))
    os.environ.setdefault("SATQUERY_MODEL_CACHE", str(cache_root / "public-models"))
    os.environ.setdefault("SATQUERY_WORKSPACE_ROOT", str(APP_ROOT))
    os.environ.setdefault("SATQUERY_PREVIEW_DIR", str(cache_root / "previews"))
    os.environ.setdefault("SATQUERY_ARTIFACT_DIR", str(cache_root / "artifacts"))
    os.environ.setdefault("SATQUERY_GROUNDER_LOCAL_MODEL", str(APP_ROOT / "models" / "grounding_dino_tiny"))
    os.environ.setdefault("SATQUERY_GROUNDER_LOCAL_FILES_ONLY", "true")
    os.environ.setdefault("SATQUERY_GROUNDER_ALLOW_DOWNLOAD", "false")
    os.environ.setdefault("SATQUERY_GROUNDER_LOAD_ON_STARTUP", "false")
    os.environ.setdefault("SVE_LOAD_ON_STARTUP", "false")
    os.environ.setdefault("RSVQA_LOAD_ON_STARTUP", "false")
    os.environ.setdefault("TTP_ENABLED", "false")
    os.environ.setdefault("SATQUERY_CHANGEREX_ENABLED", "true")
    os.environ.setdefault("SATQUERY_CHANGE_ENGINE", "changerex")
    os.environ.setdefault("SATQUERY_CHANGEREX_DEVICE", "cpu")


configure_environment()

from satquery_agent.api import agent_health, agent_image_query  # noqa: E402
from satquery_agent.image_ingestion import preview_file_path  # noqa: E402
from satquery_agent.models import (  # noqa: E402
    ImageModality,
    InputMode,
    Modality,
    ObservationRole,
)


_BACKEND: Any | None = None


def _model_dump(value: Any) -> dict[str, Any]:
    if hasattr(value, "model_dump"):
        return value.model_dump(mode="json")
    return value.dict()


def preload_models() -> None:
    """Load the six private artifacts and public GPU models once per replica.

    ZeroGPU emulates CUDA during module initialization, which is the supported
    place to construct model singletons and move their weights to CUDA.  No
    inference or startup smoke request is performed here.
    """

    global _BACKEND
    if _BACKEND is not None:
        return

    import backend
    from satquery_agent.services.sve_service import get_sve_service
    from satquery_agent.specialists.captioner import get_captioner
    from satquery_agent.specialists.grounder import get_grounder
    from satquery_agent.specialists.rsvqa_specialist import get_rsvqa_specialist

    # Importing backend constructs and registers the three SAR translation
    # models. The remaining managers are process-wide singletons.
    _BACKEND = backend
    if any(
        model is None
        for model in (backend.PIX2PIX_MODEL, backend.SARFUSIONFORMER_MODEL, backend.COLOR_CORRECTOR)
    ):
        raise RuntimeError("One or more verified SAR translation artifacts could not be loaded.")
    get_sve_service().load()
    get_rsvqa_specialist().load()
    grounder = get_grounder()
    grounder.load()
    if grounder.grounding_specialist_health().status != "ready":
        raise RuntimeError("The verified Grounding Specialist artifact could not be loaded.")
    get_captioner().load()


def _upload(path_value: str | Path | None) -> UploadFile:
    if not path_value:
        raise ValueError("An image file is required.")
    path = Path(path_value).resolve()
    if not path.is_file():
        raise ValueError("The uploaded image is unavailable.")
    mime_type = mimetypes.guess_type(path.name)[0] or "application/octet-stream"
    return UploadFile(
        path.open("rb"),
        size=path.stat().st_size,
        filename=path.name,
        headers=Headers({"content-type": mime_type}),
    )


def _preview_references(value: Any) -> Iterable[str]:
    if isinstance(value, dict):
        for nested in value.values():
            yield from _preview_references(nested)
    elif isinstance(value, list):
        for nested in value:
            yield from _preview_references(nested)
    elif isinstance(value, str) and value.startswith("/api/agent/previews/"):
        yield value


def _published_artifacts(response: dict[str, Any]) -> tuple[list[dict[str, Any]], list[str]]:
    manifest: list[dict[str, Any]] = []
    paths: list[str] = []
    seen: set[str] = set()
    for reference in _preview_references(response):
        if reference in seen:
            continue
        seen.add(reference)
        path = preview_file_path(reference.rsplit("/", 1)[-1])
        if path is None:
            continue
        index = len(paths)
        paths.append(str(path))
        manifest.append(
            {
                "index": index,
                "kind": "preview",
                "reference": reference,
                "filename": path.name,
            }
        )
    return manifest, paths


async def _execute_agent_query(
    *,
    query: str,
    input_mode: InputMode,
    primary_modality: Modality,
    primary_image_modality: ImageModality | None,
    primary_image_path: str | Path | None,
    secondary_modality: Modality | None = None,
    secondary_image_path: str | Path | None = None,
    primary_observation_role: ObservationRole | None = None,
    secondary_observation_role: ObservationRole | None = None,
    primary_date: str | None = None,
    secondary_date: str | None = None,
    use_cache: bool = True,
    force_rerun: bool = False,
) -> tuple[dict[str, Any], list[dict[str, Any]], list[str]]:
    uploads: list[UploadFile] = []
    try:
        primary_upload = _upload(primary_image_path)
        uploads.append(primary_upload)
        secondary_upload = None
        if secondary_image_path:
            secondary_upload = _upload(secondary_image_path)
            uploads.append(secondary_upload)
        response = await agent_image_query(
            query=query,
            input_mode=input_mode,
            primary_modality=primary_modality,
            primary_image_modality=primary_image_modality,
            primary_image=primary_upload,
            secondary_modality=secondary_modality,
            secondary_image=secondary_upload,
            primary_observation_role=primary_observation_role,
            secondary_observation_role=secondary_observation_role,
            primary_date=primary_date,
            secondary_date=secondary_date,
            use_cache=use_cache,
            force_rerun=force_rerun,
        )
        serialized = _model_dump(response)
        manifest, paths = _published_artifacts(serialized)
        return {"ok": True, "status_code": 200, "response": serialized}, manifest, paths
    except HTTPException as error:
        return {
            "ok": False,
            "status_code": error.status_code,
            "error": error.detail,
        }, [], []
    except (OSError, ValueError) as error:
        return {
            "ok": False,
            "status_code": 400,
            "error": {"code": "INVALID_INPUT", "message": str(error)},
        }, [], []
    except Exception:
        return {
            "ok": False,
            "status_code": 500,
            "error": {
                "code": "INFERENCE_FAILED",
                "message": "SatQuery could not complete this inference request.",
            },
        }, [], []
    finally:
        for upload in uploads:
            await upload.close()


def _run(**kwargs: Any) -> tuple[dict[str, Any], list[dict[str, Any]], list[str]]:
    return asyncio.run(_execute_agent_query(**kwargs))


def run_single_image(
    primary_image: str | None,
    query: str,
    primary_modality: str = "optical",
    primary_image_modality: str = "auto",
    use_cache: bool = True,
    force_rerun: bool = False,
) -> tuple[dict[str, Any], list[dict[str, Any]], list[str]]:
    """Execute the existing single-image agent workflow."""

    try:
        return _run(
            query=query,
            input_mode=InputMode.SINGLE,
            primary_modality=Modality(primary_modality),
            primary_image_modality=ImageModality(primary_image_modality),
            primary_image_path=primary_image,
            use_cache=use_cache,
            force_rerun=force_rerun,
        )
    except ValueError as error:
        return _invalid_parameter(error)


def run_optical_sar(
    optical_image: str | None,
    sar_image: str | None,
    query: str,
    use_cache: bool = True,
    force_rerun: bool = False,
) -> tuple[dict[str, Any], list[dict[str, Any]], list[str]]:
    """Execute the existing optical-plus-SAR cross-modal workflow."""

    return _run(
        query=query,
        input_mode=InputMode.CROSS_MODAL,
        primary_modality=Modality.OPTICAL,
        primary_image_modality=ImageModality.AUTO,
        primary_image_path=optical_image,
        secondary_modality=Modality.SAR,
        secondary_image_path=sar_image,
        primary_observation_role=ObservationRole.OPTICAL,
        secondary_observation_role=ObservationRole.SAR,
        use_cache=use_cache,
        force_rerun=force_rerun,
    )


def run_bitemporal(
    earlier_image: str | None,
    later_image: str | None,
    query: str,
    primary_date: str,
    secondary_date: str,
    primary_modality: str = "optical",
    use_cache: bool = True,
    force_rerun: bool = False,
) -> tuple[dict[str, Any], list[dict[str, Any]], list[str]]:
    """Execute the existing bi-temporal workflow."""

    try:
        modality = Modality(primary_modality)
        return _run(
            query=query,
            input_mode=InputMode.BI_TEMPORAL,
            primary_modality=modality,
            primary_image_modality=ImageModality.AUTO,
            primary_image_path=earlier_image,
            secondary_modality=modality,
            secondary_image_path=later_image,
            primary_date=primary_date,
            secondary_date=secondary_date,
            use_cache=use_cache,
            force_rerun=force_rerun,
        )
    except ValueError as error:
        return _invalid_parameter(error)


def _invalid_parameter(error: ValueError) -> tuple[dict[str, Any], list[dict[str, Any]], list[str]]:
    return {
        "ok": False,
        "status_code": 422,
        "error": {"code": "INVALID_PARAMETER", "message": str(error)},
    }, [], []


def space_health() -> dict[str, Any]:
    """Return model lifecycle state without performing inference or allocating a GPU."""

    backend_health: dict[str, Any] | None = None
    if _BACKEND is not None:
        backend_health = _BACKEND.health()
    return {
        "status": "ready" if _BACKEND is not None else "initializing",
        "runtime": "hugging-face-zerogpu",
        "backend": backend_health,
        "agent": _model_dump(agent_health()),
    }


def api_contract() -> dict[str, Any]:
    """Describe the stable bridge contract without touching model execution."""

    return {
        "schema_version": "hf-space-bridge-v1",
        "endpoints": {
            "single_image": {"gpu_duration_seconds": 90},
            "optical_sar": {"gpu_duration_seconds": 75},
            "bitemporal": {"gpu_duration_seconds": 90},
            "health": {"gpu": False},
            "contract": {"gpu": False},
        },
        "result_outputs": ["result_envelope", "artifact_manifest", "artifact_files"],
        "private_artifacts": {
            "repository": "AnirudhShashikumar/SatQuery_AI_Models",
            "count": 6,
            "authentication": "HF_TOKEN Space secret",
            "integrity": "pinned SHA-256 allowlist",
        },
    }