88 lines
3.2 KiB
Python
88 lines
3.2 KiB
Python
"""Row-based splitting for spreadsheet documents."""
|
|
|
|
from __future__ import annotations
|
|
|
|
from rag_cut.models import Block, BlockType, SplitConfig
|
|
from rag_cut.parsers.pdf.tables import (
|
|
build_table_embedding_text,
|
|
detect_spreadsheet_layout,
|
|
extract_table_keywords,
|
|
rows_to_markdown,
|
|
)
|
|
|
|
|
|
def _chunk_meta(table_block: Block, chunk_rows: list[list[str]], md: str, row_range: list[int]) -> dict:
|
|
"""Build per-chunk metadata without leaking the full source table."""
|
|
header_rows = table_block.meta.get("header_rows") or 1
|
|
keywords = extract_table_keywords(
|
|
table_block.meta.get("table_title") or "",
|
|
md,
|
|
table_block.meta.get("table_description") or "",
|
|
)
|
|
embedding_text = build_table_embedding_text(
|
|
table_title=table_block.meta.get("table_title") or "",
|
|
markdown=md,
|
|
description=table_block.meta.get("table_description") or "",
|
|
keywords=keywords,
|
|
)
|
|
return {
|
|
"sheet": table_block.meta.get("sheet"),
|
|
"table_title": table_block.meta.get("table_title"),
|
|
"table_description": table_block.meta.get("table_description"),
|
|
"header_rows": header_rows,
|
|
"header_row_start": table_block.meta.get("header_row_start"),
|
|
"header_row_end": table_block.meta.get("header_row_end"),
|
|
"data_start_row": table_block.meta.get("data_start_row"),
|
|
"preamble_rows": table_block.meta.get("preamble_rows", 0),
|
|
"row_count": len(chunk_rows),
|
|
"col_count": max((len(r) for r in chunk_rows), default=0),
|
|
"chunk_rows": chunk_rows,
|
|
"row_range": row_range,
|
|
"keywords": keywords,
|
|
"embedding_text": embedding_text,
|
|
}
|
|
|
|
|
|
def split_by_row(blocks: list[Block], config: SplitConfig) -> list[list[Block]]:
|
|
table_block = next((b for b in blocks if b.type == BlockType.TABLE and b.meta.get("rows")), None)
|
|
if not table_block:
|
|
return [blocks] if blocks else []
|
|
|
|
rows: list[list[str]] = table_block.meta["rows"]
|
|
layout = {
|
|
"header_row_start": table_block.meta.get("header_row_start"),
|
|
"header_row_end": table_block.meta.get("header_row_end"),
|
|
"data_start_row": table_block.meta.get("data_start_row"),
|
|
"header_rows": table_block.meta.get("header_rows"),
|
|
}
|
|
if not layout["header_row_start"]:
|
|
layout = detect_spreadsheet_layout(rows)
|
|
|
|
h_start = max(1, config.header_row_start) - 1
|
|
h_end = max(h_start + 1, config.header_row_end)
|
|
header = rows[h_start:h_end]
|
|
header_rows = layout.get("header_rows") or len(header)
|
|
data_start = max(config.start_row - 1, h_end)
|
|
data_rows = rows[data_start:]
|
|
rows_per = max(1, config.rows_per_chunk)
|
|
|
|
groups: list[list[Block]] = []
|
|
for i in range(0, len(data_rows), rows_per):
|
|
slice_rows = data_rows[i : i + rows_per]
|
|
chunk_rows = header + slice_rows
|
|
md = rows_to_markdown(chunk_rows, header_rows=header_rows)
|
|
row_range = [
|
|
data_start + i + 1,
|
|
data_start + i + len(slice_rows),
|
|
]
|
|
groups.append(
|
|
[
|
|
Block(
|
|
type=BlockType.TABLE,
|
|
markdown=md,
|
|
meta=_chunk_meta(table_block, chunk_rows, md, row_range),
|
|
)
|
|
]
|
|
)
|
|
return groups
|