479 lines
16 KiB
Python
479 lines
16 KiB
Python
import pytest
|
|
|
|
from haiku.rag.store.engine import Store
|
|
from haiku.rag.store.models.document import Document
|
|
from haiku.rag.store.models.document_item import DocumentItem
|
|
from haiku.rag.store.repositories.document import DocumentRepository
|
|
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
|
|
):
|
|
"""list_all excludes content and docling_document by default."""
|
|
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.txt",
|
|
title="Test Document",
|
|
metadata={"key": "value"},
|
|
)
|
|
created = await doc_repo.create(doc)
|
|
|
|
docs = await doc_repo.list_all()
|
|
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].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."""
|
|
async with Store(temp_db_path, create=True) as store:
|
|
doc_repo = DocumentRepository(store)
|
|
|
|
first_doc = qa_corpus[0]
|
|
document_text = first_doc["document_extracted"]
|
|
|
|
doc1 = Document(
|
|
content=document_text,
|
|
uri="https://example.com/doc1.txt",
|
|
metadata={"source": "test", "category": "A"},
|
|
)
|
|
doc2 = Document(
|
|
content=document_text,
|
|
uri="https://arxiv.org/paper.pdf",
|
|
metadata={"source": "test", "category": "B"},
|
|
)
|
|
doc3 = Document(
|
|
content=document_text,
|
|
uri="https://example.com/doc3.txt",
|
|
metadata={"source": "test", "category": "A"},
|
|
)
|
|
|
|
created_doc1 = await doc_repo.create(doc1)
|
|
created_doc2 = await doc_repo.create(doc2)
|
|
created_doc3 = await doc_repo.create(doc3)
|
|
|
|
all_documents = await doc_repo.list_all()
|
|
assert len(all_documents) == 3
|
|
|
|
arxiv_documents = await doc_repo.list_all(filter="uri LIKE '%arxiv%'")
|
|
assert len(arxiv_documents) == 1
|
|
assert arxiv_documents[0].id == created_doc2.id
|
|
|
|
example_documents = await doc_repo.list_all(filter="uri LIKE '%example.com%'")
|
|
assert len(example_documents) == 2
|
|
assert {doc.id for doc in example_documents} == {
|
|
created_doc1.id,
|
|
created_doc3.id,
|
|
}
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_document_create_batch(qa_corpus: list[dict[str, str]], temp_db_path):
|
|
"""create accepts a list of documents and writes them in a single version."""
|
|
async with Store(temp_db_path, create=True) as store:
|
|
doc_repo = DocumentRepository(store)
|
|
|
|
content = qa_corpus[0]["document_extracted"]
|
|
doc_a = Document(content=content, uri="https://example.com/a.txt", title="A")
|
|
doc_b = Document(content=content, uri="https://example.com/b.txt", title="B")
|
|
|
|
before = await store.documents_table.version()
|
|
created = await doc_repo.create([doc_a, doc_b])
|
|
after = await store.documents_table.version()
|
|
|
|
assert isinstance(created, list)
|
|
assert len(created) == 2
|
|
assert created[0].id is not None
|
|
assert created[1].id is not None
|
|
assert created[0].id != created[1].id
|
|
assert after - before == 1
|
|
|
|
round_a = await doc_repo.get_by_id(created[0].id)
|
|
round_b = await doc_repo.get_by_id(created[1].id)
|
|
assert round_a is not None and round_a.title == "A"
|
|
assert round_b is not None and round_b.title == "B"
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_document_create_empty_batch(temp_db_path):
|
|
"""create([]) is a no-op returning an empty list with no version bump."""
|
|
async with Store(temp_db_path, create=True) as store:
|
|
doc_repo = DocumentRepository(store)
|
|
|
|
before = await store.documents_table.version()
|
|
created = await doc_repo.create([])
|
|
after = await store.documents_table.version()
|
|
|
|
assert created == []
|
|
assert after == before
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_document_item_create_all(temp_db_path):
|
|
"""create_all writes items spanning multiple documents in a single version."""
|
|
async with Store(temp_db_path, create=True) as store:
|
|
item_repo = DocumentItemRepository(store)
|
|
|
|
items = [
|
|
DocumentItem(
|
|
document_id="doc-1", position=0, self_ref="#/texts/0", text="a"
|
|
),
|
|
DocumentItem(
|
|
document_id="doc-1", position=1, self_ref="#/texts/1", text="b"
|
|
),
|
|
DocumentItem(
|
|
document_id="doc-2", position=0, self_ref="#/texts/0", text="c"
|
|
),
|
|
]
|
|
|
|
before = await store.document_items_table.version()
|
|
await item_repo.create_all(items)
|
|
after = await store.document_items_table.version()
|
|
|
|
assert after - before == 1
|
|
|
|
doc1_items = await item_repo.get_all_items("doc-1")
|
|
doc2_items = await item_repo.get_all_items("doc-2")
|
|
assert [i.text for i in doc1_items] == ["a", "b"]
|
|
assert [i.text for i in doc2_items] == ["c"]
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_document_item_create_all_empty(temp_db_path):
|
|
"""create_all([]) is a no-op with no version bump."""
|
|
async with Store(temp_db_path, create=True) as store:
|
|
item_repo = DocumentItemRepository(store)
|
|
|
|
before = await store.document_items_table.version()
|
|
await item_repo.create_all([])
|
|
after = await store.document_items_table.version()
|
|
|
|
assert after == before
|
|
|
|
|
|
def test_document_get_docling_document():
|
|
"""Test parsing stored DoclingDocument JSON."""
|
|
doc_json = {
|
|
"name": "test_doc",
|
|
"texts": [
|
|
{
|
|
"self_ref": "#/texts/0",
|
|
"text": "Test text",
|
|
"orig": "Test text",
|
|
"label": "paragraph",
|
|
},
|
|
],
|
|
"tables": [],
|
|
"pictures": [],
|
|
"groups": [],
|
|
"body": {"self_ref": "#/body", "children": []},
|
|
"furniture": {"self_ref": "#/furniture", "children": []},
|
|
}
|
|
|
|
import json
|
|
|
|
from haiku.rag.store.compression import compress_json
|
|
|
|
document = Document(
|
|
content="Test content",
|
|
docling_document=compress_json(json.dumps(doc_json)),
|
|
docling_version="1.3.0",
|
|
)
|
|
|
|
docling_doc = document.get_docling_document()
|
|
|
|
assert docling_doc is not None
|
|
assert docling_doc.name == "test_doc"
|
|
assert len(docling_doc.texts) == 1
|
|
assert docling_doc.texts[0].text == "Test text"
|
|
|
|
|
|
def test_document_get_docling_document_none():
|
|
"""Test get_docling_document returns None when not stored."""
|
|
document = Document(content="Test content")
|
|
|
|
assert document.docling_document is None
|
|
assert document.get_docling_document() is None
|
|
|
|
|
|
def test_set_docling_splits_structure_and_pages():
|
|
"""set_docling stores structure and pages separately."""
|
|
import json
|
|
|
|
from docling_core.types.doc.document import DoclingDocument
|
|
from docling_core.types.doc.labels import DocItemLabel
|
|
|
|
from haiku.rag.store.compression import decompress_json
|
|
|
|
docling_doc = DoclingDocument(name="split_test")
|
|
docling_doc.add_text(label=DocItemLabel.PARAGRAPH, text="Hello world")
|
|
|
|
document = Document(content="test")
|
|
document.set_docling(docling_doc)
|
|
|
|
assert document.docling_document is not None
|
|
assert document.docling_version == docling_doc.version
|
|
|
|
# Structure should not contain pages
|
|
structure = json.loads(decompress_json(document.docling_document))
|
|
assert "pages" not in structure
|
|
assert structure["name"] == "split_test"
|
|
|
|
# get_docling_document should work from the split structure
|
|
parsed = document.get_docling_document()
|
|
assert parsed is not None
|
|
assert parsed.name == "split_test"
|
|
assert len(list(parsed.iterate_items())) > 0
|
|
|
|
|
|
def test_set_docling_with_page_images():
|
|
"""set_docling stores page images in docling_pages."""
|
|
import json
|
|
|
|
from docling_core.types.doc.base import Size
|
|
from docling_core.types.doc.document import DoclingDocument, PageItem
|
|
from docling_core.types.doc.labels import DocItemLabel
|
|
|
|
from haiku.rag.store.compression import decompress_json
|
|
|
|
docling_doc = DoclingDocument(name="pages_test")
|
|
docling_doc.add_text(label=DocItemLabel.PARAGRAPH, text="Content")
|
|
docling_doc.pages[1] = PageItem(
|
|
size=Size(width=612, height=792),
|
|
page_no=1,
|
|
)
|
|
|
|
document = Document(content="test")
|
|
document.set_docling(docling_doc)
|
|
|
|
assert document.docling_pages is not None
|
|
|
|
# Pages blob should contain page data
|
|
pages = json.loads(decompress_json(document.docling_pages))
|
|
assert "1" in pages
|
|
|
|
|
|
def test_set_docling_dict_path_matches_string_path():
|
|
"""The ingestion dict path must produce the same stored bytes as the
|
|
legacy string path.
|
|
|
|
``set_docling`` feeds ``compress_docling_data`` the dict from
|
|
``model_dump(mode="json")`` instead of round-tripping through
|
|
``model_dump_json()`` + ``json.loads``. The on-disk format must not change
|
|
— in particular int-keyed ``pages`` must still serialize to string keys.
|
|
"""
|
|
import json
|
|
|
|
from docling_core.types.doc.base import Size
|
|
from docling_core.types.doc.document import DoclingDocument, PageItem
|
|
from docling_core.types.doc.labels import DocItemLabel
|
|
|
|
from haiku.rag.store.compression import (
|
|
compress_docling_data,
|
|
compress_docling_split,
|
|
decompress_json,
|
|
)
|
|
|
|
docling_doc = DoclingDocument(name="equivalence_test")
|
|
docling_doc.add_text(label=DocItemLabel.PARAGRAPH, text="Hello world")
|
|
docling_doc.pages[1] = PageItem(size=Size(width=612, height=792), page_no=1)
|
|
|
|
struct_dict, pages_dict = compress_docling_data(
|
|
docling_doc.model_dump(mode="json")
|
|
)
|
|
struct_str, pages_str = compress_docling_split(docling_doc.model_dump_json())
|
|
|
|
assert json.loads(decompress_json(struct_dict)) == json.loads(
|
|
decompress_json(struct_str)
|
|
)
|
|
assert pages_dict is not None and pages_str is not None
|
|
assert json.loads(decompress_json(pages_dict)) == json.loads(
|
|
decompress_json(pages_str)
|
|
)
|
|
|
|
|
|
def test_get_page_images():
|
|
"""get_page_images returns requested pages from docling_pages blob."""
|
|
import json
|
|
|
|
from haiku.rag.store.compression import compress_json
|
|
|
|
pages_data = {
|
|
"1": {"size": {"width": 612, "height": 792}, "page_no": 1},
|
|
"2": {"size": {"width": 612, "height": 792}, "page_no": 2},
|
|
"3": {"size": {"width": 612, "height": 792}, "page_no": 3},
|
|
}
|
|
document = Document(
|
|
content="test",
|
|
docling_pages=compress_json(json.dumps(pages_data)),
|
|
)
|
|
|
|
result = document.get_page_images([1, 3])
|
|
assert len(result) == 2
|
|
assert 1 in result
|
|
assert 3 in result
|
|
assert 2 not in result
|
|
|
|
# Missing pages are skipped
|
|
result = document.get_page_images([99])
|
|
assert len(result) == 0
|
|
|
|
# None docling_pages returns empty
|
|
doc_no_pages = Document(content="test")
|
|
assert doc_no_pages.get_page_images([1]) == {}
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_get_docling_data_loads_only_docling_columns(
|
|
qa_corpus: list[dict[str, str]], temp_db_path
|
|
):
|
|
"""get_docling_data returns docling blob without loading content."""
|
|
import json
|
|
|
|
from haiku.rag.store.compression import compress_json
|
|
|
|
async with Store(temp_db_path, create=True) as store:
|
|
doc_repo = DocumentRepository(store)
|
|
|
|
doc_json = {
|
|
"name": "test_doc",
|
|
"texts": [],
|
|
"tables": [],
|
|
"pictures": [],
|
|
"groups": [],
|
|
"body": {"self_ref": "#/body", "children": []},
|
|
"furniture": {"self_ref": "#/furniture", "children": []},
|
|
}
|
|
compressed = compress_json(json.dumps(doc_json))
|
|
|
|
doc = Document(
|
|
content=qa_corpus[0]["document_extracted"],
|
|
uri="https://example.com/doc.txt",
|
|
docling_document=compressed,
|
|
docling_version="2.1.0",
|
|
)
|
|
created = await doc_repo.create(doc)
|
|
assert created.id is not None
|
|
|
|
result = await doc_repo.get_docling_data(created.id)
|
|
assert result is not None
|
|
assert result.id == created.id
|
|
assert result.content == ""
|
|
assert result.docling_document == compressed
|
|
assert result.docling_version == "2.1.0"
|
|
|
|
# Verify docling document can be parsed
|
|
docling_doc = result.get_docling_document()
|
|
assert docling_doc is not None
|
|
assert docling_doc.name == "test_doc"
|
|
|
|
# Non-existent ID returns None
|
|
assert await doc_repo.get_docling_data("nonexistent-id") is None
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_get_pages_data_loads_only_pages_column(
|
|
qa_corpus: list[dict[str, str]], temp_db_path
|
|
):
|
|
"""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}})
|
|
)
|
|
|
|
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.txt",
|
|
docling_pages=pages_blob,
|
|
)
|
|
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.content == ""
|
|
assert result.docling_pages == pages_blob
|
|
|
|
# Non-existent ID returns None
|
|
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
|
|
):
|
|
"""Test get_by_uri handles URIs with special characters like single quotes."""
|
|
async with Store(temp_db_path, create=True) as store:
|
|
doc_repo = DocumentRepository(store)
|
|
|
|
first_doc = qa_corpus[0]
|
|
document_text = first_doc["document_extracted"]
|
|
|
|
doc_with_quote = Document(
|
|
content=document_text,
|
|
uri="Hamish and Andy's Gap Year",
|
|
metadata={"source": "test"},
|
|
)
|
|
|
|
created_doc = await doc_repo.create(doc_with_quote)
|
|
|
|
retrieved = await doc_repo.get_by_uri("Hamish and Andy's Gap Year")
|
|
assert retrieved is not None
|
|
assert retrieved.id == created_doc.id
|
|
assert retrieved.uri == "Hamish and Andy's Gap Year"
|