diff --git a/haiku_rag_slim/haiku/rag/store/engine.py b/haiku_rag_slim/haiku/rag/store/engine.py index 060bba7a..5a7cb79d 100644 --- a/haiku_rag_slim/haiku/rag/store/engine.py +++ b/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: "") diff --git a/haiku_rag_slim/haiku/rag/store/models/chunk.py b/haiku_rag_slim/haiku/rag/store/models/chunk.py index 42ea1368..57069cd5 100644 --- a/haiku_rag_slim/haiku/rag/store/models/chunk.py +++ b/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) diff --git a/haiku_rag_slim/haiku/rag/store/models/document.py b/haiku_rag_slim/haiku/rag/store/models/document.py index e23d5d3c..1430f1c5 100644 --- a/haiku_rag_slim/haiku/rag/store/models/document.py +++ b/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) diff --git a/haiku_rag_slim/haiku/rag/store/repositories/document.py b/haiku_rag_slim/haiku/rag/store/repositories/document.py index 82062058..3320c3ec 100644 --- a/haiku_rag_slim/haiku/rag/store/repositories/document.py +++ b/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, }, ) diff --git a/tests/test_chunk.py b/tests/test_chunk.py index cbb341fa..78e74c51 100644 --- a/tests/test_chunk.py +++ b/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 == []