Fix defects found reviewing the coverage work
check_source_accessible narrowed its handler to ValueError, but Path.exists re-raises errno values outside its ignored set (EACCES, ENAMETOOLONG). Those were swallowed before and now escaped into the rebuild sweep the guard exists to protect. Catch OSError too. Restore the arity guard in _common_path_prefix: without it an empty list raises from min() and a single label yields a prefix covering the whole path. Two tests would have hung rather than failed on regression (the vacuum skip and the protected-wait cancellation); both are now bounded. The import vacuum test raced against the done-callback that discards the task, and now spies on the call instead, with a negative control. Replace assertions that could not fail: blank-query search against an empty corpus, a batch flush counted against an empty table, a picture description asserting its own input state, and an FS scheme check with nothing on disk to resolve. The get_model matrix asserted only the returned type across 26 cases and now pins the per-provider settings. The three batching tests now count flushes, which revealed embed-only writes through chunks_table.add rather than _flush_rebuild_batch.
This commit is contained in:
Yiorgis Gozadinos
parent
7ba78fdde3
commit
f96a428ef1
22 changed files with 518 additions and 97 deletions
|
|
@ -908,8 +908,10 @@ async def update_document(
|
|||
def check_source_accessible(uri: str) -> bool:
|
||||
"""Check if a document's source URI is accessible.
|
||||
|
||||
A stored URI that no longer parses (``urlparse`` rejects malformed IPv6
|
||||
hosts) counts as inaccessible rather than aborting the caller's sweep.
|
||||
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)
|
||||
|
|
@ -918,5 +920,5 @@ def check_source_accessible(uri: str) -> bool:
|
|||
elif parsed_url.scheme in ("http", "https", "s3"):
|
||||
return True
|
||||
return False
|
||||
except ValueError:
|
||||
except (ValueError, OSError):
|
||||
return False
|
||||
|
|
|
|||
|
|
@ -281,6 +281,9 @@ 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
|
||||
|
|
|
|||
|
|
@ -330,6 +330,8 @@ 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: # pragma: no cover - families always have >=2 members
|
||||
return ""
|
||||
lo, hi = min(labels), max(labels)
|
||||
end = 0
|
||||
while end < len(lo) and lo[end] == hi[end]:
|
||||
|
|
|
|||
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
File diff suppressed because one or more lines are too long
|
|
@ -338,6 +338,11 @@ async def test_fetch_rejects_foreign_scheme(tmp_path):
|
|||
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")
|
||||
|
||||
|
|
|
|||
|
|
@ -713,19 +713,16 @@ async def test_head_returns_none_for_empty_multistatus():
|
|||
assert await src.head("https://nc.example.com/dav/a.md") is None
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_entry_with_status_but_no_prop_has_no_revision():
|
||||
"""A 200 propstat carrying no <prop> still yields an entry, without a revision."""
|
||||
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
|
||||
|
||||
def handler(request: httpx.Request) -> httpx.Response:
|
||||
return httpx.Response(207, content=_raw_multistatus(_STATUS_WITHOUT_PROP))
|
||||
entries = _parse_multistatus(_raw_multistatus(_STATUS_WITHOUT_PROP))
|
||||
|
||||
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
|
||||
assert len(entries) == 1
|
||||
assert entries[0].revision is None
|
||||
assert _parse_multistatus(_raw_multistatus(_NO_HREF)) == []
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
|
|
|
|||
|
|
@ -380,9 +380,20 @@ class TestVfsReadPaths:
|
|||
|
||||
sandbox = Sandbox(temp_db_path, AppConfig(), AnalysisContext())
|
||||
|
||||
content = await _read_vfs_text(sandbox, f"/documents/{doc_id}/content.txt")
|
||||
# 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()
|
||||
|
||||
assert content == "the stored body"
|
||||
async with HaikuRAG(temp_db_path, create=False) as client:
|
||||
await client.update_document(doc_id, content="the rewritten body")
|
||||
|
||||
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:
|
||||
|
|
|
|||
|
|
@ -427,5 +427,7 @@ async def test_wait_protected_reraises_when_recovery_itself_is_cancelled():
|
|||
|
||||
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 outer
|
||||
await asyncio.wait_for(outer, timeout=5)
|
||||
|
|
|
|||
|
|
@ -457,12 +457,14 @@ async def test_ensure_fts_index_warns_on_failure(temp_db_path):
|
|||
|
||||
|
||||
@pytest.mark.vcr()
|
||||
async def test_chunk_repository_get_by_id_and_list_all_pagination(temp_db_path):
|
||||
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:
|
||||
doc = await client.create_document(
|
||||
content="First paragraph.\n\nSecond paragraph.\n\nThird paragraph."
|
||||
)
|
||||
# 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)
|
||||
|
|
@ -482,18 +484,23 @@ 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)
|
||||
# Fail loudly if the fixture stops producing enough chunks to page.
|
||||
assert len(everything) >= 2
|
||||
|
||||
if len(everything) > 1:
|
||||
second = await client.chunk_repository.list_all(limit=1, offset=1)
|
||||
assert len(second) == 1
|
||||
assert second[0].id == everything[1].id
|
||||
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 has nothing to search for."""
|
||||
"""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(" ") == []
|
||||
|
||||
|
||||
|
|
|
|||
|
|
@ -862,7 +862,10 @@ class TestDoclingServeChunkerRefResolution:
|
|||
return [
|
||||
{
|
||||
"raw_text": "body",
|
||||
"doc_items": [{"self_ref": "#/texts/0", "label": "paragraph"}],
|
||||
# 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"}],
|
||||
}
|
||||
]
|
||||
|
||||
|
|
@ -872,4 +875,4 @@ class TestDoclingServeChunkerRefResolution:
|
|||
|
||||
assert len(chunks) == 1
|
||||
assert chunks[0].metadata["doc_item_refs"] == ["#/texts/0"]
|
||||
assert chunks[0].metadata["labels"] == ["paragraph"]
|
||||
assert chunks[0].metadata["labels"] == ["caption"]
|
||||
|
|
|
|||
|
|
@ -1,3 +1,4 @@
|
|||
import asyncio
|
||||
import json
|
||||
import tempfile
|
||||
import threading
|
||||
|
|
@ -2521,32 +2522,40 @@ async def test_visualize_chunk_falls_back_when_expansion_drops_refs(temp_db_path
|
|||
|
||||
|
||||
@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."""
|
||||
@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 = True
|
||||
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)
|
||||
await client.import_documents(
|
||||
[
|
||||
DocumentImport(
|
||||
docling_document=docling_doc,
|
||||
chunks=chunks,
|
||||
uri="test://batch-vacuum",
|
||||
)
|
||||
]
|
||||
)
|
||||
|
||||
assert client._vacuum_tasks
|
||||
# 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 an unchanged source with a new title updates just the title."""
|
||||
"""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"
|
||||
|
|
@ -2563,6 +2572,7 @@ async def test_reingesting_a_source_applies_an_explicit_title(temp_db_path):
|
|||
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):
|
||||
|
|
|
|||
|
|
@ -540,4 +540,4 @@ def test_load_default_config_falls_back_to_builtin_defaults(monkeypatch):
|
|||
|
||||
config = _load_default_config()
|
||||
|
||||
assert config.environment == AppConfig().environment
|
||||
assert config.model_dump() == AppConfig().model_dump()
|
||||
|
|
|
|||
|
|
@ -1466,7 +1466,13 @@ class TestExpandWithItemsWindowEdges:
|
|||
rag.document_item_repository, doc.id, [resolvable, unmatched], 5000
|
||||
)
|
||||
|
||||
assert "untouched" in [r.content for r in expanded]
|
||||
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():
|
||||
|
|
|
|||
|
|
@ -45,8 +45,7 @@ async def test_operations_work_after_database_created(tmp_path):
|
|||
assert doc.content == "Test content"
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_default_db_path_comes_from_storage_data_dir(tmp_path):
|
||||
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
|
||||
|
|
@ -62,7 +61,12 @@ async def test_default_db_path_comes_from_storage_data_dir(tmp_path):
|
|||
@pytest.mark.asyncio
|
||||
async def test_vacuum_is_callable_on_the_client(temp_db_path):
|
||||
"""The public vacuum() delegates to the store."""
|
||||
from unittest.mock import AsyncMock, patch
|
||||
|
||||
from haiku.rag.client import HaikuRAG
|
||||
|
||||
async with HaikuRAG(temp_db_path, create=True) as client:
|
||||
await client.vacuum()
|
||||
with patch.object(client.store, "vacuum", new=AsyncMock()) as store_vacuum:
|
||||
await client.vacuum()
|
||||
|
||||
store_vacuum.assert_awaited_once()
|
||||
|
|
|
|||
|
|
@ -160,8 +160,8 @@ def test_vector_dim_property_reports_configured_dimension():
|
|||
"provider,env_var",
|
||||
[("voyageai", "VOYAGE_API_KEY"), ("cohere", "CO_API_KEY")],
|
||||
)
|
||||
def test_saas_providers_build_offline(monkeypatch, provider, env_var):
|
||||
"""Construction only wires the SDK; no request is made."""
|
||||
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(
|
||||
|
|
@ -174,6 +174,9 @@ def test_saas_providers_build_offline(monkeypatch, provider, env_var):
|
|||
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():
|
||||
|
|
|
|||
|
|
@ -344,6 +344,11 @@ class TestVectorIndexCreation:
|
|||
|
||||
@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)
|
||||
|
||||
|
|
@ -351,8 +356,10 @@ class TestVectorIndexCreation:
|
|||
raise RuntimeError("index build failed")
|
||||
|
||||
with patch.object(store.chunks_table, "create_index", boom):
|
||||
await store._ensure_vector_index()
|
||||
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)
|
||||
|
||||
|
|
@ -378,10 +385,13 @@ class TestStoreMiscellany:
|
|||
|
||||
@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:
|
||||
# Returns immediately rather than blocking on the held lock.
|
||||
await store.vacuum()
|
||||
# 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):
|
||||
|
|
|
|||
|
|
@ -1071,18 +1071,38 @@ async def test_rebuild_blocks_tag_operations(temp_db_path, monkeypatch):
|
|||
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):
|
||||
|
|
@ -1130,8 +1150,16 @@ async def test_hydrate_skips_documents_deleted_mid_rebuild(temp_db_path):
|
|||
assert [doc async for doc in _hydrate(client, [stored])] == []
|
||||
|
||||
|
||||
async def test_apply_descriptions_skips_pictures_without_text():
|
||||
"""A picture whose generated description is empty is left untouched."""
|
||||
@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
|
||||
|
|
@ -1140,9 +1168,13 @@ async def test_apply_descriptions_skips_pictures_without_text():
|
|||
ref = docling_doc.pictures[0].self_ref
|
||||
document = Document(content="x", uri="test://doc")
|
||||
|
||||
_apply_descriptions_sync(docling_doc, document, {ref: ""})
|
||||
_apply_descriptions_sync(docling_doc, document, {ref: description})
|
||||
|
||||
assert docling_doc.pictures[0].meta is None
|
||||
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
|
||||
|
|
@ -1234,6 +1266,22 @@ async def test_rebuild_embed_only_flushes_in_batches(temp_db_path, monkeypatch):
|
|||
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)
|
||||
|
|
@ -1243,6 +1291,9 @@ async def test_rebuild_embed_only_flushes_in_batches(temp_db_path, monkeypatch):
|
|||
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):
|
||||
|
|
@ -1266,6 +1317,7 @@ 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 = []
|
||||
|
|
@ -1280,6 +1332,8 @@ async def test_rebuild_full_flushes_in_batches(temp_db_path, monkeypatch):
|
|||
|
||||
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):
|
||||
|
|
@ -1324,7 +1378,13 @@ async def test_rebuild_full_flushes_pending_before_source_rebuild(temp_db_path):
|
|||
async for doc_id in client.rebuild_database(mode=RebuildMode.FULL)
|
||||
]
|
||||
|
||||
assert content_doc.id in processed
|
||||
# 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()
|
||||
|
|
@ -1341,6 +1401,7 @@ async def test_rebuild_descriptions_flushes_in_batches(temp_db_path, monkeypatch
|
|||
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}
|
||||
|
|
@ -1366,6 +1427,9 @@ async def test_rebuild_descriptions_flushes_in_batches(temp_db_path, monkeypatch
|
|||
|
||||
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):
|
||||
|
|
@ -1408,7 +1472,10 @@ async def test_rebuild_embed_only_recovers_picture_bytes(
|
|||
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)
|
||||
|
|
@ -1449,3 +1516,11 @@ async def test_rebuild_embed_only_recovers_picture_bytes(
|
|||
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")
|
||||
|
|
|
|||
|
|
@ -140,22 +140,46 @@ Emoji test: 🚀 ✅ 📝"""
|
|||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"kwargs",
|
||||
"kwargs,expected_settings",
|
||||
[
|
||||
{"provider": "ollama", "name": "llama3"},
|
||||
{"provider": "ollama", "name": "gpt-oss", "enable_thinking": False},
|
||||
{"provider": "ollama", "name": "gpt-oss", "enable_thinking": True},
|
||||
{"provider": "ollama", "name": "llama3", "temperature": 0.5, "max_tokens": 100},
|
||||
{"provider": "openai", "name": "gpt-4o"},
|
||||
{"provider": "openai", "name": "o1", "enable_thinking": True},
|
||||
{"provider": "openai", "name": "o1", "enable_thinking": False},
|
||||
{
|
||||
"provider": "openai",
|
||||
"name": "gpt-4o",
|
||||
"enable_thinking": False,
|
||||
"temperature": 0.7,
|
||||
"max_tokens": 500,
|
||||
},
|
||||
({"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",
|
||||
|
|
@ -168,9 +192,17 @@ Emoji test: 🚀 ✅ 📝"""
|
|||
"openai_all_settings",
|
||||
],
|
||||
)
|
||||
def test_get_model_returns_openai_chat_model(kwargs):
|
||||
"""Every ollama and openai configuration resolves to an OpenAIChatModel."""
|
||||
assert isinstance(get_model(ModelConfig(**kwargs)), OpenAIChatModel)
|
||||
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():
|
||||
|
|
@ -296,8 +328,14 @@ def test_get_model_anthropic():
|
|||
|
||||
|
||||
@pytest.mark.skipif(not HAS_ANTHROPIC, reason="Anthropic not installed")
|
||||
@pytest.mark.parametrize("enable_thinking", [True, False])
|
||||
def test_get_model_anthropic_with_thinking(enable_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
|
||||
|
||||
|
|
@ -307,7 +345,10 @@ def test_get_model_anthropic_with_thinking(enable_thinking):
|
|||
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")
|
||||
|
|
@ -343,8 +384,10 @@ def test_get_model_groq():
|
|||
|
||||
|
||||
@pytest.mark.skipif(not HAS_GROQ, reason="Groq not installed")
|
||||
@pytest.mark.parametrize("enable_thinking", [True, False])
|
||||
def test_get_model_groq_with_thinking(enable_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
|
||||
|
||||
|
|
@ -354,7 +397,10 @@ def test_get_model_groq_with_thinking(enable_thinking):
|
|||
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")
|
||||
|
|
@ -370,19 +416,39 @@ def test_get_model_bedrock():
|
|||
|
||||
|
||||
@pytest.mark.skipif(not HAS_BEDROCK, reason="Bedrock not installed")
|
||||
@pytest.mark.parametrize("enable_thinking", [True, False])
|
||||
@pytest.mark.parametrize(
|
||||
"name",
|
||||
"name,enable_thinking,expected_fields",
|
||||
[
|
||||
"anthropic.claude-3-5-sonnet-20241022-v2:0",
|
||||
"openai.o3-mini-v1:0",
|
||||
"qwen.qwen3-32b-v1:0",
|
||||
(
|
||||
"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",
|
||||
("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",
|
||||
],
|
||||
ids=["claude", "o_series", "qwen", "unmapped_family"],
|
||||
)
|
||||
def test_get_model_bedrock_with_thinking(name, enable_thinking):
|
||||
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
|
||||
|
||||
|
|
@ -392,7 +458,16 @@ def test_get_model_bedrock_with_thinking(name, enable_thinking):
|
|||
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():
|
||||
|
|
@ -803,7 +878,12 @@ async def test_render_picture_handles_stored_bytes(stored, renders):
|
|||
|
||||
result = await _render_picture(client, "doc1", "#/pictures/0")
|
||||
|
||||
assert (result is not None) is renders
|
||||
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():
|
||||
|
|
|
|||
Loading…
Reference in a new issue