陈辅元
Cursor
@@ -0,0 +1,164 @@
|
||||
"""Render block groups into chunk markdown strings with layout metadata."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from rag_cut.models import Block, BlockType, Chunk
|
||||
from rag_cut.parsers.pdf.tables import table_meta_summary
|
||||
|
||||
|
||||
def block_to_layout_dict(block: Block) -> dict:
|
||||
"""Serialize a block for chunk metadata and frontend positional rendering."""
|
||||
entry: dict = {
|
||||
"type": block.type.value,
|
||||
"order_index": block.meta.get("order_index"),
|
||||
"page": block.meta.get("page"),
|
||||
"pages": block.meta.get("pages"),
|
||||
"bbox": block.meta.get("bbox"),
|
||||
"bboxes": block.meta.get("bboxes"),
|
||||
"parent_heading": block.meta.get("parent_heading"),
|
||||
"nearest_heading": block.meta.get("nearest_heading"),
|
||||
"bound_heading": block.meta.get("bound_heading"),
|
||||
"preceding_text": block.meta.get("preceding_text"),
|
||||
"following_text": block.meta.get("following_text"),
|
||||
}
|
||||
if block.type == BlockType.HEADING:
|
||||
entry["text"] = block.text
|
||||
entry["level"] = block.level
|
||||
elif block.type == BlockType.IMAGE:
|
||||
entry["text"] = block.text
|
||||
entry["image_path"] = block.image_path
|
||||
entry["image_id"] = block.image_id
|
||||
entry["ocr_text"] = block.ocr_text
|
||||
elif block.type == BlockType.TABLE:
|
||||
entry["text"] = block.markdown or block.text
|
||||
entry["markdown"] = block.markdown or block.text
|
||||
entry["ocr_text"] = block.ocr_text
|
||||
entry["image_path"] = block.image_path
|
||||
entry["image_id"] = block.image_id
|
||||
entry["crop_path"] = block.meta.get("crop_path")
|
||||
entry.update(table_meta_summary(block.meta))
|
||||
entry["embedding_text"] = block.meta.get("embedding_text")
|
||||
else:
|
||||
entry["text"] = block.text or block.markdown
|
||||
return entry
|
||||
|
||||
|
||||
def collect_chunk_layout_meta(blocks: list[Block]) -> dict:
|
||||
"""Aggregate page/bbox/image/table metadata for a chunk group."""
|
||||
pages = sorted({b.meta.get("page") for b in blocks if b.meta.get("page") is not None})
|
||||
for b in blocks:
|
||||
for p in b.meta.get("pages") or []:
|
||||
if p is not None:
|
||||
pages.append(p)
|
||||
pages = sorted(set(pages))
|
||||
|
||||
bboxes = [
|
||||
{
|
||||
"order_index": b.meta.get("order_index"),
|
||||
"page": b.meta.get("page"),
|
||||
"bbox": b.meta.get("bbox"),
|
||||
"type": b.type.value,
|
||||
}
|
||||
for b in blocks
|
||||
if b.meta.get("bbox")
|
||||
]
|
||||
images = [
|
||||
{
|
||||
"order_index": b.meta.get("order_index"),
|
||||
"image_id": b.image_id,
|
||||
"image_path": b.image_path,
|
||||
"page": b.meta.get("page"),
|
||||
"bbox": b.meta.get("bbox"),
|
||||
"ocr_text": b.ocr_text,
|
||||
"bound_heading": b.meta.get("bound_heading"),
|
||||
"preceding_text": b.meta.get("preceding_text"),
|
||||
"following_text": b.meta.get("following_text"),
|
||||
}
|
||||
for b in blocks
|
||||
if b.type == BlockType.IMAGE
|
||||
]
|
||||
tables = [
|
||||
{
|
||||
"order_index": b.meta.get("order_index"),
|
||||
"table_title": b.meta.get("table_title"),
|
||||
"image_path": b.image_path or b.meta.get("crop_path"),
|
||||
"crop_path": b.meta.get("crop_path"),
|
||||
"page": b.meta.get("page"),
|
||||
"pages": b.meta.get("pages"),
|
||||
"bbox": b.meta.get("bbox"),
|
||||
"bboxes": b.meta.get("bboxes"),
|
||||
"markdown": b.markdown,
|
||||
"ocr_text": b.ocr_text,
|
||||
"footnotes": b.meta.get("footnotes"),
|
||||
"keywords": b.meta.get("keywords"),
|
||||
"nearest_heading": b.meta.get("nearest_heading"),
|
||||
"chapter": b.meta.get("chapter"),
|
||||
"row_count": b.meta.get("row_count"),
|
||||
"col_count": b.meta.get("col_count"),
|
||||
"header_rows": b.meta.get("header_rows"),
|
||||
"table_source": b.meta.get("table_source"),
|
||||
"preceding_text": b.meta.get("preceding_text"),
|
||||
"following_text": b.meta.get("following_text"),
|
||||
"cross_page": b.meta.get("cross_page"),
|
||||
"embedding_text": b.meta.get("embedding_text"),
|
||||
}
|
||||
for b in blocks
|
||||
if b.type == BlockType.TABLE
|
||||
]
|
||||
|
||||
meta: dict = {
|
||||
"blocks": [block_to_layout_dict(b) for b in blocks],
|
||||
"bboxes": bboxes,
|
||||
"images": images,
|
||||
"tables": tables,
|
||||
}
|
||||
if pages:
|
||||
meta["pages"] = pages
|
||||
if len(pages) == 1:
|
||||
meta["page"] = pages[0]
|
||||
|
||||
headings = [b.text for b in blocks if b.type == BlockType.HEADING]
|
||||
if headings:
|
||||
meta["heading"] = headings[0]
|
||||
meta["nearest_heading"] = headings[0]
|
||||
elif blocks and blocks[0].meta.get("nearest_heading"):
|
||||
meta["nearest_heading"] = blocks[0].meta.get("nearest_heading")
|
||||
|
||||
table_blocks = [b for b in blocks if b.type == BlockType.TABLE]
|
||||
if table_blocks:
|
||||
primary = table_blocks[0]
|
||||
meta["table_title"] = primary.meta.get("table_title") or meta.get("nearest_heading")
|
||||
if not meta.get("chunk_strategy"):
|
||||
meta["chunk_strategy"] = "table_with_context"
|
||||
meta["retrieval"] = meta.get("retrieval", True)
|
||||
if primary.meta.get("embedding_text"):
|
||||
meta["embedding_text"] = primary.meta["embedding_text"]
|
||||
if primary.meta.get("keywords"):
|
||||
meta["keywords"] = primary.meta["keywords"]
|
||||
|
||||
order_indices = [b.meta.get("order_index") for b in blocks if b.meta.get("order_index") is not None]
|
||||
if order_indices:
|
||||
meta["order_range"] = meta.get("order_range") or [min(order_indices), max(order_indices)]
|
||||
|
||||
return meta
|
||||
|
||||
|
||||
def render_blocks(blocks: list[Block], index: int, meta: dict | None = None) -> Chunk:
|
||||
"""Render blocks in reading order: heading → body → table/image → notes."""
|
||||
parts: list[str] = []
|
||||
for block in blocks:
|
||||
rendered = block.render().strip()
|
||||
if rendered:
|
||||
parts.append(rendered)
|
||||
content = "\n\n".join(parts)
|
||||
|
||||
layout_meta = collect_chunk_layout_meta(blocks)
|
||||
chunk_meta = {**(meta or {}), **layout_meta}
|
||||
|
||||
return Chunk(
|
||||
index=index,
|
||||
content=content,
|
||||
char_count=len(content),
|
||||
block_types=[b.type.value for b in blocks],
|
||||
meta=chunk_meta,
|
||||
)
|
||||
Reference in New Issue
Block a user