tests/rag/test_dedup_mechanism.py

"""Behavioral test for the `rag-dedup-v2` over-fetch + dedup + truncate
mechanism (Task 8, Step 2/3). Written and confirmed FAILING before
`ChromaDBClient.query_observed` gained the mechanism.

Fixture design (deliberately avoids the "just returns fewer candidates" trap
the plan preamble's acceptance gate calls out): 5 documents ("alpha".."epsilon"),
each contributing 3 near-duplicate chunks (adjacent-chunker-overlap style --
differ only in a trailing variant token), ranked by embedding distance so
that document "alpha"'s 3 duplicates are the 3 nearest neighbors, "beta"'s 3
duplicates are next, and so on -- 15 raw candidates total, ranked in 5 clean
groups of 3.

Requesting top-5:
- `rag-naive-v1` (default, no over-fetch) can only ever see the nearest 5
  raw candidates -- all 3 of "alpha"'s duplicates plus the first 2 of
  "beta"'s -- so it returns a low-diversity, high-duplication set.
- `rag-dedup-v2` (over-fetch factor 3 -> fetches all 15, dedups, truncates
  back to 5) can see every group, dedup down to one representative per
  group, and return exactly 5 candidates again -- but now one from each of
  the 5 distinct documents. Same final count (5), strictly higher diversity,
  strictly lower duplication -- the fair comparison the plan requires.

v1's own behavior (default/unspecified pipeline_version) must be completely
unchanged by any of this -- asserted directly against the same fixture.
"""

from __future__ import annotations

import os

os.environ.setdefault("CEREBRAS_API_KEY", "test-key")

import chromadb
import pytest

from app.observability import bootstrap
from app.observability.config import ObservabilityConfig
from app.observability.retrieval import candidate_stats


# 12+-token base phrases, no shared vocabulary across groups, matching the
# Jaccard-clearing design in tests/observability/test_candidate_stats.py.
_GROUP_BASES = {
    "alpha": "alpha membrane transport channel proteins regulate ion flow across the lipid bilayer",
    "beta": "beta ribosome translation machinery assembles polypeptide chains from messenger rna transcripts today",
    "gamma": "gamma mitochondrial electron transport chain produces cellular energy via oxidative phosphorylation",
    "delta": "delta cytoskeleton microtubule dynamics control cell shape and intracellular cargo transport",
    "epsilon": "epsilon nuclear pore complex regulates macromolecule trafficking between nucleus and cytoplasm",
}
_GROUP_ORDER = ["alpha", "beta", "gamma", "delta", "epsilon"]
_VARIANTS = ["one", "two", "three"]


def _text(group: str, variant: str) -> str:
    return f"{_GROUP_BASES[group]} variant {variant}"


# Rank 1..15, strictly increasing distance from the query origin -> Chroma's
# L2-nearest-neighbor order is exactly _GROUP_ORDER x _VARIANTS, in order.
_ALL_TEXTS: list[str] = [
    _text(group, variant) for group in _GROUP_ORDER for variant in _VARIANTS
]
_QUERY_TEXT = "find the most relevant passages"

_VECTORS: dict[str, list[float]] = {_QUERY_TEXT: [0.0, 0.0]}
for _rank, _text_value in enumerate(_ALL_TEXTS, start=1):
    _VECTORS[_text_value] = [float(_rank), 0.0]


class _FakeEmbedder:
    """Deterministic embedder: exact string -> fixed vector (see _VECTORS)."""

    def embed(self, texts: list[str]) -> list[list[float]]:
        return [list(_VECTORS[t]) for t in texts]


@pytest.fixture
def isolated_db(monkeypatch):
    from chromadb.api.client import SharedSystemClient

    SharedSystemClient.clear_system_cache()
    monkeypatch.setattr(
        chromadb, "PersistentClient", lambda *a, **k: chromadb.EphemeralClient()
    )
    from app.rag.chromadb_client import ChromaDBClient

    return ChromaDBClient(embedder=_FakeEmbedder())


def _seed(db) -> None:
    documents = list(_ALL_TEXTS)
    metadatas = [
        {
            "source": f"paper-{group}.pdf",
            "document_id": f"doc-{group}",
            "chunk_index": i,
            "type": "content",
            "project_id": "proj-1",
        }
        for i, group in enumerate(
            g for g in _GROUP_ORDER for _ in _VARIANTS
        )
    ]
    ids = [f"chunk-{i}" for i in range(len(documents))]
    db.upsert("library", documents, metadatas, ids)


def _query_vec(db):
    return db.embedder.embed([_QUERY_TEXT])[0]


# --- v1: unchanged, low-diversity, high-duplication top-5 -------------------


def test_naive_v1_default_returns_low_diversity_top5(isolated_db):
    _seed(isolated_db)
    out = isolated_db.query(
        "library", _query_vec(isolated_db), n_results=5, where={"project_id": "proj-1"}
    )
    assert len(out) == 5
    stats = candidate_stats(out)
    # Nearest 5 raw candidates = all 3 "alpha" duplicates + first 2 "beta" duplicates.
    assert stats["document_diversity"] == 2
    assert stats["duplicate_candidate_ratio"] == pytest.approx(3 / 5)


