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

Merge pull request #518 from ggozad/chore/coverage-100-and-test-consolidation

Browse source 
Reach and enforce 100% test coverage
This commit is contained in:
Yiorgis Gozadinos 2026-07-27 14:09:21 +03:00 committed by GitHub
commit 9168bc15ef
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
77 changed files with 4708 additions and 849 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

2
.github/workflows/build-docs.yml vendored
View file

@ -15,7 +15,7 @@ jobs:
runs-on: ubuntu-latest
steps:
- uses: actions/checkout@v4
- uses: astral-sh/setup-uv@v5
- uses: astral-sh/setup-uv@v9.0.0
- run: uv sync --group dev
- run: uv run zensical build
- uses: actions/configure-pages@v5

2
.github/workflows/build-publish-slim.yml vendored
View file

@ -9,7 +9,7 @@ jobs:
runs-on: ubuntu-latest
steps:
- uses: actions/checkout@v4
- uses: astral-sh/setup-uv@v4
- uses: astral-sh/setup-uv@v9.0.0
with:
enable-cache: true
- name: Set up Python

2
.github/workflows/build-publish.yml vendored
View file

@ -12,7 +12,7 @@ jobs:
if: ${{ github.event_name == 'workflow_dispatch' || github.event.workflow_run.conclusion == 'success' }}
steps:
- uses: actions/checkout@v4
- uses: astral-sh/setup-uv@v4
- uses: astral-sh/setup-uv@v9.0.0
with:
enable-cache: true
- name: Set up Python

10
.github/workflows/test.yml vendored
View file

@ -11,13 +11,13 @@ jobs:
runs-on: ubuntu-latest
steps:
- uses: actions/checkout@v4
- uses: astral-sh/setup-uv@v4
- uses: astral-sh/setup-uv@v9.0.0
with:
enable-cache: true
- name: Set up Python
uses: actions/setup-python@v5
with:
python-version-file: "pyproject.toml"
python-version-file: ".python-version"
- name: Install dependencies
run: uv sync --all-extras
- name: Lint
@ -52,13 +52,13 @@ jobs:
HF_TOKEN: ${{ secrets.HF_TOKEN }}
steps:
- uses: actions/checkout@v4
- uses: astral-sh/setup-uv@v4
- uses: astral-sh/setup-uv@v9.0.0
with:
enable-cache: true
- name: Set up Python
uses: actions/setup-python@v5
with:
python-version-file: "pyproject.toml"
python-version-file: ".python-version"
- name: Install dependencies
run: uv sync --all-extras
- name: Cache HuggingFace models
@ -77,7 +77,7 @@ jobs:
env:
HF_HUB_OFFLINE: ${{ steps.hf-cache.outputs.cache-hit == 'true' && '1' || '0' }}
TRANSFORMERS_OFFLINE: ${{ steps.hf-cache.outputs.cache-hit == 'true' && '1' || '0' }}
run: uv run pytest -m "not integration" --cov=haiku --cov-report=xml
run: uv run pytest -m "not integration" --cov=haiku --cov-report=xml --cov-report=term-missing:skip-covered
- name: Upload coverage to Codecov
uses: codecov/codecov-action@v5
with:

9
CHANGELOG.md
View file

@ -1,6 +1,15 @@
# Changelog
## [Unreleased]
### Fixed
- `Store.set_haiku_version` stamps the store's own config into a recreated settings row instead of the process-global `Config`.
- `check_source_accessible` returns `False` for a URI it cannot resolve (unparseable host, unreadable path) instead of raising and aborting a full rebuild.
### Removed
- `SettingsRepository.create`, `get_by_id`, `update`, `delete`, `list_all` and `ChunkRepository.update`, `delete`, `get_chunks_in_range`.
## [0.70.0] - 2026-07-25
### Added

17
haiku_rag_slim/haiku/rag/client/documents.py
View file

