mirror of
https://github.com/infiniflow/ragflow.git
synced 2026-08-14 04:36:52 +08:00
fix(tests): isolate token chunker imports (#18018)
This commit is contained in:
@@ -22,11 +22,8 @@ import numpy as np
|
||||
import pdfplumber
|
||||
from PIL import Image
|
||||
|
||||
from api.db.services.file2document_service import File2DocumentService
|
||||
from api.db.services.file_service import FileService
|
||||
from common import settings
|
||||
from common.misc_utils import get_uuid
|
||||
from deepdoc.parser.pdf_parser import LOCK_KEY_pdfplumber, RAGFlowPdfParser
|
||||
from rag.utils.base64_image import image2id
|
||||
|
||||
PDF_PREVIEW_GAP = 6
|
||||
@@ -47,6 +44,8 @@ def _extract_raw_positions(item):
|
||||
|
||||
position_tag = item.get("position_tag")
|
||||
if isinstance(position_tag, str) and position_tag:
|
||||
from deepdoc.parser.pdf_parser import RAGFlowPdfParser
|
||||
|
||||
return [[pos[0][-1], *pos[1:]] for pos in RAGFlowPdfParser.extract_positions(position_tag)]
|
||||
|
||||
position_int = item.get("position_int")
|
||||
@@ -187,6 +186,9 @@ def finalize_pdf_chunk(chunk):
|
||||
|
||||
|
||||
def _fetch_source_blob(from_upstream, canvas):
|
||||
from api.db.services.file2document_service import File2DocumentService
|
||||
from api.db.services.file_service import FileService
|
||||
|
||||
if canvas._doc_id:
|
||||
bucket, name = File2DocumentService.get_storage_address(doc_id=canvas._doc_id)
|
||||
return settings.STORAGE_IMPL.get(bucket, name)
|
||||
@@ -196,6 +198,8 @@ def _fetch_source_blob(from_upstream, canvas):
|
||||
|
||||
|
||||
def _load_pdf_page_images(blob, zoom=PDF_PREVIEW_ZOOM):
|
||||
from deepdoc.parser.pdf_parser import LOCK_KEY_pdfplumber
|
||||
|
||||
with sys.modules[LOCK_KEY_pdfplumber]:
|
||||
with pdfplumber.open(io.BytesIO(blob)) as pdf:
|
||||
return [page.to_image(resolution=72 * zoom, antialias=True).annotated for page in pdf.pages]
|
||||
|
||||
46
test/unit_test/rag/test_flow_chunker_import.py
Normal file
46
test/unit_test/rag/test_flow_chunker_import.py
Normal file
@@ -0,0 +1,46 @@
|
||||
#
|
||||
# Copyright 2026 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 importlib
|
||||
import sys
|
||||
|
||||
|
||||
def test_token_chunker_imports_without_pdf_parser_symbols(pdf_parser_stub, monkeypatch):
|
||||
monkeypatch.delattr(pdf_parser_stub, "RAGFlowPdfParser")
|
||||
for module_name in [
|
||||
"rag.flow.chunker",
|
||||
"rag.flow.chunker.token_chunker",
|
||||
"rag.flow.parser.pdf_chunk_metadata",
|
||||
]:
|
||||
sys.modules.pop(module_name, None)
|
||||
|
||||
token_chunker = importlib.import_module("rag.flow.chunker.token_chunker")
|
||||
|
||||
assert callable(token_chunker._merge_text_chunks_by_token_size)
|
||||
|
||||
position_tag = "@@1\t10\t20\t30\t40##"
|
||||
extracted_tags = []
|
||||
|
||||
class RuntimePdfParser:
|
||||
@staticmethod
|
||||
def extract_positions(value):
|
||||
extracted_tags.append(value)
|
||||
return [([0], 10.0, 20.0, 30.0, 40.0)]
|
||||
|
||||
monkeypatch.setattr(pdf_parser_stub, "RAGFlowPdfParser", RuntimePdfParser, raising=False)
|
||||
|
||||
assert token_chunker.extract_pdf_positions({"position_tag": position_tag}) == [[1, 10.0, 20.0, 30.0, 40.0]]
|
||||
assert extracted_tags == [position_tag]
|
||||
Reference in New Issue
Block a user