def test_naive_v1_explicit_matches_default(isolated_db):
    _seed(isolated_db)
    q = _query_vec(isolated_db)
    default_out = isolated_db.query_observed(
        "library", q, n_results=5, where={"project_id": "proj-1"}
    )
    explicit_out = isolated_db.query_observed(
        "library", q, n_results=5, where={"project_id": "proj-1"}, pipeline_version="rag-naive-v1"
    )
    assert default_out.candidates == explicit_out.candidates


# --- v2: same final count, strictly more diverse, strictly less duplicated --


def test_dedup_v2_returns_same_count_more_diverse_top5(isolated_db):
    _seed(isolated_db)
    out = isolated_db.query(
        "library",
        _query_vec(isolated_db),
        n_results=5,
        where={"project_id": "proj-1"},
        pipeline_version="rag-dedup-v2",
    )
    # Same final count as v1 -- the acceptance-gate-critical assertion: this
    # is NOT just "fewer candidates returned."
    assert len(out) == 5
    stats = candidate_stats(out)
    assert stats["document_diversity"] == 5
    assert stats["duplicate_candidate_ratio"] == 0.0
    # One representative from each of the 5 distinct documents.
    assert {c["document_id"] for c in out} == {f"doc-{g}" for g in _GROUP_ORDER}


def test_dedup_v2_keeps_the_most_relevant_member_of_each_group(isolated_db):
    _seed(isolated_db)
    out = isolated_db.query(
        "library",
        _query_vec(isolated_db),
        n_results=5,
        where={"project_id": "proj-1"},
        pipeline_version="rag-dedup-v2",
    )
    # Within each kept document, the surviving chunk is the nearest-ranked
    # ("variant one") member of that group, not an arbitrary one.
    for c in out:
        assert c["text"].endswith("variant one")


# --- `duplicates_removed_from_pool`: direct pool-level evidence (final-review fix) --
#
# `duplicate_candidate_ratio` (see test_candidate_stats.py) only ever sees the
# final, post-truncate candidate list -- it reads ~0.0 whether the dedup step
# removed nothing from the over-fetched pool or removed several duplicates
# that then got truncated away, and can't tell those two cases apart on its
# own. `duplicates_removed_from_pool` is a direct count taken from the raw
# pool *before* truncation, closing that gap. This fixture (5 documents x 3
# near-duplicate variants each) gives the over-fetched v2 pool of 15 real,
# removable duplicates, so the count below is provably non-zero -- unlike the
# real corpus analyzed in `docs/observability/dedup-v2-evaluation.md`, where
# it reads 0.


def _local_observability_config(tmp_path) -> ObservabilityConfig:
    return ObservabilityConfig(
        enabled=True,
        mode="local",
        otlp_endpoint="http://127.0.0.1:4318",
        signoz_ui_url="http://127.0.0.1:8080",
        artifact_root=tmp_path,
    )


def _latest_result_prepare_row(store) -> dict:
    rows = [r for r in store.all() if r["operation"] == "rag.result_prepare"]
    assert rows, "expected at least one rag.result_prepare observation"
    return rows[-1]


def test_dedup_v2_reports_nonzero_duplicates_removed_from_pool(isolated_db, tmp_path):
    bootstrap.initialize_observability(config=_local_observability_config(tmp_path))
    try:
        _seed(isolated_db)
        isolated_db.query_observed(
            "library",
            _query_vec(isolated_db),
            n_results=5,
            where={"project_id": "proj-1"},
            pipeline_version="rag-dedup-v2",
        )
        row = _latest_result_prepare_row(bootstrap.get_state().local_store)
        # Pool = fetch_n = requested_n(5) * overfetch_factor(3) = 15, in 5
        # groups of 3 near-duplicates each. Dedup keeps the first (highest-
        # ranked) member of each group -- 5 survive, 10 are removed -- before
        # the survivors are truncated back to requested_n(5).
        assert row["attributes"]["duplicates_removed_from_pool"] == 10
    finally:
        bootstrap.shutdown_observability(timeout_seconds=2.0)


def test_naive_v1_reports_zero_duplicates_removed_from_pool(isolated_db, tmp_path):
    """v1 never over-fetches or dedups, so `duplicates_removed_from_pool` is
    always exactly 0 -- by construction (the dedup branch never runs), not
    because this fixture's naive top-5 pool happens to be duplicate-free (it
    isn't: see test_naive_v1_default_returns_low_diversity_top5's
    duplicate_candidate_ratio=3/5 on that same top-5)."""
    bootstrap.initialize_observability(config=_local_observability_config(tmp_path))
    try:
        _seed(isolated_db)
        isolated_db.query_observed(
            "library", _query_vec(isolated_db), n_results=5, where={"project_id": "proj-1"}
        )
        row = _latest_result_prepare_row(bootstrap.get_state().local_store)
        assert row["attributes"]["duplicates_removed_from_pool"] == 0
    finally:
        bootstrap.shutdown_observability(timeout_seconds=2.0)