Add provider connectivity probes to haiku-rag doctor

This commit is contained in:
Yiorgis Gozadinos 2026-06-23 10:34:27 +03:00
parent c2cb3cedf3
commit cc1d8d1e4c
No known key found for this signature in database
5 changed files with 467 additions and 25 deletions

View file

@ -3,7 +3,7 @@
### Added ### Added
- `haiku-rag doctor` checks a database for consistency (orphaned chunks/items, documents without chunks, dangling `doc_item_refs`, vector-dimension mismatch, unembedded chunks, missing picture data, settings/embedding drift, pending migrations, vector-index coverage, provider API keys) and exits 1 when any check fails. - `haiku-rag doctor` checks a database for consistency (orphaned chunks/items, documents without chunks, dangling `doc_item_refs`, vector-dimension mismatch, unembedded chunks, missing picture data, settings/embedding drift, pending migrations, vector-index coverage, provider API keys) and probes configured provider endpoints (Ollama `/api/tags` with model presence, docling-serve `/health`, OpenAI-compatible/vLLM `/models`); exits 1 when any check fails.
## [0.60.0] - 2026-06-22 ## [0.60.0] - 2026-06-22

View file

@ -308,6 +308,14 @@ Checks include:
- the vector index covers all chunks - the vector index covers all chunks
- API keys are set for configured providers - API keys are set for configured providers
It also probes the external endpoints the config uses and reports them under a Providers section:
- Ollama is reachable and the configured models are installed (`{base_url}/api/tags`)
- docling-serve is reachable when used as the converter or chunker (`{base_url}/health`)
- custom OpenAI-compatible and vLLM endpoints respond (`{base_url}/models`)
SaaS providers (OpenAI, Anthropic, Cohere, Jina, ZeroEntropy, Voyage) are covered by the API-key check rather than a network probe. In-process local models (sentence-transformers, cross-encoder, mxbai, jina-local) have no endpoint and are reported as such.
Each failure prints the command that fixes it (`rebuild`, `create-index`, `migrate`, `rebuild --set-embedder`). `doctor` makes no changes. It exits with status 1 when any check fails, so it can gate CI or monitoring. Each failure prints the command that fixes it (`rebuild`, `create-index`, `migrate`, `rebuild --set-embedder`). `doctor` makes no changes. It exits with status 1 when any check fails, so it can gate CI or monitoring.
### Migrate Database ### Migrate Database

View file

