""" 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")