928 lines
28 KiB
Python
928 lines
28 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,
|
|
ConversionOptions,
|
|
DoclingServeConfig,
|
|
EmbeddingModelConfig,
|
|
EmbeddingsConfig,
|
|
ModelConfig,
|
|
PictureDescriptionConfig,
|
|
ProcessingConfig,
|
|
ProvidersConfig,
|
|
)
|
|
from haiku.rag.doctor import (
|
|
CheckResult,
|
|
DoctorReport,
|
|
Severity,
|
|
_active_models,
|
|
_check_api_keys,
|
|
_check_embedding_drift,
|
|
_check_vector_index,
|
|
_model_present,
|
|
_probe_endpoint,
|
|
_provider_targets,
|
|
_resolve_endpoint,
|
|
_sample,
|
|
run_doctor,
|
|
run_provider_checks,
|
|
)
|
|
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,
|
|
multimodal: bool = False,
|
|
):
|
|
return AppConfig(
|
|
embeddings=EmbeddingsConfig(
|
|
model=EmbeddingModelConfig(
|
|
provider=provider,
|
|
name=name,
|
|
vector_dim=vector_dim,
|
|
multimodal=multimodal,
|
|
)
|
|
)
|
|
)
|
|
|
|
|
|
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.fixture(autouse=True)
|
|
def _stub_provider_probe(monkeypatch):
|
|
"""Default every provider probe to reachable with the test models present,
|
|
so database-integrity tests don't depend on a live Ollama. Provider tests
|
|
re-patch this with their own behavior."""
|
|
|
|
async def probe(_client, _url):
|
|
return (
|
|
True,
|
|
None,
|
|
{
|
|
"models": [
|
|
{"name": "test"},
|
|
{"name": "gpt-oss:latest"},
|
|
{"name": "qwen3-embedding:4b"},
|
|
]
|
|
},
|
|
)
|
|
|
|
monkeypatch.setattr("haiku.rag.doctor._probe_endpoint", probe)
|
|
|
|
|
|
@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
|
|
|
|
|
|
async def _add_doc(db, doc_id, *, items, metadata=None, chunks=None):
|
|
docs_tbl = await db.open_table("documents")
|
|
meta_tbl = await db.open_table("document_meta")
|
|
await docs_tbl.add([DocumentRecord(id=doc_id, content="x")])
|
|
await meta_tbl.add(
|
|
[
|
|
DocumentMetaRecord(
|
|
document_id=doc_id,
|
|
uri=f"test://{doc_id}",
|
|
metadata=json.dumps(metadata or {}),
|
|
)
|
|
]
|
|
)
|
|
if items:
|
|
items_tbl = await db.open_table("document_items")
|
|
await items_tbl.add(items)
|
|
if chunks:
|
|
chunks_tbl = await db.open_table("chunks")
|
|
await chunks_tbl.add(chunks)
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_document_with_text_but_no_chunks_warns(temp_db_path):
|
|
db = await _build_db(temp_db_path)
|
|
await _add_doc(
|
|
db,
|
|
"d2",
|
|
items=[
|
|
DocumentItemRecord(
|
|
document_id="d2",
|
|
position=0,
|
|
self_ref="#/texts/0",
|
|
label="text",
|
|
text="real content",
|
|
)
|
|
],
|
|
)
|
|
report = await run_doctor(_config(), temp_db_path, {})
|
|
result = _result(report, "documents_text_no_chunks")
|
|
assert result.severity is Severity.WARN
|
|
assert "d2" in result.details
|
|
assert report.count(Severity.FAIL) == 0
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_empty_document_no_chunks_is_ok(temp_db_path):
|
|
db = await _build_db(temp_db_path)
|
|
await _add_doc(db, "d2", items=[])
|
|
report = await run_doctor(_config(), temp_db_path, {})
|
|
assert _result(report, "documents_without_chunks").severity is Severity.OK
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_heading_only_document_no_chunks_is_ok(temp_db_path):
|
|
db = await _build_db(temp_db_path)
|
|
await _add_doc(
|
|
db,
|
|
"d2",
|
|
items=[
|
|
DocumentItemRecord(
|
|
document_id="d2",
|
|
position=0,
|
|
self_ref="#/texts/0",
|
|
label="section_header",
|
|
text="title: haiku.rag",
|
|
)
|
|
],
|
|
)
|
|
report = await run_doctor(_config(), temp_db_path, {})
|
|
assert _result(report, "documents_without_chunks").severity is Severity.OK
|
|
assert all(r.name != "documents_text_no_chunks" for r in report.results)
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_image_only_document_text_embedder_warns(temp_db_path):
|
|
db = await _build_db(temp_db_path)
|
|
await _add_doc(
|
|
db,
|
|
"d2",
|
|
items=[
|
|
DocumentItemRecord(
|
|
document_id="d2", position=0, self_ref="#/pictures/0", label="picture"
|
|
)
|
|
],
|
|
)
|
|
report = await run_doctor(_config(), temp_db_path, {})
|
|
result = _result(report, "documents_images_unsearchable")
|
|
assert result.severity is Severity.WARN
|
|
assert "d2" in result.details
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_image_only_document_multimodal_embedder_warns(temp_db_path):
|
|
db = await _build_db(temp_db_path, provider="vllm", name="qwen-vl")
|
|
await _add_doc(
|
|
db,
|
|
"d2",
|
|
items=[
|
|
DocumentItemRecord(
|
|
document_id="d2", position=0, self_ref="#/pictures/0", label="picture"
|
|
)
|
|
],
|
|
)
|
|
report = await run_doctor(
|
|
_config(provider="vllm", name="qwen-vl", multimodal=True), temp_db_path, {}
|
|
)
|
|
result = _result(report, "documents_pictures_no_chunks")
|
|
assert result.severity is Severity.WARN
|
|
assert "d2" in result.details
|
|
|
|
|
|
@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_chunked_document_without_items_warns(temp_db_path):
|
|
db = await _build_db(temp_db_path)
|
|
await _add_doc(
|
|
db,
|
|
"d2",
|
|
items=[],
|
|
chunks=[
|
|
ChunkRecord(
|
|
id="c2", document_id="d2", content="x", vector=[0.1] * VECTOR_DIM
|
|
)
|
|
],
|
|
)
|
|
report = await run_doctor(_config(), temp_db_path, {})
|
|
result = _result(report, "documents_without_items")
|
|
assert result.severity is Severity.WARN
|
|
assert "d2" in result.details
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_empty_document_without_items_is_ok(temp_db_path):
|
|
db = await _build_db(temp_db_path)
|
|
await _add_doc(db, "d2", items=[])
|
|
report = await run_doctor(_config(), temp_db_path, {})
|
|
assert _result(report, "documents_without_items").severity is Severity.OK
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_missing_picture_data_in_text_document_is_ok(temp_db_path):
|
|
db = await _build_db(temp_db_path)
|
|
await _add_doc(
|
|
db,
|
|
"d2",
|
|
metadata={"content_type": "text/markdown"},
|
|
items=[
|
|
DocumentItemRecord(
|
|
document_id="d2",
|
|
position=0,
|
|
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.OK
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_missing_picture_data_in_pdf_document_warns(temp_db_path):
|
|
db = await _build_db(temp_db_path)
|
|
await _add_doc(
|
|
db,
|
|
"d2",
|
|
metadata={"content_type": "application/pdf"},
|
|
items=[
|
|
DocumentItemRecord(
|
|
document_id="d2",
|
|
position=0,
|
|
self_ref="#/pictures/0",
|
|
label="picture",
|
|
picture_data=None,
|
|
)
|
|
],
|
|
)
|
|
report = await run_doctor(_config(), temp_db_path, {})
|
|
result = _result(report, "picture_data")
|
|
assert result.severity is Severity.WARN
|
|
assert "d2" in result.details
|
|
|
|
|
|
@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
|
|
|
|
|
|
# --- Active models / API keys ---
|
|
|
|
|
|
def test_api_key_not_required_for_custom_openai_base_url():
|
|
config = AppConfig(
|
|
embeddings=EmbeddingsConfig(
|
|
model=EmbeddingModelConfig(
|
|
provider="openai",
|
|
name="x",
|
|
vector_dim=4,
|
|
base_url="http://localhost:1234/v1",
|
|
)
|
|
)
|
|
)
|
|
assert _check_api_keys(config, {}).severity is Severity.OK
|
|
|
|
|
|
def test_api_key_required_for_openai_without_base_url():
|
|
config = AppConfig(
|
|
embeddings=EmbeddingsConfig(
|
|
model=EmbeddingModelConfig(provider="openai", name="x", vector_dim=4)
|
|
)
|
|
)
|
|
result = _check_api_keys(config, {})
|
|
assert result.severity is Severity.FAIL
|
|
assert any("OPENAI_API_KEY" in d for d in result.details)
|
|
|
|
|
|
def test_active_models_includes_picture_description_when_enabled():
|
|
config = AppConfig(processing=ProcessingConfig(pictures="description"))
|
|
names = [name for _p, name, _b in _active_models(config)]
|
|
assert "ministral-3" in names
|
|
|
|
|
|
def test_active_models_excludes_picture_description_by_default():
|
|
names = [name for _p, name, _b in _active_models(AppConfig())]
|
|
assert "ministral-3" not in names
|
|
|
|
|
|
def test_active_models_includes_title_model_when_auto_title():
|
|
base = _active_models(AppConfig())
|
|
with_title = _active_models(AppConfig(processing=ProcessingConfig(auto_title=True)))
|
|
assert len(with_title) == len(base) + 1
|
|
|
|
|
|
def test_picture_description_model_checked_for_api_key():
|
|
config = AppConfig(
|
|
processing=ProcessingConfig(
|
|
pictures="description",
|
|
conversion_options=ConversionOptions(
|
|
picture_description=PictureDescriptionConfig(
|
|
model=ModelConfig(provider="openai", name="gpt-4o")
|
|
)
|
|
),
|
|
)
|
|
)
|
|
result = _check_api_keys(config, {})
|
|
assert result.severity is Severity.FAIL
|
|
assert any("OPENAI_API_KEY" in d for d in result.details)
|
|
|
|
|
|
# --- Provider connectivity ---
|
|
|
|
|
|
def test_resolve_endpoint_ollama_strips_v1():
|
|
assert _resolve_endpoint("ollama", "http://h:1/v1", "http://fallback") == (
|
|
"http://h:1/api/tags",
|
|
"ollama",
|
|
"http://h:1",
|
|
)
|
|
|
|
|
|
def test_resolve_endpoint_ollama_uses_provider_fallback():
|
|
assert _resolve_endpoint("ollama", None, "http://fallback:11434") == (
|
|
"http://fallback:11434/api/tags",
|
|
"ollama",
|
|
"http://fallback:11434",
|
|
)
|
|
|
|
|
|
def test_resolve_endpoint_vllm_default_and_models_path():
|
|
assert _resolve_endpoint("vllm", None, "http://o") == (
|
|
"http://localhost:8000/v1/models",
|
|
"openai",
|
|
"http://localhost:8000/v1",
|
|
)
|
|
|
|
|
|
def test_resolve_endpoint_vllm_appends_v1():
|
|
assert _resolve_endpoint("vllm", "http://vllm:8000", "http://o") == (
|
|
"http://vllm:8000/v1/models",
|
|
"openai",
|
|
"http://vllm:8000/v1",
|
|
)
|
|
|
|
|
|
def test_resolve_endpoint_openai_saas_is_skipped():
|
|
assert _resolve_endpoint("openai", None, "http://o") is None
|
|
|
|
|
|
def test_resolve_endpoint_openai_with_base_url():
|
|
assert _resolve_endpoint("openai", "http://lmstudio:1234/v1", "http://o") == (
|
|
"http://lmstudio:1234/v1/models",
|
|
"openai",
|
|
"http://lmstudio:1234/v1",
|
|
)
|
|
|
|
|
|
def test_resolve_endpoint_local_provider():
|
|
assert _resolve_endpoint("sentence-transformers", None, "http://o") == "local"
|
|
|
|
|
|
def test_model_present_tag_insensitive():
|
|
assert _model_present("gpt-oss", {"gpt-oss:latest"})
|
|
assert _model_present("qwen:4b", {"qwen:4b"})
|
|
assert not _model_present("qwen:4b", {"qwen:8b"})
|
|
|
|
|
|
def test_provider_targets_default_groups_ollama_models():
|
|
targets, local = _provider_targets(AppConfig())
|
|
assert not local
|
|
assert len(targets) == 1
|
|
entry = next(iter(targets.values()))
|
|
assert entry["kind"] == "ollama"
|
|
assert {"qwen3-embedding:4b", "gpt-oss"} <= entry["models"]
|
|
|
|
|
|
def test_provider_targets_includes_docling_serve():
|
|
config = AppConfig(
|
|
processing=ProcessingConfig(converter="docling-serve"),
|
|
providers=ProvidersConfig(
|
|
docling_serve=DoclingServeConfig(base_url="http://docling:5001")
|
|
),
|
|
)
|
|
targets, _ = _provider_targets(config)
|
|
assert "http://docling:5001/health" in targets
|
|
assert targets["http://docling:5001/health"]["kind"] == "docling-serve"
|
|
|
|
|
|
def test_provider_targets_collects_local_providers():
|
|
config = AppConfig(
|
|
embeddings=EmbeddingsConfig(
|
|
model=EmbeddingModelConfig(
|
|
provider="sentence-transformers", name="x", vector_dim=4
|
|
)
|
|
)
|
|
)
|
|
_, local = _provider_targets(config)
|
|
assert "sentence-transformers" in local
|
|
|
|
|
|
def _fake_probe(result):
|
|
async def probe(_client, _url):
|
|
return result
|
|
|
|
return probe
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_provider_check_ok_when_models_present(monkeypatch):
|
|
monkeypatch.setattr(
|
|
"haiku.rag.doctor._probe_endpoint",
|
|
_fake_probe(
|
|
(
|
|
True,
|
|
None,
|
|
{
|
|
"models": [
|
|
{"name": "qwen3-embedding:4b"},
|
|
{"name": "gpt-oss:latest"},
|
|
]
|
|
},
|
|
)
|
|
),
|
|
)
|
|
results = await run_provider_checks(AppConfig())
|
|
assert all(r.severity is Severity.OK for r in results)
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_provider_check_warns_on_missing_model(monkeypatch):
|
|
monkeypatch.setattr(
|
|
"haiku.rag.doctor._probe_endpoint",
|
|
_fake_probe((True, None, {"models": [{"name": "something-else"}]})),
|
|
)
|
|
results = await run_provider_checks(AppConfig())
|
|
result = next(r for r in results if r.name.startswith("provider:"))
|
|
assert result.severity is Severity.WARN
|
|
assert result.details
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_provider_check_fails_when_unreachable(monkeypatch):
|
|
monkeypatch.setattr(
|
|
"haiku.rag.doctor._probe_endpoint",
|
|
_fake_probe((False, "Connection refused", None)),
|
|
)
|
|
results = await run_provider_checks(AppConfig())
|
|
result = next(r for r in results if r.name.startswith("provider:"))
|
|
assert result.severity is Severity.FAIL
|
|
assert "Connection refused" in result.details
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_provider_check_reports_local_provider(monkeypatch):
|
|
monkeypatch.setattr(
|
|
"haiku.rag.doctor._probe_endpoint",
|
|
_fake_probe((True, None, {"models": [{"name": "gpt-oss:latest"}]})),
|
|
)
|
|
config = AppConfig(
|
|
embeddings=EmbeddingsConfig(
|
|
model=EmbeddingModelConfig(
|
|
provider="sentence-transformers", name="x", vector_dim=4
|
|
)
|
|
)
|
|
)
|
|
results = await run_provider_checks(config)
|
|
local = next(r for r in results if r.name == "provider:sentence-transformers")
|
|
assert local.severity is Severity.OK
|
|
assert "local" in local.message
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_run_doctor_includes_provider_results(temp_db_path, monkeypatch):
|
|
await _build_db(temp_db_path)
|
|
monkeypatch.setattr(
|
|
"haiku.rag.doctor._probe_endpoint",
|
|
_fake_probe(
|
|
(True, None, {"models": [{"name": "test"}, {"name": "gpt-oss:latest"}]})
|
|
),
|
|
)
|
|
report = await run_doctor(_config(), temp_db_path, {})
|
|
assert any(r.name.startswith("provider:") for r in report.results)
|
|
assert not report.failed
|
|
|
|
|
|
async def _probe_with_handler(handler):
|
|
import httpx
|
|
|
|
transport = httpx.MockTransport(handler)
|
|
async with httpx.AsyncClient(transport=transport) as client:
|
|
return await _probe_endpoint(client, "http://x")
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_probe_endpoint_success_with_json():
|
|
import httpx
|
|
|
|
reachable, error, payload = await _probe_with_handler(
|
|
lambda _request: httpx.Response(200, json={"models": []})
|
|
)
|
|
assert reachable and error is None and payload == {"models": []}
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_probe_endpoint_success_non_json():
|
|
import httpx
|
|
|
|
reachable, _, payload = await _probe_with_handler(
|
|
lambda _request: httpx.Response(200, content=b"not json")
|
|
)
|
|
assert reachable and payload is None
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_probe_endpoint_http_error_status():
|
|
import httpx
|
|
|
|
reachable, error, _ = await _probe_with_handler(
|
|
lambda _request: httpx.Response(503)
|
|
)
|
|
assert not reachable
|
|
assert error is not None and "503" in error
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_probe_endpoint_connection_error():
|
|
import httpx
|
|
|
|
def handler(_request):
|
|
raise httpx.ConnectError("refused")
|
|
|
|
reachable, error, _ = await _probe_with_handler(handler)
|
|
assert not reachable
|
|
assert error is not None and "refused" in error
|