1047 lines
37 KiB
Python
1047 lines
37 KiB
Python
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,
|
||
connect_lancedb,
|
||
get_database_stats,
|
||
)
|
||
from haiku.rag.store.repositories.settings import SettingsRepository
|
||
from haiku.rag.store.upgrades import get_pending_upgrades
|
||
|
||
# Cap how many offending ids we collect per check; doctor is a summary, not a dump.
|
||
_SAMPLE_LIMIT = 5
|
||
|
||
# API providers and the environment variable that carries their key.
|
||
_PROVIDER_ENV_VARS: dict[str, str] = {
|
||
"openai": "OPENAI_API_KEY",
|
||
"anthropic": "ANTHROPIC_API_KEY",
|
||
"cohere": "CO_API_KEY",
|
||
"voyageai": "VOYAGE_API_KEY",
|
||
"jina": "JINA_API_KEY",
|
||
"zeroentropy": "ZEROENTROPY_API_KEY",
|
||
}
|
||
|
||
# Providers backed by in-process local models — no endpoint to probe.
|
||
_LOCAL_PROVIDERS = {"sentence-transformers", "mxbai", "cross-encoder", "jina-local"}
|
||
|
||
# Item labels that never yield a standalone chunk: pictures (handled via the
|
||
# image path), headings (folded into chunk context, not embedded alone), and
|
||
# page furniture. A document whose only items carry these labels is expected to
|
||
# have no chunks.
|
||
_NON_BODY_LABELS = {
|
||
"picture",
|
||
"section_header",
|
||
"title",
|
||
"page_header",
|
||
"page_footer",
|
||
"caption",
|
||
}
|
||
|
||
# Operators care whether an endpoint answers now, not eventually.
|
||
_PROBE_TIMEOUT_S = 2.0
|
||
|
||
|
||
class Severity(StrEnum):
|
||
OK = "ok"
|
||
WARN = "warn"
|
||
FAIL = "fail"
|
||
|
||
|
||
class CheckResult(BaseModel):
|
||
name: str
|
||
severity: Severity
|
||
message: str
|
||
remediation: str | None = None
|
||
details: list[str] = Field(default_factory=list)
|
||
|
||
|
||
class DoctorReport(BaseModel):
|
||
results: list[CheckResult] = Field(default_factory=list)
|
||
|
||
@property
|
||
def failed(self) -> bool:
|
||
return any(r.severity is Severity.FAIL for r in self.results)
|
||
|
||
def count(self, severity: Severity) -> int:
|
||
return sum(1 for r in self.results if r.severity is severity)
|
||
|
||
|
||
def _sample(ids: list[str]) -> list[str]:
|
||
"""Cap a list of offending ids for display, noting how many were elided."""
|
||
if len(ids) <= _SAMPLE_LIMIT:
|
||
return list(ids)
|
||
extra = len(ids) - _SAMPLE_LIMIT
|
||
return [*ids[:_SAMPLE_LIMIT], f"... (+{extra} more)"]
|
||
|
||
|
||
def _active_models(config: AppConfig) -> list[tuple[str, str, str | None]]:
|
||
"""(provider, name, base_url) for every model role the config activates.
|
||
|
||
Picture-description and title models are only included when their feature
|
||
is enabled (``processing.pictures == "description"`` / ``auto_title``), so
|
||
doctor checks exactly the providers the next ingest will use.
|
||
"""
|
||
models = [
|
||
(
|
||
config.embeddings.model.provider,
|
||
config.embeddings.model.name,
|
||
config.embeddings.model.base_url,
|
||
)
|
||
]
|
||
for model in (config.reranking.model, config.qa.model, config.analysis.model):
|
||
if model is not None:
|
||
models.append((model.provider, model.name, model.base_url))
|
||
|
||
proc = config.processing
|
||
if proc.pictures == "description":
|
||
pd = proc.conversion_options.picture_description.model
|
||
models.append((pd.provider, pd.name, pd.base_url))
|
||
if proc.auto_title:
|
||
tm = proc.title_model
|
||
models.append((tm.provider, tm.name, tm.base_url))
|
||
return models
|
||
|
||
|
||
def _check_api_keys(config: AppConfig, environ: dict[str, str]) -> CheckResult:
|
||
# A custom base_url points at a self-hosted OpenAI-compatible endpoint that
|
||
# uses a placeholder key, so the SaaS key is only required when a provider
|
||
# is used without one. Reachability of custom endpoints is the probe's job.
|
||
need_key = {
|
||
provider
|
||
for provider, _name, base_url in _active_models(config)
|
||
if not base_url and provider in _PROVIDER_ENV_VARS
|
||
}
|
||
missing = [
|
||
f"{provider} ({_PROVIDER_ENV_VARS[provider]})"
|
||
for provider in sorted(need_key)
|
||
if not environ.get(_PROVIDER_ENV_VARS[provider])
|
||
]
|
||
if missing:
|
||
return CheckResult(
|
||
name="api_keys",
|
||
severity=Severity.FAIL,
|
||
message="Configured providers are missing their API key.",
|
||
remediation="Set the listed environment variables.",
|
||
details=missing,
|
||
)
|
||
return CheckResult(
|
||
name="api_keys",
|
||
severity=Severity.OK,
|
||
message="API keys present for all configured providers.",
|
||
)
|
||
|
||
|
||
def _check_tables_present(stats: dict) -> CheckResult:
|
||
missing = [name for name in REQUIRED_TABLES if not stats[name]["exists"]]
|
||
if missing:
|
||
return CheckResult(
|
||
name="tables_present",
|
||
severity=Severity.FAIL,
|
||
message="Required tables are missing.",
|
||
remediation="Run 'haiku-rag init' for a new database or 'haiku-rag migrate'.",
|
||
details=missing,
|
||
)
|
||
return CheckResult(
|
||
name="tables_present",
|
||
severity=Severity.OK,
|
||
message="All required tables are present.",
|
||
)
|
||
|
||
|
||
def _classify_unchunked(
|
||
no_chunk_ids: set[str],
|
||
labels_by_doc: dict[str, set[str]],
|
||
supports_images: bool,
|
||
) -> list[CheckResult]:
|
||
"""Classify chunk-less documents by what they hold.
|
||
|
||
A document with body-text items but no chunks is always a problem. A
|
||
picture-only document is a problem under a multimodal embedder (its picture
|
||
chunks are missing) and an indexing gap under a text-only embedder (which
|
||
cannot embed images). A document carrying only headings/furniture (or no
|
||
items at all) is expected to have no chunks.
|
||
"""
|
||
text_docs: list[str] = []
|
||
picture_docs: list[str] = []
|
||
for doc_id in no_chunk_ids:
|
||
labels = labels_by_doc.get(doc_id, set())
|
||
if any(label not in _NON_BODY_LABELS for label in labels):
|
||
text_docs.append(doc_id)
|
||
elif "picture" in labels:
|
||
picture_docs.append(doc_id)
|
||
|
||
results: list[CheckResult] = []
|
||
if text_docs:
|
||
results.append(
|
||
CheckResult(
|
||
name="documents_text_no_chunks",
|
||
severity=Severity.WARN,
|
||
message=f"{len(text_docs)} document(s) have text content but no chunks.",
|
||
remediation="haiku-rag rebuild",
|
||
details=_sample(sorted(text_docs)),
|
||
)
|
||
)
|
||
if picture_docs and supports_images:
|
||
results.append(
|
||
CheckResult(
|
||
name="documents_pictures_no_chunks",
|
||
severity=Severity.WARN,
|
||
message=f"{len(picture_docs)} document(s) with pictures have no chunks.",
|
||
remediation="haiku-rag rebuild",
|
||
details=_sample(sorted(picture_docs)),
|
||
)
|
||
)
|
||
elif picture_docs:
|
||
results.append(
|
||
CheckResult(
|
||
name="documents_images_unsearchable",
|
||
severity=Severity.WARN,
|
||
message=(
|
||
f"{len(picture_docs)} image-only document(s) have no chunks; "
|
||
"a text-only embedder cannot index images."
|
||
),
|
||
remediation=(
|
||
"Set embeddings.model.multimodal: true on a vllm, voyageai, or "
|
||
"cohere model and rebuild to index images."
|
||
),
|
||
details=_sample(sorted(picture_docs)),
|
||
)
|
||
)
|
||
if not results:
|
||
results.append(
|
||
CheckResult(
|
||
name="documents_without_chunks",
|
||
severity=Severity.OK,
|
||
message="Every document with content has chunks.",
|
||
)
|
||
)
|
||
return results
|
||
|
||
|
||
async def _column_values(table, column: str) -> list:
|
||
rows = await table.query().select([column]).to_list()
|
||
return [row[column] for row in rows]
|
||
|
||
|
||
class _DuplicateFamily(BaseModel):
|
||
members: list[str]
|
||
keep: str
|
||
similarity: dict[str, float]
|
||
sizes: dict[str, int]
|
||
|
||
|
||
def _duplicate_families(
|
||
doc_ids: list[str],
|
||
centroids: np.ndarray,
|
||
counts: np.ndarray,
|
||
cfg: DuplicateDetectionConfig,
|
||
) -> list[_DuplicateFamily]:
|
||
"""Cluster documents whose embedding centroids are nearly identical.
|
||
|
||
``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.
|
||
"""
|
||
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)
|
||
|
||
# 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))
|
||
|
||
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
|
||
for start in range(0, n, block):
|
||
sims = unit[start : start + block] @ unit.T
|
||
for row in range(sims.shape[0]):
|
||
gi = start + row
|
||
cols = (
|
||
gi + 1 + np.nonzero(sims[row, gi + 1 :] >= cfg.similarity_threshold)[0]
|
||
)
|
||
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 []
|
||
|
||
components: dict[int, list[int]] = {}
|
||
for idx in range(n):
|
||
components.setdefault(find(idx), []).append(idx)
|
||
|
||
families: list[_DuplicateFamily] = []
|
||
for indices in components.values():
|
||
if len(indices) < 2:
|
||
continue
|
||
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,
|
||
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)
|
||
|
||
|
||
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 document to keep and
|
||
``similarity`` is the highest centroid cosine to another group member."""
|
||
groups = []
|
||
for n, family in enumerate(families, start=1):
|
||
groups.append(
|
||
{
|
||
"group": n,
|
||
"keep": family.keep,
|
||
"documents": [
|
||
{
|
||
"document_id": member,
|
||
"document": label(member),
|
||
"chunks": family.sizes[member],
|
||
"similarity": family.similarity[member],
|
||
"keep_suggested": member == family.keep,
|
||
}
|
||
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_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_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
|
||
|
||
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 similarity 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.keep]}:"
|
||
)
|
||
for member in family.members:
|
||
details.append(f" #{number[member]} {short(member)}")
|
||
sims = ", ".join(
|
||
f"#{number[m]} {family.similarity[m]:.0%}" for m in family.members
|
||
)
|
||
details.append(f" similarity: {sims}")
|
||
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 near-identical documents "
|
||
f"(potential duplicates), {total_docs} documents."
|
||
),
|
||
remediation=(
|
||
"Review each group and remove redundant copies; duplication may be intentional."
|
||
),
|
||
details=details,
|
||
)
|
||
|
||
|
||
async def run_db_checks(
|
||
store: Store,
|
||
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()
|
||
.select(["document_id", "metadata", "uri", "title"])
|
||
.to_list()
|
||
)
|
||
meta_doc_ids = {row["document_id"] for row in meta_rows}
|
||
content_type_by_doc = {
|
||
row["document_id"]: json.loads(row.get("metadata") or "{}").get(
|
||
"content_type", ""
|
||
)
|
||
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}
|
||
|
||
notify("Reading chunks")
|
||
chunk_rows = (
|
||
await store.chunks_table.query()
|
||
.select(["id", "document_id", "metadata"])
|
||
.to_list()
|
||
)
|
||
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"])
|
||
.to_list()
|
||
)
|
||
item_doc_ids = {row["document_id"] for row in item_rows}
|
||
self_refs_by_doc: dict[str, set[str]] = {}
|
||
labels_by_doc: dict[str, set[str]] = {}
|
||
for row in item_rows:
|
||
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
|
||
if orphan_docs or orphan_meta:
|
||
details = [f"document with no meta: {d}" for d in _sample(sorted(orphan_docs))]
|
||
details += [f"meta with no document: {d}" for d in _sample(sorted(orphan_meta))]
|
||
results.append(
|
||
CheckResult(
|
||
name="document_meta_parity",
|
||
severity=Severity.FAIL,
|
||
message="documents and document_meta are out of sync.",
|
||
remediation="haiku-rag rebuild",
|
||
details=details,
|
||
)
|
||
)
|
||
else:
|
||
results.append(
|
||
CheckResult(
|
||
name="document_meta_parity",
|
||
severity=Severity.OK,
|
||
message="documents and document_meta are consistent.",
|
||
)
|
||
)
|
||
|
||
# Orphaned chunks / items reference a document that no longer exists.
|
||
orphan_chunk_docs = chunk_doc_ids - doc_ids
|
||
results.append(
|
||
CheckResult(
|
||
name="orphaned_chunks",
|
||
severity=Severity.FAIL if orphan_chunk_docs else Severity.OK,
|
||
message=(
|
||
"Chunks reference missing documents."
|
||
if orphan_chunk_docs
|
||
else "No orphaned chunks."
|
||
),
|
||
remediation="haiku-rag rebuild" if orphan_chunk_docs else None,
|
||
details=_sample(sorted(orphan_chunk_docs)),
|
||
)
|
||
)
|
||
|
||
orphan_item_docs = item_doc_ids - doc_ids
|
||
results.append(
|
||
CheckResult(
|
||
name="orphaned_document_items",
|
||
severity=Severity.FAIL if orphan_item_docs else Severity.OK,
|
||
message=(
|
||
"Document items reference missing documents."
|
||
if orphan_item_docs
|
||
else "No orphaned document items."
|
||
),
|
||
remediation="haiku-rag rebuild" if orphan_item_docs else None,
|
||
details=_sample(sorted(orphan_item_docs)),
|
||
)
|
||
)
|
||
|
||
notify("Checking document chunking")
|
||
# Documents with no chunks, classified by what they contain and whether the
|
||
# embedder can index images.
|
||
results += _classify_unchunked(
|
||
doc_ids - chunk_doc_ids, labels_by_doc, store.embedder.supports_images
|
||
)
|
||
|
||
# A chunked document must have items; one without them is corrupt. Empty
|
||
# documents legitimately have neither, so only flag the chunked ones.
|
||
docs_missing_items = (doc_ids & chunk_doc_ids) - item_doc_ids
|
||
results.append(
|
||
CheckResult(
|
||
name="documents_without_items",
|
||
severity=Severity.WARN if docs_missing_items else Severity.OK,
|
||
message=(
|
||
f"{len(docs_missing_items)} chunked document(s) have no document items."
|
||
if docs_missing_items
|
||
else "Every chunked document has document items."
|
||
),
|
||
remediation="haiku-rag rebuild" if docs_missing_items else None,
|
||
details=_sample(sorted(docs_missing_items)),
|
||
)
|
||
)
|
||
|
||
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:
|
||
refs = json.loads(row.get("metadata") or "{}").get("doc_item_refs") or []
|
||
known = self_refs_by_doc.get(row["document_id"], set())
|
||
if any(ref not in known for ref in refs):
|
||
dangling.append(row["id"])
|
||
results.append(
|
||
CheckResult(
|
||
name="dangling_doc_item_refs",
|
||
severity=Severity.FAIL if dangling else Severity.OK,
|
||
message=(
|
||
f"{len(dangling)} chunk(s) reference missing document items."
|
||
if dangling
|
||
else "All chunk doc_item_refs resolve."
|
||
),
|
||
remediation="haiku-rag rebuild" if dangling else None,
|
||
details=_sample(dangling),
|
||
)
|
||
)
|
||
|
||
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 = (
|
||
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
|
||
if stored_dim and stored_dim != actual_dim:
|
||
results.append(
|
||
CheckResult(
|
||
name="vector_dimension",
|
||
severity=Severity.FAIL,
|
||
message=(
|
||
f"Chunk vector size {actual_dim} does not match stored "
|
||
f"vector_dim {stored_dim}."
|
||
),
|
||
remediation="haiku-rag rebuild",
|
||
)
|
||
)
|
||
else:
|
||
results.append(
|
||
CheckResult(
|
||
name="vector_dimension",
|
||
severity=Severity.OK,
|
||
message=f"Chunk vectors are {actual_dim}-dimensional.",
|
||
)
|
||
)
|
||
|
||
# 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_count else Severity.OK,
|
||
message=(
|
||
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_count else None,
|
||
details=zero_sample,
|
||
)
|
||
)
|
||
|
||
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_ids,
|
||
centroids,
|
||
counts,
|
||
uri_by_doc,
|
||
title_by_doc,
|
||
config.doctor.duplicates,
|
||
yaml_path=duplicates_out,
|
||
)
|
||
)
|
||
|
||
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.
|
||
missing_picture_docs = [
|
||
row["document_id"]
|
||
for row in await store.document_items_table.query()
|
||
.select(["document_id"])
|
||
.where("label = 'picture' AND picture_data IS NULL")
|
||
.to_list()
|
||
]
|
||
real_missing = [
|
||
doc_id
|
||
for doc_id in missing_picture_docs
|
||
if not content_type_by_doc.get(doc_id, "").startswith("text/")
|
||
]
|
||
results.append(
|
||
CheckResult(
|
||
name="picture_data",
|
||
severity=Severity.WARN if real_missing else Severity.OK,
|
||
message=(
|
||
f"{len(real_missing)} picture item(s) in image/PDF documents "
|
||
"have no image data."
|
||
if real_missing
|
||
else "Pictures that should carry image data have it."
|
||
),
|
||
remediation="haiku-rag rebuild" if real_missing else None,
|
||
details=_sample(sorted(set(real_missing))),
|
||
)
|
||
)
|
||
|
||
notify("Checking settings and indexes")
|
||
# Settings must hold exactly one canonical row.
|
||
total_settings = await store.settings_table.count_rows()
|
||
canonical = len(
|
||
await store.settings_table.query().where("id = 'settings'").to_list()
|
||
)
|
||
if total_settings == 0 or canonical != 1:
|
||
results.append(
|
||
CheckResult(
|
||
name="settings_row",
|
||
severity=Severity.FAIL,
|
||
message=(
|
||
f"Expected exactly one 'settings' row, found {canonical} "
|
||
f"(of {total_settings} total)."
|
||
),
|
||
remediation="haiku-rag migrate",
|
||
)
|
||
)
|
||
else:
|
||
results.append(
|
||
CheckResult(
|
||
name="settings_row",
|
||
severity=Severity.OK,
|
||
message="Settings row is present.",
|
||
)
|
||
)
|
||
|
||
results.append(_check_embedding_drift(stored, config))
|
||
|
||
stored_version = str(stored.get("version", "unknown"))
|
||
pending = (
|
||
get_pending_upgrades(stored_version) if stored_version != "unknown" else []
|
||
)
|
||
results.append(
|
||
CheckResult(
|
||
name="pending_migrations",
|
||
severity=Severity.WARN if pending else Severity.OK,
|
||
message=(
|
||
f"{len(pending)} migration(s) pending (db version {stored_version})."
|
||
if pending
|
||
else f"Database is up to date (version {stored_version})."
|
||
),
|
||
remediation="haiku-rag migrate" if pending else None,
|
||
details=[f"{step.version}: {step.description or ''}" for step in pending],
|
||
)
|
||
)
|
||
|
||
results.append(_check_vector_index(stats))
|
||
|
||
return results
|
||
|
||
|
||
def _check_embedding_drift(stored: dict, config: AppConfig) -> CheckResult:
|
||
stored_model = stored.get("embeddings", {}).get("model", {})
|
||
current_model = config.embeddings.model
|
||
if not stored_model:
|
||
return CheckResult(
|
||
name="embedding_drift",
|
||
severity=Severity.OK,
|
||
message="No stored embedding identity to compare.",
|
||
)
|
||
|
||
stored_dim = stored_model.get("vector_dim")
|
||
if stored_dim and stored_dim != current_model.vector_dim:
|
||
return CheckResult(
|
||
name="embedding_drift",
|
||
severity=Severity.FAIL,
|
||
message=(
|
||
f"Embedding vector_dim differs: stored {stored_dim} -> "
|
||
f"config {current_model.vector_dim}."
|
||
),
|
||
remediation="haiku-rag rebuild",
|
||
)
|
||
|
||
drift: list[str] = []
|
||
if stored_model.get("provider") not in (None, current_model.provider):
|
||
drift.append(
|
||
f"provider: {stored_model['provider']} -> {current_model.provider}"
|
||
)
|
||
if stored_model.get("name") not in (None, current_model.name):
|
||
drift.append(f"name: {stored_model['name']} -> {current_model.name}")
|
||
if drift:
|
||
return CheckResult(
|
||
name="embedding_drift",
|
||
severity=Severity.WARN,
|
||
message="Embedding identity differs from config (vector_dim matches).",
|
||
remediation="haiku-rag rebuild --set-embedder",
|
||
details=drift,
|
||
)
|
||
return CheckResult(
|
||
name="embedding_drift",
|
||
severity=Severity.OK,
|
||
message="Embedding identity matches the stored settings.",
|
||
)
|
||
|
||
|
||
def _check_vector_index(stats: dict) -> CheckResult:
|
||
chunks = stats["chunks"]
|
||
num_chunks = chunks.get("num_rows", 0)
|
||
if not chunks.get("has_vector_index"):
|
||
if num_chunks >= 100_000:
|
||
return CheckResult(
|
||
name="vector_index",
|
||
severity=Severity.WARN,
|
||
message=(
|
||
"No vector index on a large collection; "
|
||
"similarity search scans every chunk and may be slow."
|
||
),
|
||
remediation="haiku-rag create-index",
|
||
)
|
||
return CheckResult(
|
||
name="vector_index",
|
||
severity=Severity.OK,
|
||
message="No vector index; similarity search is exact (brute-force).",
|
||
)
|
||
unindexed = chunks.get("num_unindexed_rows", 0)
|
||
if unindexed > 0:
|
||
return CheckResult(
|
||
name="vector_index",
|
||
severity=Severity.WARN,
|
||
message=f"{unindexed} chunk(s) are not in the vector index.",
|
||
remediation="haiku-rag create-index",
|
||
)
|
||
return CheckResult(
|
||
name="vector_index",
|
||
severity=Severity.OK,
|
||
message="Vector index covers all chunks.",
|
||
)
|
||
|
||
|
||
def _resolve_endpoint(
|
||
provider: str, base_url: str | None, ollama_base: str
|
||
) -> tuple[str, str, str] | str | None:
|
||
"""Map a model's provider to a probe target.
|
||
|
||
Returns ``(probe_url, kind, display)``, the literal ``"local"`` for an
|
||
in-process model, or ``None`` for a SaaS provider covered by the API-key
|
||
check.
|
||
"""
|
||
if provider == "ollama":
|
||
base = (base_url or ollama_base).rstrip("/")
|
||
if base.endswith("/v1"):
|
||
base = base[:-3].rstrip("/")
|
||
return f"{base}/api/tags", "ollama", base
|
||
if provider == "vllm":
|
||
base = (base_url or "http://localhost:8000/v1").rstrip("/")
|
||
if not base.endswith("/v1"):
|
||
base = f"{base}/v1"
|
||
return f"{base}/models", "openai", base
|
||
if provider == "openai" and base_url:
|
||
base = base_url.rstrip("/")
|
||
return f"{base}/models", "openai", base
|
||
if provider in _LOCAL_PROVIDERS:
|
||
return "local"
|
||
return None
|
||
|
||
|
||
def _provider_targets(
|
||
config: AppConfig,
|
||
) -> tuple[dict[str, dict], set[str]]:
|
||
"""Collect probe targets (keyed by probe URL) and local-only providers."""
|
||
targets: dict[str, dict] = {}
|
||
local: set[str] = set()
|
||
ollama_base = config.providers.ollama.base_url
|
||
|
||
def add_model(provider: str, name: str, base_url: str | None) -> None:
|
||
resolved = _resolve_endpoint(provider, base_url, ollama_base)
|
||
if resolved is None:
|
||
return
|
||
if resolved == "local":
|
||
local.add(provider)
|
||
return
|
||
probe_url, kind, display = resolved
|
||
entry = targets.setdefault(
|
||
probe_url, {"kind": kind, "display": display, "models": set()}
|
||
)
|
||
if name:
|
||
entry["models"].add(name)
|
||
|
||
proc = config.processing
|
||
if proc.converter == "docling-serve" or proc.chunker == "docling-serve":
|
||
for url in config.providers.docling_serve.base_urls:
|
||
base = url.rstrip("/")
|
||
targets.setdefault(
|
||
f"{base}/health",
|
||
{"kind": "docling-serve", "display": base, "models": set()},
|
||
)
|
||
|
||
for provider, name, base_url in _active_models(config):
|
||
add_model(provider, name, base_url)
|
||
|
||
return targets, local
|
||
|
||
|
||
def _model_present(expected: str, available: set[str]) -> bool:
|
||
if expected in available:
|
||
return True
|
||
if ":" not in expected:
|
||
return any(a.split(":", 1)[0] == expected for a in available)
|
||
return False
|
||
|
||
|
||
async def _probe_endpoint(
|
||
client: httpx.AsyncClient, url: str
|
||
) -> tuple[bool, str | None, dict | None]:
|
||
try:
|
||
response = await client.get(url)
|
||
except httpx.HTTPError as exc:
|
||
return False, str(exc), None
|
||
if not response.is_success:
|
||
return False, f"HTTP {response.status_code}", None
|
||
try:
|
||
return True, None, response.json()
|
||
except ValueError:
|
||
return True, None, None
|
||
|
||
|
||
def _endpoint_result(
|
||
entry: dict, reachable: bool, error: str | None, payload: dict | None
|
||
) -> CheckResult:
|
||
kind = entry["kind"]
|
||
display = entry["display"]
|
||
name = f"provider:{display}"
|
||
if not reachable:
|
||
return CheckResult(
|
||
name=name,
|
||
severity=Severity.FAIL,
|
||
message=f"{kind} at {display} is unreachable.",
|
||
remediation="Start the service or fix the configured base_url.",
|
||
details=[error] if error else [],
|
||
)
|
||
if kind == "ollama":
|
||
available = {m.get("name", "") for m in (payload or {}).get("models", [])}
|
||
missing = [
|
||
model
|
||
for model in sorted(entry["models"])
|
||
if not _model_present(model, available)
|
||
]
|
||
if missing:
|
||
return CheckResult(
|
||
name=name,
|
||
severity=Severity.WARN,
|
||
message=f"ollama at {display} is reachable but missing model(s).",
|
||
remediation="ollama pull <model>",
|
||
details=missing,
|
||
)
|
||
return CheckResult(
|
||
name=name,
|
||
severity=Severity.OK,
|
||
message=f"{kind} at {display} is reachable.",
|
||
)
|
||
|
||
|
||
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)
|
||
)
|
||
for url, (reachable, error, payload) in zip(targets, probes):
|
||
results.append(_endpoint_result(targets[url], reachable, error, payload))
|
||
|
||
for provider in sorted(local):
|
||
results.append(
|
||
CheckResult(
|
||
name=f"provider:{provider}",
|
||
severity=Severity.OK,
|
||
message=f"{provider}: local model, nothing to probe.",
|
||
)
|
||
)
|
||
return results
|
||
|
||
|
||
async def run_doctor(
|
||
config: AppConfig,
|
||
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)
|
||
|
||
results: list[CheckResult] = []
|
||
if not any(entry["exists"] for entry in stats.values()):
|
||
results.append(
|
||
CheckResult(
|
||
name="tables_present",
|
||
severity=Severity.FAIL,
|
||
message="Database is empty.",
|
||
remediation="haiku-rag init",
|
||
)
|
||
)
|
||
else:
|
||
results.append(_check_tables_present(stats))
|
||
missing = [name for name in REQUIRED_TABLES if not stats[name]["exists"]]
|
||
if not missing:
|
||
async with Store(
|
||
db_path,
|
||
config=config,
|
||
skip_validation=True,
|
||
read_only=True,
|
||
skip_migration_check=True,
|
||
) as store:
|
||
results += await run_db_checks(
|
||
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, on_progress=on_progress)
|
||
return DoctorReport(results=results)
|