425 lines
14 KiB
Python
425 lines
14 KiB
Python
import json
|
|
from importlib import metadata
|
|
from unittest.mock import AsyncMock, MagicMock
|
|
|
|
import lancedb
|
|
import pytest
|
|
from typer.testing import CliRunner
|
|
|
|
from haiku.rag.cli import _cli as cli
|
|
from haiku.rag.config.models import AppConfig, EmbeddingModelConfig, EmbeddingsConfig
|
|
from haiku.rag.doctor import (
|
|
CheckResult,
|
|
DoctorReport,
|
|
Severity,
|
|
_check_embedding_drift,
|
|
_check_vector_index,
|
|
_sample,
|
|
run_doctor,
|
|
)
|
|
from haiku.rag.store.engine import (
|
|
DocumentItemRecord,
|
|
DocumentMetaRecord,
|
|
DocumentRecord,
|
|
SettingsRecord,
|
|
create_chunk_model,
|
|
)
|
|
|
|
runner = CliRunner()
|
|
|
|
CURRENT_VERSION = metadata.version("haiku.rag-slim")
|
|
VECTOR_DIM = 4
|
|
ChunkRecord = create_chunk_model(VECTOR_DIM)
|
|
|
|
|
|
def _config(provider: str = "ollama", name: str = "test", vector_dim: int = VECTOR_DIM):
|
|
return AppConfig(
|
|
embeddings=EmbeddingsConfig(
|
|
model=EmbeddingModelConfig(
|
|
provider=provider, name=name, vector_dim=vector_dim
|
|
)
|
|
)
|
|
)
|
|
|
|
|
|
async def _build_db(
|
|
path,
|
|
*,
|
|
version: str = CURRENT_VERSION,
|
|
provider: str = "ollama",
|
|
name: str = "test",
|
|
vector_dim: int = VECTOR_DIM,
|
|
stored_vector_dim: int | None = None,
|
|
):
|
|
"""Create a consistent single-document database without touching an embedder.
|
|
|
|
``stored_vector_dim`` records a different dimension in settings than the
|
|
chunks table actually uses, to exercise the vector-dimension check.
|
|
"""
|
|
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)
|
|
chunks_tbl = await db.create_table("chunks", schema=create_chunk_model(vector_dim))
|
|
items_tbl = await db.create_table("document_items", schema=DocumentItemRecord)
|
|
|
|
await settings_tbl.add(
|
|
[
|
|
SettingsRecord(
|
|
id="settings",
|
|
settings=json.dumps(
|
|
{
|
|
"version": version,
|
|
"embeddings": {
|
|
"model": {
|
|
"provider": provider,
|
|
"name": name,
|
|
"vector_dim": stored_vector_dim or vector_dim,
|
|
}
|
|
},
|
|
}
|
|
),
|
|
)
|
|
]
|
|
)
|
|
await docs_tbl.add([DocumentRecord(id="d1", content="hello")])
|
|
await meta_tbl.add([DocumentMetaRecord(document_id="d1", uri="test://d1")])
|
|
await items_tbl.add(
|
|
[
|
|
DocumentItemRecord(
|
|
document_id="d1", position=0, self_ref="#/texts/0", text="x"
|
|
)
|
|
]
|
|
)
|
|
chunk_model = create_chunk_model(vector_dim)
|
|
await chunks_tbl.add(
|
|
[
|
|
chunk_model(
|
|
id="c1",
|
|
document_id="d1",
|
|
content="hello",
|
|
metadata=json.dumps({"doc_item_refs": ["#/texts/0"]}),
|
|
vector=[0.1] * vector_dim,
|
|
)
|
|
]
|
|
)
|
|
return db
|
|
|
|
|
|
def _result(report: DoctorReport, name: str) -> CheckResult:
|
|
return next(r for r in report.results if r.name == name)
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_healthy_db_all_ok(temp_db_path):
|
|
await _build_db(temp_db_path)
|
|
report = await run_doctor(_config(), temp_db_path, {})
|
|
assert not report.failed
|
|
assert report.count(Severity.WARN) == 0
|
|
assert all(r.severity is Severity.OK for r in report.results)
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_empty_db_fails(temp_db_path):
|
|
report = await run_doctor(_config(), temp_db_path, {})
|
|
assert report.failed
|
|
assert _result(report, "tables_present").message == "Database is empty."
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_missing_table_fails_without_opening_store(temp_db_path):
|
|
db = await lancedb.connect_async(temp_db_path)
|
|
await db.create_table("settings", schema=SettingsRecord)
|
|
report = await run_doctor(_config(), temp_db_path, {})
|
|
assert report.failed
|
|
tables = _result(report, "tables_present")
|
|
assert tables.severity is Severity.FAIL
|
|
assert "documents" in tables.details
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_orphaned_chunk_fails(temp_db_path):
|
|
db = await _build_db(temp_db_path)
|
|
chunks_tbl = await db.open_table("chunks")
|
|
await chunks_tbl.add(
|
|
[
|
|
ChunkRecord(
|
|
id="orphan",
|
|
document_id="ghost",
|
|
content="x",
|
|
vector=[0.2] * VECTOR_DIM,
|
|
)
|
|
]
|
|
)
|
|
report = await run_doctor(_config(), temp_db_path, {})
|
|
result = _result(report, "orphaned_chunks")
|
|
assert result.severity is Severity.FAIL
|
|
assert "ghost" in result.details
|
|
assert report.failed
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_orphaned_document_item_fails(temp_db_path):
|
|
db = await _build_db(temp_db_path)
|
|
items_tbl = await db.open_table("document_items")
|
|
await items_tbl.add(
|
|
[DocumentItemRecord(document_id="ghost", position=0, self_ref="#/texts/0")]
|
|
)
|
|
report = await run_doctor(_config(), temp_db_path, {})
|
|
assert _result(report, "orphaned_document_items").severity is Severity.FAIL
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_document_without_chunks_warns(temp_db_path):
|
|
db = await _build_db(temp_db_path)
|
|
docs_tbl = await db.open_table("documents")
|
|
meta_tbl = await db.open_table("document_meta")
|
|
await docs_tbl.add([DocumentRecord(id="d2", content="no chunks")])
|
|
await meta_tbl.add([DocumentMetaRecord(document_id="d2", uri="test://d2")])
|
|
report = await run_doctor(_config(), temp_db_path, {})
|
|
assert _result(report, "documents_without_chunks").severity is Severity.WARN
|
|
assert _result(report, "documents_without_items").severity is Severity.WARN
|
|
assert not report.failed
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_document_meta_parity_fails(temp_db_path):
|
|
db = await _build_db(temp_db_path)
|
|
docs_tbl = await db.open_table("documents")
|
|
await docs_tbl.add([DocumentRecord(id="d2", content="no meta")])
|
|
report = await run_doctor(_config(), temp_db_path, {})
|
|
result = _result(report, "document_meta_parity")
|
|
assert result.severity is Severity.FAIL
|
|
assert any("d2" in d for d in result.details)
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_dangling_doc_item_ref_fails(temp_db_path):
|
|
db = await _build_db(temp_db_path)
|
|
chunks_tbl = await db.open_table("chunks")
|
|
await chunks_tbl.add(
|
|
[
|
|
ChunkRecord(
|
|
id="c2",
|
|
document_id="d1",
|
|
content="x",
|
|
metadata=json.dumps({"doc_item_refs": ["#/texts/999"]}),
|
|
vector=[0.3] * VECTOR_DIM,
|
|
)
|
|
]
|
|
)
|
|
report = await run_doctor(_config(), temp_db_path, {})
|
|
result = _result(report, "dangling_doc_item_refs")
|
|
assert result.severity is Severity.FAIL
|
|
assert "c2" in result.details
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_unembedded_chunk_warns(temp_db_path):
|
|
db = await _build_db(temp_db_path)
|
|
chunks_tbl = await db.open_table("chunks")
|
|
await chunks_tbl.add(
|
|
[
|
|
ChunkRecord(
|
|
id="zero",
|
|
document_id="d1",
|
|
content="x",
|
|
metadata=json.dumps({"doc_item_refs": ["#/texts/0"]}),
|
|
vector=[0.0] * VECTOR_DIM,
|
|
)
|
|
]
|
|
)
|
|
report = await run_doctor(_config(), temp_db_path, {})
|
|
result = _result(report, "unembedded_chunks")
|
|
assert result.severity is Severity.WARN
|
|
assert "zero" in result.details
|
|
assert not report.failed
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_missing_picture_data_warns(temp_db_path):
|
|
db = await _build_db(temp_db_path)
|
|
items_tbl = await db.open_table("document_items")
|
|
await items_tbl.add(
|
|
[
|
|
DocumentItemRecord(
|
|
document_id="d1",
|
|
position=1,
|
|
self_ref="#/pictures/0",
|
|
label="picture",
|
|
picture_data=None,
|
|
)
|
|
]
|
|
)
|
|
report = await run_doctor(_config(), temp_db_path, {})
|
|
assert _result(report, "picture_data").severity is Severity.WARN
|
|
assert not report.failed
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_picture_with_data_ok(temp_db_path):
|
|
db = await _build_db(temp_db_path)
|
|
items_tbl = await db.open_table("document_items")
|
|
await items_tbl.add(
|
|
[
|
|
DocumentItemRecord(
|
|
document_id="d1",
|
|
position=1,
|
|
self_ref="#/pictures/0",
|
|
label="picture",
|
|
picture_data=b"\x89PNG",
|
|
)
|
|
]
|
|
)
|
|
report = await run_doctor(_config(), temp_db_path, {})
|
|
assert _result(report, "picture_data").severity is Severity.OK
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_embedding_name_drift_warns(temp_db_path):
|
|
await _build_db(temp_db_path, name="test")
|
|
report = await run_doctor(_config(name="different"), temp_db_path, {})
|
|
result = _result(report, "embedding_drift")
|
|
assert result.severity is Severity.WARN
|
|
assert not report.failed
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_embedding_dim_drift_fails(temp_db_path):
|
|
await _build_db(temp_db_path, vector_dim=VECTOR_DIM)
|
|
report = await run_doctor(_config(vector_dim=VECTOR_DIM + 1), temp_db_path, {})
|
|
assert _result(report, "embedding_drift").severity is Severity.FAIL
|
|
assert report.failed
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_embedding_provider_drift_warns(temp_db_path):
|
|
await _build_db(temp_db_path, provider="ollama")
|
|
report = await run_doctor(_config(provider="vllm"), temp_db_path, {})
|
|
result = _result(report, "embedding_drift")
|
|
assert result.severity is Severity.WARN
|
|
assert any("provider" in d for d in result.details)
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_vector_dimension_mismatch_fails(temp_db_path):
|
|
await _build_db(
|
|
temp_db_path, vector_dim=VECTOR_DIM, stored_vector_dim=VECTOR_DIM + 1
|
|
)
|
|
report = await run_doctor(_config(vector_dim=VECTOR_DIM + 1), temp_db_path, {})
|
|
result = _result(report, "vector_dimension")
|
|
assert result.severity is Severity.FAIL
|
|
assert report.failed
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_pending_migration_warns(temp_db_path):
|
|
await _build_db(temp_db_path, version="0.40.0")
|
|
report = await run_doctor(_config(), temp_db_path, {})
|
|
assert _result(report, "pending_migrations").severity is Severity.WARN
|
|
assert not report.failed
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_missing_api_key_fails(temp_db_path):
|
|
await _build_db(temp_db_path, provider="openai", name="text-embedding-3-small")
|
|
config = _config(provider="openai", name="text-embedding-3-small")
|
|
report = await run_doctor(config, temp_db_path, environ={})
|
|
result = _result(report, "api_keys")
|
|
assert result.severity is Severity.FAIL
|
|
assert any("OPENAI_API_KEY" in d for d in result.details)
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_present_api_key_ok(temp_db_path):
|
|
await _build_db(temp_db_path, provider="openai", name="text-embedding-3-small")
|
|
config = _config(provider="openai", name="text-embedding-3-small")
|
|
report = await run_doctor(config, temp_db_path, environ={"OPENAI_API_KEY": "sk-x"})
|
|
assert _result(report, "api_keys").severity is Severity.OK
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_settings_row_missing_fails(temp_db_path):
|
|
db = await _build_db(temp_db_path)
|
|
settings_tbl = await db.open_table("settings")
|
|
await settings_tbl.delete("id = 'settings'")
|
|
report = await run_doctor(_config(), temp_db_path, {})
|
|
assert _result(report, "settings_row").severity is Severity.FAIL
|
|
assert report.failed
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_many_orphans_are_sampled(temp_db_path):
|
|
db = await _build_db(temp_db_path)
|
|
chunks_tbl = await db.open_table("chunks")
|
|
await chunks_tbl.add(
|
|
[
|
|
ChunkRecord(
|
|
id=f"o{i}",
|
|
document_id=f"ghost{i}",
|
|
content="x",
|
|
vector=[0.2] * VECTOR_DIM,
|
|
)
|
|
for i in range(8)
|
|
]
|
|
)
|
|
report = await run_doctor(_config(), temp_db_path, {})
|
|
details = _result(report, "orphaned_chunks").details
|
|
assert len(details) == 6
|
|
assert details[-1] == "... (+3 more)"
|
|
|
|
|
|
def test_sample_returns_all_within_limit():
|
|
assert _sample(["a", "b"]) == ["a", "b"]
|
|
|
|
|
|
def test_embedding_drift_ok_without_stored_identity():
|
|
assert _check_embedding_drift({}, _config()).severity is Severity.OK
|
|
|
|
|
|
def test_vector_index_ok_below_threshold():
|
|
stats = {"chunks": {"num_rows": 10, "has_vector_index": False}}
|
|
assert _check_vector_index(stats).severity is Severity.OK
|
|
|
|
|
|
def test_vector_index_warns_when_missing_above_threshold():
|
|
stats = {"chunks": {"num_rows": 300, "has_vector_index": False}}
|
|
result = _check_vector_index(stats)
|
|
assert result.severity is Severity.WARN
|
|
assert result.remediation == "haiku-rag create-index"
|
|
|
|
|
|
def test_vector_index_warns_on_unindexed_backlog():
|
|
stats = {
|
|
"chunks": {"num_rows": 300, "has_vector_index": True, "num_unindexed_rows": 5}
|
|
}
|
|
assert _check_vector_index(stats).severity is Severity.WARN
|
|
|
|
|
|
def test_vector_index_ok_when_fully_indexed():
|
|
stats = {
|
|
"chunks": {"num_rows": 300, "has_vector_index": True, "num_unindexed_rows": 0}
|
|
}
|
|
assert _check_vector_index(stats).severity is Severity.OK
|
|
|
|
|
|
def test_cli_doctor_nonexistent_db_exits_1(tmp_path):
|
|
result = runner.invoke(cli, ["doctor", "--db", str(tmp_path / "nope.lancedb")])
|
|
assert result.exit_code == 1
|
|
assert "does not exist" in result.output
|
|
|
|
|
|
def test_cli_doctor_exits_0_when_healthy(monkeypatch):
|
|
app = MagicMock()
|
|
app.doctor = AsyncMock(return_value=False)
|
|
monkeypatch.setattr("haiku.rag.cli.create_app", lambda *_a, **_k: app)
|
|
result = runner.invoke(cli, ["doctor", "--db", "/tmp/whatever.lancedb"])
|
|
assert result.exit_code == 0
|
|
|
|
|
|
def test_cli_doctor_exits_1_on_failure(monkeypatch):
|
|
app = MagicMock()
|
|
app.doctor = AsyncMock(return_value=True)
|
|
monkeypatch.setattr("haiku.rag.cli.create_app", lambda *_a, **_k: app)
|
|
result = runner.invoke(cli, ["doctor", "--db", "/tmp/whatever.lancedb"])
|
|
assert result.exit_code == 1
|