Update Document & Chunk schemas to include DoclingDocument references.

haiku_rag_slim/haiku/rag/store/engine.py

@@ -22,6 +22,8 @@ class DocumentRecord(LanceModel):
     uri: str | None = None
     title: str | None = None
     metadata: str = Field(default="{}")
+    docling_document_json: str | None = None
+    docling_version: str | None = None
     created_at: str = Field(default_factory=lambda: "")
     updated_at: str = Field(default_factory=lambda: "")
 

haiku_rag_slim/haiku/rag/store/models/chunk.py

@@ -1,6 +1,25 @@
 from pydantic import BaseModel
 
 
+class ChunkMetadata(BaseModel):
+    """
+    Structured metadata for a chunk, including DoclingDocument references.
+
+    Attributes:
+        doc_item_refs: JSON pointer references to DocItems in the parent DoclingDocument
+                       (e.g., ["#/texts/5", "#/texts/6", "#/tables/0"])
+        headings: Section heading hierarchy for this chunk
+                  (e.g., ["Chapter 1", "Section 1.1"])
+        labels: Semantic labels for each doc_item (e.g., ["paragraph", "table"])
+        page_numbers: Page numbers where the chunk content appears
+    """
+
+    doc_item_refs: list[str] = []
+    headings: list[str] | None = None
+    labels: list[str] = []
+    page_numbers: list[int] = []
+
+
 class Chunk(BaseModel):
     """
     Represents a chunk with content, metadata, and optional document information.
@@ -15,3 +34,7 @@ class Chunk(BaseModel):
     document_title: str | None = None
     document_meta: dict = {}
     embedding: list[float] | None = None
+
+    def get_chunk_metadata(self) -> ChunkMetadata:
+        """Parse metadata dict into structured ChunkMetadata."""
+        return ChunkMetadata.model_validate(self.metadata)

haiku_rag_slim/haiku/rag/store/models/document.py

@@ -13,5 +13,7 @@ class Document(BaseModel):
     uri: str | None = None
     title: str | None = None
     metadata: dict = {}
+    docling_document_json: str | None = None
+    docling_version: str | None = None
     created_at: datetime = Field(default_factory=datetime.now)
     updated_at: datetime = Field(default_factory=datetime.now)

haiku_rag_slim/haiku/rag/store/repositories/document.py

@@ -37,6 +37,8 @@ class DocumentRepository:
             uri=record.uri,
             title=record.title,
             metadata=json.loads(record.metadata),
+            docling_document_json=record.docling_document_json,
+            docling_version=record.docling_version,
             created_at=datetime.fromisoformat(record.created_at)
             if record.created_at
             else datetime.now(),
@@ -60,6 +62,8 @@ class DocumentRepository:
             uri=entity.uri,
             title=entity.title,
             metadata=json.dumps(entity.metadata),
+            docling_document_json=entity.docling_document_json,
+            docling_version=entity.docling_version,
             created_at=now,
             updated_at=now,
         )
@@ -102,6 +106,8 @@ class DocumentRepository:
                 "uri": entity.uri,
                 "title": entity.title,
                 "metadata": json.dumps(entity.metadata),
+                "docling_document_json": entity.docling_document_json,
+                "docling_version": entity.docling_version,
                 "updated_at": now,
             },
         )

tests/test_chunk.py

@@ -5,7 +5,7 @@ from haiku.rag.client import HaikuRAG
 from haiku.rag.config import Config
 from haiku.rag.converters import get_converter
 from haiku.rag.store.engine import Store
-from haiku.rag.store.models.chunk import Chunk
+from haiku.rag.store.models.chunk import Chunk, ChunkMetadata
 from haiku.rag.store.models.document import Document
 from haiku.rag.store.repositories.chunk import ChunkRepository
 from haiku.rag.store.repositories.document import DocumentRepository
@@ -197,3 +197,37 @@ async def test_adjacent_chunks(temp_db_path):
         assert chunk.document_id == created_document.id
 
     store.close()
+
+
+def test_chunk_metadata_parsing():
+    """Test ChunkMetadata parsing from chunk metadata dict."""
+    metadata_dict = {
+        "doc_item_refs": ["#/texts/0", "#/texts/1", "#/tables/0"],
+        "headings": ["Chapter 1", "Section 1.1"],
+        "labels": ["paragraph", "paragraph", "table"],
+        "page_numbers": [1, 1, 2],
+    }
+
+    chunk = Chunk(
+        content="Test content",
+        metadata=metadata_dict,
+    )
+
+    chunk_meta = chunk.get_chunk_metadata()
+
+    assert isinstance(chunk_meta, ChunkMetadata)
+    assert chunk_meta.doc_item_refs == ["#/texts/0", "#/texts/1", "#/tables/0"]
+    assert chunk_meta.headings == ["Chapter 1", "Section 1.1"]
+    assert chunk_meta.labels == ["paragraph", "paragraph", "table"]
+    assert chunk_meta.page_numbers == [1, 1, 2]
+
+
+def test_chunk_metadata_defaults():
+    """Test ChunkMetadata with empty/default values."""
+    chunk = Chunk(content="Test content", metadata={})
+    chunk_meta = chunk.get_chunk_metadata()
+
+    assert chunk_meta.doc_item_refs == []
+    assert chunk_meta.headings is None
+    assert chunk_meta.labels == []
+    assert chunk_meta.page_numbers == []