java-topology/defects/langchain-0002/test/LangChainMultiQueryUniqueDedupTest.py

"""
LangChainMultiQueryUniqueDedupTest.py

Unit test for langchain-0002: MultiQueryRetriever _unique_documents O(D^2) fix.

Tests both the defective (O(D^2)) and fixed (O(D)) implementations to confirm:
1. Correctness: fixed returns same result as defective.
2. Performance: fixed is significantly faster at scale.

No install required. Stubs Document with a minimal dataclass.
"""

import time
import sys
from typing import Any, Optional, Sequence


class Document:
    """Minimal Document stub matching LangChain's Document interface."""
    def __init__(self, page_content: str, metadata: dict = None, id: Optional[str] = None):
        self.page_content = page_content
        self.metadata = metadata if metadata is not None else {}
        self.id = id

    def __eq__(self, other):
        if not isinstance(other, Document):
            return False
        return (self.id == other.id and
                self.page_content == other.page_content and
                self.metadata == other.metadata)

    def __repr__(self):
        return f"Document(id={self.id!r}, page_content={self.page_content[:20]!r})"


# ---- Defective implementation (O(D^2)) ----
def _unique_documents_defective(documents: Sequence[Document]) -> list[Document]:
    """Original O(D^2): slice-in-loop + linear scan."""
    return [doc for i, doc in enumerate(documents) if doc not in documents[:i]]


# ---- Fixed implementation (O(D)) ----
def _unique_documents_fixed(documents: Sequence[Document]) -> list[Document]:
    """Fixed O(D): hashable proxy key in a seen set.

    metadata values may be unhashable (e.g. list), so we always stringify values.
    This keeps the key hashable regardless of metadata value types.
    """
    seen: set = set()
    result: list = []
    for doc in documents:
        meta_key: tuple = tuple(sorted((k, str(v)) for k, v in doc.metadata.items()))
        key = (doc.id, doc.page_content, meta_key)
        if key not in seen:
            seen.add(key)
            result.append(doc)
    return result


# ---- Tests ----

def test_empty():
    assert _unique_documents_fixed([]) == []
    print("PASS test_empty")


def test_no_duplicates():
    docs = [Document(f"content {i}", {"src": f"src{i}"}) for i in range(5)]
    result = _unique_documents_fixed(docs)
    assert len(result) == 5
    assert result == docs
    print("PASS test_no_duplicates")


def test_all_duplicates():
    doc = Document("same content", {"src": "a"})
    docs = [doc, doc, doc, doc]
    result = _unique_documents_fixed(docs)
    assert len(result) == 1
    assert result[0] == doc
    print("PASS test_all_duplicates")


def test_preserves_order():
    """First occurrence is kept."""
    docs = [
        Document("alpha", {"src": "a"}),
        Document("beta", {"src": "b"}),
        Document("alpha", {"src": "a"}),   # duplicate of first
        Document("gamma", {"src": "c"}),
    ]
    result = _unique_documents_fixed(docs)
    assert len(result) == 3
    assert result[0].page_content == "alpha"
    assert result[1].page_content == "beta"
    assert result[2].page_content == "gamma"
    print("PASS test_preserves_order")


def test_same_content_different_metadata():
    """Same page_content but different metadata -> both kept."""
    docs = [
        Document("shared content", {"src": "a"}),
        Document("shared content", {"src": "b"}),
    ]
    result = _unique_documents_fixed(docs)
    assert len(result) == 2
    print("PASS test_same_content_different_metadata")


def test_with_ids():
    """Documents with distinct ids are deduped by id+content+meta."""
    docs = [
        Document("content", {"src": "a"}, id="doc-1"),
        Document("content", {"src": "a"}, id="doc-1"),  # exact duplicate
        Document("content", {"src": "a"}, id="doc-2"),  # different id -> kept
    ]
    result = _unique_documents_fixed(docs)
    assert len(result) == 2
    print("PASS test_with_ids")


def test_matches_defective_on_small():
    """Fixed produces same result as defective on small inputs."""
    docs = []
    for i in range(20):
        # Every other doc is a duplicate
        idx = i if i % 2 == 0 else i - 1
        docs.append(Document(f"content {idx}", {"src": f"src{idx}"}))

    expected = _unique_documents_defective(docs)
    result = _unique_documents_fixed(docs)
    assert result == expected, f"Mismatch:\nexpected={expected}\ngot={result}"
    print("PASS test_matches_defective_on_small")


def test_performance():
    """Fixed is at least 10x faster than defective at D=500."""
    # Q=10 queries, k=50 results each, 30% duplicates
    import random
    random.seed(42)
    unique_count = 350
    unique_docs = [Document(f"content {i} " + "x" * 200, {"src": f"src{i}", "idx": i})
                   for i in range(unique_count)]
    # Build D=500 docs with 150 duplicates
    docs = unique_docs[:]
    for _ in range(150):
        docs.append(random.choice(unique_docs))
    random.shuffle(docs)
    D = len(docs)

    # Warmup
    _unique_documents_defective(docs[:50])
    _unique_documents_fixed(docs[:50])

    t0 = time.perf_counter()
    expected = _unique_documents_defective(docs)
    t_defective = time.perf_counter() - t0

    t0 = time.perf_counter()
    result = _unique_documents_fixed(docs)
    t_fixed = time.perf_counter() - t0

    # Correctness: same set of unique documents (order may differ slightly due to
    # key construction, but set equality should hold)
    assert set(d.page_content for d in result) == set(d.page_content for d in expected)
    assert len(result) == len(expected)

    speedup = t_defective / t_fixed if t_fixed > 0 else float("inf")
    print(f"PASS test_performance: D={D}, defective={t_defective*1000:.1f}ms, "
          f"fixed={t_fixed*1000:.1f}ms, speedup={speedup:.1f}x")

    assert speedup >= 5.0, f"Expected >=5x speedup, got {speedup:.1f}x"


def test_unhashable_metadata_values():
    """Docs with unhashable metadata values (e.g. list) are handled via str fallback."""
    docs = [
        Document("content", {"tags": ["a", "b"]}),
        Document("content", {"tags": ["a", "b"]}),  # duplicate
        Document("content", {"tags": ["c", "d"]}),  # different -> kept
    ]
    result = _unique_documents_fixed(docs)
    assert len(result) == 2
    print("PASS test_unhashable_metadata_values")


if __name__ == "__main__":
    tests = [
        test_empty,
        test_no_duplicates,
        test_all_duplicates,
        test_preserves_order,
        test_same_content_different_metadata,
        test_with_ids,
        test_matches_defective_on_small,
        test_performance,
        test_unhashable_metadata_values,
    ]
    failed = 0
    for t in tests:
        try:
            t()
        except Exception as e:
            print(f"FAIL {t.__name__}: {e}")
            failed += 1
    if failed:
        print(f"\n{failed}/{len(tests)} FAILED")
        sys.exit(1)
    else:
        print(f"\n{len(tests)}/{len(tests)} PASSED")