110adf23ee
build_toc stored item positions, while the sandbox instructions describe item_range as a slice into items.jsonl; a gap in positions pulled the next heading into a section. Ranges are now indices into the position-ordered items, and get_document_section slices by index too. docs/mcp.md names the tools that take sources.
513 lines
20 KiB
Python
513 lines
20 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.chunk import Chunk
|
|
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 TestMetadataJson:
|
|
"""metadata.json is served by a reader callback, like the other VFS files."""
|
|
|
|
async def test_metadata_reader_returns_document_fields(self, temp_db_path):
|
|
async with HaikuRAG(temp_db_path, create=True) as client:
|
|
doc_id = await _empty_doc(client, uri="test://meta", title="Meta Doc")
|
|
|
|
sandbox = Sandbox(temp_db_path, AppConfig(), AnalysisContext())
|
|
raw = await _read_vfs_text(sandbox, f"/documents/{doc_id}/metadata.json")
|
|
meta = json.loads(raw)
|
|
|
|
assert meta["id"] == doc_id
|
|
assert meta["title"] == "Meta Doc"
|
|
assert meta["uri"] == "test://meta"
|
|
|
|
|
|
@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()
|
|
|
|
# Rewrite via the repository rather than client.update_document: the
|
|
# latter re-chunks and re-embeds, which this file deliberately avoids
|
|
# so these tests need no embedding endpoint.
|
|
async with HaikuRAG(temp_db_path, create=False) as client:
|
|
doc.content = "the rewritten body"
|
|
await client.document_repository.update(doc)
|
|
|
|
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_chunks_jsonl_lists_chunks_in_order_with_their_metadata(
|
|
self, temp_db_path
|
|
):
|
|
"""One row per chunk, in chunk order, carrying the stored metadata as
|
|
is; the second read of a document is served from the sandbox's cache."""
|
|
config = AppConfig()
|
|
dim = config.embeddings.model.vector_dim
|
|
async with HaikuRAG(temp_db_path, create=True) as client:
|
|
doc_id = await _empty_doc(client, uri="test://paras", title="Paras")
|
|
for order, para_no in enumerate(["13", "14"]):
|
|
await client.chunk_repository.create(
|
|
Chunk(
|
|
document_id=doc_id,
|
|
content=f"Paragraph {para_no}.",
|
|
embedding=[0.1] * dim,
|
|
order=order,
|
|
metadata={"para_no": para_no, "doc_item_refs": []},
|
|
)
|
|
)
|
|
|
|
sandbox = Sandbox(temp_db_path, config, AnalysisContext())
|
|
first = await _read_vfs_text(sandbox, f"/documents/{doc_id}/chunks.jsonl")
|
|
rows = [json.loads(line) for line in first.split("\n")]
|
|
|
|
assert [row["metadata"]["para_no"] for row in rows] == ["13", "14"]
|
|
assert all(set(row) == {"chunk_id", "metadata"} for row in rows)
|
|
assert rows[0]["metadata"] == {"para_no": "13", "doc_item_refs": []}
|
|
assert sandbox._chunks_jsonl_cache[doc_id] == first
|
|
assert (
|
|
await _read_vfs_text(sandbox, f"/documents/{doc_id}/chunks.jsonl") == first
|
|
)
|
|
|
|
async def test_item_range_is_a_line_slice_into_items_jsonl(self, temp_db_path):
|
|
"""`item_range` indexes lines of items.jsonl, as documented, not item
|
|
positions: a gap in positions must not pull the next heading into a
|
|
section."""
|
|
async with HaikuRAG(temp_db_path, create=True) as client:
|
|
doc_id = await _empty_doc(client, uri="test://slice", title="Slice")
|
|
items = [
|
|
_header(doc_id, 0, 1, "Intro"),
|
|
_para(doc_id, 1),
|
|
_header(doc_id, 3, 1, "Methods"),
|
|
_para(doc_id, 4),
|
|
]
|
|
await client.document_item_repository.create_items(doc_id, items)
|
|
|
|
sandbox = Sandbox(temp_db_path, AppConfig(), AnalysisContext())
|
|
toc = await _read_toc(sandbox, doc_id)
|
|
raw = await _read_vfs_text(sandbox, f"/documents/{doc_id}/items.jsonl")
|
|
lines = raw.split("\n")
|
|
|
|
intro, methods = toc["tree"]
|
|
assert intro["item_range"] == [0, 2]
|
|
assert methods["item_range"] == [2, 4]
|
|
start, end = intro["item_range"]
|
|
assert [json.loads(line)["self_ref"] for line in lines[start:end]] == [
|
|
"#/texts/0",
|
|
"#/texts/1",
|
|
]
|
|
|
|
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"]
|