diff --git a/backend/package/yuxi/knowledge/chunking/ragflow_like/nlp.py b/backend/package/yuxi/knowledge/chunking/ragflow_like/nlp.py index 32eb948f..5d3375d6 100644 --- a/backend/package/yuxi/knowledge/chunking/ragflow_like/nlp.py +++ b/backend/package/yuxi/knowledge/chunking/ragflow_like/nlp.py @@ -55,17 +55,28 @@ def count_tokens(text: str) -> int: return max(1, len(parts)) if text.strip() else 0 -def hard_split_by_token_limit(text: str, chunk_token_num: int) -> list[str]: - """将文本按 token 上限硬切，用于 naive_merge 之后的兜底保护。""" +def hard_split_by_token_limit(text: str, chunk_token_num: int, hard_limit_token_num: int | None = None) -> list[str]: + """将文本按 token 上限硬切，用于 naive_merge 之后的兜底保护。 + + hard_limit_token_num 只在调用方显式传入时生效，用于允许略超目标长度的块 + 保持完整；默认保持严格不超过 chunk_token_num 的历史行为。 + """ token_iter = list(re.finditer(r"[A-Za-z0-9_]+|[一-鿿]", text or "")) if not token_iter: cleaned = (text or "").strip() return [cleaned] if cleaned else [] - chunks: list[str] = [] + max_tokens = max(int(chunk_token_num or 0), 1) + hard_limit = None + if hard_limit_token_num is not None: + hard_limit = max(int(hard_limit_token_num or 0), max_tokens) + if len(token_iter) <= hard_limit: + cleaned = (text or "").strip() + return [cleaned] if cleaned else [] + + spans: list[tuple[int, int]] = [] start = 0 index = 0 - max_tokens = max(int(chunk_token_num or 0), 1) while index < len(token_iter): next_index = min(index + max_tokens, len(token_iter)) @@ -73,16 +84,24 @@ def hard_split_by_token_limit(text: str, chunk_token_num: int) -> list[str]: end = token_iter[next_index].start() else: end = len(text) - piece = text[start:end].strip() - if piece: - chunks.append(piece) + if text[start:end].strip(): + spans.append((start, end)) start = end index = next_index tail = text[start:].strip() if tail: - chunks.append(tail) - return chunks + spans.append((start, len(text))) + + if hard_limit is not None and len(spans) >= 2: + prev_start, _ = spans[-2] + _, tail_end = spans[-1] + candidate = text[prev_start:tail_end].strip() + if count_tokens(candidate) <= hard_limit: + spans[-2] = (prev_start, tail_end) + spans.pop() + + return [text[start:end].strip() for start, end in spans if text[start:end].strip()] def random_choices(arr: list[str], k: int) -> list[str]: diff --git a/backend/package/yuxi/knowledge/chunking/ragflow_like/parsers/book.py b/backend/package/yuxi/knowledge/chunking/ragflow_like/parsers/book.py index 1c1621f5..c71ed47e 100644 --- a/backend/package/yuxi/knowledge/chunking/ragflow_like/parsers/book.py +++ b/backend/package/yuxi/knowledge/chunking/ragflow_like/parsers/book.py @@ -23,6 +23,21 @@ def _iter_sections(markdown_content: str) -> list[tuple[str, str]]: return sections +def _ensure_chunk_token_limit(chunks: list[str], chunk_token_num: int) -> list[str]: + max_tokens = int(chunk_token_num or 0) + normalized = [chunk.strip() for chunk in chunks if chunk and chunk.strip()] + if max_tokens <= 0: + return normalized + + protected: list[str] = [] + for chunk in normalized: + if nlp.count_tokens(chunk) <= max_tokens: + protected.append(chunk) + else: + protected.extend(nlp.hard_split_by_token_limit(chunk, max_tokens)) + return protected + + def chunk_markdown(markdown_content: str, parser_config: dict[str, Any] | None = None) -> list[str]: parser_config = parser_config or {} @@ -51,11 +66,14 @@ def chunk_markdown(markdown_content: str, parser_config: dict[str, Any] | None = ) if chunks: - return chunks + return _ensure_chunk_token_limit(chunks, chunk_token_num) - return nlp.naive_merge( - sections, - chunk_token_num=chunk_token_num, - delimiter=delimiter, - overlapped_percent=overlapped_percent, + return _ensure_chunk_token_limit( + nlp.naive_merge( + sections, + chunk_token_num=chunk_token_num, + delimiter=delimiter, + overlapped_percent=overlapped_percent, + ), + chunk_token_num, ) diff --git a/backend/package/yuxi/knowledge/chunking/ragflow_like/parsers/general.py b/backend/package/yuxi/knowledge/chunking/ragflow_like/parsers/general.py index 2ef9e621..d7e719f8 100644 --- a/backend/package/yuxi/knowledge/chunking/ragflow_like/parsers/general.py +++ b/backend/package/yuxi/knowledge/chunking/ragflow_like/parsers/general.py @@ -4,6 +4,8 @@ from typing import Any from yuxi.knowledge.chunking.ragflow_like import nlp +GENERAL_HARD_LIMIT_RATIO = 1.5 + def _unescape_delimiter(delimiter: str) -> str: return delimiter.replace("\\n", "\n").replace("\\r", "\r").replace("\\t", "\t").replace("\\\\", "\\") @@ -31,11 +33,12 @@ def _iter_sections(markdown_content: str, delimiter: str) -> list[tuple[str, str def _ensure_chunk_token_limit(chunks: list[str], chunk_token_num: int) -> list[str]: - """对输出 chunk 做 token 上限保护：超长的直接硬切。""" + """对输出 chunk 做 token 上限保护：默认 512 token 时允许到 768 再硬切。""" max_tokens = int(chunk_token_num or 0) if max_tokens <= 0: return [c.strip() for c in chunks if c and c.strip()] + hard_limit = max(max_tokens, int(max_tokens * GENERAL_HARD_LIMIT_RATIO)) protected: list[str] = [] for chunk in chunks: cleaned = (chunk or "").strip() @@ -44,7 +47,7 @@ def _ensure_chunk_token_limit(chunks: list[str], chunk_token_num: int) -> list[s if nlp.count_tokens(cleaned) <= max_tokens: protected.append(cleaned) else: - protected.extend(nlp.hard_split_by_token_limit(cleaned, max_tokens)) + protected.extend(nlp.hard_split_by_token_limit(cleaned, max_tokens, hard_limit_token_num=hard_limit)) return protected diff --git a/backend/test/unit/plugins/test_ragflow_like_chunking.py b/backend/test/unit/plugins/test_ragflow_like_chunking.py index 2efbedda..d9d4c535 100644 --- a/backend/test/unit/plugins/test_ragflow_like_chunking.py +++ b/backend/test/unit/plugins/test_ragflow_like_chunking.py @@ -122,6 +122,32 @@ def test_book_chunking_hierarchical_merge() -> None: assert any("第一章" in ck["content"] for ck in chunks) +def test_book_chunking_should_apply_overlength_protection() -> None: + content = "\n".join( + [ + "第一章 总则", + "第一节 适用范围", + "超长正文" * 1200, + "第二节 基本原则", + "应当遵循最小改动原则。", + ] + ) + max_chunk_tokens = 180 + + chunks = chunk_markdown( + markdown_content=content, + file_id="file_book_long", + filename="book.txt", + processing_params={ + "chunk_preset_id": "book", + "chunk_parser_config": {"chunk_token_num": max_chunk_tokens, "delimiter": "\\n"}, + }, + ) + + assert len(chunks) > 1 + assert max(count_tokens(ck["content"]) for ck in chunks) <= max_chunk_tokens + + def test_split_sentences_chinese_should_keep_quote_boundary() -> None: text = '他说：“你好。”然后问：“你在吗？”最后结束！' sentences = split_sentences_chinese(text) diff --git a/backend/test/unit/test_chunking_token_limit.py b/backend/test/unit/test_chunking_token_limit.py index de700e6d..8c128e1d 100644 --- a/backend/test/unit/test_chunking_token_limit.py +++ b/backend/test/unit/test_chunking_token_limit.py @@ -86,6 +86,17 @@ class TestHardSplitByTokenLimit: for chunk in result: assert nlp.count_tokens(chunk) <= 512 + def test_optional_hard_limit_keeps_slightly_oversized_text(self): + text = "内容" * 300 # 600 CJK tokens + result = nlp.hard_split_by_token_limit(text, 512, hard_limit_token_num=768) + assert result == [text] + + def test_optional_hard_limit_merges_short_tail(self): + text = "内容" * 635 # 1270 CJK tokens -> 512 + 512 + 246 + result = nlp.hard_split_by_token_limit(text, 512, hard_limit_token_num=768) + assert len(result) == 2 + assert [nlp.count_tokens(chunk) for chunk in result] == [512, 758] + def test_splits_long_english_text(self): text = "hello world " * 1000 # ~2000 word tokens result = nlp.hard_split_by_token_limit(text, 512) @@ -120,16 +131,21 @@ class TestEnsureChunkTokenLimit: result = general._ensure_chunk_token_limit(chunks, 512) assert result == ["短文本一", "短文本二", "短文本三"] - def test_oversized_chunk_gets_split(self): + def test_slightly_oversized_chunk_passes_through(self): long_text = "内容" * 300 # ~600 CJK tokens chunks = ["短文本", long_text, "短文本二"] result = general._ensure_chunk_token_limit(chunks, 512) + assert result == ["短文本", long_text, "短文本二"] + + def test_oversized_chunk_gets_split_with_merged_tail(self): + long_text = "内容" * 635 # 1270 CJK tokens -> 512 + 758 + chunks = ["短文本", long_text, "短文本二"] + result = general._ensure_chunk_token_limit(chunks, 512) assert result[0] == "短文本" assert result[-1] == "短文本二" middle_chunks = result[1:-1] - assert len(middle_chunks) > 1 - for chunk in middle_chunks: - assert nlp.count_tokens(chunk) <= 512 + assert len(middle_chunks) == 2 + assert [nlp.count_tokens(chunk) for chunk in middle_chunks] == [512, 758] def test_empty_chunks_filtered(self): chunks = ["有效文本", "", " ", "另一段"] @@ -159,7 +175,7 @@ class TestGeneralChunkMarkdown: chunks = general.chunk_markdown(doc, {"chunk_token_num": 512}) assert len(chunks) > 1 for chunk in chunks: - assert nlp.count_tokens(chunk) <= 512 + assert nlp.count_tokens(chunk) <= 768 def test_empty_document_returns_empty(self): assert general.chunk_markdown("", {"chunk_token_num": 512}) == []