scripts/test_visual_upgrades_stub.py

"""Stub tests for the Sept-2026 upgrade options:
1. env-overridable model picks (RERANK_MODEL / VISUAL_EMBED_MODEL)
2. visual_rerank.rerank_keys ordering (fail-open, png gating)
3. dim-aware pages collection params
4. zerank id-alias fallback in the reranker factory"""
import importlib
import os
import sys
import tempfile
from pathlib import Path

REPO = Path(__file__).resolve().parents[1]
sys.path.insert(0, str(REPO))
os.environ["FACTORY_RAG_DATA"] = tempfile.mkdtemp(prefix="fr-upg-")

import config  # noqa: E402
import ingest_lib  # noqa: E402
import visual_rerank  # noqa: E402


def main():
    failures = []

    # 1. env overrides + swap tiers
    os.environ["RERANK_MODEL"] = "zeroentropy/zerank-2-reranker"
    os.environ["VISUAL_EMBED_MODEL"] = "TomoroAI/tomoro-ai-colqwen3-embed-8b-awq"
    cfg2 = importlib.reload(config)
    if cfg2.RERANK_MODEL != "zeroentropy/zerank-2-reranker":
        failures.append("RERANK_MODEL env override broken")
    if cfg2.VISUAL_EMBED_DIM != 320 or cfg2.VISUAL_EMBED_MODEL != "TomoroAI/tomoro-ai-colqwen3-embed-8b-awq":
        failures.append("VISUAL_EMBED_MODEL/DIM swap broken")
    if "visual_embed_model" not in cfg2.snapshot():
        failures.append("snapshot must expose visual_embed_model (advisor)")
    os.environ["ENABLE_VISUAL_RERANK"] = "1"
    cfg2 = importlib.reload(config)
    if cfg2.ENABLE_VISUAL_RERANK is not True:
        failures.append("ENABLE_VISUAL_RERANK env override broken")
    del os.environ["RERANK_MODEL"], os.environ["VISUAL_EMBED_MODEL"], os.environ["ENABLE_VISUAL_RERANK"]
    importlib.reload(config)

    # 2. visual_rerank.rerank_keys ordering
    sha = "abc123"
    png_dir = ingest_lib.PAGES_DIR / sha
    png_dir.mkdir(parents=True, exist_ok=True)
    from PIL import Image
    for pg in (1, 2):
        Image.new("RGB", (4, 4)).save(png_dir / f"p{pg:04d}.png")

    class FakeRE:
        def __init__(self, scores):
            self.scores = scores
            self.seen = None

        def predict(self, pairs, prompt=None, **kw):
            self.seen = [p[1]["image"] for p in pairs]
            return self.scores

    fake = FakeRE([0.1, 0.9])  # page 2 more relevant
    vis_flag_orig = config.ENABLE_VISUAL_RERANK
    config.ENABLE_VISUAL_RERANK = True
    visual_rerank.get_visual_reranker = lambda: fake
    try:
        keys = [(sha, 1), (sha, 2), (sha, None)]  # third has no page -> rest
        out = visual_rerank.rerank_keys("torque?", keys)
        if out[:2] != [(sha, 2), (sha, 1)]:
            failures.append(f"rerank_keys order wrong: {out}")
        if out[2] != (sha, None):
            failures.append("no-png entry must land after scored ones")
        if not any(p.endswith("p0002.png") for p in fake.seen):
            failures.append("reranker did not receive page PNG path")
        # fail-open: reranker exception -> original order
        visual_rerank.get_visual_reranker = lambda: None
        if visual_rerank.rerank_keys("q", keys) != keys:
            failures.append("rerank_keys must be fail-open with no reranker")
        config.ENABLE_VISUAL_RERANK = False
        if visual_rerank.rerank_keys("q", keys) != keys:
            failures.append("rerank_keys must no-op when flag off")
    finally:
        visual_rerank.get_visual_reranker = None
        config.ENABLE_VISUAL_RERANK = vis_flag_orig if False else vis_flag_orig
    config.ENABLE_VISUAL_RERANK = vis_flag_orig

    # 3. dim-aware pages collection params
    import types
    _tm = types.ModuleType("torch")
    _tm.float32 = "f32"; _tm.float16 = "f16"; _tm.bfloat16 = "bf16"
    sys.modules.setdefault("torch", _tm)  # CPU sandbox: no torch; stubbed attrs suffice
    import visual_lane
    dim_orig = config.VISUAL_EMBED_DIM
    try:
        config.VISUAL_EMBED_DIM = 320
        params = visual_lane.config_pages_vector_params()
        if params.size != 320:
            failures.append("tomoro dim (320) not used for pages collection")
        config.VISUAL_EMBED_DIM = 128
        params = visual_lane.config_pages_vector_params()
        if params.size != 128:
            failures.append("colqwen2.5 dim (128) not used")
    finally:
        config.VISUAL_EMBED_DIM = dim_orig

    # 4. zerank id-alias fallback (fake transformers in sys.modules)
    saved_tf = sys.modules.get("transformers")
    # 4. zerank id-alias fallback (fake transformers module in sys.modules)
    import types
    import rerankers

    class FakeTok:
        def __init__(self):
            self.name = None

        @classmethod
        def from_pretrained(cls, name, padding_side=None):
            if name == "zeroentropy/zerank-2-reranker":
                raise OSError("alias missing")
            obj = cls()
            obj.name = name
            return obj

        def encode(self, s, add_special_tokens=False):
            return [1]

    class FakeCausalLM:
        @classmethod
        def from_pretrained(cls, name, **kw):
            obj = cls()
            return obj

        def to(self, device):
            return self

        def eval(self):
            return self

    fake_tf = types.ModuleType("transformers")
    fake_tf.AutoTokenizer = FakeTok
    fake_tf.AutoModelForCausalLM = FakeCausalLM
    saved_tf = sys.modules.get("transformers")
    sys.modules["transformers"] = fake_tf
    try:
        r = rerankers._QwenPairwiseReranker("zeroentropy/zerank-2-reranker")
        if r.tok is None or r.tok.name != "zeroentropy/zerank-2":
            failures.append("zerank id-alias fallback broken")
    finally:
        if saved_tf is not None:
            sys.modules["transformers"] = saved_tf
        else:
            sys.modules.pop("transformers", None)

    if failures:
        print("FAIL:")
        for f in failures:
            print("  -", f)
        sys.exit(1)
    print("PASS: env overrides + swap tiers, rerank_keys ordering + fail-open, "
          "dim-aware collection, zerank alias fallback")


if __name__ == "__main__":
    main()