test_rag_integration.py

#!/usr/bin/env python3
"""
End-to-end RAG integration test (NO OpenAI key required).

Drives the REAL src/cdms/rag_search.py pipeline over the REAL 2,109 chunks from
data/cdms_metadata.db, using an in-memory Qdrant and a local hashing embedder in
place of OpenAI. Demonstrates the before/after for the ISA-reported bug:

  * OLD path (raw top-k, threshold 0.3, no diversity) -> results dominated by the
    product that owns most of the index (Roundup = 71% of chunks).
  * NEW path (rag_search.search: threshold 0.4 + product diversity + abstention)
    -> varied products, and an honest "no results" for products we don't have.

Requires: qdrant-client, numpy  (pip install qdrant-client numpy)
Run:      python test_rag_integration.py
"""

import hashlib
import re
import sqlite3
import sys
from collections import Counter

import numpy as np

sys.path.insert(0, ".")

from src.rag.vector_store import QdrantVectorStore
from src.cdms.rag_search import CDMSRAGSearch
from src.cdms.product_catalog import normalize_filename

DIM = 1536  # match QdrantVectorStore's configured vector size
_TOKEN = re.compile(r"[a-z0-9]+")


def embed(text: str):
    """Deterministic lexical embedding (hashed bag-of-words, L2-normalized).

    Not semantic, but faithful for this test: it reproduces the index-dominance
    effect (more Roundup chunks -> Roundup fills the top-k) so we can show the
    fix changing the outcome. COSINE distance in Qdrant matches the normalization.
    """
    v = np.zeros(DIM, dtype=np.float32)
    for tok in _TOKEN.findall(text.lower()):
        idx = int(hashlib.md5(tok.encode()).hexdigest(), 16) % DIM
        v[idx] += 1.0
    n = np.linalg.norm(v)
    if n > 0:
        v /= n
    return v.tolist()


class LocalEmbedder:
    """Drop-in stand-in for OpenAIEmbeddingService."""

    def generate_embedding(self, text: str):
        return embed(text)


def products_of(results):
    return dict(Counter(normalize_filename(r.get("source_file", "")) for r in results))


def main() -> bool:
    # 1) Load real chunks from the shipped DB.
    conn = sqlite3.connect("data/cdms_metadata.db")
    rows = conn.execute(
        """SELECT dc.id, dc.content, dc.page_number, dc.document_id, d.filename
           FROM document_chunks dc JOIN documents d ON dc.document_id = d.id"""
    ).fetchall()
    print(f"Loaded {len(rows)} real chunks from data/cdms_metadata.db")

    # 2) Build an in-memory Qdrant and index every chunk with local embeddings.
    store = QdrantVectorStore()  # Docker unavailable -> in-memory mode
    for cid, content, page, docid, filename in rows:
        if not content:
            continue
        payload = {
            "content": content,
            "source_file": filename,
            "page_number": page or 0,
            "document_id": docid,
        }
        store.add_document_chunk(str(cid), embed(content), payload)
    print("Indexed chunks into in-memory Qdrant.\n")

    # 3) Wire the REAL search pipeline to the local store + embedder.
    searcher = CDMSRAGSearch()
    searcher.vector_store = store
    searcher.embedding_service = LocalEmbedder()

    # Helper: OLD behavior = raw top-5 at the old 0.3 threshold, no diversity.
    def old_top5(query):
        return store.search_documents(embed(query), limit=5, score_threshold=0.05)

    passed, failed = 0, []

    def check(name, cond, detail=""):
        nonlocal passed
        if cond:
            passed += 1
            print(f"  ✅ {name} {detail}")
        else:
            failed.append(name)
            print(f"  ❌ {name} {detail}")

    # --- Case 1: general query should not be monopolized by one product -----
    print("=" * 70)
    print("CASE 1: general query — dominance vs diversity")
    print("=" * 70)
    q1 = "What is the application rate and mixing instructions?"
    old = old_top5(q1)
    new = searcher.search(q1, score_threshold=0.05)  # low thr = lexical scale
    old_p, new_p = products_of(old), products_of(new)
    print(f"  query: {q1!r}")
    print(f"  OLD top-5 products: {old_p}")
    print(f"  NEW top-5 products: {new_p}")
    check("NEW caps any single product at <=2", max(new_p.values()) <= 2 if new_p else False)
    check("NEW returns >=2 distinct products", len(new_p) >= 2 if new_p else False)

    # --- Case 2: specific product we HAVE resolves to that product ----------
    print("\n" + "=" * 70)
    print("CASE 2: specific indexed product (Sevin)")
    print("=" * 70)
    q2 = "Is Sevin safe to use on vegetables?"
    res2 = searcher.search(q2, product_name="sevin", score_threshold=0.0)
    p2 = products_of(res2)
    print(f"  query: {q2!r}  ->  products: {p2}")
    check("returns only Sevin chunks", set(p2.keys()) <= {"sevin"} and bool(p2))

    # --- Case 3: product we DON'T have -> abstain (the core bug) ------------
    print("\n" + "=" * 70)
    print("CASE 3: un-indexed product (Trust) — must ABSTAIN, not substitute")
    print("=" * 70)
    q3 = "What is the application rate for Trust herbicide?"
    old3 = old_top5(q3)  # OLD global search: no abstention, no product filter
    res3 = searcher.search(q3, product_name="Trust", score_threshold=0.0)
    print(f"  query: {q3!r}")
    print(f"  OLD path would answer from: {products_of(old3)}  <-- THE BUG (wrong products)")
    print(f"  NEW path returns: {len(res3)} results (abstains)")
    check("OLD path substitutes a different product", len(old3) > 0 and "trust" not in products_of(old3))
    check("abstains on un-indexed product", len(res3) == 0)

    # --- Case 4: 0-chunk product (ACQUIT) also abstains --------------------
    print("\n" + "=" * 70)
    print("CASE 4: present-but-unprocessed product (ACQUIT, 0 chunks)")
    print("=" * 70)
    res4 = searcher.search("ACQUIT label safety", product_name="ACQUIT", score_threshold=0.0)
    print(f"  -> {len(res4)} results (expected 0)")
    check("abstains on 0-chunk product", len(res4) == 0)

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


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