src/compliment_forest/server.py

from __future__ import annotations

from collections.abc import Callable
from pathlib import Path
from typing import Any

import gradio as gr
import httpx
from fastapi import HTTPException
from fastapi.responses import FileResponse
from fastapi.staticfiles import StaticFiles
from pydantic import BaseModel, ConfigDict, Field
from starlette.responses import StreamingResponse

from .backends.image import (
    DemoImageBackend,
    FluxImageBackend,
    HfInferenceImageBackend,
    ModalImageBackend,
    ZeroGpuImageBackend,
)
from .backends.music import ModalMusicBackend, NoMusicBackend
from .backends.text import (
    DemoTextBackend,
    HfInferenceTextBackend,
    LlamaCppTextBackend,
    ModalTextBackend,
    TransformersTextBackend,
)
from .config import AppConfig
from .orchestrator import ForestOrchestrator, build_guided_situation
from .schema import ForestStyle, IntakeQuestion, IntakeTurn, StreamEvent
from .trace import TraceRecorder


class ForestRequest(BaseModel):
    model_config = ConfigDict(extra="forbid", str_strip_whitespace=True)

    name: str = Field(min_length=1, max_length=80)
    situation: str = Field(min_length=1, max_length=1200)
    seed: int | None = Field(default=None, ge=0, le=2_147_483_647)
    style: ForestStyle | None = None
    intake: list[IntakeTurn] = Field(default_factory=list, max_length=5)


class IntakeNextRequest(BaseModel):
    model_config = ConfigDict(extra="forbid", str_strip_whitespace=True)

    name: str = Field(min_length=1, max_length=80)
    situation: str = Field(min_length=1, max_length=1200)
    history: list[IntakeTurn] = Field(default_factory=list, max_length=5)
    seed: int | None = Field(default=None, ge=0, le=2_147_483_647)


def build_orchestrator(
    config: AppConfig,
    *,
    gpu_image_generator: Callable[[str, int, str], str] | None = None,
    gpu_text_generator: Callable[[list[dict[str, str]], dict[str, object]], str] | None = None,
) -> ForestOrchestrator:
    if config.text_backend == "llama_cpp":
        text_backend = LlamaCppTextBackend(
            base_url=config.llama_base_url,
            model=config.llama_model,
        )
    elif config.text_backend == "hf_inference":
        text_backend = HfInferenceTextBackend(model=config.hf_text_model)
    elif config.text_backend == "transformers":
        if gpu_text_generator is None:
            raise ValueError("transformers text backend requires a GPU text generator")
        text_backend = TransformersTextBackend(
            model=config.transformers_text_model,
            generator=gpu_text_generator,
        )
    elif config.text_backend == "modal":
        assert config.modal_text_endpoint is not None
        assert config.modal_signing_key is not None
        text_backend = ModalTextBackend(
            endpoint=config.modal_text_endpoint,
            signing_key=config.modal_signing_key.get_secret_value(),
        )
    else:
        text_backend = DemoTextBackend()

    if config.image_backend == "modal":
        assert config.modal_image_endpoint is not None
        assert config.modal_signing_key is not None
        image_backend = ModalImageBackend(
            endpoint=config.modal_image_endpoint,
            signing_key=config.modal_signing_key.get_secret_value(),
            fallback=HfInferenceImageBackend(model=config.hf_image_model),
        )
    elif config.image_backend == "zerogpu":
        if gpu_image_generator is None:
            raise ValueError("zerogpu image backend requires a GPU image generator")
        image_backend = ZeroGpuImageBackend(
            gpu_image_generator,
            fallback=HfInferenceImageBackend(model=config.hf_image_model),
        )
    elif config.image_backend == "hf_inference":
        image_backend = HfInferenceImageBackend(model=config.hf_image_model)
    elif config.image_backend == "flux":
        image_backend = FluxImageBackend(
            model_id=config.flux_model_id,
            lora_id=config.flux_lora_id,
            local_files_only=config.local_files_only,
        )
    else:
        image_backend = DemoImageBackend()

    if config.music_backend == "modal":
        assert config.modal_music_endpoint is not None
        assert config.modal_signing_key is not None
        music_backend = ModalMusicBackend(
            endpoint=config.modal_music_endpoint,
            signing_key=config.modal_signing_key.get_secret_value(),
        )
    else:
        music_backend = NoMusicBackend()
    trace_recorder = TraceRecorder(config.trace_path) if config.trace_path else None
    return ForestOrchestrator(
        text_backend=text_backend,
        image_backend=image_backend,
        music_backend=music_backend,
        trace_recorder=trace_recorder,
    )


