mirror of
https://github.com/infiniflow/ragflow.git
synced 2026-07-08 12:24:48 +08:00
1e80419c21
fix: restore TitleChunker output for json/chunks upstream formats
## Summary
The refactor commit e194027b (#14247) introduced two regressions that
caused `TitleChunker` to produce zero chunks when the upstream Parser
node outputs `json` or `chunks` format (e.g. PDF parsing).
## Root Cause
### 1. Dead code in `extract_line_records` (critical)
After refactor, when `payload` is `None` (which is the case for `json`
and `chunks` output formats), the method returns an empty list
immediately via `return []`, so no records are ever extracted from
structured upstream output. The original `json`/`chunks` handling code
became unreachable dead code.
### 2. Unconditional overwrite in `build_chunks_from_record_groups`
The `chunks` variable assigned in the `if` branch for markdown/text/html
formats was unconditionally overwritten by the statement below it, due
to a missing `else` keyword.
## Fix
- Remove the premature `return []` so the `json`/`chunks` branch is
reachable again.
- Add `else` branch in `build_chunks_from_record_groups` so the two
format families are handled independently.
## Test Plan
- [x] Verified no lint errors on the changed file
- [ ] Tested with a PDF document parsed via DeepDOC → TitleChunker
pipeline
- [ ] Tested with markdown input through TitleChunker
- [ ] Tested hierarchy and group chunking modes
## Impact
- Fixes the regression where documents parsed with `json`/`chunks`
output format produced no chunks from `TitleChunker`.
- No API or configuration changes. Fully backward compatible.
Signed-off-by: noob <yixiao121314@outlook.com>
340 lines
12 KiB
Python
340 lines
12 KiB
Python
#
|
|
# Copyright 2025 The InfiniFlow Authors. All Rights Reserved.
|
|
#
|
|
# Licensed under the Apache License, Version 2.0 (the "License");
|
|
# you may not use this file except in compliance with the License.
|
|
# You may obtain a copy of the License at
|
|
#
|
|
# http://www.apache.org/licenses/LICENSE-2.0
|
|
#
|
|
# Unless required by applicable law or agreed to in writing, software
|
|
# distributed under the License is distributed on an "AS IS" BASIS,
|
|
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
|
# See the License for the specific language governing permissions and
|
|
# limitations under the License.
|
|
|
|
import random
|
|
import re
|
|
import sys
|
|
from abc import ABC, abstractmethod
|
|
from collections import Counter
|
|
from copy import deepcopy
|
|
|
|
from deepdoc.parser.pdf_parser import RAGFlowPdfParser
|
|
from deepdoc.parser.utils import extract_pdf_outlines
|
|
from rag.flow.base import ProcessBase, ProcessParamBase
|
|
from rag.flow.parser.pdf_chunk_metadata import (
|
|
PDF_POSITIONS_KEY,
|
|
extract_pdf_positions,
|
|
finalize_pdf_chunk,
|
|
merge_pdf_positions,
|
|
restore_pdf_text_previews,
|
|
)
|
|
from rag.nlp import not_bullet, not_title
|
|
|
|
BODY_LEVEL = sys.maxsize - 1
|
|
|
|
|
|
class TitleChunkerParam(ProcessParamBase):
|
|
def __init__(self):
|
|
super().__init__()
|
|
self.levels = []
|
|
self.hierarchy = None
|
|
self.include_heading_content = False
|
|
self.root_chunk_as_heading = False
|
|
|
|
def check(self):
|
|
if self.method in {"hierarchy", "group"}:
|
|
self.check_empty(self.levels, "Hierarchical setups.")
|
|
if self.method == "hierarchy":
|
|
self.check_empty(self.hierarchy, "Hierarchy number.")
|
|
|
|
def get_input_form(self) -> dict[str, dict]:
|
|
return {}
|
|
|
|
|
|
class BaseTitleChunker(ABC):
|
|
start_message = "Start to chunk by title."
|
|
|
|
def __init__(self, process: ProcessBase, from_upstream):
|
|
self.process = process
|
|
self.param = process._param
|
|
self.from_upstream = from_upstream
|
|
|
|
|
|
async def invoke(self):
|
|
self.process.set_output("output_format", "chunks")
|
|
self.process.callback(random.randint(1, 5) / 100.0, self.start_message)
|
|
line_records = self.extract_line_records()
|
|
resolved = self.resolve_levels(line_records)
|
|
chunks = self.build_chunks(line_records, resolved)
|
|
await self.set_chunks(chunks)
|
|
self.process.callback(1, "Done.")
|
|
|
|
|
|
def extract_line_records(self):
|
|
"""
|
|
Normalize all upstream input payloads into a unified ordered record stream.
|
|
All level resolution and chunk construction logic operates on this standard stream,
|
|
decoupling downstream chunking strategies from different upstream output formats.
|
|
"""
|
|
import logging
|
|
logger = logging.getLogger(__name__)
|
|
|
|
payload = None
|
|
# Extract raw content payload based on upstream output format type
|
|
if self.from_upstream.output_format == "markdown":
|
|
payload = self.from_upstream.markdown_result or ""
|
|
elif self.from_upstream.output_format == "text":
|
|
payload = self.from_upstream.text_result or ""
|
|
elif self.from_upstream.output_format == "html":
|
|
payload = self.from_upstream.html_result or ""
|
|
|
|
# Boundary robustness fix: explicit None check to distinguish `None` and empty string ""
|
|
# Prevents empty payload from unexpectedly falling through to structured chunk branch
|
|
if payload is not None:
|
|
lines = payload.split("\n")
|
|
input_line_count = len(lines)
|
|
|
|
# Format-branched text processing to preserve original document semantics
|
|
# Plain text: perform full whitespace stripping and invalid empty line filtering
|
|
if self.from_upstream.output_format == "text":
|
|
clean_lines = [line.strip() for line in lines if line.strip()]
|
|
# Markdown & HTML: retain original indentation/spacing, only filter pure blank lines
|
|
else:
|
|
clean_lines = [line for line in lines if line.strip()]
|
|
|
|
output_line_count = len(clean_lines)
|
|
# Production observability log: added format dimension per project coding guidelines
|
|
logger.info(
|
|
f"payload filter: format={self.from_upstream.output_format} before={input_line_count} after={output_line_count}"
|
|
)
|
|
|
|
return [
|
|
{
|
|
"text": line,
|
|
"doc_type_kwd": "text",
|
|
"img_id": None,
|
|
"layout": "",
|
|
PDF_POSITIONS_KEY: []
|
|
}
|
|
for line in clean_lines
|
|
]
|
|
items = self.from_upstream.chunks if self.from_upstream.output_format == "chunks" else self.from_upstream.json_result
|
|
return [
|
|
{
|
|
"text": item.get("text") or "",
|
|
"doc_type_kwd": str(item.get("doc_type_kwd") or "text"),
|
|
"img_id": item.get("img_id"),
|
|
"layout": "{} {}".format(item.get("layout_type", ""), item.get("layoutno", "")).strip(),
|
|
PDF_POSITIONS_KEY: extract_pdf_positions(item),
|
|
}
|
|
for item in items or []
|
|
]
|
|
|
|
def extract_outlines(self):
|
|
file = self.from_upstream.file or {}
|
|
source = (
|
|
file.get("blob")
|
|
or file.get("binary")
|
|
or file.get("path")
|
|
or file.get("name")
|
|
)
|
|
if not source:
|
|
return []
|
|
return extract_pdf_outlines(source)
|
|
|
|
|
|
@staticmethod
|
|
def match_regex_level(text, level_group):
|
|
stripped = text.strip()
|
|
for level, pattern in enumerate(level_group, start=1):
|
|
if re.match(pattern, stripped) and not not_bullet(stripped):
|
|
return level
|
|
return None
|
|
|
|
|
|
@staticmethod
|
|
def select_level_group(lines, raw_levels):
|
|
if not raw_levels:
|
|
return []
|
|
|
|
# Select one regex family before assigning numeric levels. Mixing
|
|
# patterns across families would make the level numbers ambiguous and
|
|
# break downstream comparisons.
|
|
hits = [0] * len(raw_levels)
|
|
for i, group in enumerate(raw_levels):
|
|
for sec in lines:
|
|
sec = sec.strip()
|
|
if not sec:
|
|
continue
|
|
for pattern in group:
|
|
if re.match(pattern, sec) and not not_bullet(sec):
|
|
hits[i] += 1
|
|
break
|
|
|
|
maximum = 0
|
|
selected = -1
|
|
for i, hit in enumerate(hits):
|
|
if hit <= maximum:
|
|
continue
|
|
selected = i
|
|
maximum = hit
|
|
|
|
if selected < 0:
|
|
return []
|
|
return [pattern for pattern in raw_levels[selected] if pattern]
|
|
|
|
|
|
@staticmethod
|
|
def match_layout_level(text, layout, fallback_level):
|
|
if re.search(r"(section|title|head)", layout, re.I) and not not_title(text.split("@")[0].strip()):
|
|
return fallback_level
|
|
return BODY_LEVEL
|
|
|
|
|
|
@staticmethod
|
|
def _outline_similarity(left, right):
|
|
left_pairs = {left[i] + left[i + 1] for i in range(len(left) - 1)}
|
|
right_pairs = {right[i] + right[i + 1] for i in range(min(len(left), len(right) - 1))}
|
|
return len(left_pairs & right_pairs) / max(len(left_pairs), len(right_pairs), 1)
|
|
|
|
|
|
def resolve_outline_levels(self, line_records):
|
|
outlines = self.extract_outlines()
|
|
if not line_records or len(outlines) / len(line_records) <= 0.03:
|
|
return None
|
|
|
|
max_level = max(level for _, level, _ in outlines) + 1
|
|
levels = []
|
|
for record in line_records:
|
|
if record["doc_type_kwd"] != "text":
|
|
levels.append(BODY_LEVEL)
|
|
continue
|
|
text = record["text"]
|
|
for outline_text, level, _ in outlines:
|
|
if self._outline_similarity(outline_text, text) > 0.8:
|
|
levels.append(level + 1)
|
|
break
|
|
else:
|
|
levels.append(BODY_LEVEL)
|
|
|
|
return {
|
|
"levels": levels,
|
|
"most_level": max(1, max_level - 1),
|
|
"source": "outline",
|
|
}
|
|
|
|
|
|
def resolve_frequency_levels(self, line_records):
|
|
level_group = self.select_level_group(
|
|
[record["text"] for record in line_records],
|
|
self.param.levels,
|
|
)
|
|
fallback_level = len(level_group) + 1
|
|
levels = []
|
|
for record in line_records:
|
|
if record["doc_type_kwd"] != "text":
|
|
levels.append(BODY_LEVEL)
|
|
continue
|
|
level = self.match_regex_level(record["text"], level_group)
|
|
if level is not None:
|
|
levels.append(level)
|
|
continue
|
|
levels.append(
|
|
self.match_layout_level(
|
|
record["text"],
|
|
record["layout"],
|
|
fallback_level,
|
|
)
|
|
)
|
|
|
|
most_level = None
|
|
for level, _ in Counter(levels).most_common():
|
|
if level < BODY_LEVEL:
|
|
most_level = level
|
|
break
|
|
|
|
return {
|
|
"levels": levels,
|
|
"most_level": most_level,
|
|
"source": "frequency",
|
|
}
|
|
|
|
|
|
def resolve_title_levels(self, line_records):
|
|
return self.resolve_outline_levels(line_records) or self.resolve_frequency_levels(line_records)
|
|
|
|
|
|
def build_chunks_from_record_groups(self, record_groups):
|
|
# Strategy code decides record grouping. This method materializes each
|
|
# group into the output chunk representation. For PDF-like inputs, the
|
|
# chunk box is defined by merged source positions and the text payload
|
|
# is normalized by removing parser tags.
|
|
if self.from_upstream.output_format in ["markdown", "text", "html"]:
|
|
chunks = [
|
|
{"text": "".join(record["text"] + "\n" for record in records)}
|
|
for records in record_groups
|
|
if records
|
|
]
|
|
else:
|
|
chunks = [
|
|
(
|
|
{
|
|
"text": RAGFlowPdfParser.remove_tag("".join(record["text"] + "\n" for record in records)),
|
|
"doc_type_kwd": "text",
|
|
PDF_POSITIONS_KEY: merge_pdf_positions(records),
|
|
}
|
|
if records[0]["doc_type_kwd"] == "text"
|
|
else {
|
|
"text": records[0]["text"],
|
|
"doc_type_kwd": records[0]["doc_type_kwd"],
|
|
"img_id": records[0]["img_id"],
|
|
PDF_POSITIONS_KEY: records[0][PDF_POSITIONS_KEY],
|
|
}
|
|
)
|
|
for records in record_groups
|
|
if records
|
|
]
|
|
|
|
if self.param.root_chunk_as_heading and len(chunks) > 1:
|
|
root_chunk = chunks[0]
|
|
root_text = root_chunk.get("text", "")
|
|
|
|
for ck in chunks[1:]:
|
|
ck['text'] = root_text + "\n" + ck.get("text", "")
|
|
|
|
return chunks[1:]
|
|
|
|
return chunks
|
|
|
|
|
|
async def set_chunks(self, chunks):
|
|
if self.from_upstream.output_format in ["markdown", "text", "html"]:
|
|
self.process.set_output("chunks", chunks)
|
|
return
|
|
|
|
# Text grouping runs before visual enrichment. Preview text and final
|
|
# box metadata are derived here from the merged PDF positions.
|
|
await restore_pdf_text_previews(chunks, self.from_upstream, self.process._canvas)
|
|
self.process.set_output("chunks", [finalize_pdf_chunk(deepcopy(chunk)) for chunk in chunks])
|
|
|
|
|
|
@abstractmethod
|
|
def resolve_levels(self, line_records):
|
|
raise NotImplementedError()
|
|
|
|
|
|
@abstractmethod
|
|
def build_chunks(self, line_records, resolved):
|
|
raise NotImplementedError()
|
|
|
|
|
|
def resolve_target_level(levels, hierarchy):
|
|
title_levels = sorted({level for level in levels if 0 < level < BODY_LEVEL})
|
|
if not title_levels:
|
|
return None
|
|
|
|
hierarchy_num = max(int(hierarchy), 1)
|
|
return title_levels[min(hierarchy_num, len(title_levels)) - 1]
|