From 7c76c3399eee4a3ecd3b57f086ca61de5e6120dc Mon Sep 17 00:00:00 2001 From: Yiorgis Gozadinos Date: Sun, 26 Jul 2026 19:23:54 +0300 Subject: [PATCH] Cover get_model provider branches and visualize_chunk edge paths Parametrize the get_model tests whose only assertion was the returned model type, extending the table to the reasoning-off, groq-parsed and Bedrock o-series/qwen/unmapped branches. Fold format_citations_no_title into the superset test that already covered its scenario. Add tests for cosine_similarity on parallel vectors, multi-citation Rich rendering, picture rendering with missing and undecodable bytes, a missing docling distribution, and the visualize_chunk short-circuits for absent documents, absent rasters and refless chunks. --- tests/test_client.py | 239 +++++++++++++++++++++++++++++++++++++++++++ tests/test_search.py | 67 +++++++++++- tests/test_utils.py | 209 ++++++++++++++++++++++++------------- 3 files changed, 441 insertions(+), 74 deletions(-) diff --git a/tests/test_client.py b/tests/test_client.py index 73012dc5..fe545765 100644 --- a/tests/test_client.py +++ b/tests/test_client.py @@ -2279,3 +2279,242 @@ def test_check_source_accessible_file_uri(tmp_path): 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 diff --git a/tests/test_search.py b/tests/test_search.py index 982b44f8..6a6e429f 100644 --- a/tests/test_search.py +++ b/tests/test_search.py @@ -480,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) @@ -525,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) == [] diff --git a/tests/test_utils.py b/tests/test_utils.py index 01809e2b..7ecbf13c 100644 --- a/tests/test_utils.py +++ b/tests/test_utils.py @@ -139,27 +139,38 @@ 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) - assert isinstance(result, OpenAIChatModel) +@pytest.mark.parametrize( + "kwargs", + [ + {"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, + }, + ], + 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_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_ollama_appends_v1_to_per_model_base_url(): @@ -183,20 +194,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,14 +296,15 @@ 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", [True, False]) +def test_get_model_anthropic_with_thinking(enable_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) @@ -345,12 +343,15 @@ 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", [True, False]) +def test_get_model_groq_with_thinking(enable_thinking): """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) @@ -369,14 +370,26 @@ 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("enable_thinking", [True, False]) +@pytest.mark.parametrize( + "name", + [ + "anthropic.claude-3-5-sonnet-20241022-v2:0", + "openai.o3-mini-v1:0", + "qwen.qwen3-32b-v1:0", + # A family with no reasoning mapping leaves the request fields untouched. + "meta.llama3-70b-instruct-v1:0", + ], + ids=["claude", "o_series", "qwen", "unmapped_family"], +) +def test_get_model_bedrock_with_thinking(name, enable_thinking): + """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) @@ -390,19 +403,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 +526,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 +745,85 @@ 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") + + assert (result is not None) is renders + + +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"