67 lines
2.2 KiB
Python
67 lines
2.2 KiB
Python
from typing import TYPE_CHECKING
|
|
|
|
from pydantic import BaseModel
|
|
|
|
if TYPE_CHECKING:
|
|
from docling_core.types.doc.document import DocItem, DoclingDocument
|
|
|
|
|
|
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] = []
|
|
|
|
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):
|
|
"""
|
|
Represents a chunk with content, metadata, and optional document information.
|
|
"""
|
|
|
|
id: str | None = None
|
|
document_id: str | None = None
|
|
content: str
|
|
metadata: dict = {}
|
|
order: int = 0
|
|
document_uri: str | None = None
|
|
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)
|