mirror of
https://github.com/infiniflow/ragflow.git
synced 2026-08-06 23:51:15 +08:00
299 lines
8.0 KiB
Python
299 lines
8.0 KiB
Python
#
|
|
# 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.util
|
|
import sys
|
|
from pathlib import Path
|
|
from types import ModuleType, SimpleNamespace
|
|
from unittest.mock import Mock
|
|
|
|
import pytest
|
|
|
|
|
|
def _package(monkeypatch, name):
|
|
package = ModuleType(name)
|
|
package.__path__ = []
|
|
monkeypatch.setitem(sys.modules, name, package)
|
|
return package
|
|
|
|
|
|
def _module(monkeypatch, name, **attributes):
|
|
module = ModuleType(name)
|
|
for key, value in attributes.items():
|
|
setattr(module, key, value)
|
|
monkeypatch.setitem(sys.modules, name, module)
|
|
return module
|
|
|
|
|
|
def _load_figure_parser(monkeypatch):
|
|
repo_root = Path(__file__).resolve().parents[4]
|
|
|
|
for package_name in (
|
|
"api",
|
|
"api.db",
|
|
"api.db.services",
|
|
"api.db.joint_services",
|
|
"common",
|
|
"rag",
|
|
"rag.app",
|
|
"rag.prompts",
|
|
"rag.utils",
|
|
):
|
|
_package(monkeypatch, package_name)
|
|
|
|
class FakeImage:
|
|
def close(self):
|
|
pass
|
|
|
|
image_module = _module(monkeypatch, "PIL.Image", Image=FakeImage)
|
|
pil_module = _package(monkeypatch, "PIL")
|
|
pil_module.Image = image_module
|
|
|
|
_module(
|
|
monkeypatch,
|
|
"common.constants",
|
|
LLMType=SimpleNamespace(VISION="vision"),
|
|
)
|
|
_module(
|
|
monkeypatch,
|
|
"api.db.services.llm_service",
|
|
LLMBundle=Mock(),
|
|
)
|
|
_module(
|
|
monkeypatch,
|
|
"api.db.joint_services.tenant_model_service",
|
|
get_tenant_default_model_by_type=Mock(),
|
|
)
|
|
|
|
def timeout(*_args, **_kwargs):
|
|
return lambda function: function
|
|
|
|
_module(monkeypatch, "common.connection_utils", timeout=timeout)
|
|
_module(
|
|
monkeypatch,
|
|
"rag.app.picture",
|
|
vision_llm_chunk=Mock(return_value="description"),
|
|
)
|
|
_module(
|
|
monkeypatch,
|
|
"rag.prompts.generator",
|
|
vision_llm_figure_describe_prompt=Mock(return_value="prompt"),
|
|
vision_llm_figure_describe_prompt_with_context=Mock(return_value="prompt"),
|
|
)
|
|
_module(
|
|
monkeypatch,
|
|
"rag.nlp",
|
|
append_context2table_image4pdf=Mock(return_value=[]),
|
|
)
|
|
_module(
|
|
monkeypatch,
|
|
"rag.utils.lazy_image",
|
|
ensure_pil_image=lambda image: image,
|
|
open_image_for_processing=lambda image, **_kwargs: (image, False),
|
|
is_image_like=lambda _image: True,
|
|
)
|
|
|
|
module_path = repo_root / "deepdoc" / "parser" / "figure_parser.py"
|
|
spec = importlib.util.spec_from_file_location(
|
|
"test_figure_parser_module",
|
|
module_path,
|
|
)
|
|
module = importlib.util.module_from_spec(spec)
|
|
monkeypatch.setitem(sys.modules, spec.name, module)
|
|
spec.loader.exec_module(module)
|
|
return module, FakeImage
|
|
|
|
|
|
@pytest.mark.p1
|
|
@pytest.mark.parametrize(
|
|
("context_above", "context_below", "prompt_name", "expected_arguments"),
|
|
[
|
|
(
|
|
"",
|
|
"",
|
|
"vision_llm_figure_describe_prompt",
|
|
{},
|
|
),
|
|
(
|
|
"Above ",
|
|
"Below",
|
|
"vision_llm_figure_describe_prompt_with_context",
|
|
{
|
|
"context_above": "Above Caption",
|
|
"context_below": "Below",
|
|
},
|
|
),
|
|
],
|
|
)
|
|
@pytest.mark.parametrize(
|
|
("language", "expected_language"),
|
|
[
|
|
("Chinese", "Chinese"),
|
|
("", "English"),
|
|
],
|
|
)
|
|
def test_docx_wrapper_passes_dataset_language_to_vision_model_and_prompt(
|
|
monkeypatch,
|
|
context_above,
|
|
context_below,
|
|
prompt_name,
|
|
expected_arguments,
|
|
language,
|
|
expected_language,
|
|
):
|
|
module, FakeImage = _load_figure_parser(monkeypatch)
|
|
model_config = {"llm_name": "vision-model"}
|
|
vision_model = object()
|
|
|
|
module.get_tenant_default_model_by_type = Mock(return_value=model_config)
|
|
module.LLMBundle = Mock(return_value=vision_model)
|
|
module.picture_vision_llm_chunk = Mock(return_value="description")
|
|
|
|
default_prompt = Mock(return_value="prompt")
|
|
contextual_prompt = Mock(return_value="prompt")
|
|
module.vision_llm_figure_describe_prompt = default_prompt
|
|
module.vision_llm_figure_describe_prompt_with_context = contextual_prompt
|
|
|
|
chunks = [
|
|
{
|
|
"image": FakeImage(),
|
|
"text": "Caption",
|
|
"context_above": context_above,
|
|
"context_below": context_below,
|
|
}
|
|
]
|
|
|
|
module.vision_figure_parser_docx_wrapper_naive(
|
|
chunks=chunks,
|
|
idx_lst=[0],
|
|
callback=lambda *_args, **_kwargs: None,
|
|
tenant_id="tenant-id",
|
|
lang=language,
|
|
)
|
|
|
|
module.LLMBundle.assert_called_once_with(
|
|
"tenant-id",
|
|
model_config,
|
|
lang=expected_language,
|
|
)
|
|
|
|
selected_prompt = getattr(module, prompt_name)
|
|
selected_prompt.assert_called_once_with(
|
|
**expected_arguments,
|
|
language=expected_language,
|
|
)
|
|
assert chunks[0]["text"].endswith("description")
|
|
|
|
|
|
@pytest.mark.p1
|
|
@pytest.mark.parametrize(
|
|
("language", "expected_language"),
|
|
[
|
|
("Chinese", "Chinese"),
|
|
("", "English"),
|
|
],
|
|
)
|
|
def test_vision_figure_parser_passes_dataset_language_to_prompt(
|
|
monkeypatch,
|
|
language,
|
|
expected_language,
|
|
):
|
|
module, FakeImage = _load_figure_parser(monkeypatch)
|
|
prompt = Mock(return_value="prompt")
|
|
module.vision_llm_figure_describe_prompt = prompt
|
|
module.picture_vision_llm_chunk = Mock(return_value="description")
|
|
|
|
parser = module.VisionFigureParser(
|
|
vision_model=object(),
|
|
figures_data=[(FakeImage(), ["caption"])],
|
|
lang=language,
|
|
)
|
|
|
|
parser(callback=lambda *_args, **_kwargs: None)
|
|
|
|
prompt.assert_called_once_with(language=expected_language)
|
|
|
|
|
|
@pytest.mark.p1
|
|
@pytest.mark.parametrize(
|
|
"wrapper_name",
|
|
[
|
|
"vision_figure_parser_docx_wrapper",
|
|
"vision_figure_parser_figure_xlsx_wrapper",
|
|
"vision_figure_parser_pdf_wrapper",
|
|
],
|
|
)
|
|
@pytest.mark.parametrize(
|
|
("language", "expected_language"),
|
|
[
|
|
("Chinese", "Chinese"),
|
|
("", "English"),
|
|
],
|
|
)
|
|
def test_figure_wrappers_pass_dataset_language_to_model_and_parser(
|
|
monkeypatch,
|
|
wrapper_name,
|
|
language,
|
|
expected_language,
|
|
):
|
|
module, FakeImage = _load_figure_parser(monkeypatch)
|
|
model_config = {"llm_name": "vision-model"}
|
|
vision_model = object()
|
|
parser_instance = Mock(return_value=[])
|
|
|
|
module.get_tenant_default_model_by_type = Mock(return_value=model_config)
|
|
module.LLMBundle = Mock(return_value=vision_model)
|
|
module.VisionFigureParser = Mock(return_value=parser_instance)
|
|
|
|
if wrapper_name == "vision_figure_parser_docx_wrapper":
|
|
arguments = {
|
|
"sections": [("caption", FakeImage())],
|
|
"tbls": [],
|
|
}
|
|
elif wrapper_name == "vision_figure_parser_figure_xlsx_wrapper":
|
|
arguments = {
|
|
"images": [
|
|
{
|
|
"image": FakeImage(),
|
|
"image_description": "caption",
|
|
}
|
|
],
|
|
}
|
|
else:
|
|
arguments = {
|
|
"tbls": [
|
|
(
|
|
(FakeImage(), ["caption"]),
|
|
[(0, 0, 0, 0, 0)],
|
|
)
|
|
],
|
|
"sections": [],
|
|
}
|
|
|
|
getattr(module, wrapper_name)(
|
|
**arguments,
|
|
callback=lambda *_args, **_kwargs: None,
|
|
tenant_id="tenant-id",
|
|
lang=language,
|
|
)
|
|
|
|
module.LLMBundle.assert_called_once_with(
|
|
"tenant-id",
|
|
model_config,
|
|
lang=expected_language,
|
|
)
|
|
assert module.VisionFigureParser.call_args.kwargs["lang"] == expected_language
|
|
parser_instance.assert_called_once()
|