@ -570,7 +570,10 @@ def _extract_pdf_attachments(
for i in range(attachment_count):
att = pdf.get_attachment(i)
name = att.get_name()
if not name:
# A malformed PDF can carry an attachment with an empty /F, so
# this is real validation on untrusted input — it just needs a
# hand-crafted file to reach, which no fixture here produces.
if not name: # pragma: no cover - needs a malformed PDF
continue
data = bytes(att.get_data())
child_uri = f"{parent_uri}#attachment={quote(name, safe='')}"
@ -906,13 +909,19 @@ async def update_document(
def check_source_accessible(uri: str) -> bool:
"""Check if a document's source URI is accessible."""
parsed_url = urlparse(uri)
"""Check if a document's source URI is accessible.
Anything the URI itself makes unanswerable counts as inaccessible rather
than aborting the caller's sweep: ``urlparse`` rejects malformed IPv6
hosts, and ``Path.exists`` re-raises errno values outside its ignored set
(an unreadable parent directory, an over-long name).
"""
try:
parsed_url = urlparse(uri)
if parsed_url.scheme == "file":
return Path(parsed_url.path).exists()
elif parsed_url.scheme in ("http", "https", "s3"):
return True
return False
except Exception:
except (ValueError, OSError):
return False

2
haiku_rag_slim/haiku/rag/client/downloads.py
View file

@ -38,7 +38,7 @@ async def download_models(
yield DownloadProgress(model="docling", status="start")
await asyncio.to_thread(download_models)
yield DownloadProgress(model="docling", status="done")
except ImportError:
except ImportError: # pragma: no cover - docling installed in test env
pass
# HuggingFace tokenizer

12
haiku_rag_slim/haiku/rag/client/rebuild.py
View file

@ -66,7 +66,7 @@ class _StagingMarkerRecord(LanceModel):
async def rebuild_database(
client: "HaikuRAG", mode: "RebuildMode | None" = None
client: "HaikuRAG", mode: "RebuildMode"
) -> AsyncGenerator[str, None]:
"""Rebuild the database with the specified mode.
@ -79,9 +79,6 @@ async def rebuild_database(
"""
from haiku.rag.client import RebuildMode
if mode is None:
mode = RebuildMode.FULL
async with client.store._rebuild_lock:
if mode == RebuildMode.SET_EMBEDDER:
await _set_embedder(client)
@ -284,11 +281,12 @@ async def _populate_staging_table(client: "HaikuRAG") -> None:
pagination drift), so peak memory stays bounded regardless of corpus
size. The vector column is omitted — the point of embed-only rebuild is
to regenerate it.
Requires ``_resolve_rebuild_recovery`` to have cleared any leftover
staging table first: ``create_table`` raises if the name is already taken.
"""
db = client.store.db
tables = (await db.list_tables()).tables
if _STAGING_TABLE_NAME in tables:
await db.drop_table(_STAGING_TABLE_NAME)
staging = await db.create_table(_STAGING_TABLE_NAME, schema=_StagingChunkRecord)
if "chunks" not in tables:
@ -301,8 +299,6 @@ async def _populate_staging_table(client: "HaikuRAG") -> None:
)
async for batch in stream:
rows = batch.to_pylist()
if not rows:
continue
records = [
_StagingChunkRecord(
id=r["id"],

4
haiku_rag_slim/haiku/rag/context.py
View file

@ -101,8 +101,6 @@ def _evidence_anchors(content: str, max_chars: int) -> list[str]:
mid = len(content) // 2
start = max(0, mid - target // 2)
anchors.append(content[start : start + target])
if not anchors:
anchors.append(content[:max_chars] if max_chars > 0 else content)
return anchors
@ -289,8 +287,6 @@ def _add_input_pages_for_surviving_refs(
) -> None:
"""Fill missing item-table pages from inputs whose own refs all survived."""
surviving = set(refs)
if not surviving:
return
for result in original_results:
if not result.page_numbers or not result.doc_item_refs:
continue

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

@ -330,7 +330,7 @@ def _common_path_prefix(labels: list[str]) -> str:
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:
if len(labels) < 2: # pragma: no cover - families always have >=2 members
return ""
lo, hi = min(labels), max(labels)
end = 0

4
haiku_rag_slim/haiku/rag/ingester/pollers/base.py
View file

@ -55,7 +55,6 @@ class BasePoller:
self._sync = sync_repo
self._breaker = breaker or CircuitBreaker(config.circuit_breaker)
self._stop = asyncio.Event()
self._task: asyncio.Task | None = None
self._last_polled_at: datetime | None = None
self._last_skip_reason: str | None = None
self._default_max_attempts = default_max_attempts
@ -84,9 +83,6 @@ class BasePoller:
async def stop(self) -> None:
self._stop.set()
if self._task is not None:
await asyncio.gather(self._task, return_exceptions=True)
self._task = None
async def _stagger_start(self) -> bool:
"""Sleep a random fraction of the interval so pollers sharing an

2
haiku_rag_slim/haiku/rag/ingester/sources/fs.py
View file

@ -143,7 +143,7 @@ class FSSource:
if path.is_symlink():
try:
resolved = path.resolve(strict=False)
except OSError:
except OSError: # pragma: no cover - strict=False absorbs these
continue
if not resolved.is_relative_to(self.root):
continue

2
haiku_rag_slim/haiku/rag/mcp.py
View file

@ -244,6 +244,6 @@ def create_mcp_server(
result = await rag.analyze(question, filter=filter, images=images)
return result.answer
except Exception as e:
return f"Error running analysis capability: {e!s}" # pragma: no cover
return f"Error running analysis capability: {e!s}"
return mcp

6
haiku_rag_slim/haiku/rag/sandbox/sandbox.py
View file

@ -356,7 +356,7 @@ class Sandbox:
return read_toc
for doc in docs:
if not doc.id:
if not doc.id: # pragma: no cover - stored rows always carry an id
continue
doc_id: str = doc.id
doc_dir = f"/documents/{doc_id}"
@ -437,7 +437,9 @@ class Sandbox:
stdout_lines: list[str] = []
def print_callback(_stream: Literal["stdout"], text: str) -> None:
def print_callback( # pragma: no cover - runs on Monty's Rust thread
_stream: Literal["stdout"], text: str
) -> None:
stdout_lines.append(text)
max_chars = self._config.analysis.max_output_chars

4
haiku_rag_slim/haiku/rag/store/engine.py
View file

@ -606,7 +606,7 @@ class Store:
for tag in tags.values()
if tag["version"] in timestamps
]
if not tagged:
if not tagged: # pragma: no cover - vacuum never cleans a tagged version
return retention
# LanceDB version timestamps are naive datetimes in local time.
@ -847,7 +847,7 @@ class Store:
)
else:
# Create new settings record
settings_data = Config.model_dump(mode="json")
settings_data = self._config.model_dump(mode="json")
settings_data["version"] = version
await self.settings_table.add(
[SettingsRecord(id="settings", settings=json.dumps(settings_data))]

2
haiku_rag_slim/haiku/rag/store/models/chunk.py
View file

@ -74,7 +74,7 @@ class ChunkMetadata(BaseModel):
continue
for prov_item in prov:
bbox = getattr(prov_item, "bbox", None)
if bbox is None:
if bbox is None: # pragma: no cover - prov always carries a bbox
continue
bounding_boxes.append(
BoundingBox(

4
haiku_rag_slim/haiku/rag/store/models/document_item.py
View file

@ -105,10 +105,6 @@ def extract_item_text(
except Exception:
pass
if caption := getattr(item, "caption", None):
if hasattr(caption, "text"):
return caption.text
return None

82
haiku_rag_slim/haiku/rag/store/repositories/chunk.py
View file

@ -4,7 +4,6 @@ from typing import TYPE_CHECKING
from uuid import uuid4
if TYPE_CHECKING:
import pandas as pd
from lancedb.query import AsyncQueryBase
from lancedb.index import FTS
@ -153,40 +152,6 @@ class ChunkRepository:
order=chunk_record.order,
)
async def update(self, entity: Chunk) -> Chunk:
"""Update an existing chunk.
Chunk must have embedding set before calling this method.
"""
self.store._assert_writable()
assert entity.id, "Chunk ID is required for update"
assert entity.embedding is not None, "Chunk must have an embedding"
await self.store.chunks_table.update(
{
"document_id": entity.document_id,
"content": entity.content,
"content_fts": self._contextualize_content(entity),
"metadata": json.dumps(
{k: v for k, v in entity.metadata.items() if k != "order"}
),
"order": int(entity.order),
"vector": entity.embedding,
},
where=f"id = '{entity.id}'",
)
return entity
async def delete(self, entity_id: str) -> bool:
"""Delete a chunk by its ID."""
self.store._assert_writable()
chunk = await self.get_by_id(entity_id)
if chunk is None:
return False
await self.store.chunks_table.delete(f"id = '{entity_id}'")
return True
async def list_all(
self, limit: int | None = None, offset: int | None = None
) -> list[Chunk]:
@ -417,46 +382,10 @@ class ChunkRepository:
)
return len(df)
async def get_chunks_in_range(
self, document_id: str, min_order: int, max_order: int
) -> list[Chunk]:
"""Get chunks for a document within an order range.
Args:
document_id: The document ID to get chunks for.
min_order: Minimum order value (inclusive).
max_order: Maximum order value (inclusive).
Returns:
List of chunks within the order range.
"""
where = (
f"document_id = '{document_id}'"
f" AND `order` >= {min_order}"
f" AND `order` <= {max_order}"
)
results = await query_to_pydantic(
self.store.chunks_table.query().where(where), self.store.ChunkRecord
)
return [
Chunk(
id=rec.id,
document_id=rec.document_id,
content=rec.content,
metadata=json.loads(rec.metadata),
order=rec.order,
)
for rec in results
]
async def _process_search_results(
self, query_result: "pd.DataFrame | AsyncQueryBase"
self, query_result: "AsyncQueryBase"
) -> list[tuple[Chunk, float]]:
"""Process search results into chunks with document info and scores.
Args:
query_result: Either a pandas DataFrame or a LanceDB async query result
"""
"""Process search results into chunks with document info and scores."""
import pandas as pd
def extract_scores(df: pd.DataFrame) -> list[float]:
@ -473,12 +402,7 @@ class ChunkRepository:
else:
raise ValueError("Unknown search result format, cannot extract scores")
# Convert everything to DataFrame for uniform processing
if isinstance(query_result, pd.DataFrame):
df = query_result
else:
# Convert LanceDB query result to DataFrame
df = await query_result.to_pandas()
df = await query_result.to_pandas()
# Extract scores
scores = extract_scores(df)

41
haiku_rag_slim/haiku/rag/store/repositories/settings.py
View file

@ -18,47 +18,6 @@ class SettingsRepository:
def __init__(self, store: Store) -> None:
self.store = store
async def create(self, entity: dict) -> dict:
"""Create settings in the database."""
settings_record = SettingsRecord(id="settings", settings=json.dumps(entity))
await self.store.settings_table.add([settings_record])
return entity
async def get_by_id(self, entity_id: str) -> dict | None:
"""Get settings by ID."""
results = await query_to_pydantic(
self.store.settings_table.query().where(f"id = '{entity_id}'").limit(1),
SettingsRecord,
)
if not results:
return None
return json.loads(results[0].settings) if results[0].settings else {}
async def update(self, entity: dict) -> dict:
"""Update existing settings."""
await self.store.settings_table.update(
{"settings": json.dumps(entity)}, where="id = 'settings'"
)
return entity
async def delete(self, entity_id: str) -> bool:
"""Delete settings by ID."""
await self.store.settings_table.delete(f"id = '{entity_id}'")
return True
async def list_all(
self, limit: int | None = None, offset: int | None = None
) -> list[dict]:
"""List all settings."""
results = await query_to_pydantic(
self.store.settings_table.query(), SettingsRecord
)
return [
json.loads(record.settings) if record.settings else {} for record in results
]
async def get_current_settings(self) -> dict:
"""Get the current settings."""
results = await query_to_pydantic(

1
pyproject.toml
View file

@ -146,6 +146,7 @@ omit = [
[tool.coverage.report]
show_missing = true
fail_under = 100
exclude_also = [
"if TYPE_CHECKING:",
"@abstractmethod",

0
tests/cassettes/test_chunk/test_chunk_content_fts_populated.yaml → tests/cassettes/test_chunk/test_chunk_content_fts[populated].yaml
View file

0
tests/cassettes/test_chunk/test_chunk_content_fts_without_headings.yaml → tests/cassettes/test_chunk/test_chunk_content_fts[without_headings].yaml
View file

80
tests/cassettes/test_chunk/test_chunk_repository_get_by_id_and_list_all_pagination.yaml Normal file
View file

File diff suppressed because one or more lines are too long

82
tests/cassettes/test_chunk/test_chunk_search_returns_empty_for_blank_query.yaml Normal file
View file

File diff suppressed because one or more lines are too long

42
tests/cassettes/test_chunk/test_chunk_search_with_precomputed_vector_skips_text_query.yaml Normal file
View file

File diff suppressed because one or more lines are too long

42
tests/cassettes/test_client/test_client_create_document_from_source_with_title.yaml
View file

File diff suppressed because one or more lines are too long

42
tests/cassettes/test_client/test_client_unchanged_file_keeps_timestamp.yaml
View file

File diff suppressed because one or more lines are too long

42
tests/cassettes/test_client/test_import_documents_schedules_vacuum_per_config[False].yaml Normal file
View file

File diff suppressed because one or more lines are too long

42
tests/cassettes/test_client/test_import_documents_schedules_vacuum_per_config[True].yaml Normal file
View file

File diff suppressed because one or more lines are too long

82
tests/cassettes/test_client/test_reingesting_a_source_applies_an_explicit_title.yaml Normal file
View file

File diff suppressed because one or more lines are too long

82
tests/cassettes/test_document_tools/TestSummarizeDocumentTool.test_summarize_document_returns_model_summary.yaml Normal file
View file

File diff suppressed because one or more lines are too long

42
tests/cassettes/test_rebuild/test_flush_rebuild_batch_is_a_noop_without_documents.yaml Normal file
View file

File diff suppressed because one or more lines are too long

86
tests/cassettes/test_rebuild/test_rebuild_descriptions_flushes_in_batches.yaml Normal file
View file

File diff suppressed because one or more lines are too long

242
tests/cassettes/test_rebuild/test_rebuild_embed_only_flushes_in_batches.yaml Normal file
View file

File diff suppressed because one or more lines are too long

42
tests/cassettes/test_rebuild/test_rebuild_embed_only_recovers_picture_bytes[recoverable].yaml Normal file
View file

File diff suppressed because one or more lines are too long

46
tests/cassettes/test_rebuild/test_rebuild_embed_only_recovers_picture_bytes[wiped].yaml Normal file
View file

File diff suppressed because one or more lines are too long

42
tests/cassettes/test_rebuild/test_rebuild_embed_only_yields_documents_without_chunks.yaml Normal file
View file

File diff suppressed because one or more lines are too long

242
tests/cassettes/test_rebuild/test_rebuild_full_flushes_in_batches.yaml Normal file
View file

File diff suppressed because one or more lines are too long

162
tests/cassettes/test_rebuild/test_rebuild_full_flushes_pending_before_source_rebuild.yaml Normal file
View file

File diff suppressed because one or more lines are too long

42
tests/cassettes/test_rebuild/test_rebuild_full_skips_document_deleted_mid_rebuild.yaml Normal file
View file

File diff suppressed because one or more lines are too long

82
tests/cassettes/test_rebuild/test_rebuild_full_warns_when_source_is_missing.yaml Normal file
View file

File diff suppressed because one or more lines are too long

82
tests/cassettes/test_rebuild/test_rebuild_warns_when_post_rebuild_vacuum_fails.yaml Normal file
View file

File diff suppressed because one or more lines are too long

42
tests/cassettes/test_rebuild/test_rechunk_raises_when_docling_blob_is_missing.yaml Normal file
View file

File diff suppressed because one or more lines are too long

27
tests/conftest.py
View file

@ -2,6 +2,8 @@ import json
import logging
import os
import tempfile
from collections.abc import Iterator
from contextlib import contextmanager
from pathlib import Path
from typing import TYPE_CHECKING, Any
@ -32,6 +34,31 @@ setattr(pydantic_ai.models, "ALLOW_MODEL_REQUESTS", False)
logging.getLogger("vcr.cassette").setLevel(logging.WARNING)
@contextmanager
def capture_logs(
logger: logging.Logger, level: int
) -> Iterator[list[logging.LogRecord]]:
"""Collect records emitted by ``logger`` at or above ``level``.
Attaches directly to the given logger instead of using ``caplog``:
``haiku.rag.logging.get_logger()`` sets ``propagate=False`` on the
``haiku.rag`` logger, so records never reach caplog's root handler once
any test in the session has called it.
"""
records: list[logging.LogRecord] = []
class _ListHandler(logging.Handler):
def emit(self, record: logging.LogRecord) -> None:
records.append(record)
handler = _ListHandler(level=level)
logger.addHandler(handler)
try:
yield records
finally:
logger.removeHandler(handler)
@pytest.fixture(scope="session")
def qa_corpus() -> list[dict[str, str]]:
corpus_path = Path(__file__).parent / "data" / "qa_corpus.json"

43
tests/ingester/test_fs_source.py
View file

@ -330,3 +330,46 @@ async def test_fs_source_fetch_reads_off_event_loop_thread(fs_root: Path):
"FSSource._read_body ran on the event-loop thread; the read+hash must "
"be dispatched via asyncio.to_thread"
)
@pytest.mark.asyncio
async def test_fetch_rejects_foreign_scheme(tmp_path):
"""`supports()` short-circuits on scheme, but fetch/head resolve directly,
so the unsupported-scheme path must be handled there too."""
src = FSSource(root=tmp_path, supported_extensions=[".md"], source_id="local")
# A same-named file under the root exists, so a scheme-blind implementation
# would happily resolve it — None/raise here really is the scheme check.
(tmp_path / "key.md").write_text("local copy")
assert await src.head((tmp_path / "key.md").as_uri()) is not None
with pytest.raises(UnsupportedSourceError):
await src.fetch("s3://bucket/key.md")
assert await src.head("s3://bucket/key.md") is None
@pytest.mark.asyncio
async def test_fetch_falls_back_to_octet_stream_for_unknown_extension(tmp_path):
target = tmp_path / "data.unknownext"
target.write_bytes(b"payload")
src = FSSource(
root=tmp_path, supported_extensions=[".unknownext"], source_id="local"
)
result = await src.fetch(target.as_uri())
assert result.content_type == "application/octet-stream"
assert result.body == b"payload"
@pytest.mark.asyncio
async def test_discover_skips_symlink_to_missing_in_root_target(tmp_path):
"""A broken symlink inside the root resolves to a path that is not a file."""
(tmp_path / "real.md").write_text("real")
(tmp_path / "broken.md").symlink_to(tmp_path / "absent.md")
src = FSSource(root=tmp_path, supported_extensions=[".md"], source_id="local")
events = [e async for e in src.discover()]
assert {e.uri for e in events} == {(tmp_path / "real.md").as_uri()}

11
tests/ingester/test_http_source.py
View file

@ -422,3 +422,14 @@ async def test_fetch_skips_head_when_no_max_size():
)
await src.fetch("https://example.com/a.md")
assert calls == ["GET"]
@pytest.mark.asyncio
async def test_aclose_closes_the_http_client():
src = HTTPSource(
source_id="urls",
urls=[],
transport=httpx.MockTransport(lambda r: httpx.Response(200)),
)
await src.aclose()
assert src._http.is_closed

156
tests/ingester/test_pollers.py
View file

@ -602,3 +602,159 @@ async def test_fs_poller_enqueues_initial_files(tmp_path, jobs, sync):
queued = await jobs.list_jobs(source_id="local")
assert {Path(j.uri).name for j in queued} == {"a.md", "b.md"}
assert all(j.status is JobStatus.QUEUED for j in queued)
# --- _dry_run_once ---
@pytest.mark.asyncio
async def test_dry_run_collects_changes_without_writing(fs_config, jobs, sync):
source = _StubSource(
"src",
[
[
_event("file:///a.md"),
_event("file:///b.md", kind=SourceEventKind.UNCHANGED),
_event("file:///c.md", kind=SourceEventKind.DELETE),
]
],
)
poller = _periodic(source, fs_config, jobs, sync)
ok, summary, changes = await poller._dry_run_once()
assert ok is True
assert summary.upsert_count == 1
assert summary.unchanged_count == 1
assert summary.delete_count == 1
assert {c.op for c in changes} == {JobOp.UPSERT, JobOp.DELETE}
# A dry run must not touch the queue.
assert await jobs.list_jobs(source_id="src") == []
@pytest.mark.asyncio
async def test_dry_run_ignores_deletes_when_delete_orphans_false(
fs_config, jobs, sync, tmp_path
):
config = FSSourceConfig(
type="fs",
id="src",
root=tmp_path,
delete_orphans=False,
poll_interval_s=0.05,
)
source = _StubSource("src", [[_event("file:///c.md", kind=SourceEventKind.DELETE)]])
poller = _periodic(source, config, jobs, sync)
ok, summary, changes = await poller._dry_run_once()
assert ok is True
assert summary.delete_count == 0
assert summary.ignored_delete_count == 1
assert changes == []
@pytest.mark.asyncio
async def test_dry_run_skipped_when_circuit_open(fs_config, jobs, sync):
class _Clock:
now = 0.0
def __call__(self):
return self.now
breaker = CircuitBreaker(
CircuitBreakerConfig(failure_threshold=1, cooldown_s=30.0),
now_fn=_Clock(),
)
source = _StubSource("src", [])
source.fail_with = RuntimeError("upstream down")
poller = _periodic(source, fs_config, jobs, sync, breaker=breaker)
assert await poller._sweep_once() is False
assert breaker.is_open is True
before = source.discover_calls
ok, summary, changes = await poller._dry_run_once()
assert ok is False
assert changes == []
assert source.discover_calls == before
assert poller.last_skip_reason == "circuit_open"
@pytest.mark.asyncio
async def test_dry_run_records_failure_when_discover_raises(fs_config, jobs, sync):
source = _StubSource("src", [])
source.fail_with = RuntimeError("upstream down")
poller = _periodic(source, fs_config, jobs, sync)
ok, summary, changes = await poller._dry_run_once()
assert ok is False
assert changes == []
assert poller._breaker.consecutive_failures == 1
@pytest.mark.asyncio
async def test_dry_run_skipped_when_queue_has_pending_work(fs_config, jobs, sync):
source = _StubSource("src", [[_event("file:///a.md")]])
poller = _periodic(source, fs_config, jobs, sync)
await jobs.enqueue("src", "file:///pending.md", JobOp.UPSERT)
ok, _summary, changes = await poller._dry_run_once()
assert ok is False
assert changes == []
assert poller.last_skip_reason == "pending_work"
@pytest.mark.asyncio
async def test_watch_deleted_skipped_when_delete_orphans_false(tmp_path, jobs, sync):
from watchfiles import Change
from haiku.rag.ingester.pollers.fs import FSPoller
from haiku.rag.ingester.sources.fs import FSSource
cfg = FSSourceConfig(
type="fs",
id="local",
root=tmp_path,
delete_orphans=False,
poll_interval_s=60.0,
)
poller = FSPoller(
source=FSSource(root=tmp_path, supported_extensions=[".md"], source_id="local"),
config=cfg,
job_repo=jobs,
sync_repo=sync,
)
await poller._handle_watch_change(Change.deleted, tmp_path / "gone.md")
assert await jobs.list_jobs(source_id="local") == []
@pytest.mark.asyncio
async def test_dry_run_manifest_reports_failed_sources(tmp_path, jobs, sync):
"""A source whose discover() raises is named in the failed list while the
manifest still carries the sources that succeeded."""
manager = PollerManager(
configs=[FSSourceConfig(type="fs", id="ok", root=tmp_path)],
job_repo=jobs,
sync_repo=sync,
)
broken = _StubSource("broken", [])
broken.fail_with = RuntimeError("upstream down")
manager._pollers.append(
_periodic(
broken,
FSSourceConfig(type="fs", id="broken", root=tmp_path, poll_interval_s=60.0),
jobs,
sync,
)
)
manifest, failed = await manager.dry_run_manifest()
assert failed == ["broken"]
assert {s.source_id for s in manifest.sources} == {"ok", "broken"}

168
tests/ingester/test_webdav_source.py
View file

@ -7,26 +7,29 @@ from haiku.rag.ingester.sources.base import FileTooLargeError, SourceEventKind
from haiku.rag.ingester.sources.webdav import WebDAVSource, _strip_etag
def test_strip_etag_strong_quoted():
assert _strip_etag('"abc123"') == "abc123"
def test_strip_etag_weak_marker():
assert _strip_etag('W/"abc123"') == "abc123"
def test_strip_etag_unquoted():
assert _strip_etag("abc123") == "abc123"
def test_strip_etag_whitespace():
assert _strip_etag(' W/"abc" ') == "abc"
def test_strip_etag_empty_returns_none():
assert _strip_etag("") is None
assert _strip_etag('""') is None
assert _strip_etag(None) is None
@pytest.mark.parametrize(
"raw,expected",
[
('"abc123"', "abc123"),
('W/"abc123"', "abc123"),
("abc123", "abc123"),
(' W/"abc" ', "abc"),
("", None),
('""', None),
(None, None),
],
ids=[
"strong_quoted",
"weak_marker",
"unquoted",
"whitespace",
"empty",
"empty_quotes",
"none",
],
)
def test_strip_etag(raw, expected):
assert _strip_etag(raw) == expected
def _transport(handler) -> httpx.MockTransport:
@ -628,3 +631,128 @@ async def test_fetch_skips_head_when_no_max_size():
)
await src.fetch("https://nc.example.com/dav/a.txt")
assert calls == ["GET"]
# Malformed multistatus bodies: a <response> that can't be decoded is dropped
# rather than aborting the whole listing.
def _raw_multistatus(*response_blocks: str) -> bytes:
body = ['<?xml version="1.0" encoding="utf-8"?>', '<d:multistatus xmlns:d="DAV:">']
body.extend(response_blocks)
body.append("</d:multistatus>")
return "\n".join(body).encode()
_NO_HREF = """ <d:response>
<d:propstat>
<d:status>HTTP/1.1 200 OK</d:status>
<d:prop><d:getetag>"r"</d:getetag></d:prop>
</d:propstat>
</d:response>"""
_EMPTY_HREF = """ <d:response>
<d:href></d:href>
<d:propstat>
<d:status>HTTP/1.1 200 OK</d:status>
<d:prop><d:getetag>"r"</d:getetag></d:prop>
</d:propstat>
</d:response>"""
_NO_STATUS = """ <d:response>
<d:href>/dav/a.md</d:href>
<d:propstat>
<d:prop><d:getetag>"r"</d:getetag></d:prop>
</d:propstat>
</d:response>"""
_NOT_FOUND_STATUS = """ <d:response>
<d:href>/dav/a.md</d:href>
<d:propstat>
<d:status>HTTP/1.1 404 Not Found</d:status>
<d:prop><d:getetag>"r"</d:getetag></d:prop>
</d:propstat>
</d:response>"""
_STATUS_WITHOUT_PROP = """ <d:response>
<d:href>/dav/a.md</d:href>
<d:propstat>
<d:status>HTTP/1.1 200 OK</d:status>
</d:propstat>
</d:response>"""
@pytest.mark.asyncio
@pytest.mark.parametrize(
"block",
[_NO_HREF, _EMPTY_HREF, _NO_STATUS, _NOT_FOUND_STATUS],
ids=["no_href", "empty_href", "propstat_without_status", "propstat_404"],
)
async def test_head_returns_none_for_undecodable_response(block):
def handler(request: httpx.Request) -> httpx.Response:
return httpx.Response(207, content=_raw_multistatus(block))
src = WebDAVSource(
source_id="nc",
base_url="https://nc.example.com/dav/",
transport=_transport(handler),
)
assert await src.head("https://nc.example.com/dav/a.md") is None
@pytest.mark.asyncio
async def test_head_returns_none_for_empty_multistatus():
def handler(request: httpx.Request) -> httpx.Response:
return httpx.Response(207, content=_raw_multistatus())
src = WebDAVSource(
source_id="nc",
base_url="https://nc.example.com/dav/",
transport=_transport(handler),
)
assert await src.head("https://nc.example.com/dav/a.md") is None
def test_entry_with_status_but_no_prop_has_no_revision():
"""A 200 propstat carrying no <prop> still yields an entry, without a
revision — distinct from the malformed bodies that yield no entry at all."""
from haiku.rag.ingester.sources.webdav import _parse_multistatus
entries = _parse_multistatus(_raw_multistatus(_STATUS_WITHOUT_PROP))
assert len(entries) == 1
assert entries[0].revision is None
assert _parse_multistatus(_raw_multistatus(_NO_HREF)) == []
@pytest.mark.asyncio
async def test_discover_skips_base_url_reported_as_file():
"""Broken servers list the base URL itself as a non-collection; it and any
href outside the base are skipped."""
body = _multistatus(
{"href": "/dav/", "etag": '"base"'},
{"href": "/outside/x.md", "etag": '"out"'},
{"href": "/dav/keep.md", "etag": '"keep"'},
)
def handler(request: httpx.Request) -> httpx.Response:
return httpx.Response(207, content=body)
src = WebDAVSource(
source_id="nc",
base_url="https://nc.example.com/dav/",
transport=_transport(handler),
)
events = [event async for event in src.discover()]
assert {e.uri for e in events} == {"https://nc.example.com/dav/keep.md"}
@pytest.mark.asyncio
async def test_aclose_closes_the_http_client():
src = WebDAVSource(
source_id="nc",
base_url="https://nc.example.com/dav/",
transport=_transport(lambda r: httpx.Response(200)),
)
await src.aclose()
assert src._http.is_closed

69
tests/sandbox/test_sandbox_toc.py
View file

@ -364,3 +364,72 @@ class TestItemsJsonlSurfacesNewFields:
assert expected <= set(r)
assert "position" not in r
assert "tree_depth" not in r
@pytest.mark.asyncio
class TestVfsReadPaths:
"""The synchronous VFS readers bridge back to the event loop; drive them
through a worker thread the way execute() does."""
async def test_content_txt_is_read_lazily_per_document(self, temp_db_path):
async with HaikuRAG(temp_db_path, create=True) as client:
doc = await client.document_repository.create(
Document(content="the stored body", uri="test://body", title="Body")
)
doc_id = doc.id
sandbox = Sandbox(temp_db_path, AppConfig(), AnalysisContext())
# Build the VFS first, then change the stored content. A lazy
# CallbackFile reads through at access time and sees the new body;
# an eager MemoryFile mount would have captured the old one.
vfs = await sandbox._build_vfs()
sandbox._loop = asyncio.get_running_loop()
# Rewrite via the repository rather than client.update_document: the
# latter re-chunks and re-embeds, which this file deliberately avoids
# so these tests need no embedding endpoint.
async with HaikuRAG(temp_db_path, create=False) as client:
doc.content = "the rewritten body"
await client.document_repository.update(doc)
content = await asyncio.to_thread(
vfs.path_read_text, PurePosixPath(f"/documents/{doc_id}/content.txt")
)
assert content == "the rewritten body"
async def test_document_files_are_read_only(self, temp_db_path):
async with HaikuRAG(temp_db_path, create=True) as client:
doc = await client.document_repository.create(
Document(content="x", uri="test://ro", title="RO")
)
doc_id = doc.id
sandbox = Sandbox(temp_db_path, AppConfig(), AnalysisContext())
vfs = await sandbox._build_vfs()
sandbox._loop = asyncio.get_running_loop()
with pytest.raises(PermissionError, match="read-only"):
await asyncio.to_thread(
vfs.path_write_text,
PurePosixPath(f"/documents/{doc_id}/content.txt"),
"nope",
)
async def test_toc_skips_gaps_in_item_positions(self, temp_db_path):
"""Positions need not be contiguous — a heading's span may cover
positions that carry no item."""
async with HaikuRAG(temp_db_path, create=True) as client:
doc_id = await _empty_doc(client, uri="test://gaps", title="Gaps")
# Positions 1 and 2 are absent between the header and the paragraph.
items = [
_header(doc_id, 0, 1, "Intro"),
_para(doc_id, 3),
]
await client.document_item_repository.create_items(doc_id, items)
sandbox = Sandbox(temp_db_path, AppConfig(), AnalysisContext())
toc = await _read_toc(sandbox, doc_id)
assert [n["title"] for n in toc["tree"]] == ["Intro"]

49
tests/store/test_document_items.py
View file

@ -986,3 +986,52 @@ class TestPictureDataPreservedThroughRoundTrip:
after = await rag.document_item_repository.get_all_picture_data(created.id)
assert after.get("#/pictures/0") == original.get("#/pictures/0")
@pytest.mark.asyncio
async def test_replace_for_document_with_no_items_deletes_existing(temp_db_path):
"""Replacing with an empty list clears the document's items."""
from haiku.rag.store.engine import Store
from haiku.rag.store.repositories.document_item import DocumentItemRepository
async with Store(temp_db_path, create=True) as store:
repo = DocumentItemRepository(store)
docling_doc = _make_docling_doc()
await repo.replace_for_document("doc-1", extract_items("doc-1", docling_doc))
assert await repo.get_all_items("doc-1")
await repo.replace_for_document("doc-1", [])
assert await repo.get_all_items("doc-1") == []
class TestExtractItemTextFallbacks:
def test_table_returns_none_when_serialization_fails(self):
"""A serializer that raises leaves the table with no extractable text
rather than aborting the extraction pass."""
doc = _doc_with_tables(1)
class _Boom:
def serialize(self, item):
raise RuntimeError("serializer exploded")
assert extract_item_text(doc.tables[0], doc, get_serializer=_Boom) is None
def test_file_backed_picture_has_no_inline_bytes(self):
"""A picture whose ImageRef points at a file rather than a data: URI
carries nothing to decode."""
from docling_core.types.doc.document import ImageRef
from haiku.rag.store.models.document_item import _decode_picture_bytes
doc, pic = _doc_with_captioned_picture("caption")
pic.image = ImageRef.model_validate(
{
"mimetype": "image/png",
"dpi": 72,
"size": {"width": 1, "height": 1},
"uri": "file:///tmp/picture.png",
}
)
assert _decode_picture_bytes(pic) is None

27
tests/store/test_restore.py
View file

@ -112,8 +112,10 @@ async def test_restore_safety_tag_name_collision(temp_db_path, monkeypatch):
await store.create_tag("release-1")
await store.create_tag("before-restore-20260715T143012Z")
await store.create_tag("before-restore-20260715T143012Z-2")
safety_tag = await store.restore_tag("release-1")
assert safety_tag == "before-restore-20260715T143012Z-2"
assert safety_tag == "before-restore-20260715T143012Z-3"
@pytest.mark.asyncio
@ -406,3 +408,26 @@ async def test_wait_protected_returns_result_on_same_tick_cancellation():
result, cancelled = await outer
assert result == "done"
assert cancelled is True
@pytest.mark.asyncio
async def test_wait_protected_reraises_when_recovery_itself_is_cancelled():
"""If the recovery coroutine ends cancelled there is nothing to wait for,
so the cancellation propagates instead of looping forever."""
import asyncio
from haiku.rag.store.engine import _wait_protected
async def self_cancelling_recovery() -> str:
current = asyncio.current_task()
assert current is not None
current.cancel()
await asyncio.sleep(0)
return "unreachable"
outer = asyncio.create_task(_wait_protected(self_cancelling_recovery()))
# Bounded: without the re-raise the retry loop spins forever on an
# already-cancelled task, and the timeout turns that into a clean failure.
with pytest.raises(asyncio.CancelledError):
await asyncio.wait_for(outer, timeout=5)

17
tests/store/test_v0_50_0_migration.py
View file

@ -126,6 +126,19 @@ class TestV0_50_0Migration:
}
),
),
# Matches the LIKE on the quoted-key form, but only nested —
# there is no top-level key to rename.
LegacyDocumentRecord(
id="nested-only",
content="x",
uri="u3",
metadata=json.dumps(
{
"raw_headers": {"etag": "abc"},
"source_revision": "v3",
}
),
),
],
)
@ -141,6 +154,10 @@ class TestV0_50_0Migration:
"my_etag_key": "v",
"source_revision": "v2",
}
assert by_id["nested-only"] == {
"raw_headers": {"etag": "abc"},
"source_revision": "v3",
}
async def test_unparseable_metadata_skipped_without_crashing(self, temp_db_path):
"""A row with malformed JSON in `metadata` must not abort the whole

511
tests/test_chunk.py
View file

@ -3,6 +3,7 @@ import pytest
from haiku.rag.client import HaikuRAG
from haiku.rag.config import Config
from haiku.rag.store.models.chunk import Chunk, ChunkMetadata, SearchResult
from tests.conftest import capture_logs
@pytest.mark.vcr()
@ -123,136 +124,86 @@ async def test_chunking_pipeline(qa_corpus: list[dict[str, str]], temp_db_path):
assert chunk.order == i
def test_chunk_metadata_parsing():
@pytest.mark.parametrize(
"metadata,refs,headings,labels,page_numbers",
[
(
{
"doc_item_refs": ["#/texts/0", "#/texts/1", "#/tables/0"],
"headings": ["Chapter 1", "Section 1.1"],
"labels": ["paragraph", "paragraph", "table"],
"page_numbers": [1, 1, 2],
},
["#/texts/0", "#/texts/1", "#/tables/0"],
["Chapter 1", "Section 1.1"],
["paragraph", "paragraph", "table"],
[1, 1, 2],
),
({}, [], None, [], []),
],
ids=["populated", "defaults"],
)
def test_chunk_metadata_parsing(metadata, refs, headings, labels, page_numbers):
"""Test ChunkMetadata parsing from chunk metadata dict."""
metadata_dict = {
"doc_item_refs": ["#/texts/0", "#/texts/1", "#/tables/0"],
"headings": ["Chapter 1", "Section 1.1"],
"labels": ["paragraph", "paragraph", "table"],
"page_numbers": [1, 1, 2],
}
chunk = Chunk(
content="Test content",
metadata=metadata_dict,
)
chunk = Chunk(content="Test content", metadata=metadata)
chunk_meta = chunk.get_chunk_metadata()
assert isinstance(chunk_meta, ChunkMetadata)
assert chunk_meta.doc_item_refs == ["#/texts/0", "#/texts/1", "#/tables/0"]
assert chunk_meta.headings == ["Chapter 1", "Section 1.1"]
assert chunk_meta.labels == ["paragraph", "paragraph", "table"]
assert chunk_meta.page_numbers == [1, 1, 2]
assert chunk_meta.doc_item_refs == refs
assert chunk_meta.headings == headings
assert chunk_meta.labels == labels
assert chunk_meta.page_numbers == page_numbers
def test_chunk_metadata_defaults():
"""Test ChunkMetadata with empty/default values."""
chunk = Chunk(content="Test content", metadata={})
chunk_meta = chunk.get_chunk_metadata()
@pytest.fixture
def two_text_docling_doc():
"""Minimal DoclingDocument with two resolvable text items."""
from docling_core.types.doc.document import DoclingDocument
assert chunk_meta.doc_item_refs == []
assert chunk_meta.headings is None
assert chunk_meta.labels == []
assert chunk_meta.page_numbers == []
return DoclingDocument.model_validate(
{
"name": "test_doc",
"texts": [
{
"self_ref": "#/texts/0",
"text": "First text",
"orig": "First text",
"label": "paragraph",
},
{
"self_ref": "#/texts/1",
"text": "Second text",
"orig": "Second text",
"label": "title",
},
],
"tables": [],
"pictures": [],
"groups": [],
"body": {"self_ref": "#/body", "children": []},
"furniture": {"self_ref": "#/furniture", "children": []},
}
)
def test_chunk_metadata_resolve_doc_items():
@pytest.mark.parametrize(
"refs,expected_texts",
[
(["#/texts/0", "#/texts/1"], ["First text", "Second text"]),
# Out-of-range and malformed refs are skipped rather than raising.
(["#/texts/0", "#/texts/999", "#/invalid/path"], ["First text"]),
([], []),
],
ids=["all_valid", "graceful_degradation", "empty_refs"],
)
def test_chunk_metadata_resolve_doc_items(two_text_docling_doc, refs, expected_texts):
"""Test resolving doc_item_refs to actual DocItem objects."""
from docling_core.types.doc.document import DoclingDocument
chunk_meta = ChunkMetadata(doc_item_refs=refs)
# Create a minimal DoclingDocument with some text items
doc_json = {
"name": "test_doc",
"texts": [
{
"self_ref": "#/texts/0",
"text": "First text",
"orig": "First text",
"label": "paragraph",
},
{
"self_ref": "#/texts/1",
"text": "Second text",
"orig": "Second text",
"label": "title",
},
],
"tables": [],
"pictures": [],
"groups": [],
"body": {"self_ref": "#/body", "children": []},
"furniture": {"self_ref": "#/furniture", "children": []},
}
docling_doc = DoclingDocument.model_validate(doc_json)
doc_items = chunk_meta.resolve_doc_items(two_text_docling_doc)
# Create chunk metadata with refs
chunk_meta = ChunkMetadata(
doc_item_refs=["#/texts/0", "#/texts/1"],
labels=["paragraph", "title"],
)
# Resolve refs
doc_items = chunk_meta.resolve_doc_items(docling_doc)
assert len(doc_items) == 2
assert getattr(doc_items[0], "text") == "First text"
assert getattr(doc_items[1], "text") == "Second text"
def test_chunk_metadata_resolve_doc_items_graceful_degradation():
"""Test that invalid refs are skipped gracefully."""
from docling_core.types.doc.document import DoclingDocument
doc_json = {
"name": "test_doc",
"texts": [
{
"self_ref": "#/texts/0",
"text": "Only text",
"orig": "Only text",
"label": "paragraph",
},
],
"tables": [],
"pictures": [],
"groups": [],
"body": {"self_ref": "#/body", "children": []},
"furniture": {"self_ref": "#/furniture", "children": []},
}
docling_doc = DoclingDocument.model_validate(doc_json)
# Create chunk metadata with one valid and one invalid ref
chunk_meta = ChunkMetadata(
doc_item_refs=["#/texts/0", "#/texts/999", "#/invalid/path"],
)
# Resolve refs - invalid ones should be skipped
doc_items = chunk_meta.resolve_doc_items(docling_doc)
assert len(doc_items) == 1
assert getattr(doc_items[0], "text") == "Only text"
def test_chunk_metadata_resolve_empty_refs():
"""Test resolving with no refs returns empty list."""
from docling_core.types.doc.document import DoclingDocument
doc_json = {
"name": "test_doc",
"texts": [],
"tables": [],
"pictures": [],
"groups": [],
"body": {"self_ref": "#/body", "children": []},
"furniture": {"self_ref": "#/furniture", "children": []},
}
docling_doc = DoclingDocument.model_validate(doc_json)
chunk_meta = ChunkMetadata()
doc_items = chunk_meta.resolve_doc_items(docling_doc)
assert doc_items == []
assert [getattr(item, "text") for item in doc_items] == expected_texts
def test_search_result_from_chunk_preserves_document_meta():
@ -286,11 +237,12 @@ def test_search_result_format_for_agent_omits_document_meta():
assert "https://example.org/report/view" not in formatted
def test_search_result_format_for_agent_with_rank():
"""Test format_for_agent with rank and total parameters."""
result = SearchResult(
@pytest.fixture
def rich_search_result():
"""SearchResult with every optional field populated."""
return SearchResult(
content="This is the chunk content about elections.",
score=0.02, # Low RRF score that would confuse agents
score=0.85,
chunk_id="chunk-123",
document_id="doc-456",
document_uri="file:///docs/report.pdf",
@ -300,16 +252,30 @@ def test_search_result_format_for_agent_with_rank():
page_numbers=[1, 2],
)
formatted = result.format_for_agent(rank=1, total=5)
@pytest.mark.parametrize(
"kwargs,present,absent",
[
# A rank is supplied, so the raw RRF score is withheld from the agent.
({"rank": 1, "total": 5}, "[rank 1 of 5]", "score:"),
({}, "(score: 0.85)", "[rank"),
],
ids=["with_rank", "score_fallback"],
)
def test_search_result_format_for_agent_rank_vs_score(
rich_search_result, kwargs, present, absent
):
"""format_for_agent shows a rank when given one, else falls back to score."""
formatted = rich_search_result.format_for_agent(**kwargs)
assert present in formatted
assert absent not in formatted
assert "[chunk-123]" in formatted
assert "[rank 1 of 5]" in formatted
assert "score:" not in formatted # Score should NOT appear when rank is provided
assert (
'Source: "Annual Report 2024" > Chapter 1 > Section 1.1 > Elections'
in formatted
)
assert "Type: table" in formatted
assert "Type: table" in formatted # table has higher priority than paragraph
assert "Content:\nThis is the chunk content about elections." in formatted
@ -357,134 +323,104 @@ def test_search_result_format_for_agent_no_captions_no_line():
assert "Figure caption" not in formatted
def test_search_result_format_for_agent_rank_only():
"""Test format_for_agent with rank but no total."""
result = SearchResult(
content="Some content.",
score=0.03,
chunk_id="chunk-abc",
)
formatted = result.format_for_agent(rank=2)
assert "[chunk-abc]" in formatted
assert "[rank 2]" in formatted
assert "score:" not in formatted
def test_search_result_format_for_agent_fallback():
"""Test format_for_agent falls back to score when no rank provided."""
result = SearchResult(
content="This is the chunk content about elections.",
score=0.85,
chunk_id="chunk-123",
document_id="doc-456",
document_uri="file:///docs/report.pdf",
document_title="Annual Report 2024",
headings=["Chapter 1", "Section 1.1", "Elections"],
labels=["paragraph", "table"],
page_numbers=[1, 2],
)
formatted = result.format_for_agent()
assert "[chunk-123]" in formatted
assert "(score: 0.85)" in formatted
assert (
'Source: "Annual Report 2024" > Chapter 1 > Section 1.1 > Elections'
in formatted
)
assert "Type: table" in formatted # table has higher priority than paragraph
assert "Content:\nThis is the chunk content about elections." in formatted
def test_search_result_format_for_agent_minimal():
"""Test format_for_agent with minimal metadata."""
@pytest.mark.parametrize(
"kwargs,present,absent",
[
({"rank": 2}, "[rank 2]", ["score:"]),
# No structural metadata at all, so no Source:/Type: lines are emitted.
({}, "(score: 0.72)", ["[rank", "Source:", "Type:"]),
],
ids=["rank_only", "minimal"],
)
def test_search_result_format_for_agent_minimal(kwargs, present, absent):
"""A result carrying only content/score/chunk_id formats without metadata lines."""
result = SearchResult(
content="Some content here.",
score=0.72,
chunk_id="chunk-abc",
)
formatted = result.format_for_agent()
formatted = result.format_for_agent(**kwargs)
assert "[chunk-abc]" in formatted
assert "(score: 0.72)" in formatted
assert "Source:" not in formatted # No title or headings
assert "Type:" not in formatted # No labels
assert present in formatted
for token in absent:
assert token not in formatted
assert "Content:\nSome content here." in formatted
def test_search_result_format_for_agent_title_only():
"""Test format_for_agent with only document title."""
@pytest.mark.parametrize(
"fields,expected_source",
[
({"document_title": "My Document"}, 'Source: "My Document"'),
(
{"headings": ["Introduction", "Background"]},
"Source: Introduction > Background",
),
],
ids=["title_only", "headings_only"],
)
def test_search_result_format_for_agent_source_line(fields, expected_source):
"""The Source: line is built from the title, the headings, or both."""
result = SearchResult(
content="Content text.",
score=0.60,
chunk_id="chunk-xyz",
document_title="My Document",
**fields,
)
formatted = result.format_for_agent()
assert 'Source: "My Document"' in formatted
assert expected_source in result.format_for_agent()
def test_search_result_format_for_agent_headings_only():
"""Test format_for_agent with only headings (no title)."""
result = SearchResult(
content="Content text.",
score=0.60,
chunk_id="chunk-xyz",
headings=["Introduction", "Background"],
)
formatted = result.format_for_agent()
assert "Source: Introduction > Background" in formatted
def test_search_result_get_primary_label():
@pytest.mark.parametrize(
"labels,expected",
[
(["paragraph", "table", "text"], "table"),
(["paragraph", "code"], "code"),
(["list_item", "code"], "code"),
(["text", "list_item"], "list_item"),
# No structural label: falls through to the first label.
(["paragraph", "text"], "paragraph"),
([], None),
],
)
def test_search_result_get_primary_label(labels, expected):
"""Test _get_primary_label prioritization."""
# Table takes priority over text labels
result = SearchResult(content="x", score=0.5, labels=["paragraph", "table", "text"])
assert result._get_primary_label() == "table"
# Code takes priority over list_item
result = SearchResult(content="x", score=0.5, labels=["list_item", "code"])
assert result._get_primary_label() == "code"
# Text labels fall through to first
result = SearchResult(content="x", score=0.5, labels=["paragraph", "text"])
assert result._get_primary_label() == "paragraph"
# Empty labels
result = SearchResult(content="x", score=0.5, labels=[])
assert result._get_primary_label() is None
result = SearchResult(content="x", score=0.5, labels=labels)
assert result._get_primary_label() == expected
@pytest.mark.vcr()
async def test_chunk_content_fts_populated(temp_db_path):
"""Test that content_fts column is populated with contextualized content."""
@pytest.mark.parametrize(
"metadata,content,expected_content_fts",
[
(
{"headings": ["Chapter 1", "Section 1.1"]},
"This is the raw chunk content.",
"Chapter 1\nSection 1.1\nThis is the raw chunk content.",
),
({}, "Plain content without headings.", "Plain content without headings."),
],
ids=["populated", "without_headings"],
)
async def test_chunk_content_fts(temp_db_path, metadata, content, expected_content_fts):
"""content_fts holds the contextualized content while content stays raw."""
from haiku.rag.embeddings import get_embedder
async with HaikuRAG(db_path=temp_db_path, config=Config, create=True) as client:
# Create a chunk with headings
chunk = Chunk(
document_id="test-doc",
content="This is the raw chunk content.",
metadata={"headings": ["Chapter 1", "Section 1.1"]},
content=content,
metadata=metadata,
order=0,
)
# Generate embedding
embedder = get_embedder(Config)
embedding = (await embedder.embed_documents([chunk.content]))[0]
chunk.embedding = embedding
# Store the chunk
await client.chunk_repository.create(chunk)
# Read the raw record from the database
records = (
await client.store.chunks_table.query()
.where(f"id = '{chunk.id}'")
@ -495,52 +431,8 @@ async def test_chunk_content_fts_populated(temp_db_path):
assert len(records) == 1
record = records[0]
# Verify content is raw (no headings)
assert record["content"] == "This is the raw chunk content."
# Verify content_fts is contextualized (headings + content)
assert (
record["content_fts"]
== "Chapter 1\nSection 1.1\nThis is the raw chunk content."
)
@pytest.mark.vcr()
async def test_chunk_content_fts_without_headings(temp_db_path):
"""Test that content_fts equals content when no headings present."""
from haiku.rag.embeddings import get_embedder
async with HaikuRAG(db_path=temp_db_path, config=Config, create=True) as client:
# Create a chunk without headings
chunk = Chunk(
document_id="test-doc",
content="Plain content without headings.",
metadata={},
order=0,
)
# Generate embedding
embedder = get_embedder(Config)
embedding = (await embedder.embed_documents([chunk.content]))[0]
chunk.embedding = embedding
# Store the chunk
await client.chunk_repository.create(chunk)
# Read the raw record from the database
records = (
await client.store.chunks_table.query()
.where(f"id = '{chunk.id}'")
.limit(1)
.to_arrow()
).to_pylist()
assert len(records) == 1
record = records[0]
# Both should be the same when no headings
assert record["content"] == "Plain content without headings."
assert record["content_fts"] == "Plain content without headings."
assert record["content"] == content
assert record["content_fts"] == expected_content_fts
async def test_ensure_fts_index_warns_on_failure(temp_db_path):
@ -557,18 +449,91 @@ async def test_ensure_fts_index_warns_on_failure(temp_db_path):
repo.store.chunks_table.create_index = _boom
records: list[logging.LogRecord] = []
class _Capture(logging.Handler):
def emit(self, record: logging.LogRecord) -> None:
records.append(record)
handler = _Capture(level=logging.WARNING)
chunk_module.logger.addHandler(handler)
try:
with capture_logs(chunk_module.logger, logging.WARNING) as records:
await repo._ensure_fts_index()
finally:
chunk_module.logger.removeHandler(handler)
assert [r for r in records if r.levelno == logging.WARNING]
assert any("index build failed" in r.getMessage() for r in records)
@pytest.mark.vcr()
async def test_chunk_repository_get_by_id_and_list_all_pagination(
qa_corpus: list[dict[str, str]], temp_db_path
):
"""get_by_id resolves a stored chunk; list_all honours limit and offset."""
async with HaikuRAG(db_path=temp_db_path, config=Config, create=True) as client:
# A corpus document is long enough to chunk more than once, which is
# what makes the offset assertion below meaningful.
doc = await client.create_document(content=qa_corpus[0]["document_extracted"])
assert doc.id is not None
stored = await client.chunk_repository.get_by_document_id(doc.id)
assert stored
fetched = await client.get_chunk_by_id(stored[0].id)
assert fetched is not None
assert fetched.id == stored[0].id
assert fetched.content == stored[0].content
assert await client.get_chunk_by_id("no-such-chunk") is None
everything = await client.chunk_repository.list_all()
assert len(everything) == len(stored)
first = await client.chunk_repository.list_all(limit=1)
assert len(first) == 1
assert first[0].id == everything[0].id
# Fail loudly if the fixture stops producing enough chunks to page.
assert len(everything) >= 2
second = await client.chunk_repository.list_all(limit=1, offset=1)
assert len(second) == 1
assert second[0].id == everything[1].id
@pytest.mark.vcr()
async def test_chunk_search_returns_empty_for_blank_query(temp_db_path):
"""A blank query with no precomputed vector short-circuits before searching."""
async with HaikuRAG(db_path=temp_db_path, config=Config, create=True) as client:
await client.create_document(content="Searchable body about elections.")
# Positive control: the corpus is non-empty, so [] is a real decision
# rather than the answer to every query.
assert await client.chunk_repository.search("elections")
assert await client.chunk_repository.search(" ") == []
@pytest.mark.vcr()
async def test_chunk_search_with_precomputed_vector_skips_text_query(temp_db_path):
"""The image-as-query path searches vector-only using a stored embedding."""
async with HaikuRAG(db_path=temp_db_path, config=Config, create=True) as client:
doc = await client.create_document(content="Vector-only search target.")
assert doc.id is not None
rows = (await client.store.chunks_table.query().limit(1).to_arrow()).to_pylist()
stored_vector = list(rows[0]["vector"])
results = await client.chunk_repository.search("", query_vector=stored_vector)
assert results
assert any(c.document_id == doc.id for c, _ in results)
async def test_get_chunk_ids_by_self_ref_grouped_without_documents(temp_db_path):
async with HaikuRAG(db_path=temp_db_path, config=Config, create=True) as client:
assert await client.chunk_repository.get_chunk_ids_by_self_ref_grouped([]) == {}
async def test_process_search_results_rejects_unknown_score_column(temp_db_path):
"""A result frame with no recognised score column is a programming error."""
import pandas as pd
async with HaikuRAG(db_path=temp_db_path, config=Config, create=True) as client:
class _Frame:
async def to_pandas(self):
return pd.DataFrame([{"id": "c1", "content": "x", "metadata": "{}"}])
with pytest.raises(ValueError, match="Unknown search result format"):
await client.chunk_repository._process_search_results(_Frame())

78
tests/test_chunker.py
View file

@ -798,3 +798,81 @@ async def test_serve_chunker_accepts_picture_laden_docling(doclaynet_first_page_
chunks = await serve_chunker.chunk(doc)
assert len(chunks) > 0, "docling-serve chunker returned 0 chunks"
class TestDoclingServeChunkerRefResolution:
"""_resolve_label_from_document and the dict-shaped doc_items branch."""
@pytest.fixture
def chunker(self):
config = AppConfig()
config.providers.docling_serve.base_url = "http://localhost:5001"
config.processing.chunk_size = 256
config.processing.chunking_tokenizer = "Qwen/Qwen3-Embedding-0.6B"
return DoclingServeChunker(config)
@pytest.fixture
def document(self):
from docling_core.types.doc.document import DoclingDocument
return DoclingDocument.model_validate(
{
"name": "doc",
"texts": [
{
"self_ref": "#/texts/0",
"text": "body",
"orig": "body",
"label": "paragraph",
}
],
"tables": [],
"pictures": [],
"groups": [],
"body": {"self_ref": "#/body", "children": []},
"furniture": {"self_ref": "#/furniture", "children": []},
}
)
@pytest.mark.parametrize(
"ref",
["not-a-ref", "#/texts/999", "#/nope/0"],
ids=["unparseable", "index_out_of_range", "unknown_collection"],
)
def test_unresolvable_ref_yields_no_label(self, document, ref):
from haiku.rag.chunkers.docling_serve import _resolve_label_from_document
assert _resolve_label_from_document(ref, document) is None
def test_resolvable_ref_yields_label(self, document):
from haiku.rag.chunkers.docling_serve import _resolve_label_from_document
assert _resolve_label_from_document("#/texts/0", document) == "paragraph"
@pytest.mark.asyncio
async def test_chunk_of_none_returns_empty(self, chunker):
assert await chunker.chunk(None) == []
@pytest.mark.asyncio
async def test_dict_shaped_doc_items_are_decoded(self, chunker, document):
"""docling-serve returns refs as strings today; the dict shape is
accepted in case the API changes."""
async def fake_chunk_api(_document):
return [
{
"raw_text": "body",
# A label the document does NOT carry, so the assertion
# proves the dict's own label was used rather than a
# lookup against the document.
"doc_items": [{"self_ref": "#/texts/0", "label": "caption"}],
}
]
chunker._call_chunk_api = fake_chunk_api # type: ignore[method-assign]
chunks = await chunker.chunk(document)
assert len(chunks) == 1
assert chunks[0].metadata["doc_item_refs"] == ["#/texts/0"]
assert chunks[0].metadata["labels"] == ["caption"]

367
tests/test_client.py
View file

@ -1,3 +1,4 @@
import asyncio
import json
import tempfile
import threading
@ -14,6 +15,7 @@ from haiku.rag.client.documents import (
DocumentImport,
_prepare_document_from_docling,
_write_fetch_body,
check_source_accessible,
)
from haiku.rag.config import Config
from haiku.rag.store.compression import decompress_json
@ -298,23 +300,6 @@ async def test_client_create_document_from_source(temp_db_path):
assert "md5" in doc2.metadata
@pytest.mark.vcr()
async def test_client_create_document_from_source_with_title(temp_db_path):
"""Test creating a document from a file source with a title."""
async with HaikuRAG(temp_db_path, create=True) as client:
with tempfile.TemporaryDirectory() as temp_dir:
test_content = "This is test content from a file."
temp_path = Path(temp_dir) / "test_title.txt"
temp_path.write_text(test_content)
doc = await client.create_document_from_source(
source=temp_path, title="My Doc"
)
assert isinstance(doc, Document)
assert doc.id is not None
assert doc.title == "My Doc"
@pytest.mark.vcr()
async def test_client_update_title_noop_behavior(temp_db_path):
"""When content is unchanged, updating title should update document without re-chunking."""
@ -326,6 +311,7 @@ async def test_client_update_title_noop_behavior(temp_db_path):
doc1 = await client.create_document_from_source(temp_path, title="Title A")
assert isinstance(doc1, Document)
assert doc1.id is not None
assert doc1.title == "Title A"
# Re-add with same content but new title
doc2 = await client.create_document_from_source(temp_path, title="Title B")
@ -646,12 +632,14 @@ async def test_client_create_update_no_op_behavior(temp_db_path):
assert doc1.id is not None
assert doc1.content == test_content
original_id = doc1.id
original_updated_at = doc1.updated_at
# Second call with same content - should return existing document (no-op)
doc2 = await client.create_document_from_source(temp_path)
assert isinstance(doc2, Document)
assert doc2.id == original_id # Same document
assert doc2.content == test_content
assert doc2.updated_at == original_updated_at # No-op leaves it untouched
# Modify file content
updated_content = "Updated content for testing."
@ -669,28 +657,6 @@ async def test_client_create_update_no_op_behavior(temp_db_path):
assert retrieved_doc.content == updated_content
@pytest.mark.vcr()
async def test_client_unchanged_file_keeps_timestamp(temp_db_path):
"""Test that unchanged files don't update the updated_at timestamp."""
async with HaikuRAG(temp_db_path, create=True) as client:
# Create a temporary file
test_content = "Test content for timestamp check."
with tempfile.TemporaryDirectory() as temp_dir:
temp_path = Path(temp_dir) / "test.txt"
temp_path.write_text(test_content)
# First call - create document
doc1 = await client.create_document_from_source(temp_path)
assert isinstance(doc1, Document)
original_updated_at = doc1.updated_at
# Second call with same content - should not update timestamp
doc2 = await client.create_document_from_source(temp_path)
assert isinstance(doc2, Document)
assert doc2.id == doc1.id
assert doc2.updated_at == original_updated_at # Timestamp should not change
@pytest.mark.vcr()
async def test_client_url_create_update_no_op_behavior(temp_db_path):
"""Test create/update/no-op behavior for URLs based on MD5 changes."""
@ -2290,3 +2256,326 @@ async def test_metadata_only_update_waits_for_write_lock(temp_db_path):
assert not task.done()
updated = await task
assert updated.metadata == {"k": "v"}
@pytest.mark.parametrize(
"uri,expected",
[
("https://example.com/doc.pdf", True),
("s3://bucket/key", True),
("mem://not-a-source", False),
# urlparse rejects a malformed IPv6 host; a stored URI that no longer
# parses must not abort the caller's rebuild sweep.
("http://[::1", False),
],
ids=["https", "s3", "unknown_scheme", "unparseable"],
)
def test_check_source_accessible(uri, expected):
assert check_source_accessible(uri) is expected
def test_check_source_accessible_file_uri(tmp_path):
existing = tmp_path / "there.txt"
existing.write_text("x")
assert check_source_accessible(existing.as_uri()) is True
assert check_source_accessible((tmp_path / "gone.txt").as_uri()) is False
def _bbox_doc(*, with_page_image: bool, pages: tuple[int, ...] = (1,)):
"""DoclingDocument with one paragraph per page, each carrying a bbox.
``with_page_image=False`` produces pages with no raster, so bounding boxes
resolve but there is nothing to draw them on.
"""
from docling_core.types.doc.base import BoundingBox, Size
from docling_core.types.doc.document import ImageRef, ProvenanceItem
from PIL import Image as PilImageModule
doc = DoclingDocument(name="bbox-doc")
size = Size(width=612.0, height=792.0)
for page_no in pages:
image = (
ImageRef.from_pil(
PilImageModule.new("RGB", (612, 792), color="white"), dpi=72
)
if with_page_image
else None
)
doc.add_page(page_no=page_no, size=size, image=image)
doc.add_text(
label=DocItemLabel.PARAGRAPH,
text=f"Content on page {page_no}.",
prov=ProvenanceItem(
page_no=page_no,
bbox=BoundingBox(l=50, t=700, r=550, b=650),
charspan=(0, 20),
),
)
return doc
@pytest.mark.asyncio
async def test_visualize_chunk_returns_empty_without_page_rasters(temp_db_path):
"""Boxes resolve, but a document ingested without page images has nothing
to render them onto."""
docling_doc = _bbox_doc(with_page_image=False)
chunks = [
Chunk(
content="Content on page 1.",
metadata={
"doc_item_refs": ["#/texts/0"],
"page_numbers": [1],
"labels": ["paragraph"],
},
order=0,
embedding=[0.1] * 2560,
)
]
async with HaikuRAG(temp_db_path, create=True) as client:
doc = await client.import_document(docling_doc, chunks, uri="test://no-raster")
stored = await client.chunk_repository.get_by_document_id(doc.id)
assert await client.visualize_chunk(stored[0]) == []
@pytest.mark.asyncio
async def test_visualize_chunk_skips_pages_without_a_raster(temp_db_path):
"""A document where only some pages carry a raster renders just those."""
from docling_core.types.doc.base import BoundingBox, Size
from docling_core.types.doc.document import ImageRef, ProvenanceItem
from PIL import Image as PilImageModule
docling_doc = DoclingDocument(name="mixed-rasters")
size = Size(width=612.0, height=792.0)
docling_doc.add_page(
page_no=1,
size=size,
image=ImageRef.from_pil(
PilImageModule.new("RGB", (612, 792), color="white"), dpi=72
),
)
docling_doc.add_page(page_no=2, size=size, image=None)
for page_no in (1, 2):
docling_doc.add_text(
label=DocItemLabel.PARAGRAPH,
text=f"Content on page {page_no}.",
prov=ProvenanceItem(
page_no=page_no,
bbox=BoundingBox(l=50, t=700, r=550, b=650),
charspan=(0, 20),
),
)
chunks = [
Chunk(
content="Content on page 1.\nContent on page 2.",
metadata={
"doc_item_refs": ["#/texts/0", "#/texts/1"],
"page_numbers": [1, 2],
"labels": ["paragraph", "paragraph"],
},
order=0,
embedding=[0.1] * 2560,
)
]
async with HaikuRAG(temp_db_path, create=True) as client:
doc = await client.import_document(
docling_doc, chunks, uri="test://mixed-rasters"
)
stored = await client.chunk_repository.get_by_document_id(doc.id)
images = await client.visualize_chunk(stored[0])
assert len(images) == 1
@pytest.mark.asyncio
async def test_visualize_chunk_returns_empty_when_pages_row_missing(temp_db_path):
docling_doc = _bbox_doc(with_page_image=True)
chunks = [
Chunk(
content="Content on page 1.",
metadata={
"doc_item_refs": ["#/texts/0"],
"page_numbers": [1],
"labels": ["paragraph"],
},
order=0,
embedding=[0.1] * 2560,
)
]
async with HaikuRAG(temp_db_path, create=True) as client:
doc = await client.import_document(docling_doc, chunks, uri="test://no-row")
stored = await client.chunk_repository.get_by_document_id(doc.id)
async def no_pages_row(document_id):
return None
client.document_repository.get_pages_data = no_pages_row # type: ignore[method-assign]
assert await client.visualize_chunk(stored[0]) == []
@pytest.mark.asyncio
async def test_visualize_chunk_skips_box_on_unstored_page(temp_db_path):
"""A bounding box referencing a page the document never registered is
skipped rather than raising."""
from docling_core.types.doc.base import BoundingBox, Size
from docling_core.types.doc.document import ImageRef, ProvenanceItem
from PIL import Image as PilImageModule
docling_doc = DoclingDocument(name="orphan-page-box")
docling_doc.add_page(
page_no=1,
size=Size(width=612.0, height=792.0),
image=ImageRef.from_pil(
PilImageModule.new("RGB", (612, 792), color="white"), dpi=72
),
)
docling_doc.add_text(
label=DocItemLabel.PARAGRAPH,
text="Content attributed to a page with no raster.",
prov=ProvenanceItem(
page_no=3,
bbox=BoundingBox(l=50, t=700, r=550, b=650),
charspan=(0, 20),
),
)
chunks = [
Chunk(
content="Content attributed to a page with no raster.",
metadata={
"doc_item_refs": ["#/texts/0"],
"page_numbers": [3],
"labels": ["paragraph"],
},
order=0,
embedding=[0.1] * 2560,
)
]
async with HaikuRAG(temp_db_path, create=True) as client:
doc = await client.import_document(
docling_doc, chunks, uri="test://orphan-page"
)
stored = await client.chunk_repository.get_by_document_id(doc.id)
assert await client.visualize_chunk(stored[0]) == []
@pytest.mark.asyncio
async def test_visualize_chunk_without_refs_falls_back_to_chunk_metadata(temp_db_path):
"""A chunk carrying no doc_item_refs has nothing to expand from."""
docling_doc = _bbox_doc(with_page_image=True)
chunks = [
Chunk(
content="Content on page 1.",
metadata={"page_numbers": [1], "labels": ["paragraph"]},
order=0,
embedding=[0.1] * 2560,
)
]
async with HaikuRAG(temp_db_path, create=True) as client:
doc = await client.import_document(docling_doc, chunks, uri="test://no-refs")
stored = await client.chunk_repository.get_by_document_id(doc.id)
assert await client.visualize_chunk(stored[0]) == []
@pytest.mark.asyncio
async def test_visualize_chunk_falls_back_when_expansion_drops_refs(temp_db_path):
"""If expansion returns results carrying no refs, the original search
results' refs are used instead."""
from haiku.rag.client import search as search_module
docling_doc = _bbox_doc(with_page_image=True)
chunks = [
Chunk(
content="Content on page 1.",
metadata={
"doc_item_refs": ["#/texts/0"],
"page_numbers": [1],
"labels": ["paragraph"],
},
order=0,
embedding=[0.1] * 2560,
)
]
async with HaikuRAG(temp_db_path, create=True) as client:
doc = await client.import_document(docling_doc, chunks, uri="test://drops-refs")
stored = await client.chunk_repository.get_by_document_id(doc.id)
async def expansion_without_refs(_client, results):
return [r.model_copy(update={"doc_item_refs": []}) for r in results]
with patch.object(search_module, "expand_context", expansion_without_refs):
images = await client.visualize_chunk(stored[0])
assert len(images) == 1
@pytest.mark.vcr()
@pytest.mark.parametrize("auto_vacuum", [True, False])
async def test_import_documents_schedules_vacuum_per_config(temp_db_path, auto_vacuum):
"""A batch import runs a background vacuum only when auto_vacuum is on."""
from haiku.rag.config import AppConfig
config = AppConfig()
config.storage.auto_vacuum = auto_vacuum
async with HaikuRAG(temp_db_path, config=config, create=True) as client:
docling_doc = await client.convert("Batch imported body.")
chunks = await client.chunk(docling_doc)
# Spy rather than inspecting _vacuum_tasks: the scheduling code
# discards each task on completion, so the set races to empty. The
# spy's count does not race, and draining the scheduled task keeps
# the assertion deterministic without pulling in the close-time pass.
with patch.object(client.store, "vacuum", new=AsyncMock()) as vacuum:
await client.import_documents(
[
DocumentImport(
docling_document=docling_doc,
chunks=chunks,
uri="test://batch-vacuum",
)
]
)
await asyncio.gather(*client._vacuum_tasks)
assert vacuum.await_count == (1 if auto_vacuum else 0)
@pytest.mark.vcr()
async def test_reingesting_a_source_applies_an_explicit_title(temp_db_path):
"""Re-adding a changed source with a title updates both, in place."""
async with HaikuRAG(temp_db_path, create=True) as client:
with tempfile.TemporaryDirectory() as temp_dir:
source = Path(temp_dir) / "retitled.txt"
source.write_text("stable content")
first = await client.create_document_from_source(source)
assert not isinstance(first, list)
source.write_text("changed content")
second = await client.create_document_from_source(
source, title="Explicit Title"
)
assert not isinstance(second, list)
assert second.id == first.id
assert second.title == "Explicit Title"
assert second.content == "changed content"
async def test_update_document_rejects_unknown_id(temp_db_path):
async with HaikuRAG(temp_db_path, create=True) as client:
with pytest.raises(ValueError, match="not found"):
await client.update_document("no-such-document", content="x")

