danvics/haiku.rag
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions

Merge pull request #475 from ggozad/feat/document-similarity

Browse source 
Detect near-duplicate documents in `doctor`
This commit is contained in:
Yiorgis Gozadinos 2026-06-26 17:05:33 +03:00 committed by GitHub
commit 7cdf93e23a
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
8 changed files with 611 additions and 9 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

4
CHANGELOG.md
View file

@ -1,6 +1,10 @@
# Changelog
## [Unreleased]
### Added
- `doctor` reports groups of near-duplicate documents (revisions sharing most of their chunks), flagging the largest member as the likely one to keep; `--duplicates-out PATH` writes the groups to a YAML file. Tuned via the `doctor.duplicates` config block (`containment_threshold`, `candidate_threshold`, `twin_similarity`, `min_chunks`).
### Security
- Bumped dependencies in `uv.lock` to patched versions for known advisories: `aiohttp` 3.14.1, `cryptography` 49.0.0, `idna` 3.18, `langchain-core` 1.4.8, `langchain-text-splitters` 1.1.2, `langsmith` 0.9.1, `lxml` 6.1.1, `pillow` 12.2.0, `pydantic-settings` 2.14.2, `pyjwt` 2.13.0, `pytest` 9.1.1, `python-multipart` 0.0.32, `requests` 2.34.2, `starlette` 1.3.1, `urllib3` 2.7.0, `vcrpy` 8.2.1.

5
docs/cli.md
View file

