backend/factory.py

"""Resolves the active LLM (per model_id) + image backends based on runtime
config.

LLM backends are cached per-model_id so a character that gets its model
swapped doesn't trigger a reload, and so multiple characters on the same
model share one instance. The first call for a given model_id triggers the
GGUF download/load (Real) or instantiates a canned-style cycle (Mock).

The image backend is mock-only: Flux 2 was dropped from the pipeline (the
frontend never called /generate_background and the cold-load was eating
into the ZeroGPU quota). MockImageBackend renders a Pillow gradient — fine
as a placeholder if the endpoint ever gets called.
"""

from __future__ import annotations

import os

from game.models import DEFAULT_MODEL_ID, get_spec

from . import config
from .interfaces import ImageBackend, LLMBackend

_llms: dict[str, LLMBackend] = {}
_image: ImageBackend | None = None


def _cache_key(model_id: str) -> str:
    """GPU and CPU Llama instances are loaded with different `n_gpu_layers`
    and can't be reused across modes. Cache them under separate keys so the
    CPU-fallback path (parent process, TOWNLET_FORCE_CPU=1) and the GPU
    path (forked workers, no env var) don't collide.
    """
    mode = "cpu" if os.getenv("TOWNLET_FORCE_CPU") == "1" else "gpu"
    return f"{model_id}::{mode}"


def get_llm(model_id: str = DEFAULT_MODEL_ID) -> LLMBackend:
    key = _cache_key(model_id)
    if key not in _llms:
        get_spec(model_id)  # validates id
        if config.USE_REAL_MODELS:
            from .real import RealLLMBackend

            _llms[key] = RealLLMBackend(model_id)
        else:
            from .mock import MockLLMBackend

            _llms[key] = MockLLMBackend(model_id)
    return _llms[key]


def get_image() -> ImageBackend:
    global _image
    if _image is None:
        from .mock import MockImageBackend

        _image = MockImageBackend()
    return _image