57
tests/test_config.py
View file

@ -42,17 +42,17 @@ def test_find_config_file_cwd(tmp_path, monkeypatch):
def test_find_config_file_user_config(tmp_path, monkeypatch):
"""Test finding config in user config directory."""
monkeypatch.delenv("HAIKU_RAG_CONFIG_PATH", raising=False)
monkeypatch.chdir(tmp_path)
# Mock get_default_data_dir to return tmp_path
def mock_get_default_data_dir():
return tmp_path
# The data dir must differ from the cwd, or the cwd branch answers first
# and this never reaches the user-directory lookup.
cwd = tmp_path / "cwd"
cwd.mkdir()
data_dir = tmp_path / "data"
data_dir.mkdir()
monkeypatch.chdir(cwd)
monkeypatch.setattr("haiku.rag.utils.get_default_data_dir", lambda: data_dir)
monkeypatch.setattr(
"haiku.rag.utils.get_default_data_dir", mock_get_default_data_dir
)
config_file = tmp_path / "haiku.rag.yaml"
config_file = data_dir / "haiku.rag.yaml"
config_file.write_text("environment: production")
found = find_config_file()
@ -518,3 +518,42 @@ def test_expand_env_var_plain_string_unchanged(tmp_path):
config = load_yaml_config(config_file)
assert config["environment"] == "production"
def test_find_config_file_returns_none_when_nothing_exists(tmp_path, monkeypatch):
"""No env var, no file in cwd, none in the data dir."""
monkeypatch.delenv("HAIKU_RAG_CONFIG_PATH", raising=False)
monkeypatch.chdir(tmp_path)
empty_data_dir = tmp_path / "data"
empty_data_dir.mkdir()
monkeypatch.setattr("haiku.rag.utils.get_default_data_dir", lambda: empty_data_dir)
assert find_config_file() is None
def test_load_default_config_falls_back_to_builtin_defaults(monkeypatch):
"""With no config file discoverable, the packaged defaults are used."""
from haiku.rag.config import _load_default_config
monkeypatch.setattr("haiku.rag.config.find_config_file", lambda _=None: None)
config = _load_default_config()
assert config.model_dump() == AppConfig().model_dump()
def test_get_config_initialises_lazily_then_reuses(monkeypatch):
"""get_config() builds the instance on first use and caches it.
Asserted explicitly rather than relying on some test happening to be the
first caller in its worker process: under xdist that depends on how cases
shard across workers, which varies with the core count.
"""
from haiku.rag import config as config_module
monkeypatch.setattr(config_module, "_config", None)
first = config_module.get_config()
assert isinstance(first, AppConfig)
assert config_module.get_config() is first

