147 lines
4.6 KiB
Python
147 lines
4.6 KiB
Python
from typing import TYPE_CHECKING
|
|
|
|
from pydantic import BaseModel
|
|
|
|
if TYPE_CHECKING:
|
|
from docling_core.types.doc.document import DocItem, DoclingDocument
|
|
|
|
|
|
class BoundingBox(BaseModel):
|
|
"""Bounding box coordinates for visual grounding."""
|
|
|
|
page_no: int
|
|
left: float
|
|
top: float
|
|
right: float
|
|
bottom: float
|
|
|
|
|
|
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
|
|
|
|
def resolve_bounding_boxes(
|
|
self, docling_document: "DoclingDocument"
|
|
) -> list[BoundingBox]:
|
|
"""Resolve doc_item_refs to bounding boxes for visual grounding.
|
|
|
|
Args:
|
|
docling_document: The parent DoclingDocument containing the items.
|
|
|
|
Returns:
|
|
List of BoundingBox objects from resolved DocItems' provenance.
|
|
"""
|
|
bounding_boxes = []
|
|
for doc_item in self.resolve_doc_items(docling_document):
|
|
prov = getattr(doc_item, "prov", None)
|
|
if not prov:
|
|
continue
|
|
for prov_item in prov:
|
|
bbox = getattr(prov_item, "bbox", None)
|
|
if bbox is None:
|
|
continue
|
|
bounding_boxes.append(
|
|
BoundingBox(
|
|
page_no=prov_item.page_no,
|
|
left=bbox.l,
|
|
top=bbox.t,
|
|
right=bbox.r,
|
|
bottom=bbox.b,
|
|
)
|
|
)
|
|
return bounding_boxes
|
|
|
|
|
|
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)
|
|
|
|
|
|
class SearchResult(BaseModel):
|
|
"""Search result with optional provenance information for citations."""
|
|
|
|
content: str
|
|
score: float
|
|
chunk_id: str | None = None
|
|
document_id: str | None = None
|
|
document_uri: str | None = None
|
|
document_title: str | None = None
|
|
doc_item_refs: list[str] = []
|
|
page_numbers: list[int] = []
|
|
headings: list[str] | None = None
|
|
labels: list[str] = []
|
|
bounding_boxes: list[BoundingBox] | None = None
|
|
|
|
@classmethod
|
|
def from_chunk(
|
|
cls,
|
|
chunk: "Chunk",
|
|
score: float,
|
|
bounding_boxes: list[BoundingBox] | None = None,
|
|
) -> "SearchResult":
|
|
"""Create from a Chunk with optional bounding boxes."""
|
|
meta = chunk.get_chunk_metadata()
|
|
return cls(
|
|
content=chunk.content,
|
|
score=score,
|
|
chunk_id=chunk.id,
|
|
document_id=chunk.document_id,
|
|
document_uri=chunk.document_uri,
|
|
document_title=chunk.document_title,
|
|
doc_item_refs=meta.doc_item_refs,
|
|
page_numbers=meta.page_numbers,
|
|
headings=meta.headings,
|
|
labels=meta.labels,
|
|
bounding_boxes=bounding_boxes,
|
|
)
|