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.
431 lines
16 KiB
Python
431 lines
16 KiB
Python
"""Tests for the per-document toc.json view and the heading_level / tree_depth
|
|
fields surfaced in items.jsonl.
|
|
|
|
The TOC is derived from `DocumentItem.heading_level` (positive only) in
|
|
position order. PDF-style corpora (all section_headers at level 1) get a flat
|
|
list of siblings; HTML/markdown corpora with real heading hierarchy get a
|
|
nested tree. Items with no section_header at all produce `tree: []`.
|
|
"""
|
|
|
|
import asyncio
|
|
import json
|
|
from pathlib import PurePosixPath
|
|
|
|
import pytest
|
|
|
|
from haiku.rag.client import HaikuRAG
|
|
from haiku.rag.config.models import AppConfig
|
|
from haiku.rag.sandbox import AnalysisContext, Sandbox
|
|
from haiku.rag.store.models.document import Document
|
|
from haiku.rag.store.models.document_item import DocumentItem
|
|
|
|
|
|
async def _empty_doc(client, *, uri: str, title: str) -> str:
|
|
"""Create a Document row directly via the repository (no chunking, no
|
|
embedder) so these tests can run without a reachable embedding endpoint.
|
|
Returns the document id."""
|
|
doc = await client.document_repository.create(
|
|
Document(content="x", uri=uri, title=title)
|
|
)
|
|
return doc.id
|
|
|
|
|
|
def _para(doc_id: str, pos: int, depth: int = 1) -> DocumentItem:
|
|
return DocumentItem(
|
|
document_id=doc_id,
|
|
position=pos,
|
|
self_ref=f"#/texts/{pos}",
|
|
label="paragraph",
|
|
text=f"para{pos}",
|
|
page_numbers=[1],
|
|
tree_depth=depth,
|
|
)
|
|
|
|
|
|
def _header(
|
|
doc_id: str, pos: int, level: int, text: str, depth: int = 1, page: int = 1
|
|
) -> DocumentItem:
|
|
return DocumentItem(
|
|
document_id=doc_id,
|
|
position=pos,
|
|
self_ref=f"#/texts/{pos}",
|
|
label="section_header",
|
|
text=text,
|
|
page_numbers=[page],
|
|
heading_level=level,
|
|
tree_depth=depth,
|
|
)
|
|
|
|
|
|
async def _read_vfs_text(sandbox: Sandbox, path: str) -> str:
|
|
"""Read a VFS file the way execute() does: the synchronous reader runs on a
|
|
worker thread and bridges DB access back to the (free) calling loop."""
|
|
vfs = await sandbox._build_vfs()
|
|
sandbox._loop = asyncio.get_running_loop()
|
|
return await asyncio.to_thread(vfs.path_read_text, PurePosixPath(path))
|
|
|
|
|
|
async def _read_toc(sandbox: Sandbox, doc_id: str) -> dict:
|
|
raw = await _read_vfs_text(sandbox, f"/documents/{doc_id}/toc.json")
|
|
return json.loads(raw)
|
|
|
|
|
|
async def _read_items_jsonl(sandbox: Sandbox, doc_id: str) -> list[dict]:
|
|
raw = await _read_vfs_text(sandbox, f"/documents/{doc_id}/items.jsonl")
|
|
return [json.loads(line) for line in raw.strip().splitlines()] if raw else []
|
|
|
|
|
|
def _flatten(tree: list[dict]) -> list[dict]:
|
|
out = []
|
|
for node in tree:
|
|
out.append(node)
|
|
out.extend(_flatten(node["children"]))
|
|
return out
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
class TestTocShape:
|
|
"""toc.json builds a section tree from heading_level + position."""
|
|
|
|
async def test_multilevel_tree(self, temp_db_path):
|
|
async with HaikuRAG(temp_db_path, create=True) as client:
|
|
doc_id = await _empty_doc(client, uri="test://multi", title="TOC Test Doc")
|
|
items = [
|
|
_header(doc_id, 0, 1, "Intro"),
|
|
_para(doc_id, 1),
|
|
_header(doc_id, 2, 2, "Background"),
|
|
_para(doc_id, 3),
|
|
_header(doc_id, 4, 3, "Prior Work"),
|
|
_para(doc_id, 5),
|
|
_header(doc_id, 6, 2, "Approach"),
|
|
_para(doc_id, 7),
|
|
_header(doc_id, 8, 1, "Methods"),
|
|
_para(doc_id, 9),
|
|
]
|
|
await client.document_item_repository.create_items(doc_id, items)
|
|
|
|
sandbox = Sandbox(temp_db_path, AppConfig(), AnalysisContext())
|
|
toc = await _read_toc(sandbox, doc_id)
|
|
|
|
assert toc["doc_id"] == doc_id
|
|
assert toc["title"] == "TOC Test Doc"
|
|
tree = toc["tree"]
|
|
# Two roots: Intro (children: Background>{Prior Work}, Approach) and Methods
|
|
assert [n["title"] for n in tree] == ["Intro", "Methods"]
|
|
intro = tree[0]
|
|
assert intro["level"] == 1
|
|
assert intro["item_range"] == [0, 8] # ends at "Methods" position
|
|
assert [c["title"] for c in intro["children"]] == ["Background", "Approach"]
|
|
|
|
background = intro["children"][0]
|
|
assert background["level"] == 2
|
|
# Background covers positions 2..5; "Approach" begins at 6 (same-level sibling)
|
|
assert background["item_range"] == [2, 6]
|
|
assert [c["title"] for c in background["children"]] == ["Prior Work"]
|
|
|
|
prior = background["children"][0]
|
|
assert prior["level"] == 3
|
|
# Prior Work has no descendants and the next same-or-shallower header is
|
|
# "Approach" at level 2, position 6.
|
|
assert prior["item_range"] == [4, 6]
|
|
assert prior["children"] == []
|
|
|
|
approach = intro["children"][1]
|
|
assert approach["item_range"] == [6, 8]
|
|
|
|
methods = tree[1]
|
|
assert methods["item_range"] == [8, 10] # to end of items
|
|
assert methods["children"] == []
|
|
|
|
async def test_flat_pdf_style(self, temp_db_path):
|
|
"""All section_headers at level 1 (PDF reality) -> flat sibling list."""
|
|
async with HaikuRAG(temp_db_path, create=True) as client:
|
|
doc_id = await _empty_doc(client, uri="test://pdf-style", title="Flat PDF")
|
|
items = [
|
|
_header(doc_id, 0, 1, "Chapter 1"),
|
|
_para(doc_id, 1),
|
|
_para(doc_id, 2),
|
|
_header(doc_id, 3, 1, "Chapter 2"),
|
|
_para(doc_id, 4),
|
|
_header(doc_id, 5, 1, "Chapter 3"),
|
|
_para(doc_id, 6),
|
|
]
|
|
await client.document_item_repository.create_items(doc_id, items)
|
|
|
|
sandbox = Sandbox(temp_db_path, AppConfig(), AnalysisContext())
|
|
toc = await _read_toc(sandbox, doc_id)
|
|
|
|
tree = toc["tree"]
|
|
assert [n["title"] for n in tree] == ["Chapter 1", "Chapter 2", "Chapter 3"]
|
|
assert all(n["level"] == 1 and n["children"] == [] for n in tree)
|
|
assert tree[0]["item_range"] == [0, 3]
|
|
assert tree[1]["item_range"] == [3, 5]
|
|
assert tree[2]["item_range"] == [5, 7]
|
|
|
|
async def test_no_headers(self, temp_db_path):
|
|
async with HaikuRAG(temp_db_path, create=True) as client:
|
|
doc_id = await _empty_doc(
|
|
client, uri="test://no-headers", title="No Headers"
|
|
)
|
|
await client.document_item_repository.create_items(
|
|
doc_id, [_para(doc_id, i) for i in range(5)]
|
|
)
|
|
|
|
sandbox = Sandbox(temp_db_path, AppConfig(), AnalysisContext())
|
|
toc = await _read_toc(sandbox, doc_id)
|
|
assert toc["tree"] == []
|
|
|
|
async def test_skip_header_with_zero_level(self, temp_db_path):
|
|
"""A section_header with ``heading_level == 0`` is skipped."""
|
|
async with HaikuRAG(temp_db_path, create=True) as client:
|
|
doc_id = await _empty_doc(client, uri="test://zero", title="Zero Level")
|
|
items = [
|
|
_header(doc_id, 0, 1, "Real H1"),
|
|
_para(doc_id, 1),
|
|
_header(doc_id, 2, 0, "Pre-migration ghost"),
|
|
_para(doc_id, 3),
|
|
]
|
|
await client.document_item_repository.create_items(doc_id, items)
|
|
|
|
sandbox = Sandbox(temp_db_path, AppConfig(), AnalysisContext())
|
|
toc = await _read_toc(sandbox, doc_id)
|
|
titles = [n["title"] for n in _flatten(toc["tree"])]
|
|
assert titles == ["Real H1"]
|
|
assert toc["tree"][0]["item_range"] == [0, 4]
|
|
|
|
async def test_node_shape_has_chunk_ids_not_position(self, temp_db_path):
|
|
async with HaikuRAG(temp_db_path, create=True) as client:
|
|
doc_id = await _empty_doc(client, uri="test://shape", title="Shape")
|
|
await client.document_item_repository.create_items(
|
|
doc_id, [_header(doc_id, 0, 1, "Only"), _para(doc_id, 1)]
|
|
)
|
|
|
|
sandbox = Sandbox(temp_db_path, AppConfig(), AnalysisContext())
|
|
toc = await _read_toc(sandbox, doc_id)
|
|
node = toc["tree"][0]
|
|
expected = {
|
|
"self_ref",
|
|
"level",
|
|
"title",
|
|
"page_numbers",
|
|
"item_range",
|
|
"chunk_ids",
|
|
"children",
|
|
}
|
|
assert expected <= set(node)
|
|
assert "position" not in node
|
|
assert node["chunk_ids"] == []
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
class TestTocChunkIdsAggregation:
|
|
"""toc.json nodes carry the union of chunk_ids covered by their item_range."""
|
|
|
|
async def test_chunk_ids_union_over_item_range(self, temp_db_path, monkeypatch):
|
|
async with HaikuRAG(temp_db_path, create=True) as client:
|
|
doc_id = await _empty_doc(client, uri="test://chunks", title="Chunks")
|
|
await client.document_item_repository.create_items(
|
|
doc_id,
|
|
[
|
|
_header(doc_id, 0, 1, "Intro"),
|
|
_para(doc_id, 1),
|
|
_para(doc_id, 2),
|
|
_header(doc_id, 3, 2, "Background"),
|
|
_para(doc_id, 4),
|
|
_header(doc_id, 5, 1, "Methods"),
|
|
_para(doc_id, 6),
|
|
],
|
|
)
|
|
|
|
from haiku.rag.store.repositories.chunk import ChunkRepository
|
|
|
|
# self_ref → list[chunk_id]. Intro covers #/texts/0..2, Background
|
|
# covers #/texts/3..4, Methods covers #/texts/5..6. Item at #/texts/2
|
|
# belongs to two chunks (cA + cB) — verifying dedup-preserving-order.
|
|
fake_index = {
|
|
"#/texts/1": ["cA"],
|
|
"#/texts/2": ["cA", "cB"],
|
|
"#/texts/3": ["cB"],
|
|
"#/texts/4": ["cC"],
|
|
"#/texts/5": ["cD"],
|
|
"#/texts/6": ["cD"],
|
|
}
|
|
|
|
async def fake_grouped(self, document_ids):
|
|
return {doc_id: fake_index}
|
|
|
|
monkeypatch.setattr(
|
|
ChunkRepository,
|
|
"get_chunk_ids_by_self_ref_grouped",
|
|
fake_grouped,
|
|
)
|
|
|
|
sandbox = Sandbox(temp_db_path, AppConfig(), AnalysisContext())
|
|
toc = await _read_toc(sandbox, doc_id)
|
|
|
|
intro = toc["tree"][0]
|
|
assert intro["title"] == "Intro"
|
|
# Intro spans positions 0..2 (header) plus its child Background's
|
|
# range — item_range is [0, 5].
|
|
assert intro["item_range"] == [0, 5]
|
|
assert intro["chunk_ids"] == ["cA", "cB", "cC"]
|
|
|
|
background = intro["children"][0]
|
|
assert background["title"] == "Background"
|
|
assert background["item_range"] == [3, 5]
|
|
assert background["chunk_ids"] == ["cB", "cC"]
|
|
|
|
methods = toc["tree"][1]
|
|
assert methods["title"] == "Methods"
|
|
assert methods["item_range"] == [5, 7]
|
|
assert methods["chunk_ids"] == ["cD"]
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
class TestTocCaching:
|
|
"""items + toc reads for a doc share one items fetch and one chunk-index fetch."""
|
|
|
|
async def test_items_and_chunk_index_fetched_once_per_doc(
|
|
self, temp_db_path, monkeypatch
|
|
):
|
|
async with HaikuRAG(temp_db_path, create=True) as client:
|
|
doc_id = await _empty_doc(client, uri="test://cache", title="Cache")
|
|
await client.document_item_repository.create_items(
|
|
doc_id,
|
|
[_header(doc_id, 0, 1, "Only"), _para(doc_id, 1), _para(doc_id, 2)],
|
|
)
|
|
|
|
sandbox = Sandbox(temp_db_path, AppConfig(), AnalysisContext())
|
|
|
|
from haiku.rag.store.repositories.chunk import ChunkRepository
|
|
from haiku.rag.store.repositories.document_item import DocumentItemRepository
|
|
|
|
items_calls = {"n": 0}
|
|
chunk_calls = {"n": 0}
|
|
original_items = DocumentItemRepository.get_all_items
|
|
original_chunks = ChunkRepository.get_chunk_ids_by_self_ref_grouped
|
|
|
|
async def counting_items(self, document_id):
|
|
items_calls["n"] += 1
|
|
return await original_items(self, document_id)
|
|
|
|
async def counting_chunks(self, document_ids):
|
|
chunk_calls["n"] += 1
|
|
return await original_chunks(self, document_ids)
|
|
|
|
monkeypatch.setattr(DocumentItemRepository, "get_all_items", counting_items)
|
|
monkeypatch.setattr(
|
|
ChunkRepository, "get_chunk_ids_by_self_ref_grouped", counting_chunks
|
|
)
|
|
|
|
# Items + toc share `_doc_items` and `_doc_chunk_index`. Four reads →
|
|
# one items fetch + one chunk-index fetch.
|
|
_ = await _read_toc(sandbox, doc_id)
|
|
_ = await _read_items_jsonl(sandbox, doc_id)
|
|
_ = await _read_toc(sandbox, doc_id)
|
|
_ = await _read_items_jsonl(sandbox, doc_id)
|
|
|
|
assert items_calls["n"] == 1
|
|
assert chunk_calls["n"] == 1
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
class TestItemsJsonlSurfacesNewFields:
|
|
"""items.jsonl row shape: heading_level is always present (0 on non-headers);
|
|
chunk_ids surfaces each item's containing chunks; position and tree_depth
|
|
are not exposed."""
|
|
|
|
async def test_jsonl_row_shape(self, temp_db_path):
|
|
async with HaikuRAG(temp_db_path, create=True) as client:
|
|
doc_id = await _empty_doc(client, uri="test://jsonl-fields", title="Fields")
|
|
items = [
|
|
_header(doc_id, 0, 1, "H1", depth=2),
|
|
_para(doc_id, 1, depth=3),
|
|
_header(doc_id, 2, 2, "H2", depth=4),
|
|
]
|
|
await client.document_item_repository.create_items(doc_id, items)
|
|
|
|
sandbox = Sandbox(temp_db_path, AppConfig(), AnalysisContext())
|
|
rows = await _read_items_jsonl(sandbox, doc_id)
|
|
|
|
assert len(rows) == 3
|
|
assert rows[0]["heading_level"] == 1
|
|
assert rows[1]["heading_level"] == 0
|
|
assert rows[2]["heading_level"] == 2
|
|
for r in rows:
|
|
expected = {
|
|
"self_ref",
|
|
"label",
|
|
"text",
|
|
"page_numbers",
|
|
"heading_level",
|
|
"chunk_ids",
|
|
}
|
|
assert expected <= set(r)
|
|
assert "position" not in r
|
|
assert "tree_depth" not in r
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
class TestVfsReadPaths:
|
|
"""The synchronous VFS readers bridge back to the event loop; drive them
|
|
through a worker thread the way execute() does."""
|
|
|
|
async def test_content_txt_is_read_lazily_per_document(self, temp_db_path):
|
|
async with HaikuRAG(temp_db_path, create=True) as client:
|
|
doc = await client.document_repository.create(
|
|
Document(content="the stored body", uri="test://body", title="Body")
|
|
)
|
|
doc_id = doc.id
|
|
|
|
sandbox = Sandbox(temp_db_path, AppConfig(), AnalysisContext())
|
|
|
|
# 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()
|
|
|
|
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:
|
|
doc = await client.document_repository.create(
|
|
Document(content="x", uri="test://ro", title="RO")
|
|
)
|
|
doc_id = doc.id
|
|
|
|
sandbox = Sandbox(temp_db_path, AppConfig(), AnalysisContext())
|
|
vfs = await sandbox._build_vfs()
|
|
sandbox._loop = asyncio.get_running_loop()
|
|
|
|
with pytest.raises(PermissionError, match="read-only"):
|
|
await asyncio.to_thread(
|
|
vfs.path_write_text,
|
|
PurePosixPath(f"/documents/{doc_id}/content.txt"),
|
|
"nope",
|
|
)
|
|
|
|
async def test_toc_skips_gaps_in_item_positions(self, temp_db_path):
|
|
"""Positions need not be contiguous — a heading's span may cover
|
|
positions that carry no item."""
|
|
async with HaikuRAG(temp_db_path, create=True) as client:
|
|
doc_id = await _empty_doc(client, uri="test://gaps", title="Gaps")
|
|
# Positions 1 and 2 are absent between the header and the paragraph.
|
|
items = [
|
|
_header(doc_id, 0, 1, "Intro"),
|
|
_para(doc_id, 3),
|
|
]
|
|
await client.document_item_repository.create_items(doc_id, items)
|
|
|
|
sandbox = Sandbox(temp_db_path, AppConfig(), AnalysisContext())
|
|
toc = await _read_toc(sandbox, doc_id)
|
|
|
|
assert [n["title"] for n in toc["tree"]] == ["Intro"]
|