mirror of
https://github.com/infiniflow/ragflow.git
synced 2026-08-10 01:11:23 +08:00
446 lines
20 KiB
Python
446 lines
20 KiB
Python
import importlib.util
|
|
import asyncio
|
|
import sys
|
|
import types
|
|
from contextlib import contextmanager
|
|
from pathlib import Path
|
|
|
|
|
|
@contextmanager
|
|
def _load_token_chunker_with_stubs():
|
|
root = Path(__file__).resolve().parents[3]
|
|
original_modules = {}
|
|
|
|
def _install(name: str, module: types.ModuleType):
|
|
original_modules.setdefault(name, sys.modules.get(name))
|
|
sys.modules[name] = module
|
|
|
|
try:
|
|
rag_pkg = types.ModuleType("rag")
|
|
rag_pkg.__path__ = [str(root / "rag")]
|
|
_install("rag", rag_pkg)
|
|
|
|
rag_flow_pkg = types.ModuleType("rag.flow")
|
|
rag_flow_pkg.__package__ = "rag"
|
|
rag_flow_pkg.__path__ = [str(root / "rag" / "flow")]
|
|
_install("rag.flow", rag_flow_pkg)
|
|
|
|
rag_flow_chunker_pkg = types.ModuleType("rag.flow.chunker")
|
|
rag_flow_chunker_pkg.__package__ = "rag.flow"
|
|
rag_flow_chunker_pkg.__path__ = [str(root / "rag" / "flow" / "chunker")]
|
|
_install("rag.flow.chunker", rag_flow_chunker_pkg)
|
|
|
|
rag_flow_parser_pkg = types.ModuleType("rag.flow.parser")
|
|
rag_flow_parser_pkg.__package__ = "rag.flow"
|
|
rag_flow_parser_pkg.__path__ = [str(root / "rag" / "flow" / "parser")]
|
|
_install("rag.flow.parser", rag_flow_parser_pkg)
|
|
|
|
common_pkg = types.ModuleType("common")
|
|
common_pkg.__path__ = [str(root / "common")]
|
|
_install("common", common_pkg)
|
|
|
|
common_float_utils = types.ModuleType("common.float_utils")
|
|
common_float_utils.normalize_overlapped_percent = lambda value: value
|
|
_install("common.float_utils", common_float_utils)
|
|
|
|
common_token_utils = types.ModuleType("common.token_utils")
|
|
common_token_utils.num_tokens_from_string = lambda text: 1
|
|
_install("common.token_utils", common_token_utils)
|
|
|
|
rag_nlp = types.ModuleType("rag.nlp")
|
|
rag_nlp.naive_merge = lambda *_args, **_kwargs: []
|
|
_install("rag.nlp", rag_nlp)
|
|
|
|
class ProcessParamBase:
|
|
def __init__(self):
|
|
pass
|
|
|
|
def check_valid_value(self, value, msg, allowed):
|
|
if value not in allowed:
|
|
raise ValueError(msg)
|
|
|
|
def check_positive_integer(self, value, msg):
|
|
pass
|
|
|
|
def check_decimal_float(self, value, msg):
|
|
pass
|
|
|
|
def check_nonnegative_number(self, value, msg):
|
|
pass
|
|
|
|
class ProcessBase:
|
|
def __init__(self, _pipeline, _id, param):
|
|
self._pipeline = _pipeline
|
|
self._id = _id
|
|
self._param = param
|
|
self._outputs = {}
|
|
self.callback = lambda *_args, **_kwargs: None
|
|
|
|
def set_output(self, key, value):
|
|
self._outputs[key] = value
|
|
|
|
rag_flow_base = types.ModuleType("rag.flow.base")
|
|
rag_flow_base.ProcessBase = ProcessBase
|
|
rag_flow_base.ProcessParamBase = ProcessParamBase
|
|
_install("rag.flow.base", rag_flow_base)
|
|
|
|
rag_flow_parser_pdf_metadata = types.ModuleType("rag.flow.parser.pdf_chunk_metadata")
|
|
rag_flow_parser_pdf_metadata.PDF_POSITIONS_KEY = "pdf_positions"
|
|
rag_flow_parser_pdf_metadata.extract_pdf_positions = lambda _item: []
|
|
rag_flow_parser_pdf_metadata.finalize_pdf_chunk = lambda chunk: chunk
|
|
|
|
async def restore_pdf_text_previews(*_args, **_kwargs):
|
|
return None
|
|
|
|
rag_flow_parser_pdf_metadata.restore_pdf_text_previews = restore_pdf_text_previews
|
|
_install("rag.flow.parser.pdf_chunk_metadata", rag_flow_parser_pdf_metadata)
|
|
|
|
try:
|
|
import pydantic # noqa: F401
|
|
|
|
schema_spec = importlib.util.spec_from_file_location(
|
|
"rag.flow.chunker.schema",
|
|
root / "rag" / "flow" / "chunker" / "schema.py",
|
|
)
|
|
if schema_spec is None or schema_spec.loader is None:
|
|
raise RuntimeError("Failed to locate rag.flow.chunker.schema stub loader.")
|
|
schema_module = importlib.util.module_from_spec(schema_spec)
|
|
_install("rag.flow.chunker.schema", schema_module)
|
|
schema_spec.loader.exec_module(schema_module)
|
|
except Exception:
|
|
schema_module = types.ModuleType("rag.flow.chunker.schema")
|
|
|
|
class TokenChunkerFromUpstream:
|
|
def __init__(
|
|
self,
|
|
name,
|
|
file=None,
|
|
chunks=None,
|
|
output_format=None,
|
|
json_result=None,
|
|
markdown_result=None,
|
|
text_result=None,
|
|
html_result=None,
|
|
_created_time=None,
|
|
_elapsed_time=None,
|
|
):
|
|
self.name = name
|
|
self.file = file
|
|
self.chunks = chunks
|
|
self.output_format = output_format
|
|
self.json_result = json_result
|
|
self.json = json_result
|
|
self.markdown_result = markdown_result
|
|
self.markdown = markdown_result
|
|
self.text_result = text_result
|
|
self.text = text_result
|
|
self.html_result = html_result
|
|
self.html = html_result
|
|
self._created_time = _created_time
|
|
self._elapsed_time = _elapsed_time
|
|
|
|
@classmethod
|
|
def model_validate(cls, data):
|
|
if isinstance(data, dict):
|
|
return cls(
|
|
name=data.get("name", ""),
|
|
file=data.get("file"),
|
|
chunks=data.get("chunks"),
|
|
output_format=data.get("output_format"),
|
|
json_result=data.get("json_result", data.get("json")),
|
|
markdown_result=data.get("markdown_result", data.get("markdown")),
|
|
text_result=data.get("text_result", data.get("text")),
|
|
html_result=data.get("html_result", data.get("html")),
|
|
_created_time=data.get("_created_time"),
|
|
_elapsed_time=data.get("_elapsed_time"),
|
|
)
|
|
raise TypeError("TokenChunkerFromUpstream expects a dict payload.")
|
|
|
|
schema_module.TokenChunkerFromUpstream = TokenChunkerFromUpstream
|
|
_install("rag.flow.chunker.schema", schema_module)
|
|
|
|
token_chunker_spec = importlib.util.spec_from_file_location(
|
|
"rag.flow.chunker.token_chunker",
|
|
root / "rag" / "flow" / "chunker" / "token_chunker.py",
|
|
)
|
|
if token_chunker_spec is None or token_chunker_spec.loader is None:
|
|
raise RuntimeError("Failed to locate rag.flow.chunker.token_chunker stub loader.")
|
|
token_chunker_module = importlib.util.module_from_spec(token_chunker_spec)
|
|
_install("rag.flow.chunker.token_chunker", token_chunker_module)
|
|
token_chunker_spec.loader.exec_module(token_chunker_module)
|
|
yield token_chunker_module
|
|
finally:
|
|
for module_name, original in original_modules.items():
|
|
if original is None:
|
|
sys.modules.pop(module_name, None)
|
|
else:
|
|
sys.modules[module_name] = original
|
|
|
|
|
|
def test_token_chunker_prefers_upstream_chunks_for_json_output_format_chunks():
|
|
# Regression for #16812: when the upstream (e.g. TitleChunker) emits
|
|
# output_format="chunks", TokenChunker must consume from_upstream.chunks and
|
|
# not fall through to the raw parser json_result. Heavy deps are stubbed so
|
|
# the real TokenChunker._invoke runs against the real schema when pydantic is
|
|
# available (see title_chunker/common.py for the same chunks-vs-json branch).
|
|
with _load_token_chunker_with_stubs() as token_chunker_module:
|
|
token_chunker = token_chunker_module.TokenChunker
|
|
param = token_chunker_module.TokenChunkerParam()
|
|
param.delimiter_mode = "one"
|
|
chunker = token_chunker(None, "token_chunker", param)
|
|
|
|
kwargs = {
|
|
"name": "token_chunker",
|
|
"output_format": "chunks",
|
|
"chunks": [{"text": "CHAPTER-AWARE"}],
|
|
"json": [{"text": "RAW-PARSER-JSON"}],
|
|
}
|
|
|
|
asyncio.run(chunker._invoke(**kwargs))
|
|
|
|
assert chunker._outputs["chunks"] == [{"text": "CHAPTER-AWARE"}]
|
|
|
|
|
|
def _build_json_chunker(param: dict, monkeypatch_positions=True):
|
|
"""Build a TokenChunker (bypassing ComponentBase.__init__) wired for the JSON
|
|
``delimiter_mode`` path, with heavy deps stubbed.
|
|
|
|
Returns ``(chunker, module)`` so callers can monkeypatch the module-global
|
|
``extract_pdf_positions`` (it is imported as a name, so rebinding the module
|
|
attribute reaches the call sites inside ``_build_json_chunks``).
|
|
"""
|
|
with _load_token_chunker_with_stubs() as token_chunker_module:
|
|
token_chunker = token_chunker_module.TokenChunker
|
|
param_obj = token_chunker_module.TokenChunkerParam()
|
|
for key, value in param.items():
|
|
setattr(param_obj, key, value)
|
|
|
|
chunker = token_chunker(None, "token_chunker", param_obj)
|
|
chunker._canvas = types.SimpleNamespace(_doc_id=None, _tenant_id="t")
|
|
if monkeypatch_positions:
|
|
# Echo per-item positions so we can assert PDF coordinates survive.
|
|
token_chunker_module.extract_pdf_positions = lambda item: item.get("positions", [])
|
|
|
|
yield token_chunker_module, chunker
|
|
|
|
|
|
def test_json_delimiter_mode_drop_delimiter_text():
|
|
# The delimiter is a boundary: its text must never appear inside a chunk.
|
|
for module, chunker in _build_json_chunker({"delimiter_mode": "delimiter", "delimiters": ["`##`"]}):
|
|
kwargs = {
|
|
"name": "token_chunker",
|
|
"output_format": "json",
|
|
"json_result": [{"text": "first part##second part##third part", "doc_type_kwd": "text"}],
|
|
}
|
|
asyncio.run(chunker._invoke(**kwargs))
|
|
chunks = chunker._outputs["chunks"]
|
|
texts = [c["text"] for c in chunks]
|
|
assert texts == ["first part", "second part", "third part"]
|
|
assert all("##" not in t for t in texts)
|
|
|
|
|
|
def test_json_delimiter_mode_newline_join_not_glued():
|
|
# Regression for #17723: JSON flush must join buffered text items with "\\n",
|
|
# never glue them. Two adjacent items "hello" + "world" must stay
|
|
# "hello\\nworld", never become "helloworld".
|
|
for module, chunker in _build_json_chunker({"delimiter_mode": "delimiter", "delimiters": []}):
|
|
kwargs = {
|
|
"name": "token_chunker",
|
|
"output_format": "json",
|
|
"json_result": [
|
|
{"text": "hello", "doc_type_kwd": "text"},
|
|
{"text": "world", "doc_type_kwd": "text"},
|
|
],
|
|
}
|
|
asyncio.run(chunker._invoke(**kwargs))
|
|
chunks = chunker._outputs["chunks"]
|
|
assert len(chunks) == 1
|
|
assert chunks[0]["text"] == "hello\nworld"
|
|
|
|
|
|
def test_json_delimiter_mode_children_delimiters_applied():
|
|
# Regression for #17723: children_delimiters (secondary split) must run before
|
|
# finalizing the JSON ``delimiter_mode`` path, or they are silently ignored.
|
|
for module, chunker in _build_json_chunker({"delimiter_mode": "delimiter", "delimiters": [], "children_delimiters": ["|"]}):
|
|
kwargs = {
|
|
"name": "token_chunker",
|
|
"output_format": "json",
|
|
"json_result": [{"text": "alpha|beta", "doc_type_kwd": "text"}],
|
|
}
|
|
asyncio.run(chunker._invoke(**kwargs))
|
|
chunks = chunker._outputs["chunks"]
|
|
texts = [c["text"] for c in chunks]
|
|
assert texts == ["alpha", "beta"]
|
|
|
|
|
|
def test_json_delimiter_mode_pdf_positions_retained():
|
|
# PDF coordinates carried on the combined chunk must survive into the output.
|
|
for module, chunker in _build_json_chunker({"delimiter_mode": "delimiter", "delimiters": []}):
|
|
kwargs = {
|
|
"name": "token_chunker",
|
|
"output_format": "json",
|
|
"json_result": [
|
|
{"text": "hello", "doc_type_kwd": "text", "positions": [[1, 0, 10, 0, 5]]},
|
|
{"text": "world", "doc_type_kwd": "text", "positions": [[2, 0, 20, 0, 8]]},
|
|
],
|
|
}
|
|
asyncio.run(chunker._invoke(**kwargs))
|
|
chunks = chunker._outputs["chunks"]
|
|
assert len(chunks) == 1
|
|
assert chunks[0].get("pdf_positions") == [[1, 0, 10, 0, 5], [2, 0, 20, 0, 8]]
|
|
|
|
|
|
def test_token_size_mode_normalized_to_delimiter():
|
|
# Backward-compat: the removed "token_size" value must still be accepted by
|
|
# check() and coerced to "delimiter" (runtime behavior is identical), so
|
|
# legacy configs / pre-fix frontends don't get rejected. Unknown values are
|
|
# still rejected.
|
|
with _load_token_chunker_with_stubs() as token_chunker_module:
|
|
param = token_chunker_module.TokenChunkerParam()
|
|
param.delimiter_mode = "token_size"
|
|
param.check()
|
|
assert param.delimiter_mode == "delimiter"
|
|
|
|
bad = token_chunker_module.TokenChunkerParam()
|
|
bad.delimiter_mode = "nope"
|
|
try:
|
|
bad.check()
|
|
raise AssertionError("expected check() to reject unknown delimiter_mode")
|
|
except Exception:
|
|
pass
|
|
|
|
|
|
def test_json_no_delimiter_mode_merges_to_token_cap():
|
|
# Regression for #17979: with no active (backtick) delimiter, the JSON path
|
|
# must merge per-item text chunks up to chunk_token_size -- mirroring the old
|
|
# "token_size" mode and the Go JSON path. Concatenating every item into a
|
|
# single chunk before the merge would defeat the cap and emit one oversized
|
|
# chunk. The per-token stub (num_tokens_from_string -> 1) makes the cap easy
|
|
# to exceed: 12 one-token items under a cap of 5 must yield several chunks.
|
|
for _module, chunker in _build_json_chunker({"delimiter_mode": "delimiter", "delimiters": [], "chunk_token_size": 5}):
|
|
kwargs = {
|
|
"name": "token_chunker",
|
|
"output_format": "json",
|
|
"json_result": [{"text": f"item{i}"} for i in range(12)],
|
|
}
|
|
asyncio.run(chunker._invoke(**kwargs))
|
|
chunks = chunker._outputs["chunks"]
|
|
# 12 one-token items under a cap of 5 must NOT collapse into one chunk.
|
|
assert len(chunks) > 1, f"cap not enforced: 12 items collapsed to {len(chunks)} chunk(s)"
|
|
# Each merged chunk holds at most one overflow unit past the cap (<= 6
|
|
# items); the final output drops tk_nums, so count item markers instead.
|
|
for c in chunks:
|
|
assert c["text"].count("item") <= 6, f"chunk exceeds cap: {c['text'].count('item')} items"
|
|
# No text lost: all 12 items must survive, joined by the "\n" glue.
|
|
joined = "\n".join(c["text"] for c in chunks)
|
|
for i in range(12):
|
|
assert f"item{i}" in joined, f"item{i} dropped from output"
|
|
|
|
|
|
def test_json_no_delimiter_mode_media_breaks_merge():
|
|
# A non-text (media) chunk interleaved between text items must stay as its
|
|
# own chunk and reset the merge, so text before/after it are sized separately.
|
|
for _module, chunker in _build_json_chunker({"delimiter_mode": "delimiter", "delimiters": [], "chunk_token_size": 5}):
|
|
kwargs = {
|
|
"name": "token_chunker",
|
|
"output_format": "json",
|
|
"json_result": [{"text": f"t{i}", "doc_type_kwd": "text"} for i in range(6)]
|
|
+ [{"text": "IMG", "doc_type_kwd": "image", "img_id": "im1"}]
|
|
+ [{"text": f"u{i}", "doc_type_kwd": "text"} for i in range(12)],
|
|
}
|
|
asyncio.run(chunker._invoke(**kwargs))
|
|
chunks = chunker._outputs["chunks"]
|
|
doc_types = [c["doc_type_kwd"] for c in chunks]
|
|
# The media chunk is preserved and breaks the text merge.
|
|
assert doc_types.count("image") == 1, doc_types
|
|
# Several text chunks on each side of the media boundary.
|
|
assert doc_types.count("text") > 2, doc_types
|
|
|
|
|
|
def test_json_delimiter_mode_pdf_positions_per_segment_not_broadcast():
|
|
# Regression for #3 (PDF coordinate leak): when consecutive text items from
|
|
# different pages are buffered and then split by a custom delimiter, each
|
|
# output segment must carry only the PDF positions of the item(s) that
|
|
# contributed to it -- not the union of every buffered item. The old code
|
|
# broadcast ``combined_pos`` to every split chunk, so a page-1 segment also
|
|
# claimed page-2 coordinates and all segments shared one PDF preview image.
|
|
for module, chunker in _build_json_chunker({"delimiter_mode": "delimiter", "delimiters": ["`。`"]}):
|
|
# Mirror production's preview-cache behaviour: a chunk's preview image is
|
|
# keyed by its position set, so chunks sharing positions share one image.
|
|
async def _restore_previews(chunks, from_upstream, canvas):
|
|
preview_cache = {}
|
|
for chunk in chunks:
|
|
positions = chunk.get("pdf_positions") or []
|
|
key = tuple(tuple(p[:5]) for p in positions)
|
|
if key in preview_cache:
|
|
chunk["img_id"] = preview_cache[key]
|
|
else:
|
|
new_id = "img-%d" % len(preview_cache)
|
|
chunk["img_id"] = new_id
|
|
preview_cache[key] = new_id
|
|
|
|
module.restore_pdf_text_previews = _restore_previews
|
|
|
|
kwargs = {
|
|
"name": "doc.pdf",
|
|
"output_format": "json",
|
|
"json_result": [
|
|
{"text": "第一章。第二段", "doc_type_kwd": "text", "positions": [[1, 0, 10, 0, 5]]},
|
|
{"text": "第三章。第四章", "doc_type_kwd": "text", "positions": [[2, 0, 20, 0, 8]]},
|
|
],
|
|
}
|
|
asyncio.run(chunker._invoke(**kwargs))
|
|
chunks = chunker._outputs["chunks"]
|
|
texts = [c["text"] for c in chunks]
|
|
# Custom "。" splits the buffered text into three segments; the "\n" join
|
|
# between the two items is NOT a split point, so the middle segment spans
|
|
# both pages.
|
|
assert texts == ["第一章", "第二段\n第三章", "第四章"], texts
|
|
|
|
positions = [c.get("pdf_positions") for c in chunks]
|
|
# Page-1-only segment must NOT carry page-2 coordinates.
|
|
assert positions[0] == [[1, 0, 10, 0, 5]], positions
|
|
# Spanning segment legitimately carries both pages.
|
|
assert positions[1] == [[1, 0, 10, 0, 5], [2, 0, 20, 0, 8]], positions
|
|
# Page-2-only segment must NOT carry page-1 coordinates.
|
|
assert positions[2] == [[2, 0, 20, 0, 8]], positions
|
|
|
|
# Previews must not be shared: distinct position sets -> distinct images.
|
|
img_ids = [c.get("img_id") for c in chunks]
|
|
assert len(set(img_ids)) == len(img_ids), img_ids
|
|
|
|
|
|
def test_json_delimiter_mode_consecutive_delimiter_keeps_boundary():
|
|
# Regression for #17723: "A####B" with pattern "##" must yield ["A", "B"],
|
|
# both boundary-adjacent segments preserved (the bug collapsed it to "A##B").
|
|
for module, chunker in _build_json_chunker({"delimiter_mode": "delimiter", "delimiters": ["`##`"]}):
|
|
kwargs = {
|
|
"name": "token_chunker",
|
|
"output_format": "json",
|
|
"json_result": [{"text": "A####B", "doc_type_kwd": "text"}],
|
|
}
|
|
asyncio.run(chunker._invoke(**kwargs))
|
|
chunks = chunker._outputs["chunks"]
|
|
texts = [c["text"] for c in chunks]
|
|
assert texts == ["A", "B"]
|
|
assert all("##" not in t for t in texts)
|
|
|
|
|
|
def test_text_delimiter_mode_one_no_atom_split():
|
|
# chunk_token_size=1 must not atom-split delimiter segments into 1-token
|
|
# chunks; the delimiter path produces delimiter-boundary chunks regardless of cap.
|
|
for _module, chunker in _build_json_chunker({"delimiter_mode": "delimiter", "delimiters": ["`|`"]}):
|
|
# text path: delimiter_mode is delimiter but a custom delimiter is
|
|
# present, so the delimiter branch (_split_text_by_pattern) is used.
|
|
kwargs = {
|
|
"name": "token_chunker",
|
|
"output_format": "text",
|
|
"text": "aaa|bbb|ccc",
|
|
}
|
|
chunk_token_size = 1
|
|
setattr(chunker._param, "chunk_token_size", chunk_token_size)
|
|
asyncio.run(chunker._invoke(**kwargs))
|
|
chunks = chunker._outputs["chunks"]
|
|
texts = [c["text"] for c in chunks]
|
|
assert texts == ["aaa", "bbb", "ccc"], f"chunk_token_size={chunk_token_size} atom-split a delimiter segment: {texts}"
|