144
tests/test_context.py
View file

@ -1360,3 +1360,147 @@ class TestExpandWithItemsPictureBytes:
e_low = by_chunk["c-low"]
assert e_low.image_data == {"#/pictures/0": "LOWBYTES"}
assert "HIGHBYTES" not in (e_low.image_data or {}).values()
class TestSpanInWindow:
def test_zero_width_span_is_inside_when_position_is_in_window(self):
from haiku.rag.context import _span_in_window
from haiku.rag.store.models.document_item import DocumentItem
item = DocumentItem(
document_id="d1", position=0, self_ref="#/pictures/0", label="picture"
)
# A picture occupies no characters, so containment is by position.
assert _span_in_window((10, 10, item), 0, 20) is True
assert _span_in_window((30, 30, item), 0, 20) is False
@pytest.mark.asyncio
class TestExpandWithItemsWindowEdges:
async def test_empty_window_returns_original_results(
self, temp_db_path, monkeypatch
):
"""Refs resolve but the surrounding window comes back empty, so there is
nothing to expand from."""
from haiku.rag.client import HaikuRAG
from haiku.rag.store.models.document import Document
async with HaikuRAG(temp_db_path, create=True) as rag:
doc = await rag.document_repository.create(
Document(content="body", uri="test://window")
)
assert doc.id is not None
await rag.document_item_repository.create_items(
doc.id,
[
DocumentItem(
document_id=doc.id,
position=0,
self_ref="#/texts/0",
label="paragraph",
text="body",
page_numbers=[1],
)
],
)
async def no_window(*_args, **_kwargs):
return []
monkeypatch.setattr(
rag.document_item_repository, "get_items_in_range", no_window
)
result = SearchResult(
content="original",
score=0.9,
document_id=doc.id,
doc_item_refs=["#/texts/0"],
)
expanded = await expand_with_items(
rag.document_item_repository, doc.id, [result], 5000
)
assert [r.content for r in expanded] == ["original"]
async def test_result_with_unmatched_refs_passes_through(self, temp_db_path):
"""Two results share a document; the one whose refs resolve is expanded
and the other is returned unchanged."""
from haiku.rag.client import HaikuRAG
from haiku.rag.store.models.document import Document
async with HaikuRAG(temp_db_path, create=True) as rag:
doc = await rag.document_repository.create(
Document(content="body", uri="test://mixed")
)
assert doc.id is not None
await rag.document_item_repository.create_items(
doc.id,
[
DocumentItem(
document_id=doc.id,
position=i,
self_ref=f"#/texts/{i}",
label="paragraph",
text=f"paragraph {i}",
page_numbers=[1],
)
for i in range(2)
],
)
resolvable = SearchResult(
content="paragraph 0",
score=0.9,
document_id=doc.id,
doc_item_refs=["#/texts/0"],
)
unmatched = SearchResult(
content="untouched",
score=0.5,
document_id=doc.id,
doc_item_refs=["#/texts/404"],
)
expanded = await expand_with_items(
rag.document_item_repository, doc.id, [resolvable, unmatched], 5000
)
assert len(expanded) == 2
by_content = {r.content for r in expanded}
# The unmatched result is passed through byte-for-byte...
assert "untouched" in by_content
# ...while the resolvable one actually grew to its neighbours.
grew = next(c for c in by_content if c != "untouched")
assert "paragraph 0" in grew and "paragraph 1" in grew
def test_build_result_skips_positions_with_no_item():
"""A sparse position map (items removed or never stored) leaves gaps in the
range; those positions contribute nothing."""
from haiku.rag.context import _build_result
original = SearchResult(content="p0", score=0.9, document_id="d1")
# Positions 1 and 2 in the 0..3 range carry no item.
pos_to_item = {
0: DocumentItem(
document_id="d1",
position=0,
self_ref="#/texts/0",
label="paragraph",
text="first",
page_numbers=[1],
),
3: DocumentItem(
document_id="d1",
position=3,
self_ref="#/texts/3",
label="paragraph",
text="last",
page_numbers=[1],
),
}
built = _build_result(0, 3, [original], pos_to_item, False, 5000)
assert built.content == "first\n\nlast"

