diff --git a/CHANGELOG.md b/CHANGELOG.md index 65a52dec..d8aa55c6 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -1,6 +1,16 @@ # Changelog ## [Unreleased] +### Added + +- `doctor` shows a spinner naming the check currently in progress while it runs. + +### Changed + +- `doctor` near-duplicate detection compares whole-document embedding centroids instead of per-chunk overlap, and no longer flags a small document contained in a larger one. The per-chunk check could take hours and tens of GB of memory on corpora with large documents. The centroid check is independent of document size and reduces each document to a centroid during the vector scan without a second copy of the matrix. +- **Breaking:** the `doctor.duplicates` keys `containment_threshold`, `candidate_threshold`, and `twin_similarity` are replaced by a single `similarity_threshold` (default `0.97`). `min_chunks` is unchanged. +- `doctor --duplicates-out` YAML reports `similarity` per document instead of `contained_fraction`. + ## [0.62.1] - 2026-06-27 ### Fixed diff --git a/docs/cli.md b/docs/cli.md index d432c94f..92e25855 100644 --- a/docs/cli.md +++ b/docs/cli.md @@ -292,7 +292,9 @@ Check the database for consistency problems and print a pass/warn/fail report: haiku-rag doctor [--db /path/to/your.lancedb] [--duplicates-out groups.yaml] ``` -`--duplicates-out PATH` additionally writes the near-duplicate document groups to a YAML file (one block per group with `keep` and a list of `documents`, each carrying `document_id`, `document`, `chunks`, `contained_fraction`, and `keep_suggested`) for offline review. +While it runs, doctor shows a spinner naming the check currently in progress. + +`--duplicates-out PATH` additionally writes the near-duplicate document groups to a YAML file (one block per group with `keep` and a list of `documents`, each carrying `document_id`, `document`, `chunks`, `similarity`, and `keep_suggested`) for offline review. Checks include: @@ -309,7 +311,7 @@ Checks include: - the configured embedding identity matches the stored settings - no database migrations are pending - the vector index covers all chunks -- near-duplicate documents (revisions sharing most of their chunks) are grouped and reported, with the largest member flagged as the likely one to keep (advisory only, never deleted; tuned via `doctor.duplicates` in config) +- near-identical documents (by embedding-centroid similarity) are grouped and reported, with the largest member flagged as the likely one to keep (advisory only, never deleted, tuned via `doctor.duplicates` in config) - API keys are set for configured providers It also probes the external endpoints the config uses and reports them under a Providers section: diff --git a/docs/configuration/index.md b/docs/configuration/index.md index 4c3c558d..2c76fd66 100644 --- a/docs/configuration/index.md +++ b/docs/configuration/index.md @@ -116,9 +116,7 @@ search: doctor: duplicates: # Near-duplicate document detection (doctor command) - containment_threshold: 0.75 # flag a group when one doc shares >= this fraction of the smaller's chunks - candidate_threshold: 0.85 # centroid similarity gate for proposing candidate pairs (recall) - twin_similarity: 0.95 # cosine at which two chunks count as the same chunk + similarity_threshold: 0.97 # cosine cutoff on document embedding centroids min_chunks: 3 # documents with fewer chunks are excluded prompts: diff --git a/haiku_rag_slim/haiku/rag/app.py b/haiku_rag_slim/haiku/rag/app.py index 07c9c691..e6a68fd0 100644 --- a/haiku_rag_slim/haiku/rag/app.py +++ b/haiku_rag_slim/haiku/rag/app.py @@ -201,6 +201,7 @@ class HaikuRAGApp: # pragma: no cover async def doctor(self, duplicates_out: Path | None = None) -> bool: """Run health checks and print a report. Returns True if any check failed.""" import os + from contextlib import nullcontext from haiku.rag.doctor import Severity, run_doctor @@ -213,10 +214,24 @@ class HaikuRAGApp: # pragma: no cover self.console.print("[red]Database path does not exist.[/red]") return True - report = await run_doctor( - self.config, self.db_path, dict(os.environ), duplicates_out=duplicates_out + status = ( + self.console.status("Running checks") if self.console.is_terminal else None ) + def on_progress(label: str) -> None: + if status is not None: + status.update(f"{label}...") + + cm = status if status is not None else nullcontext() + with cm: + report = await run_doctor( + self.config, + self.db_path, + dict(os.environ), + duplicates_out=duplicates_out, + on_progress=on_progress, + ) + glyphs = { Severity.OK: "[green]✓[/green]", Severity.WARN: "[yellow]![/yellow]", diff --git a/haiku_rag_slim/haiku/rag/config/models.py b/haiku_rag_slim/haiku/rag/config/models.py index f3a1a564..0f2f682d 100644 --- a/haiku_rag_slim/haiku/rag/config/models.py +++ b/haiku_rag_slim/haiku/rag/config/models.py @@ -106,15 +106,13 @@ class AnalysisConfig(BaseModel): class DuplicateDetectionConfig(BaseModel): """Thresholds for doctor's near-duplicate document detection. - Detection clusters documents that share most of their chunks (revisions of - one another). ``containment_threshold`` is the decision knob; the others - tune the cheap centroid pre-filter, what counts as a shared chunk, and which - tiny documents to skip. + Detection clusters whole documents whose embedding centroids are nearly + identical — the same document ingested twice, or a light revision. + ``similarity_threshold`` is the cosine cutoff; ``min_chunks`` skips + documents too small to compare meaningfully. """ - containment_threshold: float = 0.75 - candidate_threshold: float = 0.85 - twin_similarity: float = 0.95 + similarity_threshold: float = 0.97 min_chunks: int = 3 diff --git a/haiku_rag_slim/haiku/rag/doctor.py b/haiku_rag_slim/haiku/rag/doctor.py index 427dd53e..8b7c3f06 100644 --- a/haiku_rag_slim/haiku/rag/doctor.py +++ b/haiku_rag_slim/haiku/rag/doctor.py @@ -235,134 +235,90 @@ async def _column_values(table, column: str) -> list: return [row[column] for row in rows] -# Backstop on how many centroid candidate pairs we verify, bounding memory and -# runtime on a pathologically self-similar corpus. -MAX_CANDIDATE_PAIRS = 200_000 - - class _DuplicateFamily(BaseModel): members: list[str] - superset: str - pairs: list[tuple[str, str, float, float]] + keep: str + similarity: dict[str, float] sizes: dict[str, int] -def _unit(vector: np.ndarray) -> np.ndarray: - norm = np.linalg.norm(vector) - return vector / norm if norm else vector - - -def _containment(source: np.ndarray, target: np.ndarray, twin: float) -> float: - """Fraction of ``source`` chunks with a near-identical chunk in ``target``.""" - return float(((source @ target.T).max(axis=1) >= twin).mean()) - - def _duplicate_families( - doc_vectors: dict[str, np.ndarray], cfg: DuplicateDetectionConfig + doc_ids: list[str], + centroids: np.ndarray, + counts: np.ndarray, + cfg: DuplicateDetectionConfig, ) -> list[_DuplicateFamily]: - """Cluster documents that share most of their chunks (revisions of one another). + """Cluster documents whose embedding centroids are nearly identical. - Block-then-verify: cheap centroid similarity proposes candidate document - pairs, then directed chunk-overlap containment confirms them. Returns one - entry per connected component of confirmed pairs. + ``centroids`` holds one summed (unnormalized) centroid per document and + ``counts`` its embedded-chunk count. Documents below the small-document + floor are dropped; the rest are normalized and clustered by union-find over + pairwise cosine above ``similarity_threshold``. One family per component, + each carrying every member's highest cosine to another member. """ - # Drop unembedded (zero) vectors and documents below the small-document - # floor; normalize the rest to unit length. - normalized: dict[str, np.ndarray] = {} - for doc_id, matrix in doc_vectors.items(): - m = np.asarray(matrix, dtype=float) - if m.ndim != 2 or m.shape[0] == 0: - continue - m = m[np.linalg.norm(m, axis=1) > 0] - if m.shape[0] < cfg.min_chunks: - continue - normalized[doc_id] = m / np.linalg.norm(m, axis=1)[:, None] - if len(normalized) < 2: + centroids = np.asarray(centroids, dtype=np.float32) + counts = np.asarray(counts) + norms = np.linalg.norm(centroids, axis=1) + eligible = np.nonzero((counts >= cfg.min_chunks) & (norms > 0))[0] + if eligible.size < 2: return [] + unit = centroids[eligible] / norms[eligible][:, None] + ids = [doc_ids[i] for i in eligible] + sizes = {doc_ids[i]: int(counts[i]) for i in eligible} + n = len(ids) - order = sorted(normalized) - centroids = np.array([_unit(normalized[d].mean(axis=0)) for d in order]) - sizes = np.array([normalized[d].shape[0] for d in order], dtype=float) + # Pairwise cosine, block-wise to avoid a full D×D matrix at once. Each row + # only compares against higher-indexed documents (upper triangle). Cluster + # with union-find and keep only each document's best similarity to a twin — + # a self-similar corpus forms one clique, so storing every pair would be + # O(D²) objects. + parent = list(range(n)) - # Stage 1: centroid candidate pairs, block-wise to avoid a full D×D matrix. - # The cap is enforced per row (truncating each row's matches) so a - # self-similar corpus can never allocate beyond MAX_CANDIDATE_PAIRS. - candidates: list[tuple[int, int]] = [] + def find(x: int) -> int: + while parent[x] != x: + parent[x] = parent[parent[x]] + x = parent[x] + return x + + best = np.zeros(n, dtype=np.float32) + linked = False block = 512 - capped = False - for start in range(0, len(order), block): - if capped: - break - sims = centroids[start : start + block] @ centroids.T + for start in range(0, n, block): + sims = unit[start : start + block] @ unit.T for row in range(sims.shape[0]): gi = start + row - targets = np.arange(gi + 1, len(order)) - if targets.size == 0: - continue - # A smaller document can be fully contained in a larger append-only - # revision even when the fixed centroid threshold would fail: - # with orthogonal chunks, cosine falls to sqrt(small / large). - # Scale the candidate gate by that size ratio, then let directed - # containment make the actual duplicate decision. - ratios = np.minimum(sizes[gi], sizes[targets]) / np.maximum( - sizes[gi], sizes[targets] + cols = ( + gi + 1 + np.nonzero(sims[row, gi + 1 :] >= cfg.similarity_threshold)[0] ) - thresholds = cfg.candidate_threshold * np.sqrt(ratios) - above = np.nonzero(sims[row, gi + 1 :] >= thresholds)[0] - remaining = MAX_CANDIDATE_PAIRS - len(candidates) - if len(above) >= remaining: - above = above[:remaining] - capped = True - candidates.extend((gi, gi + 1 + int(j)) for j in above) - if capped: - break - - # Stage 2: confirm candidates with directed chunk-overlap containment. - adjacency: dict[int, set[int]] = {} - edges: dict[tuple[int, int], tuple[float, float]] = {} - for i, j in candidates: - a_to_b = _containment( - normalized[order[i]], normalized[order[j]], cfg.twin_similarity - ) - b_to_a = _containment( - normalized[order[j]], normalized[order[i]], cfg.twin_similarity - ) - if max(a_to_b, b_to_a) >= cfg.containment_threshold: - adjacency.setdefault(i, set()).add(j) - adjacency.setdefault(j, set()).add(i) - edges[(i, j)] = (a_to_b, b_to_a) - if not edges: + if cols.size == 0: + continue + linked = True + row_best = sims[row, cols] + best[gi] = max(best[gi], float(row_best.max())) + best[cols] = np.maximum(best[cols], row_best) + ri = find(gi) + for gj in cols.tolist(): + parent[find(gj)] = ri + if not linked: return [] - # Cluster confirmed pairs into families (connected components). + components: dict[int, list[int]] = {} + for idx in range(n): + components.setdefault(find(idx), []).append(idx) + families: list[_DuplicateFamily] = [] - seen: set[int] = set() - for node in adjacency: - if node in seen: + for indices in components.values(): + if len(indices) < 2: continue - component: set[int] = set() - stack = [node] - while stack: - cur = stack.pop() - if cur in seen: - continue - seen.add(cur) - component.add(cur) - stack.extend(adjacency[cur] - seen) - members = sorted(order[i] for i in component) - # Largest document (most chunks) is the likely superset; smallest id on a tie. - superset = min(members, key=lambda d: (-normalized[d].shape[0], d)) - pairs = sorted( - (order[i], order[j], round(ab, 3), round(ba, 3)) - for (i, j), (ab, ba) in edges.items() - if i in component and j in component - ) + members = sorted(ids[i] for i in indices) + # Largest document (most chunks) is the one to keep; smallest id on a tie. + keep = min(members, key=lambda d: (-sizes[d], d)) families.append( _DuplicateFamily( members=members, - superset=superset, - pairs=pairs, - sizes={d: int(normalized[d].shape[0]) for d in members}, + keep=keep, + similarity={ids[i]: round(float(best[i]), 3) for i in indices}, + sizes={d: sizes[d] for d in members}, ) ) return sorted(families, key=lambda f: f.members) @@ -387,25 +343,21 @@ def _common_path_prefix(labels: list[str]) -> str: def _write_duplicates_out( path: Path, families: list[_DuplicateFamily], label: Callable[[str], str] ) -> None: - """One block per group; ``keep_suggested`` marks the superset and - ``contained_fraction`` is how much of the document is covered by the rest.""" + """One block per group; ``keep_suggested`` marks the document to keep and + ``similarity`` is the highest centroid cosine to another group member.""" groups = [] for n, family in enumerate(families, start=1): - contained = dict.fromkeys(family.members, 0.0) - for a, b, a_to_b, b_to_a in family.pairs: - contained[a] = max(contained[a], a_to_b) - contained[b] = max(contained[b], b_to_a) groups.append( { "group": n, - "keep": family.superset, + "keep": family.keep, "documents": [ { "document_id": member, "document": label(member), "chunks": family.sizes[member], - "contained_fraction": round(contained[member], 3), - "keep_suggested": member == family.superset, + "similarity": family.similarity[member], + "keep_suggested": member == family.keep, } for member in family.members ], @@ -416,13 +368,15 @@ def _write_duplicates_out( def _check_duplicate_documents( - doc_vectors: dict[str, np.ndarray], + doc_ids: list[str], + centroids: np.ndarray, + counts: np.ndarray, uri_by_doc: Mapping[str, str | None], title_by_doc: Mapping[str, str | None], cfg: DuplicateDetectionConfig, yaml_path: Path | None = None, ) -> CheckResult: - families = _duplicate_families(doc_vectors, cfg) + families = _duplicate_families(doc_ids, centroids, counts, cfg) def label(doc_id: str) -> str: return uri_by_doc.get(doc_id) or title_by_doc.get(doc_id) or doc_id @@ -439,7 +393,7 @@ def _check_duplicate_documents( # The terminal report is a summary: show the first few groups whole and # point at the YAML export for the rest. One block per shown group — a - # header, each member on its own numbered line, then a compact overlap line. + # header, each member on its own numbered line, then a compact similarity line. shown = families[:_SAMPLE_LIMIT] prefix = _common_path_prefix([label(m) for f in shown for m in f.members]) @@ -453,15 +407,14 @@ def _check_duplicate_documents( for n, family in enumerate(shown, start=1): number = {member: i for i, member in enumerate(family.members, start=1)} details.append( - f"group {n} — {len(family.members)} docs, keep #{number[family.superset]}:" + f"group {n} — {len(family.members)} docs, keep #{number[family.keep]}:" ) for member in family.members: details.append(f" #{number[member]} {short(member)}") - overlaps = ", ".join( - f"#{number[a]}→#{number[b]} {ab:.0%}, #{number[b]}→#{number[a]} {ba:.0%}" - for a, b, ab, ba in family.pairs + sims = ", ".join( + f"#{number[m]} {family.similarity[m]:.0%}" for m in family.members ) - details.append(f" overlap: {overlaps}") + details.append(f" similarity: {sims}") if len(families) > len(shown): details.append( f"... (+{len(families) - len(shown)} more groups; " @@ -473,11 +426,11 @@ def _check_duplicate_documents( name="duplicate_documents", severity=Severity.WARN, message=( - f"{len(families)} group(s) of documents with substantial chunk overlap " - f"(potential duplicates/revisions), {total_docs} documents." + f"{len(families)} group(s) of near-identical documents " + f"(potential duplicates), {total_docs} documents." ), remediation=( - "Review each group and remove redundant revisions; overlap may be intentional." + "Review each group and remove redundant copies; duplication may be intentional." ), details=details, ) @@ -488,13 +441,16 @@ async def run_db_checks( config: AppConfig, stats: dict, duplicates_out: Path | None = None, + on_progress: Callable[[str], None] | None = None, ) -> list[CheckResult]: """Referential and content-integrity checks against an open read-only Store. Assumes all required tables exist (the caller short-circuits otherwise). """ + notify = on_progress or (lambda _label: None) results: list[CheckResult] = [] + notify("Reading document records") doc_ids = set(await _column_values(store.documents_table, "id")) meta_rows = ( await store.document_meta_table.query() @@ -511,6 +467,7 @@ async def run_db_checks( uri_by_doc = {row["document_id"]: row.get("uri") for row in meta_rows} title_by_doc = {row["document_id"]: row.get("title") for row in meta_rows} + notify("Reading chunks") chunk_rows = ( await store.chunks_table.query() .select(["id", "document_id", "metadata"]) @@ -518,6 +475,7 @@ async def run_db_checks( ) chunk_doc_ids = {row["document_id"] for row in chunk_rows} + notify("Reading document items") item_rows = ( await store.document_items_table.query() .select(["document_id", "self_ref", "label"]) @@ -530,6 +488,7 @@ async def run_db_checks( self_refs_by_doc.setdefault(row["document_id"], set()).add(row["self_ref"]) labels_by_doc.setdefault(row["document_id"], set()).add(row["label"]) + notify("Checking referential integrity") # documents <-> document_meta must be 1:1. orphan_docs = doc_ids - meta_doc_ids orphan_meta = meta_doc_ids - doc_ids @@ -585,6 +544,7 @@ async def run_db_checks( ) ) + notify("Checking document chunking") # Documents with no chunks, classified by what they contain and whether the # embedder can index images. results += _classify_unchunked( @@ -608,6 +568,7 @@ async def run_db_checks( ) ) + notify("Checking chunk references") # Chunk metadata may reference self_refs that do not exist for that document. dangling: list[str] = [] for row in chunk_rows: @@ -629,6 +590,7 @@ async def run_db_checks( ) ) + notify("Scanning chunk vectors") # Vector dimension consistency and unembedded (all-zero) vectors share one # scan of the vector column — the heaviest check on large corpora. arrow = ( @@ -660,35 +622,60 @@ async def run_db_checks( ) ) - ids = arrow.column("id").to_pylist() - vectors = np.asarray(arrow.column("vector").to_pylist(), dtype=float) - zero_ids: list[str] = [] - if vectors.size: - zero_ids = [ids[i] for i in np.nonzero(~vectors.any(axis=1))[0]] + # Reshape the Arrow fixed-size-list child buffer directly into an (N, dim) + # float32 matrix. Going through to_pylist() would box N*dim Python floats + # (tens of GB and most of the wall-clock on large corpora); the stored + # vectors are already float32, so this keeps the layout and the dtype. + vec_col = arrow.column("vector").combine_chunks() + vectors = vec_col.values.to_numpy(zero_copy_only=False).reshape(-1, actual_dim) + embedded = vectors.any(axis=1) if vectors.size else np.zeros(0, dtype=bool) + + # Unembedded (all-zero) chunks: report a count and a few sampled ids without + # materializing every chunk id. + zero_rows = np.nonzero(~embedded)[0] + zero_count = int(zero_rows.size) + id_col = arrow.column("id") + zero_sample = [id_col[int(i)].as_py() for i in zero_rows[:_SAMPLE_LIMIT]] + if zero_count > _SAMPLE_LIMIT: + zero_sample.append(f"... (+{zero_count - _SAMPLE_LIMIT} more)") results.append( CheckResult( name="unembedded_chunks", - severity=Severity.WARN if zero_ids else Severity.OK, + severity=Severity.WARN if zero_count else Severity.OK, message=( - f"{len(zero_ids)} chunk(s) have an all-zero (unembedded) vector." - if zero_ids + f"{zero_count} chunk(s) have an all-zero (unembedded) vector." + if zero_count else "All chunks are embedded." ), - remediation="haiku-rag rebuild --embed-only" if zero_ids else None, - details=_sample(zero_ids), + remediation="haiku-rag rebuild --embed-only" if zero_count else None, + details=zero_sample, ) ) - # Near-duplicate documents (revisions sharing most chunks), grouped from the - # same vector scan rather than a second pass. - chunk_doc_ids_ordered = arrow.column("document_id").to_pylist() - indices_by_doc: dict[str, list[int]] = {} - for index, doc_id in enumerate(chunk_doc_ids_ordered): - indices_by_doc.setdefault(doc_id, []).append(index) - doc_vectors = {doc_id: vectors[idx] for doc_id, idx in indices_by_doc.items()} + notify("Detecting near-duplicate documents") + # Near-identical documents (centroid cosine). Reduce each document's chunk + # vectors to one summed centroid during the scan: dictionary-encode the + # document ids into integer codes, then sum each document's embedded rows in + # a single pass per document — no second full copy of the vector matrix. + encoded = arrow.column("document_id").combine_chunks().dictionary_encode() + doc_ids = encoded.dictionary.to_pylist() + codes = encoded.indices.to_numpy(zero_copy_only=False) + centroids = np.zeros((len(doc_ids), actual_dim), dtype=np.float32) + counts = np.zeros(len(doc_ids), dtype=np.int64) + order = np.argsort(codes, kind="stable") + bounds = np.searchsorted(codes, np.arange(len(doc_ids) + 1), sorter=order) + for d in range(len(doc_ids)): + rows = order[bounds[d] : bounds[d + 1]] + rows = rows[embedded[rows]] + counts[d] = rows.size + if rows.size: + centroids[d] = vectors[rows].sum(axis=0) + del vectors results.append( _check_duplicate_documents( - doc_vectors, + doc_ids, + centroids, + counts, uri_by_doc, title_by_doc, config.doctor.duplicates, @@ -696,6 +683,7 @@ async def run_db_checks( ) ) + notify("Checking picture data") # Pictures from image/PDF sources should carry raster bytes. Pictures that # are external image references in a text document (markdown, HTML) have no # embedded bytes by nature, so a missing raster there is expected. @@ -726,6 +714,7 @@ async def run_db_checks( ) ) + notify("Checking settings and indexes") # Settings must hold exactly one canonical row. total_settings = await store.settings_table.count_rows() canonical = len( @@ -978,12 +967,16 @@ def _endpoint_result( ) -async def run_provider_checks(config: AppConfig) -> list[CheckResult]: +async def run_provider_checks( + config: AppConfig, on_progress: Callable[[str], None] | None = None +) -> list[CheckResult]: """Probe the external endpoints the current config actually uses.""" targets, local = _provider_targets(config) results: list[CheckResult] = [] if targets: + if on_progress is not None: + on_progress("Probing provider endpoints") async with httpx.AsyncClient(timeout=_PROBE_TIMEOUT_S) as client: probes = await asyncio.gather( *(_probe_endpoint(client, url) for url in targets) @@ -1007,12 +1000,15 @@ async def run_doctor( db_path: Path, environ: dict[str, str], duplicates_out: Path | None = None, + on_progress: Callable[[str], None] | None = None, ) -> DoctorReport: """Open the database read-only and run every diagnostic check. Opens with validation and migration checks skipped so a drifted or pre-migration database can still be diagnosed rather than refusing to open. """ + notify = on_progress or (lambda _label: None) + notify("Inspecting tables") db = await connect_lancedb(config, db_path) stats = await get_database_stats(db) @@ -1038,9 +1034,14 @@ async def run_doctor( skip_migration_check=True, ) as store: results += await run_db_checks( - store, config, stats, duplicates_out=duplicates_out + store, + config, + stats, + duplicates_out=duplicates_out, + on_progress=on_progress, ) + notify("Checking API keys") results.append(_check_api_keys(config, environ)) - results += await run_provider_checks(config) + results += await run_provider_checks(config, on_progress=on_progress) return DoctorReport(results=results) diff --git a/tests/test_doctor.py b/tests/test_doctor.py index 343b5c19..0f6ddc2f 100644 --- a/tests/test_doctor.py +++ b/tests/test_doctor.py @@ -172,6 +172,17 @@ async def test_healthy_db_all_ok(temp_db_path): assert all(r.severity is Severity.OK for r in report.results) +@pytest.mark.asyncio +async def test_doctor_reports_progress(temp_db_path): + await _build_db(temp_db_path) + labels: list[str] = [] + await run_doctor(_config(), temp_db_path, {}, on_progress=labels.append) + assert "Inspecting tables" in labels + assert "Scanning chunk vectors" in labels + assert "Detecting near-duplicate documents" in labels + assert "Probing provider endpoints" in labels + + @pytest.mark.asyncio async def test_empty_db_fails(temp_db_path): report = await run_doctor(_config(), temp_db_path, {}) @@ -944,102 +955,87 @@ async def test_probe_endpoint_connection_error(): # --- Duplicate-document detection ---------------------------------------- -def _docs(spec: dict[str, list[int]], dim: int = 8) -> dict[str, np.ndarray]: - """Build per-document chunk matrices from one-hot indices. +def _centroids( + spec: dict[str, list[int]], dim: int = 8 +) -> tuple[list[str], np.ndarray, np.ndarray]: + """Summed one-hot centroids + chunk counts per document, as + ``_duplicate_families`` consumes them. - A shared index across documents is a shared (identical) chunk; distinct - indices are orthogonal, so they never count as twins. + Orthogonal one-hot chunks make the centroid cosine of two documents equal to + ``shared / sqrt(len(a) * len(b))``: identical documents score 1.0, fully + distinct documents score 0.0. """ eye = np.eye(dim) - return { - doc: np.array([eye[i] for i in idxs], dtype=float) for doc, idxs in spec.items() - } + doc_ids = list(spec) + centroids = np.array( + [eye[idxs].sum(axis=0) for idxs in spec.values()], dtype=np.float32 + ) + counts = np.array([len(idxs) for idxs in spec.values()], dtype=np.int64) + return doc_ids, centroids, counts -# Stage-2 (containment/clustering) unit tests disable the centroid gate -# (candidate_threshold=0.0) so every pair is verified; orthogonal one-hot chunks -# would otherwise drop centroids below the default gate. The gate itself is -# exercised by the end-to-end tests below. -def _stage2_cfg(**kw) -> DuplicateDetectionConfig: - return DuplicateDetectionConfig(candidate_threshold=0.0, **kw) - - -def test_duplicate_families_revision_pair(): +def test_duplicate_families_identical_docs_flagged(): families = _duplicate_families( - _docs({"a": [0, 1, 2, 3], "b": [0, 1, 2, 4]}), _stage2_cfg() + *_centroids({"a": [0, 1, 2, 3], "b": [0, 1, 2, 3]}), DuplicateDetectionConfig() ) assert len(families) == 1 assert set(families[0].members) == {"a", "b"} + assert families[0].similarity == {"a": pytest.approx(1.0), "b": pytest.approx(1.0)} -def test_duplicate_families_append_only_is_asymmetric(): +def test_duplicate_families_append_only_not_flagged(): + # A is fully contained in the larger B, but their centroids diverge + # (cosine sqrt(3/6) ~= 0.71), so it stays below the similarity cutoff. families = _duplicate_families( - _docs({"a": [0, 1, 2], "b": [0, 1, 2, 3, 4, 5]}), _stage2_cfg() - ) - assert len(families) == 1 - fam = families[0] - assert fam.superset == "b" # the larger document - # directed containment: all of A is in B (1.0); only half of B is in A. - a_to_b = next(p for p in fam.pairs if p[:2] == ("a", "b")) - assert a_to_b[2] == pytest.approx(1.0) - assert a_to_b[3] == pytest.approx(0.5) - - -def test_duplicate_families_append_only_passes_default_centroid_gate(): - # Fixed 0.85 centroid gating misses this: centroid cosine is sqrt(3 / 6), - # but A is fully contained in B and should reach the containment verifier. - families = _duplicate_families( - _docs({"a": [0, 1, 2], "b": [0, 1, 2, 3, 4, 5]}), + *_centroids({"a": [0, 1, 2], "b": [0, 1, 2, 3, 4, 5]}), DuplicateDetectionConfig(), ) - assert len(families) == 1 - assert set(families[0].members) == {"a", "b"} + assert families == [] def test_duplicate_families_distinct_docs_none(): families = _duplicate_families( - _docs({"a": [0, 1, 2], "b": [3, 4, 5]}), _stage2_cfg() + *_centroids({"a": [0, 1, 2], "b": [3, 4, 5]}), DuplicateDetectionConfig() ) assert families == [] -def test_duplicate_families_three_way_chain_one_family(): +def test_duplicate_families_three_way_one_family(): families = _duplicate_families( - _docs( - { - "a": [0, 1, 2, 3], - "b": [0, 1, 2, 3, 4], - "c": [0, 1, 2, 3, 4, 5], - } - ), - _stage2_cfg(), + *_centroids({"a": [0, 1, 2, 3], "b": [0, 1, 2, 3], "c": [0, 1, 2, 3]}), + DuplicateDetectionConfig(), ) assert len(families) == 1 assert set(families[0].members) == {"a", "b", "c"} - assert families[0].superset == "c" + # Equal sizes -> smallest id is kept. + assert families[0].keep == "a" + + +def test_duplicate_families_clique_single_family(): + # A self-similar corpus (all identical) is one clique. Union-find collapses + # it to a single family without materializing every pair. + spec = {chr(ord("a") + k): [0, 1, 2, 3] for k in range(8)} + families = _duplicate_families(*_centroids(spec), DuplicateDetectionConfig()) + assert len(families) == 1 + assert set(families[0].members) == set(spec) + assert all(s == pytest.approx(1.0) for s in families[0].similarity.values()) def test_duplicate_families_tiny_docs_ignored(): # min_chunks = 3 excludes the one-chunk documents. - families = _duplicate_families(_docs({"a": [0], "b": [0]}), _stage2_cfg()) + families = _duplicate_families( + *_centroids({"a": [0], "b": [0]}), DuplicateDetectionConfig() + ) assert families == [] -def test_duplicate_families_caps_candidates_during_collection(monkeypatch): - # Three mutually-identical docs would yield 3 candidate pairs, but a cap of 1 - # must stop collection after the first (a,b), leaving c unconfirmed. - monkeypatch.setattr("haiku.rag.doctor.MAX_CANDIDATE_PAIRS", 1) - docs = _docs({"a": [0, 1, 2], "b": [0, 1, 2], "c": [0, 1, 2]}, dim=3) - families = _duplicate_families(docs, _stage2_cfg()) - assert len(families) == 1 - assert set(families[0].members) == {"a", "b"} - - def test_duplicate_families_threshold_is_configurable(): - # Share 3 of 5 chunks each -> containment 0.6 both ways. - spec = {"a": [0, 1, 2, 3, 4], "b": [0, 1, 2, 5, 6]} - assert _duplicate_families(_docs(spec), _stage2_cfg()) == [] - flagged = _duplicate_families(_docs(spec), _stage2_cfg(containment_threshold=0.6)) + # Share 3 of 4 chunks each -> centroid cosine 0.75. + spec = {"a": [0, 1, 2, 3], "b": [0, 1, 2, 4]} + assert _duplicate_families(*_centroids(spec), DuplicateDetectionConfig()) == [] + flagged = _duplicate_families( + *_centroids(spec), DuplicateDetectionConfig(similarity_threshold=0.7) + ) assert len(flagged) == 1 assert set(flagged[0].members) == {"a", "b"} @@ -1051,9 +1047,10 @@ def test_duplicate_documents_report_truncates_summary(): idx = [3 * k, 3 * k + 1, 3 * k + 2] spec[f"a{k}"] = idx spec[f"b{k}"] = list(idx) - docs = _docs(spec, dim=3 * pairs) uris = {d: f"file:///srv/shared/library/docs/{d}.pdf" for d in spec} - result = _check_duplicate_documents(docs, uris, {}, _stage2_cfg()) + result = _check_duplicate_documents( + *_centroids(spec, dim=3 * pairs), uris, {}, DuplicateDetectionConfig() + ) assert result.severity is Severity.WARN # The summary message still reports the full total. assert f"{pairs} group(s)" in result.message @@ -1064,10 +1061,14 @@ def test_duplicate_documents_report_truncates_summary(): def test_duplicate_documents_report_factors_common_path(): - docs = _docs({"a": [0, 1, 2], "b": [0, 1, 2]}, dim=3) base = "file:///srv/shared/library/docs/" uris = {"a": base + "alpha.pdf", "b": base + "beta.pdf"} - result = _check_duplicate_documents(docs, uris, {}, _stage2_cfg()) + result = _check_duplicate_documents( + *_centroids({"a": [0, 1, 2], "b": [0, 1, 2]}, dim=3), + uris, + {}, + DuplicateDetectionConfig(), + ) assert f"common path: {base}" in result.details member_lines = [d for d in result.details if d.lstrip().startswith("#")] assert {d.strip() for d in member_lines} == {"#1 alpha.pdf", "#2 beta.pdf"} @@ -1076,10 +1077,12 @@ def test_duplicate_documents_report_factors_common_path(): def test_duplicate_documents_writes_yaml(tmp_path): # a,b identical (a 4-chunk duplicate); c distinct and excluded. - docs = _docs({"a": [0, 1, 2, 3], "b": [0, 1, 2, 3], "c": [4, 5, 6]}, dim=8) + spec = {"a": [0, 1, 2, 3], "b": [0, 1, 2, 3], "c": [4, 5, 6]} uris = {"a": "file:///x/a.pdf", "b": "file:///x/b.pdf", "c": "file:///x/c.pdf"} out = tmp_path / "dups.yaml" - _check_duplicate_documents(docs, uris, {}, _stage2_cfg(), yaml_path=out) + _check_duplicate_documents( + *_centroids(spec, dim=8), uris, {}, DuplicateDetectionConfig(), yaml_path=out + ) data = yaml.safe_load(out.read_text()) assert len(data["groups"]) == 1 group = data["groups"][0] @@ -1088,6 +1091,7 @@ def test_duplicate_documents_writes_yaml(tmp_path): assert [d["document_id"] for d in docs_out] == ["a", "b"] assert [d["document"] for d in docs_out] == ["file:///x/a.pdf", "file:///x/b.pdf"] assert all(d["chunks"] == 4 for d in docs_out) + assert all(d["similarity"] == pytest.approx(1.0) for d in docs_out) assert {d["document_id"]: d["keep_suggested"] for d in docs_out} == { "a": True, "b": False, @@ -1095,10 +1099,14 @@ def test_duplicate_documents_writes_yaml(tmp_path): def test_duplicate_documents_writes_empty_yaml_when_none(tmp_path): - docs = _docs({"a": [0, 1, 2], "b": [3, 4, 5]}, dim=6) # distinct + spec = {"a": [0, 1, 2], "b": [3, 4, 5]} # distinct out = tmp_path / "dups.yaml" _check_duplicate_documents( - docs, {"a": "u", "b": "v"}, {}, _stage2_cfg(), yaml_path=out + *_centroids(spec, dim=6), + {"a": "u", "b": "v"}, + {}, + DuplicateDetectionConfig(), + yaml_path=out, ) assert yaml.safe_load(out.read_text()) == {"groups": []} @@ -1162,7 +1170,7 @@ async def _build_dup_db(path, docs: dict[str, list[int]], *, vector_dim: int = 8 @pytest.mark.asyncio async def test_duplicate_documents_check_warns_end_to_end(temp_db_path): - await _build_dup_db(temp_db_path, {"a": [0, 1, 2, 3], "b": [0, 1, 2, 3, 4]}) + await _build_dup_db(temp_db_path, {"a": [0, 1, 2, 3], "b": [0, 1, 2, 3]}) report = await run_doctor(_config(vector_dim=8), temp_db_path, {}) result = _result(report, "duplicate_documents") assert result.severity is Severity.WARN @@ -1179,14 +1187,10 @@ async def test_duplicate_documents_check_ok_when_distinct(temp_db_path): @pytest.mark.asyncio async def test_duplicate_documents_check_reads_config(temp_db_path): - # Share 3 of 5 -> containment 0.6. Disable the centroid gate on both runs so - # only containment_threshold decides the outcome. + # Share 3 of 5 -> centroid cosine 0.6, below the default 0.97 cutoff. await _build_dup_db(temp_db_path, {"a": [0, 1, 2, 3, 4], "b": [0, 1, 2, 5, 6]}) base = _config(vector_dim=8) - base.doctor = DoctorConfig( - duplicates=DuplicateDetectionConfig(candidate_threshold=0.0) - ) assert ( _result( await run_doctor(base, temp_db_path, {}), "duplicate_documents" @@ -1196,9 +1200,7 @@ async def test_duplicate_documents_check_reads_config(temp_db_path): tuned = _config(vector_dim=8) tuned.doctor = DoctorConfig( - duplicates=DuplicateDetectionConfig( - candidate_threshold=0.0, containment_threshold=0.6 - ) + duplicates=DuplicateDetectionConfig(similarity_threshold=0.5) ) assert ( _result(