From a5d8af07e993fd8f59dbce219e9a3e82f619b2df Mon Sep 17 00:00:00 2001 From: Yiorgis Gozadinos Date: Fri, 28 Nov 2025 11:16:15 +0200 Subject: [PATCH] Resolve DoclingDocument --- .../haiku/rag/store/models/chunk.py | 27 +++++ .../haiku/rag/store/models/document.py | 17 ++++ tests/test_chunk.py | 98 +++++++++++++++++++ tests/test_document.py | 43 ++++++++ 4 files changed, 185 insertions(+) diff --git a/haiku_rag_slim/haiku/rag/store/models/chunk.py b/haiku_rag_slim/haiku/rag/store/models/chunk.py index 57069cd5..576ffd80 100644 --- a/haiku_rag_slim/haiku/rag/store/models/chunk.py +++ b/haiku_rag_slim/haiku/rag/store/models/chunk.py @@ -1,5 +1,10 @@ +from typing import TYPE_CHECKING + from pydantic import BaseModel +if TYPE_CHECKING: + from docling_core.types.doc.document import DocItem, DoclingDocument + class ChunkMetadata(BaseModel): """ @@ -19,6 +24,28 @@ class ChunkMetadata(BaseModel): labels: list[str] = [] page_numbers: list[int] = [] + def resolve_doc_items(self, docling_document: "DoclingDocument") -> list["DocItem"]: + """Resolve doc_item_refs to actual DocItem objects. + + Args: + docling_document: The parent DoclingDocument containing the items. + + Returns: + List of resolved DocItem objects. Items that fail to resolve are skipped. + """ + from docling_core.types.doc.document import RefItem + + doc_items = [] + for ref in self.doc_item_refs: + try: + ref_item = RefItem.model_validate({"$ref": ref}) + doc_item = ref_item.resolve(docling_document) + doc_items.append(doc_item) + except Exception: + # Graceful degradation: skip refs that can't be resolved + continue + return doc_items + class Chunk(BaseModel): """ diff --git a/haiku_rag_slim/haiku/rag/store/models/document.py b/haiku_rag_slim/haiku/rag/store/models/document.py index 1430f1c5..7f73f904 100644 --- a/haiku_rag_slim/haiku/rag/store/models/document.py +++ b/haiku_rag_slim/haiku/rag/store/models/document.py @@ -1,7 +1,11 @@ from datetime import datetime +from typing import TYPE_CHECKING from pydantic import BaseModel, Field +if TYPE_CHECKING: + from docling_core.types.doc.document import DoclingDocument + class Document(BaseModel): """ @@ -17,3 +21,16 @@ class Document(BaseModel): docling_version: str | None = None created_at: datetime = Field(default_factory=datetime.now) updated_at: datetime = Field(default_factory=datetime.now) + + def get_docling_document(self) -> "DoclingDocument | None": + """Parse and return the stored DoclingDocument. + + Returns: + The parsed DoclingDocument, or None if not stored. + """ + if self.docling_document_json is None: + return None + + from docling_core.types.doc.document import DoclingDocument + + return DoclingDocument.model_validate_json(self.docling_document_json) diff --git a/tests/test_chunk.py b/tests/test_chunk.py index 78e74c51..380f8980 100644 --- a/tests/test_chunk.py +++ b/tests/test_chunk.py @@ -231,3 +231,101 @@ def test_chunk_metadata_defaults(): assert chunk_meta.headings is None assert chunk_meta.labels == [] assert chunk_meta.page_numbers == [] + + +def test_chunk_metadata_resolve_doc_items(): + """Test resolving doc_item_refs to actual DocItem objects.""" + from docling_core.types.doc.document import DoclingDocument + + # Create a minimal DoclingDocument with some text items + doc_json = { + "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": []}, + } + docling_doc = DoclingDocument.model_validate(doc_json) + + # Create chunk metadata with refs + chunk_meta = ChunkMetadata( + doc_item_refs=["#/texts/0", "#/texts/1"], + labels=["paragraph", "title"], + ) + + # Resolve refs + doc_items = chunk_meta.resolve_doc_items(docling_doc) + + assert len(doc_items) == 2 + assert getattr(doc_items[0], "text") == "First text" + assert getattr(doc_items[1], "text") == "Second text" + + +def test_chunk_metadata_resolve_doc_items_graceful_degradation(): + """Test that invalid refs are skipped gracefully.""" + from docling_core.types.doc.document import DoclingDocument + + doc_json = { + "name": "test_doc", + "texts": [ + { + "self_ref": "#/texts/0", + "text": "Only text", + "orig": "Only text", + "label": "paragraph", + }, + ], + "tables": [], + "pictures": [], + "groups": [], + "body": {"self_ref": "#/body", "children": []}, + "furniture": {"self_ref": "#/furniture", "children": []}, + } + docling_doc = DoclingDocument.model_validate(doc_json) + + # Create chunk metadata with one valid and one invalid ref + chunk_meta = ChunkMetadata( + doc_item_refs=["#/texts/0", "#/texts/999", "#/invalid/path"], + ) + + # Resolve refs - invalid ones should be skipped + doc_items = chunk_meta.resolve_doc_items(docling_doc) + + assert len(doc_items) == 1 + assert getattr(doc_items[0], "text") == "Only text" + + +def test_chunk_metadata_resolve_empty_refs(): + """Test resolving with no refs returns empty list.""" + from docling_core.types.doc.document import DoclingDocument + + doc_json = { + "name": "test_doc", + "texts": [], + "tables": [], + "pictures": [], + "groups": [], + "body": {"self_ref": "#/body", "children": []}, + "furniture": {"self_ref": "#/furniture", "children": []}, + } + docling_doc = DoclingDocument.model_validate(doc_json) + + chunk_meta = ChunkMetadata() + doc_items = chunk_meta.resolve_doc_items(docling_doc) + + assert doc_items == [] diff --git a/tests/test_document.py b/tests/test_document.py index 97855a39..f2ccb4c8 100644 --- a/tests/test_document.py +++ b/tests/test_document.py @@ -138,3 +138,46 @@ async def test_document_list_with_filter(qa_corpus: Dataset, temp_db_path): assert {doc.id for doc in example_documents} == {created_doc1.id, created_doc3.id} store.close() + + +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 + + document = Document( + content="Test content", + docling_document_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_json is None + assert document.get_docling_document() is None