backend/app/models_registry.py

"""Enabled model catalogue — small models (≤32B params) from three labs.

The "Build Small" constraint is ≤32B parameters; every model below also fits in ≤32 GB
of memory at a sensible quantization. Availability is computed from config/deps so the
UI can show which are actually live.

  • OpenBMB (MiniCPM)        — vision-language OCR + text reasoning  (the OCR/IDP engine)
  • Cohere (Aya-Vision)      — vision-language OCR / VQA
  • Black Forest Labs (FLUX) — image GENERATION → synthetic test documents (not an OCR
                               model; used to stress-test OCR quality)
"""
from __future__ import annotations

import importlib.util

MAX_PARAMS_B = 32  # hackathon "small models" guardrail


def _has(mod: str) -> bool:
    return importlib.util.find_spec(mod) is not None


def model_catalog(settings) -> dict:
    minicpm_api = bool(settings.minicpm_base_url)
    transformers = _has("transformers")
    cohere_enabled = transformers and __import__("os").getenv("COHERE_OCR_ENABLE", "").lower() in {"1", "true", "yes"}
    cohere_api = bool(getattr(settings, "cohere_api_key", None))
    diffusers = _has("diffusers")
    bfl_api = bool(settings.bfl_api_key)

    labs = [
        {
            "lab": "OpenBMB", "homepage": "https://github.com/OpenBMB/MiniCPM",
            "models": [
                {"name": "MiniCPM-V-4.6", "id": settings.minicpm_model, "params_b": 8.0,
                 "size_gb_int4": 5.5, "modality": "vision-language (OCR + reasoning)",
                 "role": "OCR backend + LLM extractor", "license": "Apache-2.0 (weights: MiniCPM Model License)",
                 "available": minicpm_api or transformers,
                 "enable": "MINICPM_BASE_URL (+ MINICPM_API_KEY) OR pip install transformers"},
                {"name": "MiniCPM-o-4.5", "id": "MiniCPM-o-4.5", "params_b": 8.0,
                 "size_gb_int4": 5.5, "modality": "omni (vision/audio) VLM",
                 "role": "alt OCR/VLM", "license": "MiniCPM Model License",
                 "available": minicpm_api, "enable": "same OpenAI-compatible endpoint"},
                {"name": "MiniCPM3-4B", "id": settings.openbmb_reasoner_model, "params_b": 4.0,
                 "size_gb_int4": 2.8, "modality": "text LLM (reasoning + function-calling, 32k ctx)",
                 "role": "ERP reasoning · NLQ→SQL · report summarization (fine-tune target)",
                 "license": "Apache-2.0 (weights: MiniCPM Model License)",
                 "available": minicpm_api or transformers,
                 "enable": "OpenAI-compatible endpoint (OPENBMB_REASONER_MODEL) OR pip install transformers"},
            ],
        },
        {
            "lab": "Cohere", "homepage": "https://huggingface.co/CohereLabs",
            "models": [
                {"name": "Aya-Vision-8B", "id": settings.cohere_ocr_model, "params_b": 8.0,
                 "size_gb_int4": 6.0, "modality": "vision-language (OCR/VQA, 23 langs)",
                 "role": "OCR backend", "license": "CC-BY-NC 4.0",
                 "available": cohere_enabled,
                 "enable": "pip install transformers torch + COHERE_OCR_ENABLE=true"},
                {"name": "Aya-Vision-32B", "id": "CohereLabs/aya-vision-32b", "params_b": 32.0,
                 "size_gb_int4": 18.0, "modality": "vision-language (OCR/VQA)",
                 "role": "alt OCR backend (max-quality small)", "license": "CC-BY-NC 4.0",
                 "available": cohere_enabled,
                 "enable": "COHERE_OCR_MODEL=CohereLabs/aya-vision-32b + COHERE_OCR_ENABLE=true"},
                {"name": "Command R7B", "id": "command-r7b-12-2024", "params_b": 7.0,
                 "size_gb_int4": 5.0, "modality": "text LLM (RAG + tool-use + reasoning, 128k ctx)",
                 "role": "ERP RAG · NLQ · grounded reasoning", "license": "CC-BY-NC 4.0",
                 "available": cohere_api,
                 "enable": "COHERE_API_KEY (api.cohere.com) OR weights CohereLabs/c4ai-command-r7b-12-2024"},
            ],
        },
        {
            "lab": "Black Forest Labs", "homepage": "https://github.com/black-forest-labs/flux",
            "models": [
                {"name": "FLUX.1 [dev]", "id": settings.bfl_model, "params_b": 12.0,
                 "size_gb_int4": 12.0, "modality": "text-to-image GENERATION",
                 "role": "synthetic test-document generator (not OCR)", "license": "FLUX.1-dev Non-Commercial",
                 "available": bfl_api or diffusers,
                 "enable": "BFL_API_KEY (api.bfl.ml) OR pip install diffusers torch"},
                {"name": "FLUX.1 [schnell]", "id": "flux-schnell", "params_b": 12.0,
                 "size_gb_int4": 12.0, "modality": "text-to-image GENERATION (fast)",
                 "role": "synthetic test-document generator", "license": "Apache-2.0",
                 "available": bfl_api or diffusers, "enable": "BFL_API_KEY OR pip install diffusers"},
            ],
        },
    ]

    # guardrail: nothing exceeds the small-model size limit
    for lab in labs:
        for m in lab["models"]:
            assert m["params_b"] <= MAX_PARAMS_B, f"{m['name']} exceeds {MAX_PARAMS_B}B"

    flat = [{"lab": lab["lab"], **m} for lab in labs for m in lab["models"]]
    return {
        "max_params_b": MAX_PARAMS_B,
        "labs": labs,
        "available": [m["name"] for m in flat if m["available"]],
        "ocr_capable": [m["name"] for m in flat if "OCR" in m["modality"] or "vision" in m["modality"]],
        "reasoning_capable": [m["name"] for m in flat
                              if "reasoning" in m["modality"] or "text LLM" in m["modality"]],
        "count": len(flat),
    }