173
tests/test_converters.py
View file

@ -405,6 +405,46 @@ class TestDoclingLocalConverter:
assert isinstance(doc, DoclingDocument)
assert doc.name == "test"
@pytest.mark.asyncio
async def test_convert_file_reads_unknown_extension_as_text(
self, converter, tmp_path
):
"""An extension in neither the docling nor the text set is read as text."""
source = tmp_path / "notes.xyz"
source.write_text("Plain body for an unknown extension.")
doc = await converter.convert_file(source)
assert isinstance(doc, DoclingDocument)
assert "Plain body for an unknown extension." in doc.export_to_markdown()
@pytest.mark.asyncio
async def test_convert_file_raises_for_undecodable_file(self, converter, tmp_path):
source = tmp_path / "binary.xyz"
source.write_bytes(b"\xff\xfe\x00\x01 not utf-8")
with pytest.raises(ValueError, match="Failed to parse file"):
await converter.convert_file(source)
@pytest.mark.asyncio
async def test_convert_text_wraps_conversion_failure(self, converter, monkeypatch):
def boom(*_args, **_kwargs):
raise RuntimeError("docling exploded")
monkeypatch.setattr(converter, "_sync_convert_docling_text", boom)
with pytest.raises(ValueError, match="Failed to convert text"):
await converter.convert_text("# Test", name="test.md")
@pytest.mark.asyncio
async def test_convert_text_falls_back_when_format_not_inferable(self, converter):
"""docling raises ConversionError for an extension it has no backend
for; the simple-document fallback keeps the text."""
doc = await converter.convert_text("just some prose", name="mystery.zzz")
assert isinstance(doc, DoclingDocument)
assert "just some prose" in doc.export_to_markdown()
@pytest.mark.asyncio
async def test_convert_code_file(self, converter):
"""Test that code files are wrapped in code blocks."""
@ -1522,3 +1562,136 @@ class TestDoclingServeConverterIntegration:
assert str(sample.image.uri).startswith("data:image/"), (
"Rehydrated picture URI should be a data: URI, not a bare artifact filename"
)
class TestDoclingServeZipParsing:
"""_parse_zip_to_docling decodes the target_type=zip payload. These drive
its branches directly — no docling-serve instance involved."""
@pytest.fixture
def converter(self):
config = AppConfig()
config.processing.converter = "docling-serve"
conv = get_converter(config)
assert isinstance(conv, DoclingServeConverter)
return conv
@staticmethod
def _zip(entries: dict[str, bytes]) -> bytes:
import io
import zipfile
buf = io.BytesIO()
with zipfile.ZipFile(buf, mode="w") as zf:
for name, blob in entries.items():
zf.writestr(name, blob)
return buf.getvalue()
@staticmethod
def _doc_json(**extra) -> dict:
base = {
"name": "document",
"texts": [],
"tables": [],
"pictures": [],
"groups": [],
"body": {"self_ref": "#/body", "children": []},
"furniture": {"self_ref": "#/furniture", "children": []},
}
base.update(extra)
return base
def test_raises_without_top_level_json(self, converter):
blob = self._zip({"artifacts/image.png": b"png"})
with pytest.raises(ValueError, match="no top-level JSON document"):
converter._parse_zip_to_docling(blob, "doc.pdf")
def test_picture_without_image_is_left_alone(self, converter):
import json as _json
doc_json = self._doc_json(
pictures=[
{
"self_ref": "#/pictures/0",
"label": "picture",
"image": None,
"prov": [],
}
]
)
blob = self._zip({"document.json": _json.dumps(doc_json).encode()})
doc = converter._parse_zip_to_docling(blob, "doc.pdf")
assert doc.pictures[0].image is None
def test_data_uri_image_is_passed_through(self, converter):
import json as _json
data_uri = "data:image/png;base64,aGVsbG8="
doc_json = self._doc_json(
pictures=[
{
"self_ref": "#/pictures/0",
"label": "picture",
"image": {
"mimetype": "image/png",
"dpi": 72,
"size": {"width": 1, "height": 1},
"uri": data_uri,
},
"prov": [],
}
]
)
blob = self._zip({"document.json": _json.dumps(doc_json).encode()})
doc = converter._parse_zip_to_docling(blob, "doc.pdf")
assert str(doc.pictures[0].image.uri) == data_uri
def test_non_dict_page_entry_is_skipped_while_inlining(self, converter):
"""A page entry that isn't an object must not blow up the image-inlining
loop with an AttributeError; it falls through to schema validation."""
import json as _json
from pydantic import ValidationError
doc_json = self._doc_json(pages={"1": "not-a-page-object"})
blob = self._zip({"document.json": _json.dumps(doc_json).encode()})
with pytest.raises(ValidationError):
converter._parse_zip_to_docling(blob, "doc.pdf")
@pytest.mark.asyncio
async def test_convert_text_rejects_unsupported_format(self, converter):
with pytest.raises(ValueError, match="Unsupported format"):
await converter.convert_text("body", format="pdf")
@pytest.mark.asyncio
async def test_convert_text_plain_builds_document_locally(self, converter):
"""format="plain" never reaches the network."""
converter.client.submit_and_poll_zip = AsyncMock(
side_effect=AssertionError("must not call docling-serve")
)
doc = await converter.convert_text("just text", format="plain")
assert isinstance(doc, DoclingDocument)
assert "just text" in doc.export_to_markdown()
@pytest.mark.asyncio
async def test_docling_serve_convert_file_wraps_text_read_failure(tmp_path):
"""An undecodable text file surfaces as a ValueError naming the path."""
config = AppConfig()
config.processing.converter = "docling-serve"
converter = get_converter(config)
assert isinstance(converter, DoclingServeConverter)
source = tmp_path / "broken.txt"
source.write_bytes(b"\xff\xfe\x00\x01 not utf-8")
with pytest.raises(ValueError, match="Failed to read text file"):
await converter.convert_file(source)

33
tests/test_database_autocreate.py
View file

@ -43,3 +43,36 @@ async def test_operations_work_after_database_created(tmp_path):
doc = await client.get_document_by_id(docs[0].id)
assert doc is not None
assert doc.content == "Test content"
def test_default_db_path_comes_from_storage_data_dir(tmp_path):
"""Omitting db_path places the database under the configured data dir."""
from haiku.rag.client import HaikuRAG
from haiku.rag.config import AppConfig
config = AppConfig()
config.storage.data_dir = tmp_path
client = HaikuRAG(config=config)
assert client._db_path == tmp_path / "haiku.rag.lancedb"
@pytest.mark.asyncio
async def test_vacuum_optimizes_tables_without_losing_rows(temp_db_path):
"""The public vacuum() runs the store's optimize pass over real rows."""
from haiku.rag.store.models.document import Document
from haiku.rag.store.repositories.document import DocumentRepository
async with HaikuRAG(temp_db_path, create=True) as client:
repo = DocumentRepository(client.store)
for i in range(3):
await repo.create(Document(content=f"body {i}", uri=f"test://doc{i}"))
before = len(await client.store.documents_table.list_versions())
await client.vacuum()
# Optimize compacts the per-document fragments into new versions; an
# unchanged count would mean nothing reached the tables.
assert len(await client.store.documents_table.list_versions()) > before
assert await client.count_documents() == 3

26
tests/test_docling_serve_client.py
View file

@ -481,3 +481,29 @@ def test_from_config_wires_retry_and_breaker():
assert client._max_attempts == 7
assert client._breaker_config.failure_threshold == 9
assert client._breaker_config.cooldown_s == 90.0
@pytest.mark.asyncio
async def test_submit_without_task_id_raises():
"""A 200 that carries no task_id is a protocol violation, not a silent pass."""
import httpx
from haiku.rag.providers.docling_serve import DoclingServeClient
def handler(request: httpx.Request) -> httpx.Response:
return httpx.Response(200, json={})
client = DoclingServeClient(base_urls="http://docling:5001")
files = {"files": ("doc.pdf", b"pdf", "application/octet-stream")}
async with httpx.AsyncClient(transport=httpx.MockTransport(handler)) as http:
with pytest.raises(ValueError, match="did not return a task_id"):
await client._submit_and_wait(
http,
"http://docling:5001",
"/v1/convert/source/async",
files,
{},
{},
"doc.pdf",
)

23
tests/test_doctor.py
View file

@ -1206,3 +1206,26 @@ async def test_duplicate_documents_check_reads_config(temp_db_path):
).severity
is Severity.WARN
)
@pytest.mark.asyncio
async def test_many_unembedded_chunks_are_sampled(temp_db_path):
"""Beyond the sample limit the detail list ends with a count of the rest."""
db = await _build_db(temp_db_path)
chunks_tbl = await db.open_table("chunks")
await chunks_tbl.add(
[
ChunkRecord(
id=f"z{i}",
document_id="d1",
content="x",
metadata=json.dumps({"doc_item_refs": ["#/texts/0"]}),
vector=[0.0] * VECTOR_DIM,
)
for i in range(8)
]
)
report = await run_doctor(_config(), temp_db_path, {})
details = _result(report, "unembedded_chunks").details
assert len(details) == 6
assert details[-1] == "... (+3 more)"

69
tests/test_document.py
View file

