Reach and enforce 100% coverage
Cover the remaining paths in the client, context, downloads, title generation, document tools and store models, and add fail_under=100 so uncovered lines fail CI. Six lines that no test can reach get a pragma with its reason: the docling import guard, the nameless PDF attachment, the FS symlink OSError guard that resolve(strict=False) absorbs, the docling bbox and LanceDB document-id shape guards, the tag-retention branch vacuum makes unreachable, and Monty's Rust-thread print callback. Fix test_find_config_file_user_config, which wrote its config into the cwd it had chdir'd to, so the cwd branch answered first and the user-directory lookup it names was never exercised.
This commit is contained in:
parent
7c120587f0
commit
f6acb65e95
25 changed files with 817 additions and 16 deletions
|
|
@ -1,6 +1,10 @@
|
|||
# Changelog
|
||||
## [Unreleased]
|
||||
|
||||
### Changed
|
||||
|
||||
- Test coverage is enforced at 100% via `fail_under` in `[tool.coverage.report]`.
|
||||
|
||||
### Fixed
|
||||
|
||||
- `Store.set_haiku_version` stamps the store's own config into a recreated settings row instead of the process-global `Config`.
|
||||
|
|
|
|||
|
|
@ -570,7 +570,8 @@ def _extract_pdf_attachments(
|
|||
for i in range(attachment_count):
|
||||
att = pdf.get_attachment(i)
|
||||
name = att.get_name()
|
||||
if not name:
|
||||
if not name: # pragma: no cover - pypdfium2 cannot produce a
|
||||
# nameless attachment, so no craftable PDF reaches this.
|
||||
continue
|
||||
data = bytes(att.get_data())
|
||||
child_uri = f"{parent_uri}#attachment={quote(name, safe='')}"
|
||||
|
|
|
|||
|
|
@ -38,7 +38,8 @@ 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 is installed in the test
|
||||
# environment, and get_package_versions depends on it.
|
||||
pass
|
||||
|
||||
# HuggingFace tokenizer
|
||||
|
|
|
|||
|
|
@ -606,7 +606,9 @@ class Store:
|
|||
for tag in tags.values()
|
||||
if tag["version"] in timestamps
|
||||
]
|
||||
if not tagged:
|
||||
if not tagged: # pragma: no cover - a tag's version is never absent from
|
||||
# list_versions: vacuum retains every version at or after the
|
||||
# oldest tag, so a tagged version is never cleaned away.
|
||||
return retention
|
||||
|
||||
# LanceDB version timestamps are naive datetimes in local time.
|
||||
|
|
|
|||
|
|
@ -74,7 +74,8 @@ 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 - docling's ProvenanceItem
|
||||
# always carries a bbox; guards against a shape change.
|
||||
continue
|
||||
bounding_boxes.append(
|
||||
BoundingBox(
|
||||
|
|
|
|||
|
|
@ -146,6 +146,7 @@ omit = [
|
|||
|
||||
[tool.coverage.report]
|
||||
show_missing = true
|
||||
fail_under = 100
|
||||
exclude_also = [
|
||||
"if TYPE_CHECKING:",
|
||||
"@abstractmethod",
|
||||
|
|
|
|||
File diff suppressed because one or more lines are too long
File diff suppressed because one or more lines are too long
File diff suppressed because one or more lines are too long
|
|
@ -1003,3 +1003,35 @@ async def test_replace_for_document_with_no_items_deletes_existing(temp_db_path)
|
|||
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
|
||||
|
|
|
|||
|
|
@ -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,24 @@ 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()))
|
||||
|
||||
with pytest.raises(asyncio.CancelledError):
|
||||
await outer
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -482,6 +482,9 @@ async def test_chunk_repository_get_by_id_and_list_all_pagination(temp_db_path):
|
|||
assert len(first) == 1
|
||||
assert first[0].id == everything[0].id
|
||||
|
||||
# offset is applied even when it selects the whole set
|
||||
assert len(await client.chunk_repository.list_all(offset=0)) == len(everything)
|
||||
|
||||
if len(everything) > 1:
|
||||
second = await client.chunk_repository.list_all(limit=1, offset=1)
|
||||
assert len(second) == 1
|
||||
|
|
@ -526,4 +529,4 @@ async def test_process_search_results_rejects_unknown_score_column(temp_db_path)
|
|||
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()) # ty: ignore[invalid-argument-type]
|
||||
await client.chunk_repository._process_search_results(_Frame())
|
||||
|
|
|
|||
|
|
@ -2518,3 +2518,54 @@ async def test_visualize_chunk_falls_back_when_expansion_drops_refs(temp_db_path
|
|||
images = await client.visualize_chunk(stored[0])
|
||||
|
||||
assert len(images) == 1
|
||||
|
||||
|
||||
@pytest.mark.vcr()
|
||||
async def test_import_documents_schedules_vacuum_when_enabled(temp_db_path):
|
||||
"""A batch import with auto_vacuum on queues a background vacuum."""
|
||||
from haiku.rag.config import AppConfig
|
||||
|
||||
config = AppConfig()
|
||||
config.storage.auto_vacuum = True
|
||||
|
||||
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)
|
||||
await client.import_documents(
|
||||
[
|
||||
DocumentImport(
|
||||
docling_document=docling_doc,
|
||||
chunks=chunks,
|
||||
uri="test://batch-vacuum",
|
||||
)
|
||||
]
|
||||
)
|
||||
|
||||
assert client._vacuum_tasks
|
||||
|
||||
|
||||
@pytest.mark.vcr()
|
||||
async def test_reingesting_a_source_applies_an_explicit_title(temp_db_path):
|
||||
"""Re-adding an unchanged source with a new title updates just the title."""
|
||||
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"
|
||||
|
||||
|
||||
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")
|
||||
|
|
|
|||
|
|
@ -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,26 @@ 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.environment == AppConfig().environment
|
||||
|
|
|
|||
|
|
@ -1373,3 +1373,128 @@ class TestSpanInWindow:
|
|||
# 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 "untouched" in [r.content for r in expanded]
|
||||
|
||||
|
||||
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"
|
||||
|
|
|
|||
|
|
@ -1672,7 +1672,7 @@ class TestDoclingServeZipParsing:
|
|||
@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( # ty: ignore[invalid-assignment]
|
||||
converter.client.submit_and_poll_zip = AsyncMock(
|
||||
side_effect=AssertionError("must not call docling-serve")
|
||||
)
|
||||
|
||||
|
|
@ -1680,3 +1680,18 @@ class TestDoclingServeZipParsing:
|
|||
|
||||
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)
|
||||
|
|
|
|||
|
|
@ -43,3 +43,26 @@ 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"
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async 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_is_callable_on_the_client(temp_db_path):
|
||||
"""The public vacuum() delegates to the store."""
|
||||
from haiku.rag.client import HaikuRAG
|
||||
|
||||
async with HaikuRAG(temp_db_path, create=True) as client:
|
||||
await client.vacuum()
|
||||
|
|
|
|||
|
|
@ -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)"
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -355,3 +355,36 @@ class TestVectorIndexCreation:
|
|||
|
||||
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):
|
||||
async with Store(temp_db_path, create=True) as store:
|
||||
async with store._vacuum_lock:
|
||||
# Returns immediately rather than blocking on the held lock.
|
||||
await store.vacuum()
|
||||
|
||||
@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")
|
||||
|
|
|
|||
|
|
@ -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]
|
||||
|
|
|
|||
|
|
@ -211,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
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize(
|
||||
"make_source,match",
|
||||
[
|
||||
(lambda d: (d / "missing.md").as_uri(), "File does not exist"),
|
||||
(lambda d: _write_unsupported(d), "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))
|
||||
|
||||
|
||||
def _write_unsupported(directory):
|
||||
target = directory / "thing.sqlite3"
|
||||
target.write_bytes(b"binary")
|
||||
return target.as_uri()
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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):
|
||||
|
|
|
|||
Loading…
Reference in a new issue