mirror of
https://github.com/infiniflow/ragflow.git
synced 2026-08-17 22:08:29 +08:00
fix: honor dataset language across VisionFigureParser paths (#17227)
This commit is contained in:
@@ -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,
|
||||
|
||||
@@ -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.")
|
||||
|
||||
@@ -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)
|
||||
|
||||
@@ -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])
|
||||
|
||||
|
||||
@@ -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.")
|
||||
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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)
|
||||
|
||||
@@ -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:
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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):
|
||||
|
||||
@@ -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.
|
||||
|
||||
@@ -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 }}
|
||||
|
||||
298
test/unit_test/deepdoc/parser/test_figure_parser.py
Normal file
298
test/unit_test/deepdoc/parser/test_figure_parser.py
Normal 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()
|
||||
50
test/unit_test/rag/app/test_one.py
Normal file
50
test/unit_test/rag/app/test_one.py
Normal 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
|
||||
100
test/unit_test/rag/app/test_vision_language_callers.py
Normal file
100
test/unit_test/rag/app/test_vision_language_callers.py
Normal 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()
|
||||
122
test/unit_test/rag/flow/parser/test_vision_language.py
Normal file
122
test/unit_test/rag/flow/parser/test_vision_language.py
Normal 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"
|
||||
117
test/unit_test/rag/prompts/test_vision_figure_prompt.py
Normal file
117
test/unit_test/rag/prompts/test_vision_figure_prompt.py
Normal 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
|
||||
Reference in New Issue
Block a user