ui/agents.py

"""
ui/agents.py
------------
Heavy model / graph loading, done once per server process and shared by every
browser session.

Streamlit gave this to us for free via @st.cache_resource. Gradio has no
equivalent, and it serves requests from a thread pool, so a naive "load if
None" would let two simultaneous first-visitors each start a ~2 GB model load.
Hence the explicit double-checked locking below: the lock is held across the
load, and a second caller blocks and then sees the finished object.

WHY LOADING NOW HAPPENS AT IMPORT, ON THE MAIN THREAD
-----------------------------------------------------
This used to load lazily: `load_agents()` from inside `demo.load`, and the
chatbot models on a daemon thread via `warm_chatbot_models_async()`. That kept
the boot splash responsive and got the Space green early — but it meant NOVA
loaded *every* model on a thread that was not the main thread, which is exactly
what ZeroGPU cannot tolerate. `spaces` fakes CUDA so nothing really initialises
it before the GPU fork, but the TorchFunctionMode/TorchDispatchMode doing that
work are thread-local to the thread that called `patch()` — the main thread, at
import — while `torch.cuda.is_available()` is patched globally to return True
for every thread. So off-main-thread model loading runs against a torch that
claims a GPU exists with none of the interception that makes the claim safe.

`preload_all()` below therefore runs at import, on the main thread, before
Gradio ever spawns a worker. See the CUDA-VIRGINITY RULE in ui/gpu.py.

The cost is honest and worth naming: the server now binds its port only after
~2 GB of weights are resident, so the Space shows "Starting" for longer instead
of showing NOVA's own splash. We trade a nicer boot animation for a GPU path
that actually works.
"""

import threading

_agents_lock = threading.Lock()
_agents = None

_chatbot_lock = threading.Lock()
_chatbot_models = None

# Set by preload_all() when a load fails, so nova_app.py can show a named error
# on the splash instead of the Space dying at import with a raw traceback.
BOOT_ERROR: str | None = None


def load_agents():
    """Import the compiled LangGraph agents. Importing the search graph also
    loads the SPECTER reranker model at module import time (by design)."""
    global _agents
    if _agents is None:
        with _agents_lock:
            if _agents is None:
                from app.modules.intent.graph import graph as intent_graph
                from app.modules.search.graph import graph as search_graph
                _agents = (intent_graph, search_graph)
    return _agents


def warm_chatbot_models():
    """Pre-load the chatbot's embedding + cross-encoder models so the first
    'Chat it out' click doesn't pay the model-load cost. We instantiate the
    exact models the chatbot uses (BAAI/bge-base-en-v1.5 + BAAI/bge-reranker-base),
    warming the weights into the HF/torch cache.

    Both are pinned to CPU: this runs at boot, outside any ZeroGPU window, and
    its whole job is to pull weights down — the GPU copy is made later, inside
    ui.gpu.vectorize_on_gpu, from the same warmed cache."""
    global _chatbot_models
    if _chatbot_models is None:
        with _chatbot_lock:
            if _chatbot_models is None:
                from vectorizeer import get_embeddings
                from langchain_community.cross_encoders import HuggingFaceCrossEncoder
                embeddings = get_embeddings(device="cpu")
                reranker = HuggingFaceCrossEncoder(
                    model_name="BAAI/bge-reranker-base",
                    model_kwargs={"device": "cpu"},
                )
                _chatbot_models = (embeddings, reranker)
    return _chatbot_models


def preload_all() -> "str | None":
    """Load every model NOW, on the calling thread. Call this at import from
    nova_app.py, which runs on the main thread.

    Replaces the old pair of lazy paths (`load_agents()` inside `demo.load`, and
    `warm_chatbot_models_async()`'s daemon thread). Both put multi-hundred-MB
    transformer loads on non-main threads, which is the one thing a ZeroGPU
    Space must not do — see the module docstring above and ui/gpu.py.

    Never raises. A failure here is almost always a missing API key, and the
    old lazy path had the nice property of surfacing that as a readable message
    on the splash rather than killing the Space at boot. That property is worth
    keeping, so the error is stashed in BOOT_ERROR and rendered by do_boot().
    """
    global BOOT_ERROR
    current = threading.current_thread()
    if current is not threading.main_thread():
        # Not fatal — but it means the caller moved this off the main thread,
        # which silently reintroduces the exact bug this function exists to
        # prevent. Say so loudly rather than fail mysteriously an hour later.
        print(
            f"[preload_all] WARNING: running on thread {current.name!r}, not the main "
            f"thread. ZeroGPU's torch patching is thread-local; models loaded here are "
            f"not covered by it and the GPU fork may fail.",
            flush=True,
        )
    try:
        load_agents()
        warm_chatbot_models()
    except Exception as e:
        import traceback
        traceback.print_exc()
        BOOT_ERROR = f"{type(e).__name__}: {e}"
    return BOOT_ERROR