Merge pull request #479 from ggozad/feat/speed-up-duplicates

Speed up doctor's duplicate-document detection

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

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:

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
+    similarity_threshold: 0.97   # cosine cutoff on document embedding centroids
-    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
     min_chunks: 3                # documents with fewer chunks are excluded
 
 prompts:

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(
+        status = (
-            self.config, self.db_path, dict(os.environ), duplicates_out=duplicates_out
+            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]",

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
+    Detection clusters whole documents whose embedding centroids are nearly
-    one another). ``containment_threshold`` is the decision knob; the others
+    identical — the same document ingested twice, or a light revision.
-    tune the cheap centroid pre-filter, what counts as a shared chunk, and which
+    ``similarity_threshold`` is the cosine cutoff; ``min_chunks`` skips
-    tiny documents to skip.
+    documents too small to compare meaningfully.
     """
 
-    containment_threshold: float = 0.75
+    similarity_threshold: float = 0.97
-    candidate_threshold: float = 0.85
-    twin_similarity: float = 0.95
     min_chunks: int = 3
 
 

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
+    keep: str
-    pairs: list[tuple[str, str, float, float]]
+    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
+    ``centroids`` holds one summed (unnormalized) centroid per document and
-    pairs, then directed chunk-overlap containment confirms them. Returns one
+    ``counts`` its embedded-chunk count. Documents below the small-document
-    entry per connected component of confirmed pairs.
+    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
+    centroids = np.asarray(centroids, dtype=np.float32)
-    # floor; normalize the rest to unit length.
+    counts = np.asarray(counts)
-    normalized: dict[str, np.ndarray] = {}
+    norms = np.linalg.norm(centroids, axis=1)
-    for doc_id, matrix in doc_vectors.items():
+    eligible = np.nonzero((counts >= cfg.min_chunks) & (norms > 0))[0]
-        m = np.asarray(matrix, dtype=float)
+    if eligible.size < 2:
-        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:
         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)
+    # Pairwise cosine, block-wise to avoid a full D×D matrix at once. Each row
-    centroids = np.array([_unit(normalized[d].mean(axis=0)) for d in order])
+    # only compares against higher-indexed documents (upper triangle). Cluster
-    sizes = np.array([normalized[d].shape[0] for d in order], dtype=float)
+    # 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.
+    def find(x: int) -> int:
-    # The cap is enforced per row (truncating each row's matches) so a
+        while parent[x] != x:
-    # self-similar corpus can never allocate beyond MAX_CANDIDATE_PAIRS.
+            parent[x] = parent[parent[x]]
-    candidates: list[tuple[int, int]] = []
+            x = parent[x]
+        return x
+
+    best = np.zeros(n, dtype=np.float32)
+    linked = False
     block = 512
-    capped = False
+    for start in range(0, n, block):
-    for start in range(0, len(order), block):
+        sims = unit[start : start + block] @ unit.T
-        if capped:
-            break
-        sims = centroids[start : start + block] @ centroids.T
         for row in range(sims.shape[0]):
             gi = start + row
-            targets = np.arange(gi + 1, len(order))
+            cols = (
-            if targets.size == 0:
+                gi + 1 + np.nonzero(sims[row, gi + 1 :] >= cfg.similarity_threshold)[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]
             )
-            thresholds = cfg.candidate_threshold * np.sqrt(ratios)
+            if cols.size == 0:
-            above = np.nonzero(sims[row, gi + 1 :] >= thresholds)[0]
+                continue
-            remaining = MAX_CANDIDATE_PAIRS - len(candidates)
+            linked = True
-            if len(above) >= remaining:
+            row_best = sims[row, cols]
-                above = above[:remaining]
+            best[gi] = max(best[gi], float(row_best.max()))
-                capped = True
+            best[cols] = np.maximum(best[cols], row_best)
-            candidates.extend((gi, gi + 1 + int(j)) for j in above)
+            ri = find(gi)
-            if capped:
+            for gj in cols.tolist():
-                break
+                parent[find(gj)] = ri
-
+    if not linked:
-    # 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:
         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 indices in components.values():
-    for node in adjacency:
+        if len(indices) < 2:
-        if node in seen:
             continue
-        component: set[int] = set()
+        members = sorted(ids[i] for i in indices)
-        stack = [node]
+        # Largest document (most chunks) is the one to keep; smallest id on a tie.
-        while stack:
+        keep = min(members, key=lambda d: (-sizes[d], d))
-            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
-        )
         families.append(
             _DuplicateFamily(
                 members=members,
-                superset=superset,
+                keep=keep,
-                pairs=pairs,
+                similarity={ids[i]: round(float(best[i]), 3) for i in indices},
-                sizes={d: int(normalized[d].shape[0]) for d in members},
+                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
+    """One block per group; ``keep_suggested`` marks the document to keep and
-    ``contained_fraction`` is how much of the document is covered by the rest."""
+    ``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),
+                        "similarity": family.similarity[member],
-                        "keep_suggested": member == family.superset,
+                        "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(
+        sims = ", ".join(
-            f"#{number[a]}→#{number[b]} {ab:.0%}, #{number[b]}→#{number[a]} {ba:.0%}"
+            f"#{number[m]} {family.similarity[m]:.0%}" for m in family.members
-            for a, b, ab, ba in family.pairs
         )
-        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"{len(families)} group(s) of near-identical documents "
-            f"(potential duplicates/revisions), {total_docs} 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()
+    # Reshape the Arrow fixed-size-list child buffer directly into an (N, dim)
-    vectors = np.asarray(arrow.column("vector").to_pylist(), dtype=float)
+    # float32 matrix. Going through to_pylist() would box N*dim Python floats
-    zero_ids: list[str] = []
+    # (tens of GB and most of the wall-clock on large corpora); the stored
-    if vectors.size:
+    # vectors are already float32, so this keeps the layout and the dtype.
-        zero_ids = [ids[i] for i in np.nonzero(~vectors.any(axis=1))[0]]
+    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."
+                f"{zero_count} chunk(s) have an all-zero (unembedded) vector."
-                if zero_ids
+                if zero_count
                 else "All chunks are embedded."
             ),
-            remediation="haiku-rag rebuild --embed-only" if zero_ids else None,
+            remediation="haiku-rag rebuild --embed-only" if zero_count else None,
-            details=_sample(zero_ids),
+            details=zero_sample,
         )
     )
 
-    # Near-duplicate documents (revisions sharing most chunks), grouped from the
+    notify("Detecting near-duplicate documents")
-    # same vector scan rather than a second pass.
+    # Near-identical documents (centroid cosine). Reduce each document's chunk
-    chunk_doc_ids_ordered = arrow.column("document_id").to_pylist()
+    # vectors to one summed centroid during the scan: dictionary-encode the
-    indices_by_doc: dict[str, list[int]] = {}
+    # document ids into integer codes, then sum each document's embedded rows in
-    for index, doc_id in enumerate(chunk_doc_ids_ordered):
+    # a single pass per document — no second full copy of the vector matrix.
-        indices_by_doc.setdefault(doc_id, []).append(index)
+    encoded = arrow.column("document_id").combine_chunks().dictionary_encode()
-    doc_vectors = {doc_id: vectors[idx] for doc_id, idx in indices_by_doc.items()}
+    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)

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]:
+def _centroids(
-    """Build per-document chunk matrices from one-hot indices.
+    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
+    Orthogonal one-hot chunks make the centroid cosine of two documents equal to
-    indices are orthogonal, so they never count as twins.
+    ``shared / sqrt(len(a) * len(b))``: identical documents score 1.0, fully
+    distinct documents score 0.0.
     """
     eye = np.eye(dim)
-    return {
+    doc_ids = list(spec)
-        doc: np.array([eye[i] for i in idxs], dtype=float) for doc, idxs in spec.items()
+    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
+def test_duplicate_families_identical_docs_flagged():
-# (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():
     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()
+        *_centroids({"a": [0, 1, 2], "b": [0, 1, 2, 3, 4, 5]}),
-    )
-    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]}),
         DuplicateDetectionConfig(),
     )
-    assert len(families) == 1
+    assert families == []
-    assert set(families[0].members) == {"a", "b"}
 
 
 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(
+        *_centroids({"a": [0, 1, 2, 3], "b": [0, 1, 2, 3], "c": [0, 1, 2, 3]}),
-            {
+        DuplicateDetectionConfig(),
-                "a": [0, 1, 2, 3],
-                "b": [0, 1, 2, 3, 4],
-                "c": [0, 1, 2, 3, 4, 5],
-            }
-        ),
-        _stage2_cfg(),
     )
     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.
+    # Share 3 of 4 chunks each -> centroid cosine 0.75.
-    spec = {"a": [0, 1, 2, 3, 4], "b": [0, 1, 2, 5, 6]}
+    spec = {"a": [0, 1, 2, 3], "b": [0, 1, 2, 4]}
-    assert _duplicate_families(_docs(spec), _stage2_cfg()) == []
+    assert _duplicate_families(*_centroids(spec), DuplicateDetectionConfig()) == []
-    flagged = _duplicate_families(_docs(spec), _stage2_cfg(containment_threshold=0.6))
+    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
+    # Share 3 of 5 -> centroid cosine 0.6, below the default 0.97 cutoff.
-    # only containment_threshold decides the outcome.
     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(
+        duplicates=DuplicateDetectionConfig(similarity_threshold=0.5)
-            candidate_threshold=0.0, containment_threshold=0.6
-        )
     )
     assert (
         _result(