fix: honor dataset language across VisionFigureParser paths (#17227)

This commit is contained in:
taek105
2026-08-05 19:40:10 +09:00
committed by GitHub
parent d2303cc46b
commit b7966c98f2
17 changed files with 808 additions and 31 deletions

View File

@@ -28,7 +28,6 @@ from rag.nlp import append_context2table_image4pdf
from rag.utils.lazy_image import ensure_pil_image, open_image_for_processing, is_image_like
# need to delete before pr
def vision_figure_parser_figure_data_wrapper(figures_data_without_positions):
if not figures_data_without_positions:
return []
@@ -46,19 +45,29 @@ def vision_figure_parser_figure_data_wrapper(figures_data_without_positions):
return res
def vision_figure_parser_docx_wrapper(sections, tbls, callback=None, **kwargs):
def _normalize_vision_language(lang):
return lang or "English"
def vision_figure_parser_docx_wrapper(sections, tbls, callback=None, lang="English", **kwargs):
lang = _normalize_vision_language(lang)
if not sections:
return tbls
try:
vision_model_config = get_tenant_default_model_by_type(kwargs["tenant_id"], LLMType.VISION)
vision_model = LLMBundle(kwargs["tenant_id"], vision_model_config)
vision_model = LLMBundle(kwargs["tenant_id"], vision_model_config, lang=lang)
callback(0.7, "Visual model detected. Attempting to enhance figure extraction...")
except Exception:
vision_model = None
if vision_model:
figures_data = vision_figure_parser_figure_data_wrapper(sections)
try:
docx_vision_parser = VisionFigureParser(vision_model=vision_model, figures_data=figures_data, **kwargs)
docx_vision_parser = VisionFigureParser(
vision_model=vision_model,
figures_data=figures_data,
lang=lang,
**kwargs,
)
boosted_figures = docx_vision_parser(callback=callback)
tbls.extend(boosted_figures)
except Exception as e:
@@ -66,13 +75,14 @@ def vision_figure_parser_docx_wrapper(sections, tbls, callback=None, **kwargs):
return tbls
def vision_figure_parser_figure_xlsx_wrapper(images, callback=None, **kwargs):
def vision_figure_parser_figure_xlsx_wrapper(images, callback=None, lang="English", **kwargs):
lang = _normalize_vision_language(lang)
tbls = []
if not images:
return []
try:
vision_model_config = get_tenant_default_model_by_type(kwargs["tenant_id"], LLMType.VISION)
vision_model = LLMBundle(kwargs["tenant_id"], vision_model_config)
vision_model = LLMBundle(kwargs["tenant_id"], vision_model_config, lang=lang)
callback(0.2, "Visual model detected. Attempting to enhance Excel image extraction...")
except Exception:
vision_model = None
@@ -90,7 +100,12 @@ def vision_figure_parser_figure_xlsx_wrapper(images, callback=None, **kwargs):
for img in images
]
try:
parser = VisionFigureParser(vision_model=vision_model, figures_data=figures_data, **kwargs)
parser = VisionFigureParser(
vision_model=vision_model,
figures_data=figures_data,
lang=lang,
**kwargs,
)
callback(0.22, "Parsing images...")
boosted_figures = parser(callback=callback)
tbls.extend(boosted_figures)
@@ -99,7 +114,8 @@ def vision_figure_parser_figure_xlsx_wrapper(images, callback=None, **kwargs):
return tbls
def vision_figure_parser_pdf_wrapper(tbls, callback=None, **kwargs):
def vision_figure_parser_pdf_wrapper(tbls, callback=None, lang="English", **kwargs):
lang = _normalize_vision_language(lang)
if not tbls:
return []
sections = kwargs.get("sections")
@@ -107,7 +123,7 @@ def vision_figure_parser_pdf_wrapper(tbls, callback=None, **kwargs):
context_size = max(0, int(parser_config.get("image_context_size", 0) or 0))
try:
vision_model_config = get_tenant_default_model_by_type(kwargs["tenant_id"], LLMType.VISION)
vision_model = LLMBundle(kwargs["tenant_id"], vision_model_config)
vision_model = LLMBundle(kwargs["tenant_id"], vision_model_config, lang=lang)
callback(0.7, "Visual model detected. Attempting to enhance figure extraction...")
except Exception:
vision_model = None
@@ -131,6 +147,7 @@ def vision_figure_parser_pdf_wrapper(tbls, callback=None, **kwargs):
figures_data=figures_data,
figure_contexts=figure_contexts,
context_size=context_size,
lang=lang,
**kwargs,
)
boosted_figures = docx_vision_parser(callback=callback)
@@ -141,12 +158,13 @@ def vision_figure_parser_pdf_wrapper(tbls, callback=None, **kwargs):
return tbls
def vision_figure_parser_docx_wrapper_naive(chunks, idx_lst, callback=None, **kwargs):
def vision_figure_parser_docx_wrapper_naive(chunks, idx_lst, callback=None, lang="English", **kwargs):
lang = _normalize_vision_language(lang)
if not chunks:
return []
try:
vision_model_config = get_tenant_default_model_by_type(kwargs["tenant_id"], LLMType.VISION)
vision_model = LLMBundle(kwargs["tenant_id"], vision_model_config)
vision_model = LLMBundle(kwargs["tenant_id"], vision_model_config, lang=lang)
callback(0.7, "Visual model detected. Attempting to enhance figure extraction...")
except Exception:
vision_model = None
@@ -164,12 +182,11 @@ def vision_figure_parser_docx_wrapper_naive(chunks, idx_lst, callback=None, **kw
# context_above + caption if any
context_above=ck.get("context_above") + ck.get("text", ""),
context_below=ck.get("context_below"),
language=lang,
)
logging.info(f"[VisionFigureParser] figure={idx} context_above_len={len(context_above)} context_below_len={len(context_below)} prompt=with_context")
logging.info(f"[VisionFigureParser] figure={idx} context_above_snippet={context_above[:512]}")
logging.info(f"[VisionFigureParser] figure={idx} context_below_snippet={context_below[:512]}")
else:
prompt = vision_llm_figure_describe_prompt()
prompt = vision_llm_figure_describe_prompt(language=lang)
logging.info(f"[VisionFigureParser] figure={idx} context_len=0 prompt=default")
try:
@@ -201,6 +218,7 @@ shared_executor = ThreadPoolExecutor(max_workers=10)
class VisionFigureParser:
def __init__(self, vision_model, figures_data, *args, **kwargs):
self.vision_model = vision_model
self.language = kwargs.get("lang") or "English"
self.figure_contexts = kwargs.get("figure_contexts") or []
self.context_size = max(0, int(kwargs.get("context_size", 0) or 0))
self._extract_figures_info(figures_data)
@@ -261,14 +279,13 @@ class VisionFigureParser:
prompt = vision_llm_figure_describe_prompt_with_context(
context_above=context_above,
context_below=context_below,
language=self.language,
)
logging.info(
f"[VisionFigureParser] figure={figure_idx} context_size={self.context_size} context_above_len={len(context_above)} context_below_len={len(context_below)} prompt=with_context"
)
logging.info(f"[VisionFigureParser] figure={figure_idx} context_above_snippet={context_above[:512]}")
logging.info(f"[VisionFigureParser] figure={figure_idx} context_below_snippet={context_below[:512]}")
else:
prompt = vision_llm_figure_describe_prompt()
prompt = vision_llm_figure_describe_prompt(language=self.language)
logging.info(f"[VisionFigureParser] figure={figure_idx} context_size={self.context_size} context_len=0 prompt=default")
description_text = picture_vision_llm_chunk(
binary=figure_binary,

View File

@@ -86,7 +86,7 @@ def chunk(filename, binary=None, from_page=0, to_page=MAXIMUM_PAGE_NUMBER, lang=
remove_contents_table(sections, eng=is_english(random_choices([t for t, _ in sections], k=200)))
tbls = vision_figure_parser_docx_wrapper(sections=sections, tbls=tbls, callback=callback, **kwargs)
tbls = vision_figure_parser_docx_wrapper(sections=sections, tbls=tbls, callback=callback, lang=lang, **kwargs)
# tbls = [((None, lns), None) for lns in tbls]
sections = [(item[0], item[1] if item[1] is not None else "") for item in sections if not isinstance(item[1], (Image.Image, LazyImage))]
callback(0.8, "Finish parsing.")

View File

@@ -268,6 +268,7 @@ def chunk(filename, binary=None, from_page=0, to_page=MAXIMUM_PAGE_NUMBER, lang=
tbls=tbls,
sections=sections,
callback=callback,
lang=lang,
**kwargs,
)
res = tokenize_table(tbls, doc, eng, language=lang)
@@ -283,7 +284,13 @@ def chunk(filename, binary=None, from_page=0, to_page=MAXIMUM_PAGE_NUMBER, lang=
elif re.search(r"\.docx?$", filename, re.IGNORECASE):
docx_parser = Docx()
ti_list, tbls = docx_parser(filename, binary, from_page=0, to_page=MAXIMUM_PAGE_NUMBER, callback=callback)
tbls = vision_figure_parser_docx_wrapper(sections=ti_list, tbls=tbls, callback=callback, **kwargs)
tbls = vision_figure_parser_docx_wrapper(
sections=ti_list,
tbls=tbls,
callback=callback,
lang=lang,
**kwargs,
)
res = tokenize_table(tbls, doc, eng, language=lang)
for text, image in ti_list:
d = copy.deepcopy(doc)

View File

@@ -123,6 +123,7 @@ def by_deepdoc(filename, binary=None, from_page=0, to_page=MAXIMUM_PAGE_NUMBER,
tbls=tables,
sections=sections,
callback=callback,
lang=lang,
**kwargs,
)
return sections, tables, pdf_parser
@@ -1027,7 +1028,13 @@ def chunk(filename, binary=None, from_page=0, to_page=MAXIMUM_PAGE_NUMBER, lang=
# images list - index of image chunk in chunks
chunks, images = naive_merge_docx(sections, int(parser_config.get("chunk_token_num", 128)), parser_config.get("delimiter", "\n!?。;!?"), table_context_size, image_context_size)
vision_figure_parser_docx_wrapper_naive(chunks=chunks, idx_lst=images, callback=callback, **kwargs)
vision_figure_parser_docx_wrapper_naive(
chunks=chunks,
idx_lst=images,
callback=callback,
lang=lang,
**kwargs,
)
callback(0.8, "Finish parsing.")
st = timer()
@@ -1147,7 +1154,7 @@ def chunk(filename, binary=None, from_page=0, to_page=MAXIMUM_PAGE_NUMBER, lang=
try:
vision_model_config = get_tenant_default_model_by_type(kwargs["tenant_id"], LLMType.VISION)
vision_model = LLMBundle(kwargs["tenant_id"], vision_model_config)
vision_model = LLMBundle(kwargs["tenant_id"], vision_model_config, lang=lang)
callback(0.2, "Visual model detected. Attempting to enhance figure extraction...")
except Exception as e:
logging.warning(f"Failed to detect figure extraction: {e}")
@@ -1168,7 +1175,12 @@ def chunk(filename, binary=None, from_page=0, to_page=MAXIMUM_PAGE_NUMBER, lang=
else:
section_images = [None] * len(sections)
section_images[idx] = combined_image
markdown_vision_parser = VisionFigureParser(vision_model=vision_model, figures_data=[((combined_image, ["markdown image"]), [(0, 0, 0, 0, 0)])], **kwargs)
markdown_vision_parser = VisionFigureParser(
vision_model=vision_model,
figures_data=[((combined_image, ["markdown image"]), [(0, 0, 0, 0, 0)])],
lang=lang,
**kwargs,
)
boosted_figures = markdown_vision_parser(callback=callback)
sections[idx] = (section_text + "\n\n" + "\n\n".join([fig[0][1] for fig in boosted_figures]), sections[idx][1])

View File

@@ -29,6 +29,9 @@ from common.constants import MAXIMUM_PAGE_NUMBER, MAXIMUM_TASK_PAGE_NUMBER
from common.parser_config_utils import normalize_layout_recognizer
logger = logging.getLogger(__name__)
class Pdf(PdfParser):
def __call__(self, filename, binary=None, from_page=0, to_page=MAXIMUM_PAGE_NUMBER, zoomin=3, callback=None):
from timeit import default_timer as timer
@@ -83,7 +86,12 @@ def chunk(filename, binary=None, from_page=0, to_page=MAXIMUM_PAGE_NUMBER, lang=
cks.append({"text": text, "image": image, "ck_type": ck_type})
vision_figure_parser_docx_wrapper_naive(cks, image_idxs, callback, **kwargs)
logger.info(
"DOCX figure vision enhancement: language=%s image_count=%d",
lang or "English",
len(image_idxs),
)
vision_figure_parser_docx_wrapper_naive(cks, image_idxs, callback, lang=lang, **kwargs)
sections = [ck["text"] for ck in cks if ck.get("text")]
callback(0.8, "Finish parsing.")

View File

@@ -184,6 +184,7 @@ def chunk(filename, binary=None, from_page=0, to_page=MAXIMUM_PAGE_NUMBER, lang=
tbls=tbls,
sections=sections,
callback=callback,
lang=lang,
**kwargs,
)
paper["tables"] = tbls

View File

@@ -63,7 +63,16 @@ def _deduplicate_column_names(columns):
class Excel(ExcelParser):
def __call__(self, fnm, binary=None, from_page=0, to_page=MAXIMUM_TASK_PAGE_NUMBER, callback=None, **kwargs):
def __call__(
self,
fnm,
binary=None,
from_page=0,
to_page=MAXIMUM_TASK_PAGE_NUMBER,
callback=None,
lang="English",
**kwargs,
):
if not binary:
wb = Excel._load_excel_to_workbook(fnm)
else:
@@ -80,7 +89,12 @@ class Excel(ExcelParser):
images = Excel._extract_images_from_worksheet(ws, sheetname=sheet_name)
pending_cell_images = []
if images:
image_descriptions = vision_figure_parser_figure_xlsx_wrapper(images=images, callback=callback, **kwargs)
image_descriptions = vision_figure_parser_figure_xlsx_wrapper(
images=images,
callback=callback,
lang=lang,
**kwargs,
)
if image_descriptions and len(image_descriptions) == len(images):
for i, bf in enumerate(image_descriptions):
desc = bf[0][1]
@@ -406,7 +420,15 @@ def chunk(filename, binary=None, from_page=0, to_page=MAXIMUM_TASK_PAGE_NUMBER,
if re.search(r"\.xlsx?$", filename, re.IGNORECASE):
callback(0.1, "Start to parse.")
excel_parser = Excel()
dfs, tbls = excel_parser(filename, binary, from_page=from_page, to_page=to_page, callback=callback, **kwargs)
dfs, tbls = excel_parser(
filename,
binary,
from_page=from_page,
to_page=to_page,
callback=callback,
lang=lang,
**kwargs,
)
elif re.search(r"\.txt$", filename, re.IGNORECASE):
callback(0.1, "Start to parse.")
txt = get_text(filename, binary)

View File

@@ -774,6 +774,7 @@ class Parser(ProcessBase):
self._canvas._tenant_id,
conf.get("vlm"),
callback=self.callback,
lang=getattr(self._canvas, "_language", None) or conf.get("lang") or "English",
)
# Emit the requested final PDF output format.
@@ -977,6 +978,7 @@ class Parser(ProcessBase):
self._canvas._tenant_id,
conf.get("vlm"),
callback=self.callback,
lang=getattr(self._canvas, "_language", None) or conf.get("lang") or "English",
)
self.set_output("json", sections)
@@ -1113,6 +1115,7 @@ class Parser(ProcessBase):
self._canvas._tenant_id,
conf.get("vlm"),
callback=self.callback,
lang=getattr(self._canvas, "_language", None) or conf.get("lang") or "English",
)
self.set_output("json", json_results)
else:

View File

@@ -164,16 +164,19 @@ def enhance_media_sections_with_vision(
tenant_id,
vlm_conf=None,
callback=None,
lang="English",
):
if not sections or not tenant_id:
return sections
lang = lang or "English"
try:
try:
vision_model_config = resolve_model_config(tenant_id, LLMType.VISION, vlm_conf["llm_id"])
except Exception:
vision_model_config = get_tenant_default_model_by_type(tenant_id, LLMType.VISION)
vision_model = LLMBundle(tenant_id, vision_model_config)
vision_model = LLMBundle(tenant_id, vision_model_config, lang=lang)
except Exception:
return sections
@@ -189,6 +192,7 @@ def enhance_media_sections_with_vision(
vision_model=vision_model,
figures_data=[((item["image"], [""]), [(0, 0, 0, 0, 0)])],
context_size=0,
lang=lang,
)(callback=callback)
except Exception:
continue

View File

@@ -360,14 +360,18 @@ def vision_llm_describe_prompt(page=None) -> str:
return template.render(page=page)
def vision_llm_figure_describe_prompt() -> str:
def vision_llm_figure_describe_prompt(language: str = "English") -> str:
template = PROMPT_JINJA_ENV.from_string(VISION_LLM_FIGURE_DESCRIBE_PROMPT)
return template.render()
return template.render(language=language)
def vision_llm_figure_describe_prompt_with_context(context_above: str, context_below: str) -> str:
def vision_llm_figure_describe_prompt_with_context(context_above: str, context_below: str, language: str = "English") -> str:
template = PROMPT_JINJA_ENV.from_string(VISION_LLM_FIGURE_DESCRIBE_PROMPT_WITH_CONTEXT)
return template.render(context_above=context_above, context_below=context_below)
return template.render(
context_above=context_above,
context_below=context_below,
language=language,
)
def tool_schema(tools_description: list[dict], complete_task=False):

View File

@@ -6,6 +6,12 @@ You are an expert visual data analyst.
Analyze the image and produce a textual representation strictly based on what is visible in the image.
## OUTPUT LANGUAGE
- Write all descriptions and field values in {{ language }}.
- Preserve all visible text verbatim in its original language; do not translate it.
- Keep the required output field names exactly as specified below.
## DECISION RULE (CRITICAL)
First, determine whether the image contains an explicit visual data representation with enumerable data units forming a coherent dataset.

View File

@@ -7,6 +7,12 @@ You are an expert visual data analyst.
Analyze the image and produce a textual representation strictly based on what is visible in the image.
Surrounding context may be used only for minimal clarification or disambiguation of terms that appear in the image, not as a source of new information.
## OUTPUT LANGUAGE
- Write all descriptions and field values in {{ language }}.
- Preserve all visible text verbatim in its original language; do not translate it.
- Keep the required output field names exactly as specified below.
## CONTEXT (ABOVE)
{{ context_above }}

View File

@@ -0,0 +1,298 @@
#
# 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()

View File

@@ -0,0 +1,50 @@
#
# 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 logging
from unittest.mock import Mock
import pytest
from rag.app import one
@pytest.mark.p1
def test_docx_chunk_forwards_language_to_vision_wrapper(monkeypatch, caplog):
docx_parser = Mock(return_value=[("caption", object(), None)])
monkeypatch.setattr(one.naive, "Docx", Mock(return_value=docx_parser))
vision_wrapper = Mock()
monkeypatch.setattr(one, "vision_figure_parser_docx_wrapper_naive", vision_wrapper)
monkeypatch.setattr(one.rag_tokenizer, "tokenize", lambda text: text)
monkeypatch.setattr(one.rag_tokenizer, "fine_grained_tokenize", lambda text: text)
monkeypatch.setattr(one, "tokenize", Mock())
with caplog.at_level(logging.INFO, logger=one.__name__):
one.chunk(
"document.docx",
binary=b"docx",
lang="Japanese",
callback=lambda *_args, **_kwargs: None,
tenant_id="tenant-id",
)
vision_wrapper.assert_called_once()
args = vision_wrapper.call_args.args
kwargs = vision_wrapper.call_args.kwargs
assert args[1] == [0]
assert kwargs["lang"] == "Japanese"
assert kwargs["tenant_id"] == "tenant-id"
assert "DOCX figure vision enhancement: language=Japanese image_count=1" in caplog.messages

View File

@@ -0,0 +1,100 @@
#
# 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 ast
from pathlib import Path
from unittest.mock import Mock
import pytest
from rag.app import naive
REPO_ROOT = Path(__file__).resolve().parents[4]
def _call_name(call):
if isinstance(call.func, ast.Name):
return call.func.id
if isinstance(call.func, ast.Attribute):
return call.func.attr
return None
@pytest.mark.p1
@pytest.mark.parametrize(
("relative_path", "expected_call_count"),
[
("rag/app/book.py", 1),
("rag/app/manual.py", 2),
("rag/app/naive.py", 2),
("rag/app/one.py", 1),
("rag/app/paper.py", 1),
("rag/app/table.py", 1),
],
)
def test_all_figure_wrapper_callers_forward_language(relative_path, expected_call_count):
tree = ast.parse((REPO_ROOT / relative_path).read_text())
calls = [node for node in ast.walk(tree) if isinstance(node, ast.Call) and (_call_name(node) or "").startswith("vision_figure_parser_")]
assert len(calls) == expected_call_count
for call in calls:
language = next((keyword.value for keyword in call.keywords if keyword.arg == "lang"), None)
assert isinstance(language, ast.Name), f"{relative_path}:{call.lineno} does not forward lang"
assert language.id == "lang"
@pytest.mark.p1
def test_markdown_chunk_forwards_language_to_model_and_figure_parser(monkeypatch):
markdown_parser = Mock(return_value=([("section", "")], [], [object()]))
monkeypatch.setattr(naive, "Markdown", Mock(return_value=markdown_parser))
monkeypatch.setattr(naive, "get_tenant_default_model_by_type", Mock(return_value={"llm_name": "vision-model"}))
vision_model = object()
llm_bundle = Mock(return_value=vision_model)
monkeypatch.setattr(naive, "LLMBundle", llm_bundle)
parser_instance = Mock(return_value=[((None, "description"), None)])
parser_factory = Mock(return_value=parser_instance)
monkeypatch.setattr(naive, "VisionFigureParser", parser_factory)
monkeypatch.setattr(naive.rag_tokenizer, "tokenize", lambda text: text)
monkeypatch.setattr(naive.rag_tokenizer, "fine_grained_tokenize", lambda text: text)
monkeypatch.setattr(naive, "num_tokens_from_string", lambda _text: 1)
monkeypatch.setattr(naive, "tokenize_table", Mock(return_value=[]))
monkeypatch.setattr(naive, "tokenize_chunks", Mock(return_value=[]))
monkeypatch.setattr(naive, "tokenize_chunks_with_images", Mock(return_value=[]))
naive.chunk(
"document.md",
binary=b"markdown",
lang="Japanese",
callback=lambda *_args, **_kwargs: None,
tenant_id="tenant-id",
is_root=False,
parser_config={
"chunk_token_num": 128,
"delimiter": "\n",
"analyze_hyperlink": False,
},
)
llm_bundle.assert_called_once_with(
"tenant-id",
{"llm_name": "vision-model"},
lang="Japanese",
)
assert parser_factory.call_args.kwargs["vision_model"] is vision_model
assert parser_factory.call_args.kwargs["lang"] == "Japanese"
parser_instance.assert_called_once()

View File

@@ -0,0 +1,122 @@
#
# 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 ast
import importlib.util
import sys
from pathlib import Path
from types import ModuleType, SimpleNamespace
from unittest.mock import Mock
import pytest
REPO_ROOT = Path(__file__).resolve().parents[5]
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_flow_utils(monkeypatch):
for package_name in (
"api",
"api.db",
"api.db.services",
"api.db.joint_services",
"common",
"deepdoc",
"deepdoc.parser",
"rag",
):
_package(monkeypatch, package_name)
_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(),
resolve_model_config=Mock(),
)
_module(monkeypatch, "common.constants", LLMType=SimpleNamespace(VISION="vision"))
_module(monkeypatch, "deepdoc.parser.figure_parser", VisionFigureParser=Mock())
_module(
monkeypatch,
"rag.nlp",
is_english=Mock(return_value=False),
random_choices=Mock(return_value=[]),
remove_contents_table=Mock(),
)
module_path = REPO_ROOT / "rag/flow/parser/utils.py"
spec = importlib.util.spec_from_file_location("test_flow_parser_utils_module", module_path)
module = importlib.util.module_from_spec(spec)
monkeypatch.setitem(sys.modules, spec.name, module)
spec.loader.exec_module(module)
return module
@pytest.mark.p1
@pytest.mark.parametrize(
("language", "expected_language"),
[
("Japanese", "Japanese"),
("", "English"),
],
)
def test_media_enhancement_forwards_language_to_model_and_parser(monkeypatch, language, expected_language):
utils = _load_flow_utils(monkeypatch)
model_config = {"llm_name": "vision-model"}
vision_model = object()
llm_bundle = Mock(return_value=vision_model)
parser_instance = Mock(return_value=[((None, "description"), None)])
parser_factory = Mock(return_value=parser_instance)
monkeypatch.setattr(utils, "resolve_model_config", Mock(return_value=model_config))
monkeypatch.setattr(utils, "LLMBundle", llm_bundle)
monkeypatch.setattr(utils, "VisionFigureParser", parser_factory)
sections = [{"text": "caption", "image": object(), "doc_type_kwd": "image"}]
result = utils.enhance_media_sections_with_vision(
sections,
"tenant-id",
{"llm_id": "vision-model"},
lang=language,
)
llm_bundle.assert_called_once_with("tenant-id", model_config, lang=expected_language)
assert parser_factory.call_args.kwargs["vision_model"] is vision_model
assert parser_factory.call_args.kwargs["lang"] == expected_language
assert result[0]["text"] == "caption\ndescription"
@pytest.mark.p1
def test_all_flow_media_enhancement_callers_forward_language():
tree = ast.parse((REPO_ROOT / "rag/flow/parser/parser.py").read_text())
calls = [node for node in ast.walk(tree) if isinstance(node, ast.Call) and isinstance(node.func, ast.Name) and node.func.id == "enhance_media_sections_with_vision"]
assert len(calls) == 3
for call in calls:
assert any(keyword.arg == "lang" for keyword in call.keywords), f"parser.py:{call.lineno} does not forward lang"

View File

@@ -0,0 +1,117 @@
#
# 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
import pytest
def _load_generator(monkeypatch):
repo_root = Path(__file__).resolve().parents[4]
json_repair = ModuleType("json_repair")
json_repair.repair_json = lambda text, **_kwargs: text
monkeypatch.setitem(sys.modules, "json_repair", json_repair)
common = ModuleType("common")
common.__path__ = [str(repo_root / "common")]
monkeypatch.setitem(sys.modules, "common", common)
misc_utils = ModuleType("common.misc_utils")
misc_utils.hash_str2int = lambda value, _mod=500: 0
monkeypatch.setitem(sys.modules, "common.misc_utils", misc_utils)
constants = ModuleType("common.constants")
constants.TAG_FLD = "tag"
monkeypatch.setitem(sys.modules, "common.constants", constants)
token_utils = ModuleType("common.token_utils")
token_utils.encoder = SimpleNamespace()
token_utils.num_tokens_from_string = len
monkeypatch.setitem(sys.modules, "common.token_utils", token_utils)
rag = ModuleType("rag")
rag.__path__ = [str(repo_root / "rag")]
monkeypatch.setitem(sys.modules, "rag", rag)
rag_nlp = ModuleType("rag.nlp")
rag_nlp.rag_tokenizer = SimpleNamespace()
monkeypatch.setitem(sys.modules, "rag.nlp", rag_nlp)
prompts = ModuleType("rag.prompts")
prompts.__path__ = [str(repo_root / "rag" / "prompts")]
monkeypatch.setitem(sys.modules, "rag.prompts", prompts)
template = ModuleType("rag.prompts.template")
template.load_prompt = lambda name: (repo_root / "rag" / "prompts" / f"{name}.md").read_text(encoding="utf-8").strip()
monkeypatch.setitem(sys.modules, "rag.prompts.template", template)
module_path = repo_root / "rag" / "prompts" / "generator.py"
spec = importlib.util.spec_from_file_location(
"test_vision_figure_prompt_generator",
module_path,
)
module = importlib.util.module_from_spec(spec)
monkeypatch.setitem(sys.modules, spec.name, module)
spec.loader.exec_module(module)
return module
@pytest.mark.p1
@pytest.mark.parametrize(
("function_name", "arguments", "expected_language"),
[
(
"vision_llm_figure_describe_prompt",
{},
"English",
),
(
"vision_llm_figure_describe_prompt",
{"language": "Chinese"},
"Chinese",
),
(
"vision_llm_figure_describe_prompt_with_context",
{"context_above": "Above", "context_below": "Below"},
"English",
),
(
"vision_llm_figure_describe_prompt_with_context",
{
"context_above": "Above",
"context_below": "Below",
"language": "Chinese",
},
"Chinese",
),
],
)
def test_figure_prompt_renders_output_language(
monkeypatch,
function_name,
arguments,
expected_language,
):
generator = _load_generator(monkeypatch)
prompt = getattr(generator, function_name)(**arguments)
assert f"Write all descriptions and field values in {expected_language}." in prompt
assert "Preserve all visible text verbatim in its original language" in prompt
assert "{{ language }}" not in prompt