@ -289,9 +289,11 @@ At the end, a separate "Versions" section lists runtime package versions:
Check the database for consistency problems and print a pass/warn/fail report:
```bash
haiku-rag doctor [--db /path/to/your.lancedb]
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.
Checks include:
- required tables are present
@ -307,6 +309,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)
- API keys are set for configured providers
It also probes the external endpoints the config uses and reports them under a Providers section:

7
docs/configuration/index.md
View file

@ -114,6 +114,13 @@ search:
vector_index_metric: cosine # cosine, l2, or dot
vector_refine_factor: 30
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
min_chunks: 3 # documents with fewer chunks are excluded
prompts:
domain_preamble: "" # Prepended to skill instructions

10
haiku_rag_slim/haiku/rag/app.py
View file

@ -198,7 +198,7 @@ class HaikuRAGApp: # pragma: no cover
f" [repr.attrib_name]docling-document schema[/repr.attrib_name]: {info.packages['docling_document_schema']}"
)
async def doctor(self) -> bool:
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
@ -213,7 +213,9 @@ 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))
report = await run_doctor(
self.config, self.db_path, dict(os.environ), duplicates_out=duplicates_out
)
glyphs = {
Severity.OK: "[green]✓[/green]",
@ -245,6 +247,10 @@ class HaikuRAGApp: # pragma: no cover
f"[yellow]{report.count(Severity.WARN)} warning(s)[/yellow], "
f"[red]{report.count(Severity.FAIL)} failure(s)[/red]"
)
if duplicates_out is not None:
self.console.print(
f"[dim]Duplicate-document groups written to {duplicates_out}[/dim]"
)
return report.failed
async def history(self, table: str | None = None, limit: int | None = None):

7
haiku_rag_slim/haiku/rag/cli.py
View file

@ -593,9 +593,14 @@ def doctor( # pragma: no cover
"--db",
help="Path to the LanceDB database file",
),
duplicates_out: Path | None = typer.Option(
None,
"--duplicates-out",
help="Write near-duplicate document groups to this YAML file",
),
):
app = create_app(db)
if asyncio.run(app.doctor()):
if asyncio.run(app.doctor(duplicates_out=duplicates_out)):
raise typer.Exit(code=1)

22
haiku_rag_slim/haiku/rag/config/models.py
View file

@ -103,6 +103,27 @@ class AnalysisConfig(BaseModel):
max_executions: int = 15
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.
"""
containment_threshold: float = 0.75
candidate_threshold: float = 0.85
twin_similarity: float = 0.95
min_chunks: int = 3
class DoctorConfig(BaseModel):
duplicates: DuplicateDetectionConfig = Field(
default_factory=DuplicateDetectionConfig
)
class PictureDescriptionConfig(BaseModel):
"""How the VLM runs over each picture when it runs at all.
@ -516,6 +537,7 @@ class AppConfig(BaseModel):
analysis: AnalysisConfig = Field(default_factory=AnalysisConfig)
processing: ProcessingConfig = Field(default_factory=ProcessingConfig)
search: SearchConfig = Field(default_factory=SearchConfig)
doctor: DoctorConfig = Field(default_factory=DoctorConfig)
providers: ProvidersConfig = Field(default_factory=ProvidersConfig)
prompts: PromptsConfig = Field(default_factory=PromptsConfig)
ingester: IngesterConfig = Field(default_factory=IngesterConfig)

292
haiku_rag_slim/haiku/rag/doctor.py
View file

@ -1,13 +1,16 @@
import asyncio
import json
from collections.abc import Callable, Mapping
from enum import StrEnum
from pathlib import Path
import httpx
import numpy as np
import yaml
from pydantic import BaseModel, Field
from haiku.rag.config import AppConfig
from haiku.rag.config.models import DuplicateDetectionConfig
from haiku.rag.store.engine import (
REQUIRED_TABLES,
Store,
@ -232,8 +235,259 @@ 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]]
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
) -> list[_DuplicateFamily]:
"""Cluster documents that share most of their chunks (revisions of one another).
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.
"""
# 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:
return []
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)
# 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]] = []
block = 512
capped = False
for start in range(0, len(order), block):
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))
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]
)
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:
return []
# Cluster confirmed pairs into families (connected components).
families: list[_DuplicateFamily] = []
seen: set[int] = set()
for node in adjacency:
if node in seen:
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
)
families.append(
_DuplicateFamily(
members=members,
superset=superset,
pairs=pairs,
sizes={d: int(normalized[d].shape[0]) for d in members},
)
)
return sorted(families, key=lambda f: f.members)
def _common_path_prefix(labels: list[str]) -> str:
"""Longest shared prefix across labels, trimmed to a path boundary.
Returns "" unless the shared prefix is long enough to be worth factoring out
of every line (deep URI trees are otherwise unreadable).
"""
if len(labels) < 2:
return ""
lo, hi = min(labels), max(labels)
end = 0
while end < len(lo) and lo[end] == hi[end]:
end += 1
cut = lo.rfind("/", 0, end)
return lo[: cut + 1] if cut > 16 else ""
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."""
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,
"documents": [
{
"document_id": member,
"document": label(member),
"chunks": family.sizes[member],
"contained_fraction": round(contained[member], 3),
"keep_suggested": member == family.superset,
}
for member in family.members
],
}
)
with open(path, "w", encoding="utf-8") as handle:
yaml.safe_dump({"groups": groups}, handle, sort_keys=False, allow_unicode=True)
def _check_duplicate_documents(
doc_vectors: dict[str, 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)
def label(doc_id: str) -> str:
return uri_by_doc.get(doc_id) or title_by_doc.get(doc_id) or doc_id
if yaml_path is not None:
_write_duplicates_out(yaml_path, families, label)
if not families:
return CheckResult(
name="duplicate_documents",
severity=Severity.OK,
message="No near-duplicate documents detected.",
)
# 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.
shown = families[:_SAMPLE_LIMIT]
prefix = _common_path_prefix([label(m) for f in shown for m in f.members])
def short(doc_id: str) -> str:
text = label(doc_id)
return text[len(prefix) :] if prefix and text.startswith(prefix) else text
details: list[str] = []
if prefix:
details.append(f"common path: {prefix}")
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]}:"
)
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
)
details.append(f" overlap: {overlaps}")
if len(families) > len(shown):
details.append(
f"... (+{len(families) - len(shown)} more groups; "
"use --duplicates-out to export all)"
)
total_docs = sum(len(f.members) for f in families)
return CheckResult(
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."
),
remediation=(
"Review each group and remove redundant revisions; overlap may be intentional."
),
details=details,
)
async def run_db_checks(
store: Store, config: AppConfig, stats: dict
store: Store,
config: AppConfig,
stats: dict,
duplicates_out: Path | None = None,
) -> list[CheckResult]:
"""Referential and content-integrity checks against an open read-only Store.
@ -244,7 +498,7 @@ async def run_db_checks(
doc_ids = set(await _column_values(store.documents_table, "id"))
meta_rows = (
await store.document_meta_table.query()
.select(["document_id", "metadata"])
.select(["document_id", "metadata", "uri", "title"])
.to_list()
)
meta_doc_ids = {row["document_id"] for row in meta_rows}
@ -254,6 +508,8 @@ async def run_db_checks(
)
for row in meta_rows
}
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}
chunk_rows = (
await store.chunks_table.query()
@ -375,7 +631,11 @@ async def run_db_checks(
# Vector dimension consistency and unembedded (all-zero) vectors share one
# scan of the vector column — the heaviest check on large corpora.
arrow = await store.chunks_table.query().select(["id", "vector"]).to_arrow()
arrow = (
await store.chunks_table.query()
.select(["id", "vector", "document_id"])
.to_arrow()
)
stored = await SettingsRepository(store).get_current_settings()
stored_dim = stored.get("embeddings", {}).get("model", {}).get("vector_dim")
actual_dim = arrow.schema.field("vector").type.list_size
@ -419,6 +679,23 @@ async def run_db_checks(
)
)
# 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()}
results.append(
_check_duplicate_documents(
doc_vectors,
uri_by_doc,
title_by_doc,
config.doctor.duplicates,
yaml_path=duplicates_out,
)
)
# 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,7 +1003,10 @@ async def run_provider_checks(config: AppConfig) -> list[CheckResult]:
async def run_doctor(
config: AppConfig, db_path: Path, environ: dict[str, str]
config: AppConfig,
db_path: Path,
environ: dict[str, str],
duplicates_out: Path | None = None,
) -> DoctorReport:
"""Open the database read-only and run every diagnostic check.
@ -757,7 +1037,9 @@ async def run_doctor(
read_only=True,
skip_migration_check=True,
) as store:
results += await run_db_checks(store, config, stats)
results += await run_db_checks(
store, config, stats, duplicates_out=duplicates_out
)
results.append(_check_api_keys(config, environ))
results += await run_provider_checks(config)

273
tests/test_doctor.py
View file

@ -3,7 +3,9 @@ from importlib import metadata
from unittest.mock import AsyncMock, MagicMock
import lancedb
import numpy as np
import pytest
import yaml
from typer.testing import CliRunner
from haiku.rag.cli import _cli as cli
@ -11,6 +13,8 @@ from haiku.rag.config.models import (
AppConfig,
ConversionOptions,
DoclingServeConfig,
DoctorConfig,
DuplicateDetectionConfig,
EmbeddingModelConfig,
EmbeddingsConfig,
ModelConfig,
@ -24,8 +28,10 @@ from haiku.rag.doctor import (
Severity,
_active_models,
_check_api_keys,
_check_duplicate_documents,
_check_embedding_drift,
_check_vector_index,
_duplicate_families,
_model_present,
_probe_endpoint,
_provider_targets,
@ -933,3 +939,270 @@ async def test_probe_endpoint_connection_error():
reachable, error, _ = await _probe_with_handler(handler)
assert not reachable
assert error is not None and "refused" in 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.
A shared index across documents is a shared (identical) chunk; distinct
indices are orthogonal, so they never count as twins.
"""
eye = np.eye(dim)
return {
doc: np.array([eye[i] for i in idxs], dtype=float) for doc, idxs in spec.items()
}
# 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():
families = _duplicate_families(
_docs({"a": [0, 1, 2, 3], "b": [0, 1, 2, 4]}), _stage2_cfg()
)
assert len(families) == 1
assert set(families[0].members) == {"a", "b"}
def test_duplicate_families_append_only_is_asymmetric():
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]}),
DuplicateDetectionConfig(),
)
assert len(families) == 1
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()
)
assert families == []
def test_duplicate_families_three_way_chain_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(),
)
assert len(families) == 1
assert set(families[0].members) == {"a", "b", "c"}
assert families[0].superset == "c"
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())
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))
assert len(flagged) == 1
assert set(flagged[0].members) == {"a", "b"}
def test_duplicate_documents_report_truncates_summary():
pairs = 7 # more than the terminal detail cap of 5
spec: dict[str, list[int]] = {}
for k in range(pairs):
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())
assert result.severity is Severity.WARN
# The summary message still reports the full total.
assert f"{pairs} group(s)" in result.message
# The terminal detail shows only the first few groups and points at export.
assert sum(1 for d in result.details if d.startswith("group ")) == 5
assert any("more groups" in d and "--duplicates-out" in d for d in result.details)
assert any("keep #" in d for d in result.details)
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())
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"}
assert not any(base in d for d in member_lines)
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)
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)
data = yaml.safe_load(out.read_text())
assert len(data["groups"]) == 1
group = data["groups"][0]
assert group["group"] == 1 and group["keep"] == "a"
docs_out = group["documents"]
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 {d["document_id"]: d["keep_suggested"] for d in docs_out} == {
"a": True,
"b": False,
}
def test_duplicate_documents_writes_empty_yaml_when_none(tmp_path):
docs = _docs({"a": [0, 1, 2], "b": [3, 4, 5]}, dim=6) # distinct
out = tmp_path / "dups.yaml"
_check_duplicate_documents(
docs, {"a": "u", "b": "v"}, {}, _stage2_cfg(), yaml_path=out
)
assert yaml.safe_load(out.read_text()) == {"groups": []}
async def _build_dup_db(path, docs: dict[str, list[int]], *, vector_dim: int = 8):
"""Build a multi-document database with one-hot chunk vectors."""
eye = np.eye(vector_dim)
db = await lancedb.connect_async(path)
settings_tbl = await db.create_table("settings", schema=SettingsRecord)
docs_tbl = await db.create_table("documents", schema=DocumentRecord)
meta_tbl = await db.create_table("document_meta", schema=DocumentMetaRecord)
chunk_model = create_chunk_model(vector_dim)
chunks_tbl = await db.create_table("chunks", schema=chunk_model)
items_tbl = await db.create_table("document_items", schema=DocumentItemRecord)
await settings_tbl.add(
[
SettingsRecord(
id="settings",
settings=json.dumps(
{
"version": CURRENT_VERSION,
"embeddings": {
"model": {
"provider": "ollama",
"name": "test",
"vector_dim": vector_dim,
}
},
}
),
)
]
)
for doc_id, idxs in docs.items():
await docs_tbl.add([DocumentRecord(id=doc_id, content="x")])
await meta_tbl.add(
[DocumentMetaRecord(document_id=doc_id, uri=f"test://{doc_id}")]
)
await items_tbl.add(
[
DocumentItemRecord(
document_id=doc_id, position=0, self_ref="#/texts/0", text="x"
)
]
)
await chunks_tbl.add(
[
chunk_model(
id=f"{doc_id}-c{n}",
document_id=doc_id,
content="x",
metadata=json.dumps({"doc_item_refs": ["#/texts/0"]}),
vector=eye[i].tolist(),
)
for n, i in enumerate(idxs)
]
)
return db
@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]})
report = await run_doctor(_config(vector_dim=8), temp_db_path, {})
result = _result(report, "duplicate_documents")
assert result.severity is Severity.WARN
blob = " ".join(result.details)
assert "test://a" in blob and "test://b" in blob
@pytest.mark.asyncio
async def test_duplicate_documents_check_ok_when_distinct(temp_db_path):
await _build_dup_db(temp_db_path, {"a": [0, 1, 2], "b": [3, 4, 5]})
report = await run_doctor(_config(vector_dim=8), temp_db_path, {})
assert _result(report, "duplicate_documents").severity is Severity.OK
@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.
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"
).severity
is Severity.OK
)
tuned = _config(vector_dim=8)
tuned.doctor = DoctorConfig(
duplicates=DuplicateDetectionConfig(
candidate_threshold=0.0, containment_threshold=0.6
)
)
assert (
_result(
await run_doctor(tuned, temp_db_path, {}), "duplicate_documents"
).severity
is Severity.WARN
)