kris/ForcePilot
1
0
Fork 0
Code Issues Pull Requests Actions 1 Packages Projects Releases Wiki Activity

fix: 保护知识库分块 token 上限

Browse Source
This commit is contained in:
Wenjie Zhang 2026-06-06 17:45:32 +08:00
parent 8d5e6f5932
commit f657edddf2
5 changed files with 104 additions and 22 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

37
backend/package/yuxi/knowledge/chunking/ragflow_like/nlp.py
View File

@ -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]:

30
backend/package/yuxi/knowledge/chunking/ragflow_like/parsers/book.py
View File

@ -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,
)

7
backend/package/yuxi/knowledge/chunking/ragflow_like/parsers/general.py
View File

@ -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

26
backend/test/unit/plugins/test_ragflow_like_chunking.py
View File

@ -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)

26
backend/test/unit/test_chunking_token_limit.py
View File

@ -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}) == []