@ -220,14 +220,25 @@ class HaikuRAGApp: # pragma: no cover
Severity.WARN: "[yellow]![/yellow]", Severity.WARN: "[yellow]![/yellow]",
Severity.FAIL: "[red]✗[/red]", Severity.FAIL: "[red]✗[/red]",
} }
self.console.rule()
for result in report.results: def render(result):
self.console.print(f"{glyphs[result.severity]} {result.message}") self.console.print(f"{glyphs[result.severity]} {result.message}")
for detail in result.details: for detail in result.details:
self.console.print(f" [dim]{detail}[/dim]") self.console.print(f" [dim]{detail}[/dim]")
if result.remediation: if result.remediation:
self.console.print(f" [dim]→ {result.remediation}[/dim]") self.console.print(f" [dim]→ {result.remediation}[/dim]")
database = [r for r in report.results if not r.name.startswith("provider:")]
providers = [r for r in report.results if r.name.startswith("provider:")]
self.console.rule("[bold]Database[/bold]")
for result in database:
render(result)
if providers:
self.console.rule("[bold]Providers[/bold]")
for result in providers:
render(result)
self.console.rule() self.console.rule()
self.console.print( self.console.print(
f"[green]{report.count(Severity.OK)} ok[/green], " f"[green]{report.count(Severity.OK)} ok[/green], "

View file

@ -1,7 +1,9 @@
import asyncio
import json import json
from enum import StrEnum from enum import StrEnum
from pathlib import Path from pathlib import Path
import httpx
import numpy as np import numpy as np
from pydantic import BaseModel, Field from pydantic import BaseModel, Field
@ -28,6 +30,12 @@ _PROVIDER_ENV_VARS: dict[str, str] = {
"zeroentropy": "ZEROENTROPY_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"}
# Operators care whether an endpoint answers now, not eventually.
_PROBE_TIMEOUT_S = 2.0
class Severity(StrEnum): class Severity(StrEnum):
OK = "ok" OK = "ok"
@ -448,6 +456,161 @@ def _check_vector_index(stats: dict) -> CheckResult:
) )
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()},
)
add_model(
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:
add_model(model.provider, model.name, model.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(
url: str, 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) -> list[CheckResult]:
"""Probe the external endpoints the current config actually uses."""
targets, local = _provider_targets(config)
results: list[CheckResult] = []
if targets:
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(url, 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( async def run_doctor(
config: AppConfig, db_path: Path, environ: dict[str, str] config: AppConfig, db_path: Path, environ: dict[str, str]
) -> DoctorReport: ) -> DoctorReport:
@ -459,29 +622,29 @@ async def run_doctor(
db = await connect_lancedb(config, db_path) db = await connect_lancedb(config, db_path)
stats = await get_database_stats(db) stats = await get_database_stats(db)
results: list[CheckResult] = []
if not any(entry["exists"] for entry in stats.values()): if not any(entry["exists"] for entry in stats.values()):
return DoctorReport( results.append(
results=[ CheckResult(
CheckResult( name="tables_present",
name="tables_present", severity=Severity.FAIL,
severity=Severity.FAIL, message="Database is empty.",
message="Database is empty.", remediation="haiku-rag init",
remediation="haiku-rag init", )
)
]
) )
else:
results = [_check_tables_present(stats)] results.append(_check_tables_present(stats))
missing = [name for name in REQUIRED_TABLES if not stats[name]["exists"]] missing = [name for name in REQUIRED_TABLES if not stats[name]["exists"]]
if not missing: if not missing:
async with Store( async with Store(
db_path, db_path,
config=config, config=config,
skip_validation=True, skip_validation=True,
read_only=True, read_only=True,
skip_migration_check=True, skip_migration_check=True,
) as store: ) as store:
results += await run_db_checks(store, config, stats) results += await run_db_checks(store, config, stats)
results.append(_check_api_keys(config, environ)) results.append(_check_api_keys(config, environ))
results += await run_provider_checks(config)
return DoctorReport(results=results) return DoctorReport(results=results)

View file

@ -7,15 +7,27 @@ import pytest
from typer.testing import CliRunner from typer.testing import CliRunner
from haiku.rag.cli import _cli as cli from haiku.rag.cli import _cli as cli
from haiku.rag.config.models import AppConfig, EmbeddingModelConfig, EmbeddingsConfig from haiku.rag.config.models import (
AppConfig,
DoclingServeConfig,
EmbeddingModelConfig,
EmbeddingsConfig,
ProcessingConfig,
ProvidersConfig,
)
from haiku.rag.doctor import ( from haiku.rag.doctor import (
CheckResult, CheckResult,
DoctorReport, DoctorReport,
Severity, Severity,
_check_embedding_drift, _check_embedding_drift,
_check_vector_index, _check_vector_index,
_model_present,
_probe_endpoint,
_provider_targets,
_resolve_endpoint,
_sample, _sample,
run_doctor, run_doctor,
run_provider_checks,
) )
from haiku.rag.store.engine import ( from haiku.rag.store.engine import (
DocumentItemRecord, DocumentItemRecord,
@ -110,6 +122,28 @@ def _result(report: DoctorReport, name: str) -> CheckResult:
return next(r for r in report.results if r.name == name) 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 @pytest.mark.asyncio
async def test_healthy_db_all_ok(temp_db_path): async def test_healthy_db_all_ok(temp_db_path):
await _build_db(temp_db_path) await _build_db(temp_db_path)
@ -423,3 +457,229 @@ def test_cli_doctor_exits_1_on_failure(monkeypatch):
monkeypatch.setattr("haiku.rag.cli.create_app", lambda *_a, **_k: app) monkeypatch.setattr("haiku.rag.cli.create_app", lambda *_a, **_k: app)
result = runner.invoke(cli, ["doctor", "--db", "/tmp/whatever.lancedb"]) result = runner.invoke(cli, ["doctor", "--db", "/tmp/whatever.lancedb"])
assert result.exit_code == 1 assert result.exit_code == 1
# --- 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