@ -8,49 +8,33 @@ from haiku.rag.store.repositories.document_item import DocumentItemRepository
@pytest.mark.asyncio
async def test_document_list_excludes_content_by_default(
qa_corpus: list[dict[str, str]], temp_db_path
@pytest.mark.parametrize("include_content", [False, True])
async def test_document_list_all(
qa_corpus: list[dict[str, str]], temp_db_path, include_content
):
"""list_all excludes content and docling_document by default."""
"""list_all excludes content and docling_document unless include_content=True."""
async with Store(temp_db_path, create=True) as store:
doc_repo = DocumentRepository(store)
content = qa_corpus[0]["document_extracted"]
doc = Document(
content=qa_corpus[0]["document_extracted"],
content=content,
uri="https://example.com/doc.txt",
title="Test Document",
metadata={"key": "value"},
)
created = await doc_repo.create(doc)
docs = await doc_repo.list_all()
docs = await doc_repo.list_all(include_content=include_content)
assert len(docs) == 1
assert docs[0].id == created.id
assert docs[0].title == "Test Document"
assert docs[0].uri == "https://example.com/doc.txt"
assert docs[0].metadata == {"key": "value"}
assert docs[0].content == ""
assert docs[0].content == (content if include_content else "")
assert docs[0].docling_document is None
@pytest.mark.asyncio
async def test_document_list_includes_content_when_requested(
qa_corpus: list[dict[str, str]], temp_db_path
):
"""list_all returns content when include_content=True."""
async with Store(temp_db_path, create=True) as store:
doc_repo = DocumentRepository(store)
content = qa_corpus[0]["document_extracted"]
doc = Document(content=content, uri="https://example.com/doc.txt")
created = await doc_repo.create(doc)
docs = await doc_repo.list_all(include_content=True)
assert len(docs) == 1
assert docs[0].id == created.id
assert docs[0].content == content
@pytest.mark.asyncio
async def test_document_list_with_filter(qa_corpus: list[dict[str, str]], temp_db_path):
"""Test listing documents with filter clause."""
@ -391,16 +375,26 @@ async def test_get_docling_data_loads_only_docling_columns(
@pytest.mark.asyncio
@pytest.mark.parametrize(
"with_pages",
# Markdown documents have no page images, so their pages blob stays None.
[True, False],
ids=["with_pages", "markdown"],
)
async def test_get_pages_data_loads_only_pages_column(
qa_corpus: list[dict[str, str]], temp_db_path
qa_corpus: list[dict[str, str]], temp_db_path, with_pages
):
"""get_pages_data returns only page image data for a document."""
import json
from haiku.rag.store.compression import compress_json
pages_blob = compress_json(
json.dumps({"1": {"size": {"width": 612, "height": 792}, "page_no": 1}})
pages_blob = (
compress_json(
json.dumps({"1": {"size": {"width": 612, "height": 792}, "page_no": 1}})
)
if with_pages
else None
)
async with Store(temp_db_path, create=True) as store:
@ -424,27 +418,6 @@ async def test_get_pages_data_loads_only_pages_column(
assert await doc_repo.get_pages_data("nonexistent-id") is None
@pytest.mark.asyncio
async def test_get_pages_data_none_for_markdown_document(
qa_corpus: list[dict[str, str]], temp_db_path
):
"""Markdown documents have no page images — get_pages_data returns None pages."""
async with Store(temp_db_path, create=True) as store:
doc_repo = DocumentRepository(store)
doc = Document(
content=qa_corpus[0]["document_extracted"],
uri="https://example.com/doc.md",
)
created = await doc_repo.create(doc)
assert created.id is not None
result = await doc_repo.get_pages_data(created.id)
assert result is not None
assert result.id == created.id
assert result.docling_pages is None
@pytest.mark.asyncio
async def test_document_get_by_uri_with_special_characters(
qa_corpus: list[dict[str, str]], temp_db_path

71
tests/test_download_models.py
View file

@ -6,6 +6,7 @@ import pytest
from haiku.rag.client.downloads import download_models
from haiku.rag.config import Config
from haiku.rag.config.models import ModelConfig
@pytest.fixture
@ -108,3 +109,73 @@ async def test_download_models_no_ollama_models(mock_to_thread):
models = {e.model for e in events}
assert "qwen3-embedding:4b" not in models
assert "gpt-oss" not in models
@pytest.mark.parametrize(
"configure,expected_model",
[
(
lambda c: setattr(
c.reranking,
"model",
ModelConfig(provider="ollama", name="rerank-model"),
),
"rerank-model",
),
(
lambda c: (
setattr(c.processing, "pictures", "description"),
setattr(
c.processing.conversion_options.picture_description.model,
"provider",
"ollama",
),
setattr(
c.processing.conversion_options.picture_description.model,
"name",
"vision-model",
),
),
"vision-model",
),
(
lambda c: (
setattr(c.processing, "auto_title", True),
setattr(c.processing.title_model, "provider", "ollama"),
setattr(c.processing.title_model, "name", "title-model"),
),
"title-model",
),
],
ids=["reranker", "picture_description", "auto_title"],
)
async def test_ollama_models_from_every_config_slot_are_pulled(
mock_to_thread, configure, expected_model
):
"""Each config slot that can name an ollama model contributes to the pull set."""
from haiku.rag.config import AppConfig
config = AppConfig()
configure(config)
@asynccontextmanager
async def mock_stream(method, url, **kwargs):
mock_resp = AsyncMock()
async def aiter_lines():
yield '{"status": "success"}'
mock_resp.aiter_lines = aiter_lines
yield mock_resp
with patch(
"haiku.rag.client.downloads.httpx.AsyncClient",
return_value=_mock_httpx_client(mock_stream),
):
pulled = {
progress.model
async for progress in download_models(config)
if progress.status == "pulling"
}
assert expected_model in pulled

47
tests/test_embedder_config.py
View file

@ -148,3 +148,50 @@ def test_vllm_embedder_does_not_double_append_v1():
base_url = embedder._base_url.rstrip("/") # type: ignore[attr-defined] # ty: ignore[unresolved-attribute]
assert base_url.endswith("/v1")
assert not base_url.endswith("/v1/v1")
def test_vector_dim_property_reports_configured_dimension():
from haiku.rag.embeddings import EmbedderWrapper
assert EmbedderWrapper(embedder=None, vector_dim=512).vector_dim == 512
@pytest.mark.parametrize(
"provider,env_var",
[("voyageai", "VOYAGE_API_KEY"), ("cohere", "CO_API_KEY")],
)
def test_saas_providers_are_wired_without_a_request(monkeypatch, provider, env_var):
"""Construction wires the SDK and reports the configured dimension."""
monkeypatch.setenv(env_var, "test-key")
config = AppConfig(
embeddings=EmbeddingsConfig(
model=EmbeddingModelConfig(
provider=provider, name="some-model", vector_dim=1024
),
),
)
embedder = get_embedder(config)
assert embedder.vector_dim == 1024
assert embedder.supports_images is False
# The provider and model reach the underlying pydantic-ai embedder.
assert embedder._embedder._model == f"{provider}:some-model" # ty: ignore[unresolved-attribute]
def test_cohere_floats_rejects_missing_embeddings():
from types import SimpleNamespace
from haiku.rag.embeddings.cohere import _floats
result = SimpleNamespace(embeddings=SimpleNamespace(float_=None))
with pytest.raises(ValueError, match="no float embeddings"):
_floats(result)
def test_voyageai_to_pil_rejects_unsupported_type():
from haiku.rag.embeddings.voyageai import _to_pil
with pytest.raises(TypeError, match="Unsupported image type"):
_to_pil("not an image") # ty: ignore[invalid-argument-type]

88
tests/test_lancedb_connection.py
View file

@ -310,3 +310,91 @@ class TestInitFailureCleanup:
pass
assert close_calls, "Store.close() was not called when _initialize raised"
class TestVectorIndexCreation:
"""_ensure_vector_index needs 256 rows of training data before it builds."""
@staticmethod
async def _seed_chunks(store, count: int) -> None:
import random
records = [
store.ChunkRecord(
document_id="doc-1",
content=f"row {i}",
content_fts=f"row {i}",
metadata="{}",
order=i,
vector=[random.random() for _ in range(store.embedder.vector_dim)],
)
for i in range(count)
]
await store.chunks_table.add(records)
@pytest.mark.asyncio
async def test_builds_index_once_enough_rows_exist(self, temp_db_path):
async with Store(temp_db_path, create=True) as store:
await self._seed_chunks(store, 256)
await store._ensure_vector_index()
indexes = await store.chunks_table.list_indices()
assert any("vector" in idx.columns for idx in indexes)
@pytest.mark.asyncio
async def test_index_failure_is_warned_not_raised(self, temp_db_path):
import logging
from haiku.rag.store import engine as engine_module
from tests.conftest import capture_logs
async with Store(temp_db_path, create=True) as store:
await self._seed_chunks(store, 256)
async def boom(*_args, **_kwargs):
raise RuntimeError("index build failed")
with patch.object(store.chunks_table, "create_index", boom):
with capture_logs(engine_module.logger, logging.WARNING) as records:
await store._ensure_vector_index()
assert any("index build failed" in r.getMessage() for r in records)
indexes = await store.chunks_table.list_indices()
assert not any("vector" in idx.columns for idx in indexes)
class TestStoreMiscellany:
@pytest.mark.asyncio
async def test_create_makes_missing_parent_directories(self, tmp_path):
nested = tmp_path / "a" / "b" / "db.lancedb"
async with Store(nested, create=True) as store:
assert store._is_new_db is True
assert nested.exists()
@pytest.mark.asyncio
async def test_stored_vector_dim_is_none_for_corrupt_settings(self, temp_db_path):
async with Store(temp_db_path, create=True) as store:
await store.settings_table.update(
{"settings": "not json at all"}, where="id = 'settings'"
)
assert await store._get_stored_vector_dim() is None
@pytest.mark.asyncio
async def test_vacuum_skips_when_already_running(self, temp_db_path):
import asyncio
async with Store(temp_db_path, create=True) as store:
async with store._vacuum_lock:
# Bounded: a regression here blocks on the held lock, and the
# timeout turns that deadlock into a clean failure.
await asyncio.wait_for(store.vacuum(), timeout=5)
@pytest.mark.asyncio
async def test_history_rejects_unknown_table(self, temp_db_path):
async with Store(temp_db_path, create=True) as store:
with pytest.raises(ValueError, match="Unknown table"):
await store.list_table_versions("not_a_table")

149
tests/test_mcp.py
View file

@ -332,3 +332,152 @@ class TestMCPImageInput:
result = await ask(question="q")
assert result == "answer"
assert captured["images"] is None
class TestMCPFileAndUrlIngestion:
@pytest.mark.asyncio
async def test_add_document_from_file(self, temp_db_path, tmp_path):
async with HaikuRAG(temp_db_path, create=True):
pass
source = tmp_path / "note.txt"
source.write_text("Ingested from a file path.")
mcp = create_mcp_server(temp_db_path, read_only=False)
add_file = await _get_tool(mcp, "add_document_from_file")
doc_id = await add_file(file_path=str(source), title="File Doc")
assert doc_id is not None
get_doc = await _get_tool(mcp, "get_document")
doc = await get_doc(document_id=doc_id)
assert doc.title == "File Doc"
@pytest.mark.asyncio
@pytest.mark.parametrize(
"tool_name,kwargs",
[
("add_document_from_file", {"file_path": "/tmp/x.txt"}),
("add_document_from_url", {"url": "https://example.com/x.txt"}),
],
)
@pytest.mark.parametrize(
"results,expected",
[
(
[Document(id="first", content="a"), Document(id="second", content="b")],
"first",
),
([], None),
],
ids=["directory_reports_first_id", "empty_directory_reports_none"],
)
async def test_add_tools_handle_multi_document_sources(
self, mcp_db, monkeypatch, tool_name, kwargs, results, expected
):
"""A source resolving to several documents reports the first id."""
async def fake_from_source(self, source, title=None, metadata=None, **kw):
return results
monkeypatch.setattr(HaikuRAG, "create_document_from_source", fake_from_source)
mcp = create_mcp_server(mcp_db, read_only=False)
add = await _get_tool(mcp, tool_name)
assert await add(**kwargs) == expected
@pytest.mark.asyncio
async def test_add_document_from_url(self, mcp_db, monkeypatch):
async def fake_from_source(self, source, title=None, metadata=None, **kwargs):
assert source == "https://example.com/doc.txt"
return Document(id="url-doc", content="fetched")
monkeypatch.setattr(HaikuRAG, "create_document_from_source", fake_from_source)
mcp = create_mcp_server(mcp_db, read_only=False)
add_url = await _get_tool(mcp, "add_document_from_url")
assert await add_url(url="https://example.com/doc.txt") == "url-doc"
class TestMCPToolsDegradeOnError:
"""Every tool swallows client failures and returns its empty value rather
than propagating an exception to the MCP transport."""
@pytest.mark.asyncio
@pytest.mark.parametrize(
"client_method,tool_name,kwargs,expected",
[
(
"create_document_from_source",
"add_document_from_file",
{"file_path": "/tmp/x.txt"},
None,
),
(
"create_document_from_source",
"add_document_from_url",
{"url": "https://example.com/x"},
None,
),
("create_document", "add_document_from_text", {"content": "x"}, None),
("delete_document", "delete_document", {"document_id": "x"}, False),
("search", "search_documents", {"query": "x"}, []),
("get_document_by_id", "get_document", {"document_id": "x"}, None),
("list_documents", "list_documents", {}, []),
],
)
async def test_tool_returns_empty_value_when_client_raises(
self, mcp_db, monkeypatch, client_method, tool_name, kwargs, expected
):
async def boom(self, *args, **kw):
raise RuntimeError("client exploded")
monkeypatch.setattr(HaikuRAG, client_method, boom)
mcp = create_mcp_server(mcp_db, read_only=False)
tool = await _get_tool(mcp, tool_name)
assert await tool(**kwargs) == expected
@pytest.mark.asyncio
async def test_list_documents_returns_empty_for_invalid_filter(self, mcp_db):
mcp = create_mcp_server(mcp_db, read_only=True)
list_docs = await _get_tool(mcp, "list_documents")
assert await list_docs(filter="no_such_column = 1") == []
@pytest.mark.asyncio
async def test_analyze_reports_the_error(self, mcp_db, monkeypatch):
async def boom(self, question, filter=None, images=None):
raise RuntimeError("sandbox exploded")
monkeypatch.setattr(HaikuRAG, "analyze", boom)
mcp = create_mcp_server(mcp_db, read_only=True)
analyze = await _get_tool(mcp, "analyze")
assert "sandbox exploded" in await analyze(question="q")
@pytest.mark.asyncio
async def test_ask_question_appends_citations_when_requested(
self, mcp_db, monkeypatch
):
from haiku.rag.store.models.citation import Citation
citation = Citation(
chunk_id="c1",
document_id="d1",
content="cited text",
document_uri="test://ai-overview",
document_title="AI Overview",
)
async def fake_ask(self, question, filter=None, images=None):
return ("the answer", [citation])
monkeypatch.setattr(HaikuRAG, "ask", fake_ask)
mcp = create_mcp_server(mcp_db, read_only=True)
ask = await _get_tool(mcp, "ask_question")
with_cite = await ask(question="q", cite=True)
assert with_cite.startswith("the answer")
assert "AI Overview" in with_cite
assert await ask(question="q", cite=False) == "the answer"

11
tests/test_pdf_split.py
View file

@ -243,3 +243,14 @@ def test_concatenate_shifts_page_nos_and_unique_self_refs():
assert page_nos == [1, 2], f"expected b's page 1 to shift to page 2, got {page_nos}"
assert sorted(merged.pages.keys()) == [1, 2]
def test_iter_pdf_slices_rejects_unopenable_pdf(tmp_path):
"""pdfium refuses non-PDF bytes; the caller sees UnsupportedSourceError."""
from haiku.rag.client.exceptions import UnsupportedSourceError
junk = tmp_path / "not-really.pdf"
junk.write_bytes(b"this is not a pdf at all")
with pytest.raises(UnsupportedSourceError, match="cannot open PDF"):
list(iter_pdf_slices(junk, slice_size=1))

29
tests/test_picture_in_context.py
View file

@ -1011,3 +1011,32 @@ async def test_rag_capability_attaches_images_for_vision_model(temp_db_path):
assert isinstance(result, ToolReturn)
assert result.content is not None
assert any(isinstance(part, BinaryContent) for part in result.content)
def test_build_picture_chunks_records_provenance_pages():
"""A picture with provenance contributes its page numbers to the chunk."""
from docling_core.types.doc.base import BoundingBox
from docling_core.types.doc.document import ProvenanceItem
from haiku.rag.client.processing import build_picture_chunks
from tests.store.test_document_items import _docling_doc_with_picture
doc = _docling_doc_with_picture()
picture = doc.pictures[0]
picture.prov = [
ProvenanceItem(
page_no=3,
bbox=BoundingBox(l=0, t=10, r=10, b=0),
charspan=(0, 0),
),
# A repeat of the same page must not be counted twice.
ProvenanceItem(
page_no=3,
bbox=BoundingBox(l=0, t=20, r=10, b=10),
charspan=(0, 0),
),
]
chunks = build_picture_chunks(doc, document_id="doc-1")
assert chunks[0].metadata["page_numbers"] == [3]

72
tests/test_processing.py
View file

@ -7,6 +7,7 @@ import pytest
from haiku.rag.client.processing import _warn_if_descriptions_missing, convert
from haiku.rag.config import AppConfig
from tests.conftest import capture_logs
def _doc_with_pictures(*, with_descriptions: bool):
@ -44,24 +45,12 @@ def _doc_without_pictures():
@pytest.fixture
def caplog_warnings(caplog):
def caplog_warnings():
"""Capture WARNING-level records from the processing logger."""
caplog.set_level(logging.WARNING, logger="haiku.rag.client.processing")
# The haiku.rag parent logger sets propagate=False after get_logger() runs,
# which can break caplog under xdist when other tests have already
# configured logging. Attach directly to the module logger.
from haiku.rag.client.processing import logger as proc_logger
records: list[logging.LogRecord] = []
class _Capture(logging.Handler):
def emit(self, record: logging.LogRecord) -> None:
records.append(record)
handler = _Capture(level=logging.WARNING)
proc_logger.addHandler(handler)
yield records
proc_logger.removeHandler(handler)
with capture_logs(proc_logger, logging.WARNING) as records:
yield records
def test_no_warning_when_picture_description_disabled(caplog_warnings):
@ -222,3 +211,56 @@ def test_merge_picture_chunks_no_pictures_returns_text_chunks():
assert result is text_chunks
assert [c.order for c in result] == [0, 1]
@pytest.mark.asyncio
async def test_convert_dispatches_large_pdfs_through_split_and_merge(
tmp_path, monkeypatch
):
"""With split_pages configured, PDF conversion routes through the
split-and-merge helper rather than the converter directly."""
from docling_core.types.doc.document import DoclingDocument
config = AppConfig()
config.processing.split_pages = 2
pdf = tmp_path / "big.pdf"
pdf.write_bytes(b"%PDF-1.4 stub")
called: dict = {}
async def fake_split(converter, path, uri, slice_size):
called["slice_size"] = slice_size
called["path"] = path
return DoclingDocument(name="merged")
monkeypatch.setattr(
"haiku.rag.converters.pdf_split.convert_pdf_with_splitting", fake_split
)
doc = await convert(config, pdf)
assert doc.name == "merged"
assert called["slice_size"] == 2
assert called["path"] == pdf
def _write_unsupported(directory):
target = directory / "thing.sqlite3"
target.write_bytes(b"binary")
return target.as_uri()
@pytest.mark.asyncio
@pytest.mark.parametrize(
"make_source,match",
[
(lambda d: (d / "missing.md").as_uri(), "File does not exist"),
(_write_unsupported, "Unsupported file extension"),
],
ids=["missing_file", "unsupported_extension"],
)
async def test_convert_rejects_bad_file_uris(tmp_path, make_source, match):
from haiku.rag.client.exceptions import UnsupportedSourceError
with pytest.raises(UnsupportedSourceError, match=match):
await convert(AppConfig(), make_source(tmp_path))

489
tests/test_rebuild.py
View file

@ -7,6 +7,7 @@ import pytest
from haiku.rag.client import HaikuRAG, RebuildMode
from haiku.rag.config import Config
from tests.conftest import capture_logs
class ChunkData(TypedDict):
@ -615,25 +616,11 @@ async def test_rebuild_full_source_failure_is_logged_and_skipped(
monkeypatch.setattr(client, "create_document_from_source", failing_create)
# Attach directly to the rebuild module's logger rather than
# relying on caplog — `haiku.rag.logging.get_logger()` (invoked
# by other tests) sets `propagate=False` on the `haiku.rag`
# logger, which breaks caplog under xdist ordering.
records: list[logging.LogRecord] = []
class _ListHandler(logging.Handler):
def emit(self, record: logging.LogRecord) -> None:
records.append(record)
handler = _ListHandler(level=logging.ERROR)
rebuild_module.logger.addHandler(handler)
try:
with capture_logs(rebuild_module.logger, logging.ERROR) as records:
processed_ids = [
doc_id
async for doc_id in client.rebuild_database(mode=RebuildMode.FULL)
]
finally:
rebuild_module.logger.removeHandler(handler)
assert processed_ids == []
assert any(
@ -843,7 +830,7 @@ async def test_patch_picture_descriptions_returns_zero_for_doc_without_pictures(
@pytest.mark.asyncio
async def test_patch_picture_descriptions_warns_on_missing_bytes(temp_db_path, caplog):
async def test_patch_picture_descriptions_warns_on_missing_bytes(temp_db_path):
"""When the docling blob has pictures but document_items.picture_data is
empty (e.g. legacy DB ingested before A2b), the helper logs a warning
and returns 0 instead of trying to drive the VLM with no input."""
@ -872,23 +859,10 @@ async def test_patch_picture_descriptions_warns_on_missing_bytes(temp_db_path, c
where=f"document_id = '{created.id}' AND label = 'picture'",
)
# Capture warnings directly off the rebuild module logger — the
# haiku.rag parent logger is configured non-propagating elsewhere
# in the suite so caplog can miss records.
from haiku.rag.client import rebuild as rebuild_module
records: list[logging.LogRecord] = []
class _ListHandler(logging.Handler):
def emit(self, record: logging.LogRecord) -> None:
records.append(record)
handler = _ListHandler(level=logging.WARNING)
rebuild_module.logger.addHandler(handler)
try:
with capture_logs(rebuild_module.logger, logging.WARNING) as records:
n = await _patch_picture_descriptions(rag, created)
finally:
rebuild_module.logger.removeHandler(handler)
assert n == 0
assert any("no stored picture bytes" in r.getMessage() for r in records)
@ -1095,3 +1069,458 @@ async def test_rebuild_blocks_tag_operations(temp_db_path, monkeypatch):
assert not client.store._rebuild_lock.locked()
await client.store.create_tag("post-rebuild")
assert set(await client.store.list_tags()) == {"post-rebuild"}
def _count_flushes(monkeypatch, rebuild_module) -> list[int]:
"""Record the size of every batch handed to _flush_rebuild_batch."""
real = rebuild_module._flush_rebuild_batch
sizes: list[int] = []
async def spy(client, documents, chunks):
sizes.append(len(documents))
return await real(client, documents, chunks)
monkeypatch.setattr(rebuild_module, "_flush_rebuild_batch", spy)
return sizes
# --- unit-level rebuild helpers (no embedder involved) ---
@pytest.mark.vcr()
async def test_flush_rebuild_batch_is_a_noop_without_documents(temp_db_path):
from haiku.rag.client.rebuild import _flush_rebuild_batch
async with HaikuRAG(temp_db_path, create=True) as client:
# A populated table makes "unchanged" distinguishable from "wiped".
existing = await client.create_document(content="keep me")
before = await client.store.documents_table.count_rows()
assert before == 1
await _flush_rebuild_batch(client, [], [])
assert await client.store.documents_table.count_rows() == before
after = await client.get_document_by_id(existing.id)
assert after is not None
assert after.updated_at == existing.updated_at
async def test_mark_phase1_complete_is_idempotent(temp_db_path):
from haiku.rag.client.rebuild import (
_STAGING_MARKER_TABLE_NAME,
_mark_phase1_complete,
)
async with HaikuRAG(temp_db_path, create=True) as client:
await _mark_phase1_complete(client)
await _mark_phase1_complete(client)
tables = (await client.store.db.list_tables()).tables
assert _STAGING_MARKER_TABLE_NAME in tables
async def test_populate_staging_returns_early_without_chunks_table(temp_db_path):
from haiku.rag.client.rebuild import _STAGING_TABLE_NAME, _populate_staging_table
async with HaikuRAG(temp_db_path, create=True) as client:
await client.store.db.drop_table("chunks")
await _populate_staging_table(client)
staging = await client.store.db.open_table(_STAGING_TABLE_NAME)
assert await staging.count_rows() == 0
async def test_hydrate_skips_documents_deleted_mid_rebuild(temp_db_path):
"""A document removed between listing and hydration is skipped."""
from haiku.rag.client.rebuild import _hydrate
from haiku.rag.store.models.document import Document
from haiku.rag.store.repositories.document import DocumentRepository
async with HaikuRAG(temp_db_path, create=True) as client:
stored = await DocumentRepository(client.store).create(
Document(content="body", uri="test://gone")
)
async def vanished(_document_id):
return None
client.get_document_by_id = vanished # type: ignore[method-assign]
assert [doc async for doc in _hydrate(client, [stored])] == []
@pytest.mark.parametrize(
"description,expected_text",
[("", None), ("a red square", "a red square")],
ids=["empty_skipped", "populated_applied"],
)
async def test_apply_descriptions_writes_only_non_empty_text(
description, expected_text
):
"""An empty generated description leaves the picture untouched; a real one
is written through to the picture meta."""
from haiku.rag.client.rebuild import _apply_descriptions_sync
from haiku.rag.store.models.document import Document
from tests.store.test_document_items import _docling_doc_with_picture
docling_doc = _docling_doc_with_picture()
ref = docling_doc.pictures[0].self_ref
document = Document(content="x", uri="test://doc")
_apply_descriptions_sync(docling_doc, document, {ref: description})
meta = docling_doc.pictures[0].meta
actual = getattr(getattr(meta, "description", None), "text", None) if meta else None
assert actual == expected_text
# The blob is re-compressed either way; page rasters must survive it.
assert document.docling_document is not None
@pytest.mark.asyncio
async def test_patch_picture_descriptions_returns_zero_without_descriptions(
temp_db_path, monkeypatch
):
"""When the VLM returns nothing, no blob rewrite is attempted."""
from haiku.rag.client.documents import _store_document_with_chunks
from haiku.rag.client.rebuild import _patch_picture_descriptions
from haiku.rag.config import AppConfig
from haiku.rag.store.models.document import Document
from tests.store.test_document_items import _docling_doc_with_picture
docling_doc = _docling_doc_with_picture()
config = AppConfig()
config.processing.pictures = "description"
async def no_descriptions(_bytes_by_ref, config=None):
return {}
monkeypatch.setattr(
"haiku.rag.providers.picture_description.describe_pictures",
no_descriptions,
)
async with HaikuRAG(temp_db_path, config=config, create=True) as rag:
document = Document(content="x", uri="test://doc")
document.set_docling(docling_doc)
created = await _store_document_with_chunks(rag, document, [], docling_doc)
assert await _patch_picture_descriptions(rag, created) == 0
@pytest.mark.vcr()
async def test_rebuild_warns_when_post_rebuild_vacuum_fails(temp_db_path, monkeypatch):
"""A failing post-rebuild vacuum is logged, not raised — the rebuild itself
already succeeded."""
import logging
from haiku.rag.client import rebuild as rebuild_module
async with HaikuRAG(temp_db_path, create=True) as client:
doc = await client.create_document(content="vacuum failure doc")
assert doc.id is not None
client._config.storage.auto_vacuum = True
async def failing_vacuum():
raise RuntimeError("vacuum exploded")
monkeypatch.setattr(client.store, "vacuum", failing_vacuum)
with capture_logs(rebuild_module.logger, logging.WARNING) as records:
processed = [
doc_id
async for doc_id in client.rebuild_database(mode=RebuildMode.RECHUNK)
]
assert doc.id in processed
assert any("vacuum failed" in r.getMessage() for r in records)
@pytest.mark.vcr()
async def test_rebuild_embed_only_yields_documents_without_chunks(temp_db_path):
"""A document whose chunks were all removed is still reported as processed."""
async with HaikuRAG(temp_db_path, create=True) as client:
doc = await client.create_document(content="doc that loses its chunks")
assert doc.id is not None
await client.chunk_repository.delete_by_document_id(doc.id)
processed = [
doc_id
async for doc_id in client.rebuild_database(mode=RebuildMode.EMBED_ONLY)
]
assert processed == [doc.id]
@pytest.mark.vcr()
async def test_rebuild_embed_only_flushes_in_batches(temp_db_path, monkeypatch):
"""Forces a tiny batch size so the mid-loop flush in phase 2 runs."""
from haiku.rag.client import rebuild as rebuild_module
monkeypatch.setattr(rebuild_module, "_REBUILD_BATCH_SIZE", 2)
async with HaikuRAG(temp_db_path, create=True) as client:
ids = []
for i in range(3):
doc = await client.create_document(content=f"embed only batch doc {i}")
assert doc.id is not None
ids.append(doc.id)
# Phase 2 writes straight to the chunks table rather than going
# through _flush_rebuild_batch, so count the adds it makes. Patch at
# class level: embed-only recreates the table, discarding any patch
# applied to the instance that exists now.
import lancedb
real_add = lancedb.AsyncTable.add
adds: list[int] = []
async def counting_add(self, records, *args, **kwargs):
if self.name == "chunks":
adds.append(len(records))
return await real_add(self, records, *args, **kwargs)
monkeypatch.setattr(lancedb.AsyncTable, "add", counting_add)
processed = [
doc_id
async for doc_id in client.rebuild_database(mode=RebuildMode.EMBED_ONLY)
]
assert sorted(processed) == sorted(ids)
for doc_id in ids:
assert await client.chunk_repository.get_by_document_id(doc_id)
# 3 docs at batch size 2: one mid-loop write plus the trailing one.
assert len(adds) == 2
@pytest.mark.vcr()
async def test_rechunk_raises_when_docling_blob_is_missing(temp_db_path):
"""RECHUNK needs the stored docling document; without it the user is told
to run a full rebuild."""
async with HaikuRAG(temp_db_path, create=True) as client:
doc = await client.create_document(content="doc with a cleared blob")
assert doc.id is not None
await client.store.documents_table.update(
{"docling_document": None}, where=f"id = '{doc.id}'"
)
with pytest.raises(ValueError, match="has no stored docling document"):
async for _ in client.rebuild_database(mode=RebuildMode.RECHUNK):
pass
@pytest.mark.vcr()
async def test_rebuild_full_flushes_in_batches(temp_db_path, monkeypatch):
from haiku.rag.client import rebuild as rebuild_module
monkeypatch.setattr(rebuild_module, "_REBUILD_BATCH_SIZE", 2)
flushes = _count_flushes(monkeypatch, rebuild_module)
async with HaikuRAG(temp_db_path, create=True) as client:
ids = []
for i in range(3):
doc = await client.create_document(content=f"full batch doc {i}")
assert doc.id is not None
ids.append(doc.id)
processed = [
doc_id async for doc_id in client.rebuild_database(mode=RebuildMode.FULL)
]
assert sorted(processed) == sorted(ids)
assert len(flushes) == 2
@pytest.mark.vcr()
async def test_rebuild_full_warns_when_source_is_missing(temp_db_path):
"""A document whose file source is gone is re-embedded from stored content."""
import logging
from haiku.rag.client import rebuild as rebuild_module
async with HaikuRAG(temp_db_path, create=True) as client:
doc = await client.create_document(
content="content whose source vanished",
uri="file:///definitely/not/here.txt",
)
assert doc.id is not None
with capture_logs(rebuild_module.logger, logging.WARNING) as records:
processed = [
doc_id
async for doc_id in client.rebuild_database(mode=RebuildMode.FULL)
]
assert doc.id in processed
assert any("Source missing" in r.getMessage() for r in records)
@pytest.mark.vcr()
async def test_rebuild_full_flushes_pending_before_source_rebuild(temp_db_path):
"""A live source is re-ingested, which creates a new document, so any
documents pending from the content path must be flushed first."""
async with HaikuRAG(temp_db_path, create=True) as client:
content_doc = await client.create_document(content="plain content doc")
assert content_doc.id is not None
with tempfile.TemporaryDirectory() as temp_dir:
source = Path(temp_dir) / "live.txt"
source.write_text("content from a source that still exists")
source_doc = await client.create_document_from_source(source)
assert not isinstance(source_doc, list)
processed = [
doc_id
async for doc_id in client.rebuild_database(mode=RebuildMode.FULL)
]
# The content-path document keeps its id and must survive the
# flush that precedes the source re-ingest; the source document
# is replaced by a freshly ingested one with a new id.
assert content_doc.id in processed
assert source_doc.id not in processed
assert len(processed) == 2
assert await client.store.documents_table.count_rows() == 2
@pytest.mark.vcr()
async def test_rebuild_descriptions_flushes_in_batches(temp_db_path, monkeypatch):
"""Two picture documents with a batch size of one exercise the mid-loop
flush in the descriptions path."""
from haiku.rag.client import rebuild as rebuild_module
from haiku.rag.client.documents import _store_document_with_chunks
from haiku.rag.config import AppConfig
from haiku.rag.store.models.document import Document
from tests.store.test_document_items import _docling_doc_with_picture
config = AppConfig()
config.processing.pictures = "description"
monkeypatch.setattr(rebuild_module, "_REBUILD_BATCH_SIZE", 1)
flushes = _count_flushes(monkeypatch, rebuild_module)
async def fake_describe(image_bytes_by_ref, *, config):
return {ref: "A red square (mocked)." for ref in image_bytes_by_ref}
monkeypatch.setattr(
"haiku.rag.providers.picture_description.describe_pictures", fake_describe
)
async with HaikuRAG(temp_db_path, config=config, create=True) as rag:
ids = []
for i in range(2):
docling_doc = _docling_doc_with_picture()
document = Document(content=f"picture doc {i}", uri=f"test://doc-{i}")
document.set_docling(docling_doc)
created = await _store_document_with_chunks(rag, document, [], docling_doc)
assert created.id is not None
ids.append(created.id)
processed = [
doc_id
async for doc_id in rag.rebuild_database(mode=RebuildMode.DESCRIPTIONS)
]
assert sorted(processed) == sorted(ids)
# 2 docs at batch size 1: one flush each, none left for the trailing pass.
assert len(flushes) == 2
@pytest.mark.vcr()
async def test_rebuild_full_skips_document_deleted_mid_rebuild(temp_db_path):
"""A document removed between listing and the content-path reload is skipped."""
async with HaikuRAG(temp_db_path, create=True) as client:
doc = await client.create_document(content="doc that disappears")
assert doc.id is not None
async def vanished(_document_id):
return None
client.get_document_by_id = vanished # type: ignore[method-assign]
processed = [
doc_id async for doc_id in client.rebuild_database(mode=RebuildMode.FULL)
]
assert processed == []
@pytest.mark.vcr()
@pytest.mark.parametrize("wipe_bytes", [True, False], ids=["wiped", "recoverable"])
async def test_rebuild_embed_only_recovers_picture_bytes(
temp_db_path, monkeypatch, wipe_bytes
):
"""Embed-only re-attaches picture bytes from document_items. When they are
gone the chunk falls back to embedding its caption as text rather than
failing the rebuild."""
import logging
from haiku.rag.client import rebuild as rebuild_module
from haiku.rag.client.documents import _store_document_with_chunks
from haiku.rag.store.models.chunk import Chunk
from haiku.rag.store.models.document import Document
from tests.store.test_document_items import _docling_doc_with_picture
docling_doc = _docling_doc_with_picture()
ref = docling_doc.pictures[0].self_ref
async with HaikuRAG(temp_db_path, create=True) as rag:
# The configured ollama embedder is text-only; stand in for a
# multimodal one so the picture-bytes recovery branch runs.
embedded_images: list[bytes] = []
async def fake_embed_image(image):
embedded_images.append(image)
return [0.2] * rag.embedder.vector_dim
monkeypatch.setattr(rag.embedder, "supports_images", True)
monkeypatch.setattr(rag.embedder, "embed_image", fake_embed_image)
document = Document(content="picture doc", uri="test://pic")
document.set_docling(docling_doc)
picture_chunk = Chunk(
content="Figure caption",
metadata={"doc_item_refs": [ref], "labels": ["picture"]},
order=0,
embedding=[0.1] * rag.embedder.vector_dim,
)
# A sibling text chunk exercises the non-picture skip in the same loop.
text_chunk = Chunk(
content="Surrounding prose",
metadata={"doc_item_refs": ["#/texts/0"], "labels": ["text"]},
order=1,
embedding=[0.1] * rag.embedder.vector_dim,
)
created = await _store_document_with_chunks(
rag, document, [picture_chunk, text_chunk], docling_doc
)
assert created.id is not None
if wipe_bytes:
await rag.store.document_items_table.update(
{"picture_data": None},
where=f"document_id = '{created.id}' AND label = 'picture'",
)
with capture_logs(rebuild_module.logger, logging.WARNING) as records:
processed = [
doc_id
async for doc_id in rag.rebuild_database(mode=RebuildMode.EMBED_ONLY)
]
assert created.id in processed
warned = any("no recoverable bytes" in r.getMessage() for r in records)
assert warned is wipe_bytes
if wipe_bytes:
# Nothing to recover, so the caption is text-embedded instead.
assert embedded_images == []
else:
# The stored PNG was re-attached and routed through embed_image.
assert len(embedded_images) == 1
assert embedded_images[0].startswith(b"\x89PNG")

27
tests/test_reranker.py
View file

@ -479,3 +479,30 @@ async def test_cross_encoder_reranker():
assert "0" in top_ids or "2" in top_ids
except ImportError:
pytest.skip("sentence-transformers not installed")
@pytest.mark.asyncio
async def test_cross_encoder_reranks_via_model_ranking(monkeypatch):
"""The rank() results map back onto the input chunks by corpus_id."""
from haiku.rag.reranking import cross_encoder as ce_module
class _StubCrossEncoder:
def __init__(self, model):
self.model = model
def rank(self, query, documents, top_k=10):
# Reverse order so the mapping back to chunks is observable.
return [
{"corpus_id": i, "score": 1.0 - (i / 10)}
for i in reversed(range(len(documents)))
][:top_k]
monkeypatch.setattr(ce_module, "CrossEncoder", _StubCrossEncoder)
reranker = ce_module.CrossEncoderReranker("stub/model")
reranked = await reranker.rerank("query", chunks, top_n=2)
assert len(reranked) == 2
last_index = len(chunks) - 1
assert reranked[0][0] is chunks[last_index]
assert reranked[0][1] == pytest.approx(1.0 - last_index / 10)

186
tests/test_search.py
View file

@ -275,72 +275,35 @@ async def test_fts_search_targets_content_fts_column(temp_db_path):
)
def test_search_result_primary_label_prioritizes_structural_types():
"""Test _get_primary_label prioritizes structural labels correctly."""
# Table should be prioritized
result = SearchResult(
content="test",
score=0.5,
chunk_id="c1",
document_id="d1",
labels=["paragraph", "table", "text"],
)
assert result._get_primary_label() == "table"
# Code should be prioritized over paragraph
result = SearchResult(
content="test",
score=0.5,
chunk_id="c2",
document_id="d2",
labels=["paragraph", "code"],
)
assert result._get_primary_label() == "code"
# list_item should be prioritized
result = SearchResult(
content="test",
score=0.5,
chunk_id="c3",
document_id="d3",
labels=["text", "list_item"],
)
assert result._get_primary_label() == "list_item"
# Returns first label when no priority match
result = SearchResult(
content="test",
score=0.5,
chunk_id="c4",
document_id="d4",
labels=["paragraph", "text"],
)
assert result._get_primary_label() == "paragraph"
# Returns None for empty labels
result = SearchResult(
content="test",
score=0.5,
chunk_id="c5",
document_id="d5",
labels=[],
)
assert result._get_primary_label() is None
# Image queries (bytes / PIL.Image)
def _png_bytes_query() -> bytes:
return b"\x89PNG\r\n\x1a\n"
def _pil_image_query():
from PIL import Image as PILImageModule
return PILImageModule.new("RGB", (8, 8), "red")
@pytest.mark.asyncio
async def test_search_with_bytes_query_uses_multimodal_embedder(
temp_db_path, monkeypatch
@pytest.mark.parametrize(
"make_query",
[_png_bytes_query, _pil_image_query],
ids=["bytes", "pil"],
)
async def test_search_with_image_query_uses_multimodal_embedder(
temp_db_path, monkeypatch, make_query
):
"""``client.search(bytes)`` embeds via ``embed_image`` and dispatches
"""``client.search(image)`` embeds via ``embed_image`` and dispatches
to vector-only chunk search (skipping FTS and reranker)."""
from haiku.rag.embeddings import EmbedderWrapper
from haiku.rag.store.models.chunk import Chunk
image_calls: list[bytes] = []
query = make_query()
image_calls: list = []
class StubMultimodal(EmbedderWrapper):
supports_images = True
@ -383,12 +346,12 @@ async def test_search_with_bytes_query_uses_multimodal_embedder(
async with HaikuRAG(temp_db_path, create=True) as rag:
rag.chunk_repository.search = fake_chunk_search # type: ignore[method-assign]
results = await rag.search(b"\x89PNG\r\n\x1a\n", limit=3, include_images=False)
results = await rag.search(query, limit=3, include_images=False)
assert len(results) == 1
assert results[0].score == 0.91
# The bytes were sent through the image embedder once.
assert image_calls == [b"\x89PNG\r\n\x1a\n"]
# The image was passed through to the image embedder untouched.
assert image_calls == [query]
# The chunk repo received a pre-computed vector and an empty text query.
assert received_kwargs["query_vector"] == [0.5, 0.5, 0.5, 0.5]
assert received_kwargs["query"] == ""
@ -493,42 +456,6 @@ async def test_search_attaches_picture_bytes_for_multimodal_reranker(
assert reranked_picture._picture_data is None
@pytest.mark.asyncio
async def test_search_with_pil_image_works_like_bytes(temp_db_path, monkeypatch):
from PIL import Image as PILImageModule
from haiku.rag.embeddings import EmbedderWrapper
from haiku.rag.store.models.chunk import Chunk
seen_types: list[type] = []
class StubMultimodal(EmbedderWrapper):
supports_images = True
def __init__(self):
super().__init__(embedder=None, vector_dim=4)
async def embed_image(self, image):
seen_types.append(type(image))
return [0.1] * 4
monkeypatch.setattr(
"haiku.rag.store.engine.get_embedder",
lambda *a, **kw: StubMultimodal(),
)
async def fake_chunk_search(**kwargs):
return [(Chunk(content="x", metadata={}), 1.0)]
async with HaikuRAG(temp_db_path, create=True) as rag:
rag.chunk_repository.search = fake_chunk_search # type: ignore[method-assign]
img = PILImageModule.new("RGB", (8, 8), "red")
results = await rag.search(img, include_images=False)
assert len(results) == 1
assert seen_types == [PILImageModule.Image]
@pytest.mark.asyncio
async def test_search_with_bytes_query_raises_for_text_only_embedder(
temp_db_path,
@ -553,20 +480,26 @@ def _picture_only_result(
)
def test_dedup_keeps_higher_scoring_picture_chunk():
@pytest.mark.parametrize(
"first_score,second_score",
# Whichever duplicate scores higher wins, regardless of arrival order.
[(0.7, 0.9), (0.9, 0.7)],
ids=["later_wins", "earlier_wins"],
)
def test_dedup_keeps_higher_scoring_picture_chunk(first_score, second_score):
"""Two results referencing the same single picture self_ref collapse
to the one with the higher score."""
from haiku.rag.client.search import _dedup_picture_chunks
text_chunk = _picture_only_result("#/pictures/0", score=0.7)
pic_chunk = _picture_only_result("#/pictures/0", score=0.9)
text_chunk = _picture_only_result("#/pictures/0", score=first_score)
pic_chunk = _picture_only_result("#/pictures/0", score=second_score)
other = _picture_only_result("#/pictures/1", score=0.6)
deduped = _dedup_picture_chunks([text_chunk, pic_chunk, other])
assert len(deduped) == 2
chosen = next(r for r in deduped if r.doc_item_refs == ["#/pictures/0"])
assert chosen.score == 0.9
assert chosen.score == max(first_score, second_score)
assert any(r.doc_item_refs == ["#/pictures/1"] for r in deduped)
@ -598,3 +531,56 @@ def test_dedup_does_not_collapse_across_documents():
deduped = _dedup_picture_chunks([a, b])
assert len(deduped) == 2
# visualize_chunk short-circuits
@pytest.mark.asyncio
async def test_expand_context_passes_through_results_without_document(temp_db_path):
"""A result with no document_id can't be expanded; it is returned as-is."""
from haiku.rag.client.search import expand_context
async with HaikuRAG(temp_db_path, create=True) as rag:
orphan = SearchResult(content="loose text", score=0.5, chunk_id="c1")
assert await expand_context(rag, [orphan]) == [orphan]
@pytest.mark.asyncio
async def test_visualize_chunk_returns_empty_for_no_chunks(temp_db_path):
async with HaikuRAG(temp_db_path, create=True) as rag:
assert await rag.visualize_chunk([]) == []
@pytest.mark.asyncio
async def test_visualize_chunk_returns_empty_without_document_id(temp_db_path):
from haiku.rag.store.models.chunk import Chunk
async with HaikuRAG(temp_db_path, create=True) as rag:
assert await rag.visualize_chunk(Chunk(content="x", metadata={})) == []
@pytest.mark.asyncio
async def test_visualize_chunk_returns_empty_when_document_missing(temp_db_path):
from haiku.rag.store.models.chunk import Chunk
async with HaikuRAG(temp_db_path, create=True) as rag:
chunk = Chunk(content="x", document_id="does-not-exist", metadata={})
assert await rag.visualize_chunk(chunk) == []
@pytest.mark.asyncio
async def test_visualize_chunk_returns_empty_when_docling_blob_absent(temp_db_path):
"""A markdown-ingested document has no docling structure to resolve boxes in."""
from haiku.rag.store.models.chunk import Chunk
from haiku.rag.store.models.document import Document
from haiku.rag.store.repositories.document import DocumentRepository
async with HaikuRAG(temp_db_path, create=True) as rag:
doc = await DocumentRepository(rag.store).create(
Document(content="plain body", uri="test://plain")
)
assert doc.id is not None
chunk = Chunk(content="plain body", document_id=doc.id, metadata={})
assert await rag.visualize_chunk(chunk) == []

45
tests/test_settings.py
View file

@ -44,6 +44,31 @@ async def test_settings_save_and_retrieve(temp_db_path):
Config.processing.chunk_size = original_chunk_size
@pytest.mark.asyncio
async def test_set_haiku_version_recreates_row_from_store_config(temp_db_path):
"""Recreating a missing settings row stamps the store's own config, not the
process-global one."""
from haiku.rag.store.engine import Store
from haiku.rag.store.repositories.settings import SettingsRepository
config = AppConfig()
config.processing.chunk_size = Config.processing.chunk_size + 512
async with Store(temp_db_path, config=config, create=True) as store:
settings_repo = SettingsRepository(store)
await store.settings_table.delete("id = 'settings'")
assert await settings_repo.get_current_settings() == {}
assert await store.get_haiku_version() == "0.0.0"
await store.set_haiku_version("1.2.3")
recreated = await settings_repo.get_current_settings()
assert recreated["version"] == "1.2.3"
assert recreated["processing"]["chunk_size"] == config.processing.chunk_size
assert await store.get_haiku_version() == "1.2.3"
class TestValidateConfigCompatibility:
"""Tests for validate_config_compatibility method."""
@ -238,3 +263,23 @@ class TestValidateConfigCompatibility:
await settings_repo.validate_config_compatibility()
assert "9999" in str(exc_info.value)
@pytest.mark.asyncio
async def test_save_current_settings_recreates_a_deleted_row(temp_db_path):
from haiku.rag.store.engine import Store
from haiku.rag.store.repositories.settings import SettingsRepository
async with Store(temp_db_path, create=True, skip_validation=True) as store:
settings_repo = SettingsRepository(store)
await store.settings_table.delete("id = 'settings'")
assert await settings_repo.get_current_settings() == {}
await settings_repo.save_current_settings()
recreated = await settings_repo.get_current_settings()
assert (
recreated["embeddings"]
== store._config.model_dump(mode="json")["embeddings"]
)

37
tests/test_title_generation.py
View file

@ -371,3 +371,40 @@ class TestRebuildTitleOnly:
# Only the second doc should have been processed
assert len(processed_ids) == 1
@pytest.mark.asyncio
async def test_generate_title_with_llm_returns_model_output(monkeypatch):
"""The agent's output is stripped and returned."""
from pydantic_ai.messages import ModelMessage, ModelResponse, TextPart
from pydantic_ai.models.function import AgentInfo, FunctionModel
from haiku.rag.client.titles import generate_title_with_llm
from haiku.rag.config import AppConfig
def respond(messages: list[ModelMessage], info: AgentInfo) -> ModelResponse:
return ModelResponse(parts=[TextPart(" A Generated Title ")])
monkeypatch.setattr(
"haiku.rag.utils.get_model", lambda *a, **kw: FunctionModel(respond)
)
assert await generate_title_with_llm(AppConfig(), "body") == "A Generated Title"
@pytest.mark.asyncio
async def test_generate_title_with_llm_returns_none_for_blank_output(monkeypatch):
from pydantic_ai.messages import ModelMessage, ModelResponse, TextPart
from pydantic_ai.models.function import AgentInfo, FunctionModel
from haiku.rag.client.titles import generate_title_with_llm
from haiku.rag.config import AppConfig
def respond(messages: list[ModelMessage], info: AgentInfo) -> ModelResponse:
return ModelResponse(parts=[TextPart(" ")])
monkeypatch.setattr(
"haiku.rag.utils.get_model", lambda *a, **kw: FunctionModel(respond)
)
assert await generate_title_with_llm(AppConfig(), "body") is None

287
tests/test_utils.py
View file

@ -139,27 +139,70 @@ Emoji test: 🚀 ✅ 📝"""
assert "🚀" in result_markdown
def test_get_model_ollama():
"""Test get_model returns OpenAIChatModel for Ollama."""
model_config = ModelConfig(provider="ollama", name="llama3")
result = get_model(model_config)
assert isinstance(result, OpenAIChatModel)
def test_get_model_ollama_without_thinking():
"""Test get_model configures thinking for gpt-oss on Ollama."""
model_config = ModelConfig(provider="ollama", name="gpt-oss", enable_thinking=False)
result = get_model(model_config)
assert isinstance(result, OpenAIChatModel)
def test_get_model_ollama_with_settings():
"""Test get_model applies temperature and max_tokens for Ollama."""
model_config = ModelConfig(
provider="ollama", name="llama3", temperature=0.5, max_tokens=100
)
result = get_model(model_config)
@pytest.mark.parametrize(
"kwargs,expected_settings",
[
({"provider": "ollama", "name": "llama3"}, None),
(
{"provider": "ollama", "name": "gpt-oss", "enable_thinking": False},
{"openai_reasoning_effort": "low"},
),
(
{"provider": "ollama", "name": "gpt-oss", "enable_thinking": True},
{"openai_reasoning_effort": "high"},
),
(
{
"provider": "ollama",
"name": "llama3",
"temperature": 0.5,
"max_tokens": 100,
},
{"temperature": 0.5, "max_tokens": 100},
),
({"provider": "openai", "name": "gpt-4o"}, None),
(
{"provider": "openai", "name": "o1", "enable_thinking": True},
{"openai_reasoning_effort": "high"},
),
(
{"provider": "openai", "name": "o1", "enable_thinking": False},
{"openai_reasoning_effort": "low"},
),
(
{
"provider": "openai",
"name": "gpt-4o",
"enable_thinking": False,
"temperature": 0.7,
"max_tokens": 500,
},
# gpt-4o is not a reasoning model, so only the common settings land.
{"temperature": 0.7, "max_tokens": 500},
),
],
ids=[
"ollama",
"ollama_thinking_off",
"ollama_thinking_on",
"ollama_with_settings",
"openai",
"openai_reasoning_thinking_on",
"openai_reasoning_thinking_off",
"openai_all_settings",
],
)
def test_get_model_openai_chat_settings(kwargs, expected_settings):
"""Each ollama/openai configuration maps onto the expected model settings."""
result = get_model(ModelConfig(**kwargs))
assert isinstance(result, OpenAIChatModel)
if expected_settings is None:
assert result.settings is None
return
assert result.settings is not None
for key, value in expected_settings.items():
assert result.settings.get(key) == value
def test_get_model_ollama_appends_v1_to_per_model_base_url():
@ -183,20 +226,6 @@ def test_get_model_ollama_does_not_double_append_v1():
assert not url.endswith("/v1/v1")
def test_get_model_openai():
"""Test get_model returns OpenAIChatModel for OpenAI."""
model_config = ModelConfig(provider="openai", name="gpt-4o")
result = get_model(model_config)
assert isinstance(result, OpenAIChatModel)
def test_get_model_openai_with_thinking():
"""Test get_model configures thinking for OpenAI reasoning models."""
model_config = ModelConfig(provider="openai", name="o1", enable_thinking=True)
result = get_model(model_config)
assert isinstance(result, OpenAIChatModel)
def test_get_model_openai_non_reasoning_model_ignores_thinking():
"""Test that non-reasoning OpenAI models don't get reasoning_effort setting."""
model_config = ModelConfig(
@ -299,17 +328,27 @@ def test_get_model_anthropic():
@pytest.mark.skipif(not HAS_ANTHROPIC, reason="Anthropic not installed")
def test_get_model_anthropic_with_thinking():
@pytest.mark.parametrize(
"enable_thinking,expected_thinking",
[
(True, {"type": "enabled", "budget_tokens": 4096}),
(False, {"type": "disabled"}),
],
)
def test_get_model_anthropic_with_thinking(enable_thinking, expected_thinking):
"""Test get_model configures thinking for Anthropic."""
from pydantic_ai.models.anthropic import AnthropicModel
model_config = ModelConfig(
provider="anthropic",
name="claude-3-5-sonnet-20241022",
enable_thinking=True,
enable_thinking=enable_thinking,
)
result = get_model(model_config)
assert isinstance(result, AnthropicModel)
assert result.settings is not None
assert result.settings.get("anthropic_thinking") == expected_thinking
@pytest.mark.skipif(not HAS_GOOGLE, reason="Google not installed")
@ -345,15 +384,23 @@ def test_get_model_groq():
@pytest.mark.skipif(not HAS_GROQ, reason="Groq not installed")
def test_get_model_groq_with_thinking():
@pytest.mark.parametrize(
"enable_thinking,expected_format", [(True, "parsed"), (False, "hidden")]
)
def test_get_model_groq_with_thinking(enable_thinking, expected_format):
"""Test get_model configures thinking format for Groq."""
from pydantic_ai.models.groq import GroqModel
model_config = ModelConfig(
provider="groq", name="llama-3.3-70b-versatile", enable_thinking=False
provider="groq",
name="llama-3.3-70b-versatile",
enable_thinking=enable_thinking,
)
result = get_model(model_config)
assert isinstance(result, GroqModel)
assert result.settings is not None
assert result.settings.get("groq_reasoning_format") == expected_format
@pytest.mark.skipif(not HAS_BEDROCK, reason="Bedrock not installed")
@ -369,17 +416,58 @@ def test_get_model_bedrock():
@pytest.mark.skipif(not HAS_BEDROCK, reason="Bedrock not installed")
def test_get_model_bedrock_with_thinking():
"""Test get_model configures thinking for Bedrock Claude models."""
@pytest.mark.parametrize(
"name,enable_thinking,expected_fields",
[
(
"anthropic.claude-3-5-sonnet-20241022-v2:0",
True,
{"thinking": {"type": "enabled", "budget_tokens": 4096}},
),
(
"anthropic.claude-3-5-sonnet-20241022-v2:0",
False,
{"thinking": {"type": "disabled"}},
),
("openai.o3-mini-v1:0", True, {"reasoning_effort": "high"}),
("openai.o3-mini-v1:0", False, {"reasoning_effort": "low"}),
("qwen.qwen3-32b-v1:0", True, {"reasoning_config": "high"}),
("qwen.qwen3-32b-v1:0", False, {"reasoning_config": "low"}),
# A family with no reasoning mapping leaves the request fields untouched.
("meta.llama3-70b-instruct-v1:0", True, None),
("meta.llama3-70b-instruct-v1:0", False, None),
],
ids=[
"claude_on",
"claude_off",
"o_series_on",
"o_series_off",
"qwen_on",
"qwen_off",
"unmapped_on",
"unmapped_off",
],
)
def test_get_model_bedrock_with_thinking(name, enable_thinking, expected_fields):
"""Each Bedrock model family maps thinking onto its own request field."""
from pydantic_ai.models.bedrock import BedrockConverseModel
model_config = ModelConfig(
provider="bedrock",
name="anthropic.claude-3-5-sonnet-20241022-v2:0",
enable_thinking=True,
name=name,
enable_thinking=enable_thinking,
)
result = get_model(model_config)
assert isinstance(result, BedrockConverseModel)
if expected_fields is None:
assert result.settings is None
return
assert result.settings is not None
assert (
result.settings.get("bedrock_additional_model_requests_fields")
== expected_fields
)
def test_get_model_unknown_provider():
@ -390,19 +478,6 @@ def test_get_model_unknown_provider():
assert result == "mistral:mistral-large-latest"
def test_get_model_with_all_settings():
"""Test get_model applies all settings together."""
model_config = ModelConfig(
provider="openai",
name="gpt-4o",
enable_thinking=False,
temperature=0.7,
max_tokens=500,
)
result = get_model(model_config)
assert isinstance(result, OpenAIChatModel)
def test_get_package_versions():
"""Test get_package_versions returns expected keys."""
from haiku.rag.utils import get_package_versions
@ -526,20 +601,7 @@ def test_format_citations_multiple_pages():
result = format_citations([citation])
assert "[1] test://doc" in result
assert "pp. 1-3" in result
def test_format_citations_no_title():
from haiku.rag.store.models.citation import Citation
from haiku.rag.utils import format_citations
citation = Citation(
document_id="doc1",
chunk_id="chunk1",
document_uri="test://doc",
content="Content",
)
result = format_citations([citation])
assert "[1] test://doc" in result
# No title: the URI stands in, and the document id never leaks.
assert "doc1" not in result
@ -758,3 +820,90 @@ def test_parse_model_option():
for bad in ["just-a-name", ":model", "provider:"]:
with pytest.raises(ValueError, match="Invalid model format"):
parse_model_option(bad)
def test_cosine_similarity_identical_vectors():
from haiku.rag.utils import cosine_similarity
assert cosine_similarity([1.0, 0.0], [1.0, 0.0]) == pytest.approx(1.0)
assert cosine_similarity([1.0, 0.0], [0.0, 1.0]) == pytest.approx(0.0)
async def test_format_citations_rich_separates_multiple_citations():
from haiku.rag.store.models.citation import Citation
from haiku.rag.utils import format_citations_rich
citations = [
Citation(
document_id=f"doc{i}",
chunk_id=f"chunk{i}",
document_uri=f"test://doc{i}",
document_title=f"Doc {i}",
content=f"Body {i}",
)
for i in (1, 2)
]
output = _render_rich(await format_citations_rich(citations))
assert "[1] Doc 1 (test://doc1)" in output
assert "[2] Doc 2 (test://doc2)" in output
@pytest.mark.parametrize(
"stored,renders",
[
(None, False),
(b"not a real image", False),
("png", True),
],
ids=["no_bytes", "undecodable_bytes", "valid_png"],
)
async def test_render_picture_handles_stored_bytes(stored, renders):
from unittest.mock import AsyncMock
from haiku.rag.utils import _render_picture
if stored == "png":
from io import BytesIO
from PIL import Image as PILImage
buf = BytesIO()
PILImage.new("RGB", (4, 4), "red").save(buf, format="PNG")
stored = buf.getvalue()
client = AsyncMock()
client.document_item_repository.get_picture_bytes = AsyncMock(return_value=stored)
result = await _render_picture(client, "doc1", "#/pictures/0")
if renders:
from textual_image.renderable import Image as RichImage
assert isinstance(result, RichImage)
else:
assert result is None
async def test_render_picture_without_client_returns_none():
from haiku.rag.utils import _render_picture
assert await _render_picture(None, "doc1", "#/pictures/0") is None
def test_get_package_versions_reports_missing_docling(monkeypatch):
from importlib import metadata as importlib_metadata
from haiku.rag.utils import get_package_versions
real_version = importlib_metadata.version
def fake_version(name):
if name == "docling":
raise importlib_metadata.PackageNotFoundError(name)
return real_version(name)
monkeypatch.setattr(importlib_metadata, "version", fake_version)
assert get_package_versions()["docling"] == "not installed"

27
tests/tools/test_document.py
View file

@ -187,6 +187,33 @@ class TestSummarizeDocumentTool:
assert "Document not found" in result
@pytest.mark.vcr()
@pytest.mark.asyncio
async def test_summarize_document_returns_model_summary(
self, doc_client, doc_config, monkeypatch
):
"""A resolvable document is summarised and labelled with its title."""
from pydantic_ai.messages import ModelMessage, ModelResponse, TextPart
from pydantic_ai.models.function import AgentInfo, FunctionModel
def respond(messages: list[ModelMessage], info: AgentInfo) -> ModelResponse:
return ModelResponse(parts=[TextPart("A concise summary.")])
monkeypatch.setattr(
"haiku.rag.tools.document.get_model",
lambda *a, **kw: FunctionModel(respond),
)
docs = await doc_client.list_documents()
assert docs and docs[0].uri
toolset = create_document_toolset(doc_config)
summarize_tool = toolset.tools["summarize_document"]
result = await summarize_tool.function(make_ctx(doc_client), docs[0].uri)
assert "A concise summary." in result
assert "Summary of" in result
@pytest.fixture
async def doc_client(temp_db_path):

46
tests/tools/test_search.py
View file

@ -214,3 +214,49 @@ def search_config():
from haiku.rag.config import Config
return Config
class TestBuildBinaryPartsFromResults:
"""Picture bytes are attached once per (document, self_ref) pair."""
def test_results_without_image_data_contribute_nothing(self):
from haiku.rag.tools.search import build_binary_parts_from_results
results = [
SearchResult(content="text only", score=0.5, chunk_id="c1", image_data=None)
]
assert build_binary_parts_from_results(results) == []
def test_duplicate_document_and_ref_is_attached_once(self):
import base64
from io import BytesIO
from PIL import Image as PILImage
from haiku.rag.tools.search import build_binary_parts_from_results
buf = BytesIO()
PILImage.new("RGB", (4, 4), "red").save(buf, format="PNG")
png = base64.b64encode(buf.getvalue()).decode()
shared = {"#/pictures/0": png}
results = [
SearchResult(
content="a",
score=0.9,
chunk_id="c1",
document_id="doc-1",
image_data=shared,
),
SearchResult(
content="b",
score=0.8,
chunk_id="c2",
document_id="doc-1",
image_data=shared,
),
]
parts = build_binary_parts_from_results(results)
assert len(parts) == 1