Batch-capable document and document-item repositories
This commit is contained in:
Yiorgis Gozadinos
parent
f83aad3cc4
commit
b43843f862
3 changed files with 152 additions and 31 deletions
|
|
@ -1,5 +1,6 @@
|
|||
import json
|
||||
from datetime import datetime
|
||||
from typing import overload
|
||||
from uuid import uuid4
|
||||
|
||||
from lancedb.index import BTree
|
||||
|
|
@ -61,17 +62,8 @@ class DocumentRepository:
|
|||
else datetime.now(),
|
||||
)
|
||||
|
||||
async def create(self, entity: Document) -> Document:
|
||||
"""Create a document in the database."""
|
||||
self.store._assert_writable()
|
||||
# Generate new UUID
|
||||
doc_id = str(uuid4())
|
||||
|
||||
# Create timestamp
|
||||
now = datetime.now().isoformat()
|
||||
|
||||
# Create document record
|
||||
doc_record = DocumentRecord(
|
||||
def _to_record(self, entity: Document, doc_id: str, now: str) -> DocumentRecord:
|
||||
return DocumentRecord(
|
||||
id=doc_id,
|
||||
content=entity.content,
|
||||
uri=entity.uri,
|
||||
|
|
@ -84,13 +76,46 @@ class DocumentRepository:
|
|||
updated_at=now,
|
||||
)
|
||||
|
||||
# Add to table
|
||||
await self.store.documents_table.add([doc_record])
|
||||
@overload
|
||||
async def create(self, entity: Document) -> Document: ...
|
||||
|
||||
entity.id = doc_id
|
||||
entity.created_at = datetime.fromisoformat(now)
|
||||
entity.updated_at = datetime.fromisoformat(now)
|
||||
return entity
|
||||
@overload
|
||||
async def create(self, entity: list[Document]) -> list[Document]: ...
|
||||
|
||||
async def create(
|
||||
self, entity: Document | list[Document]
|
||||
) -> Document | list[Document]:
|
||||
"""Create one or more documents in the database.
|
||||
|
||||
A list is written in a single table version regardless of length.
|
||||
"""
|
||||
self.store._assert_writable()
|
||||
|
||||
if isinstance(entity, Document):
|
||||
doc_id = str(uuid4())
|
||||
now = datetime.now().isoformat()
|
||||
await self.store.documents_table.add([self._to_record(entity, doc_id, now)])
|
||||
entity.id = doc_id
|
||||
entity.created_at = datetime.fromisoformat(now)
|
||||
entity.updated_at = datetime.fromisoformat(now)
|
||||
return entity
|
||||
|
||||
documents = entity
|
||||
if not documents:
|
||||
return []
|
||||
|
||||
now = datetime.now().isoformat()
|
||||
created_at = datetime.fromisoformat(now)
|
||||
records = []
|
||||
for document in documents:
|
||||
doc_id = str(uuid4())
|
||||
records.append(self._to_record(document, doc_id, now))
|
||||
document.id = doc_id
|
||||
document.created_at = created_at
|
||||
document.updated_at = created_at
|
||||
|
||||
await self.store.documents_table.add(records)
|
||||
return documents
|
||||
|
||||
async def get_by_id(self, entity_id: str) -> Document | None:
|
||||
"""Get a document by its ID."""
|
||||
|
|
|
|||
|
|
@ -37,26 +37,36 @@ class DocumentItemRepository:
|
|||
tree_depth=row.get("tree_depth", 0) or 0,
|
||||
)
|
||||
|
||||
def _to_record(self, document_id: str, item: DocumentItem) -> DocumentItemRecord:
|
||||
return DocumentItemRecord(
|
||||
document_id=document_id,
|
||||
position=item.position,
|
||||
self_ref=item.self_ref,
|
||||
label=item.label,
|
||||
text=item.text,
|
||||
page_numbers=json.dumps(item.page_numbers),
|
||||
picture_data=item.picture_data,
|
||||
heading_level=item.heading_level,
|
||||
tree_depth=item.tree_depth,
|
||||
)
|
||||
|
||||
async def create_items(self, document_id: str, items: list[DocumentItem]) -> None:
|
||||
"""Bulk insert items for a document."""
|
||||
if not items:
|
||||
return
|
||||
|
||||
self.store._assert_writable()
|
||||
records = [
|
||||
DocumentItemRecord(
|
||||
document_id=document_id,
|
||||
position=item.position,
|
||||
self_ref=item.self_ref,
|
||||
label=item.label,
|
||||
text=item.text,
|
||||
page_numbers=json.dumps(item.page_numbers),
|
||||
picture_data=item.picture_data,
|
||||
heading_level=item.heading_level,
|
||||
tree_depth=item.tree_depth,
|
||||
)
|
||||
for item in items
|
||||
]
|
||||
records = [self._to_record(document_id, item) for item in items]
|
||||
await self.store.document_items_table.add(records)
|
||||
|
||||
async def create_all(self, items: list[DocumentItem]) -> None:
|
||||
"""Bulk insert items spanning any number of documents in a single
|
||||
table version, keyed by each item's own ``document_id``."""
|
||||
if not items:
|
||||
return
|
||||
|
||||
self.store._assert_writable()
|
||||
records = [self._to_record(item.document_id, item) for item in items]
|
||||
await self.store.document_items_table.add(records)
|
||||
|
||||
async def get_all_items(self, document_id: str) -> list[DocumentItem]:
|
||||
|
|
|
|||
|
|
@ -2,7 +2,9 @@ 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
|
||||
|
|
@ -93,6 +95,90 @@ async def test_document_list_with_filter(qa_corpus: list[dict[str, str]], temp_d
|
|||
}
|
||||
|
||||
|
||||
@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 = {
|
||||
|
|
|
|||
Loading…
Reference in a new issue