f96a428ef1
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.
539 lines
19 KiB
Python
539 lines
19 KiB
Python
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()
|
|
async def test_chunk_repository_operations(
|
|
qa_corpus: list[dict[str, str]], temp_db_path
|
|
):
|
|
"""Test ChunkRepository operations."""
|
|
async with HaikuRAG(db_path=temp_db_path, config=Config, create=True) as client:
|
|
# Get the first document from the corpus
|
|
first_doc = qa_corpus[0]
|
|
document_text = first_doc["document_extracted"]
|
|
|
|
# Create a document first with chunks
|
|
created_document = await client.create_document(
|
|
content=document_text, metadata={"source": "test"}
|
|
)
|
|
assert created_document.id is not None
|
|
|
|
# Test getting chunks by document ID
|
|
chunks = await client.chunk_repository.get_by_document_id(created_document.id)
|
|
assert len(chunks) > 0
|
|
assert all(chunk.document_id == created_document.id for chunk in chunks)
|
|
|
|
# Test chunk search
|
|
results = await client.chunk_repository.search(
|
|
"election", limit=2, search_type="vector"
|
|
)
|
|
assert len(results) <= 2
|
|
assert all(hasattr(chunk, "content") for chunk, _ in results)
|
|
|
|
# Test deleting chunks by document ID
|
|
deleted = await client.chunk_repository.delete_by_document_id(
|
|
created_document.id
|
|
)
|
|
assert deleted is True
|
|
|
|
# Verify chunks are gone
|
|
chunks_after_delete = await client.chunk_repository.get_by_document_id(
|
|
created_document.id
|
|
)
|
|
assert len(chunks_after_delete) == 0
|
|
|
|
|
|
@pytest.mark.vcr()
|
|
async def test_chunk_repository_pagination(
|
|
qa_corpus: list[dict[str, str]], temp_db_path
|
|
):
|
|
"""Test ChunkRepository pagination with get_by_document_id and count_by_document_id."""
|
|
async with HaikuRAG(db_path=temp_db_path, config=Config, create=True) as client:
|
|
# Get the first document from the corpus (should produce multiple chunks)
|
|
first_doc = qa_corpus[0]
|
|
document_text = first_doc["document_extracted"]
|
|
|
|
# Create a document with chunks
|
|
created_document = await client.create_document(
|
|
content=document_text, metadata={"source": "test"}
|
|
)
|
|
assert created_document.id is not None
|
|
|
|
# Get total chunk count
|
|
total_count = await client.chunk_repository.count_by_document_id(
|
|
created_document.id
|
|
)
|
|
assert total_count > 0
|
|
|
|
# Get all chunks without pagination
|
|
all_chunks = await client.chunk_repository.get_by_document_id(
|
|
created_document.id
|
|
)
|
|
assert len(all_chunks) == total_count
|
|
|
|
# Test pagination with limit
|
|
limit = min(2, total_count)
|
|
first_batch = await client.chunk_repository.get_by_document_id(
|
|
created_document.id, limit=limit
|
|
)
|
|
assert len(first_batch) == limit
|
|
assert first_batch[0].id == all_chunks[0].id
|
|
|
|
# Test pagination with offset
|
|
if total_count > limit:
|
|
second_batch = await client.chunk_repository.get_by_document_id(
|
|
created_document.id, limit=limit, offset=limit
|
|
)
|
|
assert len(second_batch) <= limit
|
|
assert second_batch[0].id == all_chunks[limit].id
|
|
|
|
# Test offset beyond available chunks
|
|
empty_batch = await client.chunk_repository.get_by_document_id(
|
|
created_document.id, limit=10, offset=total_count + 100
|
|
)
|
|
assert len(empty_batch) == 0
|
|
|
|
|
|
@pytest.mark.vcr()
|
|
async def test_chunking_pipeline(qa_corpus: list[dict[str, str]], temp_db_path):
|
|
"""Test document chunking using client primitives."""
|
|
from haiku.rag.client import HaikuRAG
|
|
from haiku.rag.embeddings import embed_chunks
|
|
|
|
async with HaikuRAG(db_path=temp_db_path, create=True) as client:
|
|
# Get the first document from the corpus
|
|
first_doc = qa_corpus[0]
|
|
document_text = first_doc["document_extracted"]
|
|
|
|
# Use client primitives: convert → chunk → embed
|
|
docling_document = await client.convert(document_text)
|
|
chunks = await client.chunk(docling_document)
|
|
embedded_chunks = await embed_chunks(chunks, client.embedder)
|
|
|
|
# Verify chunks were created with embeddings
|
|
assert len(chunks) > 0
|
|
assert all(chunk.embedding is None for chunk in chunks) # Before embedding
|
|
assert all(chunk.embedding is not None for chunk in embedded_chunks) # After
|
|
|
|
# Verify chunk order
|
|
for i, chunk in enumerate(chunks):
|
|
assert chunk.order == i
|
|
|
|
|
|
@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."""
|
|
chunk = Chunk(content="Test content", metadata=metadata)
|
|
|
|
chunk_meta = chunk.get_chunk_metadata()
|
|
|
|
assert isinstance(chunk_meta, ChunkMetadata)
|
|
assert chunk_meta.doc_item_refs == refs
|
|
assert chunk_meta.headings == headings
|
|
assert chunk_meta.labels == labels
|
|
assert chunk_meta.page_numbers == page_numbers
|
|
|
|
|
|
@pytest.fixture
|
|
def two_text_docling_doc():
|
|
"""Minimal DoclingDocument with two resolvable text items."""
|
|
from docling_core.types.doc.document import DoclingDocument
|
|
|
|
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": []},
|
|
}
|
|
)
|
|
|
|
|
|
@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."""
|
|
chunk_meta = ChunkMetadata(doc_item_refs=refs)
|
|
|
|
doc_items = chunk_meta.resolve_doc_items(two_text_docling_doc)
|
|
|
|
assert [getattr(item, "text") for item in doc_items] == expected_texts
|
|
|
|
|
|
def test_search_result_from_chunk_preserves_document_meta():
|
|
"""Document metadata flows from Chunk to SearchResult for citation
|
|
consumers (UIs)."""
|
|
chunk = Chunk(
|
|
id="chunk-1",
|
|
document_id="doc-1",
|
|
content="Some content.",
|
|
document_uri="file:///docs/report.pdf",
|
|
document_meta={"source_url": "https://example.org/report/view"},
|
|
)
|
|
|
|
result = SearchResult.from_chunk(chunk, score=0.9)
|
|
|
|
assert result.document_meta == {"source_url": "https://example.org/report/view"}
|
|
|
|
|
|
def test_search_result_format_for_agent_omits_document_meta():
|
|
"""Document metadata is UI plumbing, never shown to the model."""
|
|
result = SearchResult(
|
|
content="Some content.",
|
|
score=0.9,
|
|
chunk_id="chunk-1",
|
|
document_meta={"source_url": "https://example.org/report/view"},
|
|
)
|
|
|
|
formatted = result.format_for_agent(rank=1, total=1)
|
|
|
|
assert "source_url" not in formatted
|
|
assert "https://example.org/report/view" not in formatted
|
|
|
|
|
|
@pytest.fixture
|
|
def rich_search_result():
|
|
"""SearchResult with every optional field populated."""
|
|
return 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],
|
|
)
|
|
|
|
|
|
@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 (
|
|
'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_picture_captions():
|
|
"""Picture captions render as labelled lines so the model can correlate them
|
|
with binary parts (BinaryContent.identifier doesn't survive serialization
|
|
to the OpenAI vision API; insertion order is the only reliable signal)."""
|
|
result = SearchResult(
|
|
content="...surrounding text...",
|
|
score=0.5,
|
|
chunk_id="chunk-xyz",
|
|
labels=["picture", "text"],
|
|
picture_captions={
|
|
"#/pictures/0": "Figure 1. Results from each model.",
|
|
"#/pictures/1": "Figure 2. Projected annual emissions.",
|
|
},
|
|
)
|
|
|
|
formatted = result.format_for_agent(rank=1, total=2)
|
|
|
|
lines = formatted.splitlines()
|
|
cap0 = next(
|
|
i for i, line in enumerate(lines) if "Figure caption (#/pictures/0)" in line
|
|
)
|
|
cap1 = next(
|
|
i for i, line in enumerate(lines) if "Figure caption (#/pictures/1)" in line
|
|
)
|
|
content_line = next(
|
|
i for i, line in enumerate(lines) if line.startswith("Content:")
|
|
)
|
|
assert cap0 < cap1 < content_line
|
|
assert "Figure 1. Results from each model." in formatted
|
|
assert "Figure 2. Projected annual emissions." in formatted
|
|
|
|
|
|
def test_search_result_format_for_agent_no_captions_no_line():
|
|
"""Without picture_captions, no caption lines appear (zero-overhead for text chunks)."""
|
|
result = SearchResult(
|
|
content="prose",
|
|
score=0.5,
|
|
chunk_id="chunk-abc",
|
|
labels=["text"],
|
|
)
|
|
formatted = result.format_for_agent(rank=1, total=1)
|
|
assert "Figure caption" not in formatted
|
|
|
|
|
|
@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(**kwargs)
|
|
|
|
assert "[chunk-abc]" in formatted
|
|
assert present in formatted
|
|
for token in absent:
|
|
assert token not in formatted
|
|
assert "Content:\nSome content here." in formatted
|
|
|
|
|
|
@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",
|
|
**fields,
|
|
)
|
|
|
|
assert expected_source in result.format_for_agent()
|
|
|
|
|
|
@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."""
|
|
result = SearchResult(content="x", score=0.5, labels=labels)
|
|
assert result._get_primary_label() == expected
|
|
|
|
|
|
@pytest.mark.vcr()
|
|
@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:
|
|
chunk = Chunk(
|
|
document_id="test-doc",
|
|
content=content,
|
|
metadata=metadata,
|
|
order=0,
|
|
)
|
|
|
|
embedder = get_embedder(Config)
|
|
embedding = (await embedder.embed_documents([chunk.content]))[0]
|
|
chunk.embedding = embedding
|
|
|
|
await client.chunk_repository.create(chunk)
|
|
|
|
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]
|
|
|
|
assert record["content"] == content
|
|
assert record["content_fts"] == expected_content_fts
|
|
|
|
|
|
async def test_ensure_fts_index_warns_on_failure(temp_db_path):
|
|
"""A failed FTS index build is surfaced at WARNING, not swallowed silently."""
|
|
import logging
|
|
|
|
from haiku.rag.store.repositories import chunk as chunk_module
|
|
|
|
async with HaikuRAG(db_path=temp_db_path, config=Config, create=True) as client:
|
|
repo = client.chunk_repository
|
|
|
|
async def _boom(*_args, **_kwargs):
|
|
raise RuntimeError("index build failed")
|
|
|
|
repo.store.chunks_table.create_index = _boom
|
|
|
|
with capture_logs(chunk_module.logger, logging.WARNING) as records:
|
|
await repo._ensure_fts_index()
|
|
|
|
assert [r for r in records if r.levelno == logging.WARNING]
|
|
assert any("index build failed" in r.getMessage() for r in records)
|
|
|
|
|
|
@pytest.mark.vcr()
|
|
async def test_chunk_repository_get_by_id_and_list_all_pagination(
|
|
qa_corpus: list[dict[str, str]], temp_db_path
|
|
):
|
|
"""get_by_id resolves a stored chunk; list_all honours limit and offset."""
|
|
async with HaikuRAG(db_path=temp_db_path, config=Config, create=True) as client:
|
|
# A corpus document is long enough to chunk more than once, which is
|
|
# what makes the offset assertion below meaningful.
|
|
doc = await client.create_document(content=qa_corpus[0]["document_extracted"])
|
|
assert doc.id is not None
|
|
|
|
stored = await client.chunk_repository.get_by_document_id(doc.id)
|
|
assert stored
|
|
|
|
fetched = await client.get_chunk_by_id(stored[0].id)
|
|
assert fetched is not None
|
|
assert fetched.id == stored[0].id
|
|
assert fetched.content == stored[0].content
|
|
|
|
assert await client.get_chunk_by_id("no-such-chunk") is None
|
|
|
|
everything = await client.chunk_repository.list_all()
|
|
assert len(everything) == len(stored)
|
|
|
|
first = await client.chunk_repository.list_all(limit=1)
|
|
assert len(first) == 1
|
|
assert first[0].id == everything[0].id
|
|
|
|
# Fail loudly if the fixture stops producing enough chunks to page.
|
|
assert len(everything) >= 2
|
|
|
|
second = await client.chunk_repository.list_all(limit=1, offset=1)
|
|
assert len(second) == 1
|
|
assert second[0].id == everything[1].id
|
|
|
|
|
|
@pytest.mark.vcr()
|
|
async def test_chunk_search_returns_empty_for_blank_query(temp_db_path):
|
|
"""A blank query with no precomputed vector short-circuits before searching."""
|
|
async with HaikuRAG(db_path=temp_db_path, config=Config, create=True) as client:
|
|
await client.create_document(content="Searchable body about elections.")
|
|
|
|
# Positive control: the corpus is non-empty, so [] is a real decision
|
|
# rather than the answer to every query.
|
|
assert await client.chunk_repository.search("elections")
|
|
assert await client.chunk_repository.search(" ") == []
|
|
|
|
|
|
@pytest.mark.vcr()
|
|
async def test_chunk_search_with_precomputed_vector_skips_text_query(temp_db_path):
|
|
"""The image-as-query path searches vector-only using a stored embedding."""
|
|
async with HaikuRAG(db_path=temp_db_path, config=Config, create=True) as client:
|
|
doc = await client.create_document(content="Vector-only search target.")
|
|
assert doc.id is not None
|
|
|
|
rows = (await client.store.chunks_table.query().limit(1).to_arrow()).to_pylist()
|
|
stored_vector = list(rows[0]["vector"])
|
|
|
|
results = await client.chunk_repository.search("", query_vector=stored_vector)
|
|
|
|
assert results
|
|
assert any(c.document_id == doc.id for c, _ in results)
|
|
|
|
|
|
async def test_get_chunk_ids_by_self_ref_grouped_without_documents(temp_db_path):
|
|
async with HaikuRAG(db_path=temp_db_path, config=Config, create=True) as client:
|
|
assert await client.chunk_repository.get_chunk_ids_by_self_ref_grouped([]) == {}
|
|
|
|
|
|
async def test_process_search_results_rejects_unknown_score_column(temp_db_path):
|
|
"""A result frame with no recognised score column is a programming error."""
|
|
import pandas as pd
|
|
|
|
async with HaikuRAG(db_path=temp_db_path, config=Config, create=True) as client:
|
|
|
|
class _Frame:
|
|
async def to_pandas(self):
|
|
return pd.DataFrame([{"id": "c1", "content": "x", "metadata": "{}"}])
|
|
|
|
with pytest.raises(ValueError, match="Unknown search result format"):
|
|
await client.chunk_repository._process_search_results(_Frame())
|