test_product_resolution.py

#!/usr/bin/env python3
"""
Offline test for ProductCatalog (no OpenAI / Qdrant required).

Verifies the fix for the "answers from a different herbicide / keeps citing the
dominant product" bug: product resolution must be driven by what is actually in
the index, and must ABSTAIN (return None) for products we don't have.

Run:  python test_product_resolution.py
"""

from src.cdms.product_catalog import (
    ProductCatalog,
    normalize_filename,
    cross_product_abstention,
    diversify_by_product,
)


def main() -> bool:
    cat = ProductCatalog()
    catalog = cat.catalog()
    available = cat.available_products()

    print("=" * 70)
    print("PRODUCT CATALOG (derived from data/cdms_metadata.db)")
    print("=" * 70)
    for product, n in sorted(catalog.items(), key=lambda kv: -kv[1]):
        flag = "" if n > 0 else "   <-- 0 chunks (unfindable / needs reprocessing)"
        print(f"  {n:5d} chunks   {product}{flag}")

    print(f"\nAvailable (has chunks): {sorted(available)}")

    # --- assertions -------------------------------------------------------
    passed, failed = 0, []

    def check(name, cond):
        nonlocal passed
        if cond:
            passed += 1
        else:
            failed.append(name)

    # Filename normalization strips the content hash and underscores.
    check("normalize roundup_hash", normalize_filename("roundup_bdc94bbee383.pdf") == "roundup")
    check("normalize Brandt_Nema_Q", normalize_filename("Brandt_Nema_Q.pdf") == "brandt nema q")
    check("normalize 24-d", normalize_filename("24-d_fa1e6bdacae6.pdf") == "24-d")

    # Products that ARE in the index resolve from a natural question.
    check("resolve roundup", cat.resolve("what is the application rate for Roundup?") == "roundup")
    check("resolve sevin", cat.resolve("Is Sevin safe for tomatoes?") == "sevin")
    check("resolve 24-d", cat.resolve("mixing instructions for 24-d") == "24-d")

    # THE BUG: a product we do NOT have must abstain (return None), NOT fall back
    # to the dominant product.
    check("abstain on Trust", cat.resolve("Tell me about the Trust herbicide") is None)
    check("abstain on Enlist", cat.resolve("What is the rate for Enlist One?") is None)
    check("abstain on gibberish", cat.resolve("qzxwv nonsense product") is None)

    # 0-chunk PDFs must NOT be considered available.
    check("acquit not available", "acquit" not in available)

    # --- cross-product abstention guard ----------------------------------
    # User asks about Sevin but retrieval only returned Roundup chunks -> ABSTAIN.
    check(
        "abstain when only wrong-product chunks returned",
        cross_product_abstention("Is Sevin safe on tomatoes?",
                                 ["roundup_bdc94bbee383.pdf", "roundup_75908642d433.pdf"],
                                 catalog=cat) == "sevin",
    )
    # User asks about Sevin and a Sevin chunk is present -> proceed (None).
    check(
        "proceed when correct-product chunk present",
        cross_product_abstention("Is Sevin safe on tomatoes?",
                                 ["sevin_0b73d6a0a2d4.pdf", "roundup_bdc94bbee383.pdf"],
                                 catalog=cat) is None,
    )
    # Generic question (no specific product) -> proceed (None), don't over-block.
    check(
        "proceed on generic question",
        cross_product_abstention("What is a pre-emergent herbicide?",
                                 ["roundup_bdc94bbee383.pdf"],
                                 catalog=cat) is None,
    )
    # Product we don't have (Trust) -> guard returns None (handled by search gate,
    # not by cross-product substitution).
    check(
        "no false abstain for un-indexed product",
        cross_product_abstention("Tell me about Trust herbicide",
                                 ["roundup_bdc94bbee383.pdf"],
                                 catalog=cat) is None,
    )

    # --- diversity re-ranking (fix for index skew) -----------------------
    # Simulate a skewed vector-search result: 5 Roundup, then 1 Sevin, 1 boron.
    # Use realistic filenames (12-hex content hash) so normalization collapses
    # the different Roundup PDFs to a single "roundup" product.
    hashes = ["bdc94bbee383", "75908642d433", "30b0dda04968", "424997b6293c", "744b26dd6d64"]
    skewed = (
        [{"source_file": f"roundup_{hashes[i]}.pdf", "score": 0.9 - i * 0.01} for i in range(5)]
        + [{"source_file": "sevin_0b73d6a0a2d4.pdf", "score": 0.62}]
        + [{"source_file": "boron_b97d7d0cfb98.pdf", "score": 0.55}]
    )
    div = diversify_by_product(skewed, limit=5, max_per_product=2)
    div_products = [normalize_filename(r["source_file"]) for r in div]
    # No product may occupy more than max_per_product slots...
    check("diversify caps per product", div_products.count("roundup") <= 2)
    # ...and the minority products must now appear in the top-5.
    check("diversify surfaces sevin", "sevin" in div_products)
    check("diversify surfaces boron", "boron" in div_products)
    # Best overall chunk (highest score) is still included.
    check("diversify keeps top hit", any(r["score"] == 0.9 for r in div))
    # When only one product exists, it still returns up to `limit` (no starvation
    # beyond the cap is fine because there's nothing else to interleave).
    only = diversify_by_product(
        [{"source_file": f"roundup_{hashes[i]}.pdf", "score": 0.9} for i in range(5)],
        limit=5, max_per_product=2,
    )
    check("single-product respects cap", len(only) == 2)

    print("\n" + "=" * 70)
    print(f"RESULT: {passed} passed, {len(failed)} failed")
    if failed:
        print("FAILED:", ", ".join(failed))
    print("=" * 70)
    return not failed


if __name__ == "__main__":
    import sys
    sys.exit(0 if main() else 1)