Consolidate duplicated client-side tests
Parametrize sibling tests that differed only in a literal value, and fold two strict-subset tests into the survivors that already covered their scenario. Every case that ran before still runs; the union of assertions is applied to each case, strengthening list_all, get_pages_data and resolve_doc_items. Replace four hand-rolled log-capture handlers with a shared capture_logs() contextmanager in conftest. 13 fewer test functions, 348 fewer lines.
This commit is contained in:
Yiorgis Gozadinos
parent
94e03d777a
commit
e2b273ee2a
11 changed files with 237 additions and 585 deletions
File diff suppressed because one or more lines are too long
File diff suppressed because one or more lines are too long
|
|
@ -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"
|
||||
|
|
|
|||
|
|
@ -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,8 @@ 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)
|
||||
|
|
|
|||
|
|
@ -298,23 +298,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 +309,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 +630,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 +655,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."""
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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):
|
||||
|
|
|
|||
|
|
@ -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)
|
||||
|
|
|
|||
|
|
@ -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,
|
||||
|
|
|
|||
Loading…
Reference in a new issue