Resolve DoclingDocument

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):
     """

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)

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 == []

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