def create_app(
    *,
    config: AppConfig | None = None,
    orchestrator: Any | None = None,
    frontend_dir: str | Path | None = None,
    gpu_image_generator: Callable[[str, int, str], str] | None = None,
    gpu_text_generator: Callable[[list[dict[str, str]], dict[str, object]], str] | None = None,
    upstream_client: httpx.Client | None = None,
) -> gr.Server:
    runtime = config or AppConfig.from_env()
    forest = None
    if runtime.upstream_space_url is None:
        forest = orchestrator or build_orchestrator(
            runtime,
            gpu_image_generator=gpu_image_generator,
            gpu_text_generator=gpu_text_generator,
        )
    proxy = upstream_client
    if runtime.upstream_space_url and proxy is None:
        proxy = httpx.Client(
            timeout=httpx.Timeout(600, connect=30),
            follow_redirects=True,
        )
    frontend = (
        Path(frontend_dir)
        if frontend_dir is not None
        else Path(__file__).resolve().parents[2] / "frontend"
    )
    app = gr.Server(
        title="The Compliment Forest",
        description="A progressive path of grounded encouragement.",
        docs_url=None,
        redoc_url=None,
    )

    # Browsers will heuristically cache static files for hours when no
    # Cache-Control header is present, and HF Spaces does not set one for
    # FastAPI-served files. Force revalidation so each Space rebuild is
    # immediately visible without a cache wipe on the user's side.
    _NO_CACHE = {"Cache-Control": "no-cache, must-revalidate"}

    @app.get("/")
    def index() -> FileResponse:
        return FileResponse(frontend / "index.html", headers=_NO_CACHE)

    @app.get("/styles.css")
    def styles() -> FileResponse:
        return FileResponse(
            frontend / "styles.css",
            media_type="text/css",
            headers=_NO_CACHE,
        )

    @app.get("/app.js")
    def javascript() -> FileResponse:
        return FileResponse(
            frontend / "app.js",
            media_type="text/javascript",
            headers=_NO_CACHE,
        )

    assets = frontend / "assets"
    if assets.exists():
        app.mount("/assets", StaticFiles(directory=assets), name="assets")

    @app.get("/health")
    def health() -> dict[str, object]:
        if runtime.upstream_space_url:
            return {
                "status": "ok",
                "runtime_mode": "upstream_proxy",
                "upstream_space_url": runtime.upstream_space_url,
                "off_grid": False,
                "fresh_images": True,
                "default_style": runtime.default_style,
                "model_parameter_budget_billions": 25,
                "phase1_model_parameter_budget_billions": 18,
            }
        hosted = bool(
            {"hf_inference", "modal", "zerogpu", "transformers"}
            & {runtime.text_backend, runtime.image_backend}
        )
        runtime_text_model = {
            "demo": "demo",
            "hf_inference": runtime.hf_text_model,
            "llama_cpp": runtime.llama_model,
            "transformers": runtime.transformers_text_model,
            "modal": "openbmb/MiniCPM4.1-8B (Modal)",
        }[runtime.text_backend]
        phase1_budget = (
            18 if runtime.text_backend == "llama_cpp" and runtime.image_backend == "flux" else None
        )
        active_budget = phase1_budget
        uses_minicpm = (
            runtime.text_backend == "modal"
            or (
                runtime.text_backend == "transformers"
                and runtime.transformers_text_model.endswith("MiniCPM4.1-8B")
            )
            or (
                runtime.text_backend == "hf_inference"
                and runtime.hf_text_model.endswith("MiniCPM4.1-8B")
            )
        )
        if uses_minicpm:
            active_budget = 25
        return {
            "status": "ok",
            "text_backend": runtime.text_backend,
            "runtime_text_model": runtime_text_model,
            "image_backend": runtime.image_backend,
            "music_backend": runtime.music_backend,
            "off_grid": not hosted,
            "fresh_images": runtime.image_backend != "demo",
            "default_style": runtime.default_style,
            "model_parameter_budget_billions": active_budget,
            "phase1_model_parameter_budget_billions": 18,
        }

    @app.post("/api/intake/next")
    def next_intake(request: IntakeNextRequest) -> IntakeQuestion:
        if runtime.upstream_space_url:
            assert proxy is not None
            try:
                response = proxy.post(
                    f"{runtime.upstream_space_url}/api/intake/next",
                    json=request.model_dump(mode="json"),
                )
                response.raise_for_status()
                return IntakeQuestion.model_validate(response.json())
            except (httpx.HTTPError, ValueError) as error:
                raise HTTPException(
                    status_code=502,
                    detail=f"The forest could not reach its generation service: {error}",
                ) from error

        from .safety import guard_input

        assert forest is not None
        guard = guard_input(request.name, request.situation)
        if not guard.allowed:
            raise HTTPException(status_code=400, detail=guard.message)
        if len(request.history) >= 5:
            raise HTTPException(status_code=400, detail="intake already complete")
        seed = (request.seed if request.seed is not None else runtime.default_seed) + len(
            request.history
        )
        try:
            return forest.next_intake_question(
                request.name,
                request.situation,
                request.history,
                seed=seed,
            )
        except ValueError as error:
            raise HTTPException(
                status_code=502,
                detail=f"The forest could not produce a question: {error}",
            ) from error

    @app.post("/api/forest")
    def generate_forest(request: ForestRequest) -> StreamingResponse:
        if runtime.upstream_space_url:

            def proxy_stream():
                assert proxy is not None
                try:
                    with proxy.stream(
                        "POST",
                        f"{runtime.upstream_space_url}/api/forest",
                        json=request.model_dump(mode="json"),
                    ) as response:
                        response.raise_for_status()
                        yield from response.iter_bytes()
                except httpx.HTTPError as error:
                    yield (
                        StreamEvent(
                            type="error",
                            message=(
                                "The forest could not reach its generation service: "
                                f"{error}"
                            ),
                        ).model_dump_json()
                        + "\n"
                    )

            return StreamingResponse(proxy_stream(), media_type="application/x-ndjson")

        def stream():
            assert forest is not None
            seed = request.seed if request.seed is not None else runtime.default_seed
            style = request.style or runtime.default_style
            model_situation = build_guided_situation(request.situation, request.intake)
            try:
                for event in forest.generate(
                    request.name,
                    request.situation,
                    seed,
                    style,
                    model_situation=model_situation,
                ):
                    yield event.model_dump_json() + "\n"
            except Exception as error:
                yield StreamEvent(
                    type="error",
                    message=f"The forest could not grow: {error}",
                ).model_dump_json() + "\n"

        return StreamingResponse(stream(), media_type="application/x-ndjson")

    return app