mirror of
https://github.com/infiniflow/ragflow.git
synced 2026-07-25 09:53:29 +08:00
Fix parser params
This commit is contained in:
@@ -16,13 +16,14 @@
|
||||
|
||||
// DOCX vision figure dispatch: enriches the parse result with
|
||||
// LLM-generated descriptions of embedded images, mirroring
|
||||
// Python's vision_figure_parser_docx_wrapper_naive in
|
||||
// deepdoc/parser/figure_parser.py.
|
||||
// Python's enhance_media_sections_with_vision in
|
||||
// rag/flow/parser/utils.py (invoked from parser.py:_doc's JSON branch).
|
||||
//
|
||||
// Unlike the PDF vision path (which replaces dispatchParse
|
||||
// entirely), DOCX vision is a post-processing step: it takes
|
||||
// the already-parsed markdown + extracted figures and augments
|
||||
// the markdown text with vision model descriptions.
|
||||
// Unlike the PDF vision path (which replaces dispatchParse entirely),
|
||||
// DOCX vision is a post-processing step. It mirrors Python exactly:
|
||||
// vision enrichment happens ONLY on the JSON output path, where each
|
||||
// item carries a doc_type_kwd and an optional image. The markdown path
|
||||
// performs no vision enrichment in Python, so it must not here either.
|
||||
|
||||
package component
|
||||
|
||||
@@ -31,7 +32,6 @@ import (
|
||||
"fmt"
|
||||
"os"
|
||||
"path/filepath"
|
||||
"sort"
|
||||
"strings"
|
||||
"sync"
|
||||
|
||||
@@ -59,18 +59,20 @@ var (
|
||||
docxVisionPromptMu sync.RWMutex
|
||||
)
|
||||
|
||||
// maybeDispatchDOCXVision checks whether the dispatch result for a
|
||||
// DOCX file contains embedded image figures and, when a vision
|
||||
// model is available, enriches the markdown with AI-generated
|
||||
// figure descriptions. It mirrors the Python flow:
|
||||
// maybeDispatchDOCXVision enriches a DOCX parse result with vision-model
|
||||
// descriptions of embedded images. It mirrors Python's
|
||||
// enhance_media_sections_with_vision (rag/flow/parser/utils.py:162), which
|
||||
// runs only in the JSON output branch of parser.py:_doc.
|
||||
//
|
||||
// 1. naive_merge_docx → chunks (text + images + context)
|
||||
// 2. vision_figure_parser_docx_wrapper_naive → LLM descriptions
|
||||
// For each JSON item whose doc_type_kwd is "image" or "table" AND that
|
||||
// carries a non-empty "image" field, the vision model describes the image
|
||||
// and the description is appended to the item's text (Python:
|
||||
// item["text"] = f"{text}\n{parsed_text}" if text else parsed_text). Items
|
||||
// without an image (e.g. DOCX tables) are left untouched, exactly as Python
|
||||
// skips them via `if item.get("image") is None: continue`.
|
||||
//
|
||||
// The function is called AFTER dispatchParse so the normal parse
|
||||
// path produces figures in dispatched.File["figures"].
|
||||
// It returns (result, handled, error). handled is true when the
|
||||
// dispatched result was modified.
|
||||
// The markdown output path receives no vision enrichment — Python's DOCX
|
||||
// markdown branch only concatenates text and never calls the vision model.
|
||||
func maybeDispatchDOCXVision(
|
||||
ctx context.Context,
|
||||
fileType utility.FileType,
|
||||
@@ -81,11 +83,9 @@ func maybeDispatchDOCXVision(
|
||||
if fileType != utility.FileTypeDOCX {
|
||||
return dispatched, false, nil
|
||||
}
|
||||
if dispatched.Err != nil || dispatched.OutputFormat != "markdown" {
|
||||
return dispatched, false, nil
|
||||
}
|
||||
figs, hasFigures := extractDOCXFiguresFromDispatch(dispatched)
|
||||
if !hasFigures {
|
||||
// Python triggers vision enrichment only on the JSON path
|
||||
// (parser.py:_doc → enhance_media_sections_with_vision).
|
||||
if dispatched.Err != nil || dispatched.OutputFormat != "json" || len(dispatched.JSON) == 0 {
|
||||
return dispatched, false, nil
|
||||
}
|
||||
|
||||
@@ -102,105 +102,74 @@ func maybeDispatchDOCXVision(
|
||||
return dispatched, false, nil
|
||||
}
|
||||
|
||||
descriptions := make([]string, len(figs))
|
||||
// Collect the indices of JSON items that carry an embeddable image.
|
||||
type target struct {
|
||||
idx int
|
||||
}
|
||||
var targets []target
|
||||
for i, item := range dispatched.JSON {
|
||||
kd, _ := item["doc_type_kwd"].(string)
|
||||
if kd != "image" && kd != "table" {
|
||||
continue
|
||||
}
|
||||
img, _ := item["image"].(string)
|
||||
if img == "" {
|
||||
continue
|
||||
}
|
||||
targets = append(targets, target{idx: i})
|
||||
}
|
||||
if len(targets) == 0 {
|
||||
return dispatched, false, nil
|
||||
}
|
||||
|
||||
descriptions := make([]string, len(targets))
|
||||
var wg sync.WaitGroup
|
||||
sem := make(chan struct{}, docxVisionConcurrency)
|
||||
|
||||
for i, fig := range figs {
|
||||
for slot, tg := range targets {
|
||||
wg.Add(1)
|
||||
go func(idx int, f map[string]any) {
|
||||
go func(slot int, itemIdx int) {
|
||||
defer wg.Done()
|
||||
sem <- struct{}{}
|
||||
defer func() { <-sem }()
|
||||
|
||||
imageB64, _ := f["image"].(string)
|
||||
ctxAbove, _ := f["context_above"].(string)
|
||||
ctxBelow, _ := f["context_below"].(string)
|
||||
|
||||
if imageB64 == "" {
|
||||
img, _ := dispatched.JSON[itemIdx]["image"].(string)
|
||||
if img == "" {
|
||||
return
|
||||
}
|
||||
|
||||
prompt, err := docxVisionPromptBuilder(ctxAbove, ctxBelow)
|
||||
if err != nil {
|
||||
// DOCX JSON items have no surrounding context (unlike the
|
||||
// former markdown path), so use the bare figure prompt —
|
||||
// matching Python's VisionFigureParser(context_size=0).
|
||||
prompt, perr := docxVisionPromptBuilder("", "")
|
||||
if perr != nil {
|
||||
return
|
||||
}
|
||||
|
||||
messages := buildVisionMessages(prompt, imageB64)
|
||||
resp, err := visionChatInvoker(ctx, driver, modelName, messages, apiConfig)
|
||||
if err != nil {
|
||||
messages := buildVisionMessages(prompt, img)
|
||||
resp, ierr := visionChatInvoker(ctx, driver, modelName, messages, apiConfig)
|
||||
if ierr != nil {
|
||||
return
|
||||
}
|
||||
descriptions[idx] = extractDOCXVisionAnswer(resp)
|
||||
}(i, fig)
|
||||
descriptions[slot] = extractDOCXVisionAnswer(resp)
|
||||
}(slot, tg.idx)
|
||||
}
|
||||
wg.Wait()
|
||||
|
||||
// Insert each description at the figure's position in the markdown,
|
||||
// matching Python's `chunks[idx]["text"] += description`.
|
||||
// Figures carry a "marker" (text immediately before the image) to
|
||||
// locate the insertion point. Process in reverse order so earlier
|
||||
// insertions don't shift later markers.
|
||||
md := dispatched.Markdown
|
||||
type indexedDesc struct {
|
||||
idx int
|
||||
desc string
|
||||
}
|
||||
var inserts []indexedDesc
|
||||
for i, d := range descriptions {
|
||||
if d = strings.TrimSpace(d); d == "" {
|
||||
modified := false
|
||||
for slot, tg := range targets {
|
||||
desc := strings.TrimSpace(descriptions[slot])
|
||||
if desc == "" {
|
||||
continue
|
||||
}
|
||||
if i >= len(figs) {
|
||||
continue
|
||||
existing, _ := dispatched.JSON[tg.idx]["text"].(string)
|
||||
if existing != "" {
|
||||
dispatched.JSON[tg.idx]["text"] = existing + "\n" + desc
|
||||
} else {
|
||||
dispatched.JSON[tg.idx]["text"] = desc
|
||||
}
|
||||
marker, _ := figs[i]["marker"].(string)
|
||||
if marker != "" {
|
||||
if pos := strings.LastIndex(md, marker); pos >= 0 {
|
||||
inserts = append(inserts, indexedDesc{idx: pos + len(marker), desc: d})
|
||||
continue
|
||||
}
|
||||
}
|
||||
// Fallback: try context_above as a search anchor.
|
||||
if ctx, _ := figs[i]["context_above"].(string); ctx != "" {
|
||||
if pos := strings.LastIndex(md, ctx); pos >= 0 {
|
||||
inserts = append(inserts, indexedDesc{idx: pos + len(ctx), desc: d})
|
||||
continue
|
||||
}
|
||||
}
|
||||
// No anchor found — append to end.
|
||||
inserts = append(inserts, indexedDesc{idx: len(md), desc: "\n\n" + d})
|
||||
modified = true
|
||||
}
|
||||
// Sort descending by position for stable insertion.
|
||||
sort.Slice(inserts, func(a, b int) bool { return inserts[a].idx > inserts[b].idx })
|
||||
for _, ins := range inserts {
|
||||
desc := ins.desc
|
||||
if !strings.HasPrefix(desc, "\n") {
|
||||
desc = "\n\n" + desc
|
||||
}
|
||||
md = md[:ins.idx] + desc + md[ins.idx:]
|
||||
}
|
||||
dispatched.Markdown = md
|
||||
|
||||
return dispatched, true, nil
|
||||
}
|
||||
|
||||
func extractDOCXFiguresFromDispatch(dispatched parserDispatchResult) ([]map[string]any, bool) {
|
||||
if dispatched.File == nil {
|
||||
return nil, false
|
||||
}
|
||||
raw, ok := dispatched.File["figures"]
|
||||
if !ok {
|
||||
return nil, false
|
||||
}
|
||||
list, ok := raw.([]map[string]any)
|
||||
if !ok {
|
||||
return nil, false
|
||||
}
|
||||
if len(list) == 0 {
|
||||
return nil, false
|
||||
}
|
||||
return list, true
|
||||
return dispatched, modified, nil
|
||||
}
|
||||
|
||||
// buildDOCXVisionPrompt loads the figure-describe prompt template
|
||||
|
||||
181
internal/ingestion/component/docx_vision_dispatch_test.go
Normal file
181
internal/ingestion/component/docx_vision_dispatch_test.go
Normal file
@@ -0,0 +1,181 @@
|
||||
//
|
||||
// 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.
|
||||
|
||||
package component
|
||||
|
||||
import (
|
||||
"context"
|
||||
"sync"
|
||||
"testing"
|
||||
|
||||
"ragflow/internal/entity"
|
||||
modelModule "ragflow/internal/entity/models"
|
||||
"ragflow/internal/utility"
|
||||
)
|
||||
|
||||
// docxVisionFakeDriver satisfies modelModule.ModelDriver but never reaches the
|
||||
// network: docxVisionCaptureInvoker intercepts the call before the driver.
|
||||
type docxVisionFakeDriver struct {
|
||||
modelModule.ModelDriver
|
||||
}
|
||||
|
||||
// docxVisionCaptureInvoker records the requested image and returns a fixed
|
||||
// description, mirroring the markdown-vision test's capture driver.
|
||||
type docxVisionCaptureInvoker struct {
|
||||
mu sync.Mutex
|
||||
images []string
|
||||
captured []modelModule.Message
|
||||
}
|
||||
|
||||
func (c *docxVisionCaptureInvoker) invoke(
|
||||
ctx context.Context,
|
||||
driver modelModule.ModelDriver,
|
||||
modelName string,
|
||||
messages []modelModule.Message,
|
||||
apiConfig *modelModule.APIConfig,
|
||||
) (*modelModule.ChatResponse, error) {
|
||||
c.mu.Lock()
|
||||
c.captured = append(c.captured, messages...)
|
||||
// Pull the data URI out of the second content part.
|
||||
if parts, ok := messages[0].Content.([]interface{}); ok && len(parts) >= 2 {
|
||||
if img, ok := parts[1].(map[string]any); ok {
|
||||
if url, ok := img["image_url"].(map[string]any); ok {
|
||||
if u, ok := url["url"].(string); ok {
|
||||
c.images = append(c.images, u)
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
c.mu.Unlock()
|
||||
ans := "a diagram of a pipeline"
|
||||
return &modelModule.ChatResponse{Answer: &ans}, nil
|
||||
}
|
||||
|
||||
// TestMaybeDispatchDOCXVision_EnhancesJSONImages verifies Diff 2.4: DOCX vision
|
||||
// enhancement must trigger on the JSON output path (like Python's
|
||||
// enhance_media_sections_with_vision in parser.py:_doc) and must NOT trigger on
|
||||
// the markdown path. Image items with a non-empty `image` field get their VLM
|
||||
// description appended to `text`; table items (no image) and text items are
|
||||
// left untouched.
|
||||
func TestMaybeDispatchDOCXVision_EnhancesJSONImages(t *testing.T) {
|
||||
origResolver := resolveTenantModelByType
|
||||
origInvoker := visionChatInvoker
|
||||
origPrompt := docxVisionPromptBuilder
|
||||
defer func() {
|
||||
resolveTenantModelByType = origResolver
|
||||
visionChatInvoker = origInvoker
|
||||
docxVisionPromptBuilder = origPrompt
|
||||
}()
|
||||
|
||||
resolveTenantModelByType = func(tenantID string, modelType entity.ModelType) (modelModule.ModelDriver, string, *modelModule.APIConfig, int, error) {
|
||||
return &docxVisionFakeDriver{}, "docx-vision-model", &modelModule.APIConfig{}, 0, nil
|
||||
}
|
||||
invoker := &docxVisionCaptureInvoker{}
|
||||
visionChatInvoker = invoker.invoke
|
||||
docxVisionPromptBuilder = func(string, string) (string, error) { return "describe the figure", nil }
|
||||
|
||||
dispatched := parserDispatchResult{
|
||||
OutputFormat: "json",
|
||||
DocType: "docx",
|
||||
JSON: []map[string]any{
|
||||
{"text": "Intro paragraph", "image": nil, "doc_type_kwd": "text"},
|
||||
{"text": "", "image": "aGVsbG8taW1hZ2U=", "doc_type_kwd": "image"},
|
||||
{"text": "<table></table>", "image": nil, "doc_type_kwd": "table"},
|
||||
},
|
||||
}
|
||||
|
||||
res, handled, err := maybeDispatchDOCXVision(
|
||||
context.Background(),
|
||||
utility.FileTypeDOCX,
|
||||
dispatched,
|
||||
map[string]any{"tenant_id": "t1"},
|
||||
defaultSetups(),
|
||||
)
|
||||
if err != nil {
|
||||
t.Fatalf("maybeDispatchDOCXVision: unexpected error: %v", err)
|
||||
}
|
||||
if !handled {
|
||||
t.Fatal("handled = false, want true (JSON image item should be enhanced)")
|
||||
}
|
||||
if len(res.JSON) != 3 {
|
||||
t.Fatalf("JSON len = %d, want 3", len(res.JSON))
|
||||
}
|
||||
if got := res.JSON[0]["text"].(string); got != "Intro paragraph" {
|
||||
t.Errorf("text item text = %q, want unchanged", got)
|
||||
}
|
||||
// image item: VLM description appended to (empty) text.
|
||||
if got, _ := res.JSON[1]["text"].(string); got != "a diagram of a pipeline" {
|
||||
t.Errorf("image item text = %q, want appended VLM description", got)
|
||||
}
|
||||
if got, _ := res.JSON[2]["text"].(string); got != "<table></table>" {
|
||||
t.Errorf("table item text = %q, want unchanged (no image)", got)
|
||||
}
|
||||
if len(invoker.images) != 1 {
|
||||
t.Fatalf("vision invoker called %d times, want 1 (only the image item)", len(invoker.images))
|
||||
}
|
||||
if want := "data:image/png;base64,aGVsbG8taW1hZ2U="; invoker.images[0] != want {
|
||||
t.Errorf("vision image data URI = %q, want %q", invoker.images[0], want)
|
||||
}
|
||||
}
|
||||
|
||||
// TestMaybeDispatchDOCXVision_JSONOnly verifies Diff 2.4: the markdown output
|
||||
// path must NOT be enhanced (Python's markdown/docx branch performs no vision
|
||||
// enrichment). A markdown result with embedded figures is returned untouched.
|
||||
func TestMaybeDispatchDOCXVision_JSONOnly(t *testing.T) {
|
||||
origResolver := resolveTenantModelByType
|
||||
origInvoker := visionChatInvoker
|
||||
defer func() {
|
||||
resolveTenantModelByType = origResolver
|
||||
visionChatInvoker = origInvoker
|
||||
}()
|
||||
|
||||
called := false
|
||||
resolveTenantModelByType = func(tenantID string, modelType entity.ModelType) (modelModule.ModelDriver, string, *modelModule.APIConfig, int, error) {
|
||||
called = true
|
||||
return &docxVisionFakeDriver{}, "m", &modelModule.APIConfig{}, 0, nil
|
||||
}
|
||||
visionChatInvoker = func(ctx context.Context, d modelModule.ModelDriver, m string, msgs []modelModule.Message, c *modelModule.APIConfig) (*modelModule.ChatResponse, error) {
|
||||
called = true
|
||||
ans := "x"
|
||||
return &modelModule.ChatResponse{Answer: &ans}, nil
|
||||
}
|
||||
|
||||
dispatched := parserDispatchResult{
|
||||
OutputFormat: "markdown",
|
||||
DocType: "docx",
|
||||
Markdown: "",
|
||||
File: map[string]any{"figures": []map[string]any{{"image": "abc", "marker": "x"}}},
|
||||
}
|
||||
|
||||
res, handled, err := maybeDispatchDOCXVision(
|
||||
context.Background(),
|
||||
utility.FileTypeDOCX,
|
||||
dispatched,
|
||||
map[string]any{"tenant_id": "t1"},
|
||||
defaultSetups(),
|
||||
)
|
||||
if err != nil {
|
||||
t.Fatalf("maybeDispatchDOCXVision: unexpected error: %v", err)
|
||||
}
|
||||
if handled {
|
||||
t.Error("handled = true, want false (markdown path must not be enhanced)")
|
||||
}
|
||||
if called {
|
||||
t.Error("vision model was resolved/invoked on the markdown path")
|
||||
}
|
||||
if res.Markdown != "" {
|
||||
t.Errorf("markdown mutated: %q", res.Markdown)
|
||||
}
|
||||
}
|
||||
@@ -169,6 +169,13 @@ func NewExtractorComponent(params map[string]any) (runtime.Component, error) {
|
||||
}
|
||||
if v, ok := params["prompt"].(string); ok {
|
||||
p.Prompt = v
|
||||
} else if v, ok := params["prompts"].(string); ok && v != "" {
|
||||
// Python agent/component/llm.py:119-120 normalizes a bare-string
|
||||
// prompts into [{"role":"user","content":prompts}]. Mirror that
|
||||
// here so a front-end/template that emits prompts as a string
|
||||
// (the graph.nodes form / dsl testdata) is not silently dropped
|
||||
// by the .([]any) assertion on the list branch below.
|
||||
p.Prompt = v
|
||||
} else if promptsRaw, ok := params["prompts"].([]any); ok && len(promptsRaw) > 0 {
|
||||
if first, ok := promptsRaw[0].(map[string]any); ok {
|
||||
if content, ok := first["content"].(string); ok {
|
||||
|
||||
@@ -570,6 +570,48 @@ func TestNewExtractorComponent_PromptsArray_PromptWins(t *testing.T) {
|
||||
}
|
||||
}
|
||||
|
||||
// TestNewExtractorComponent_PromptsString verifies that a bare-string
|
||||
// "prompts" (the shape emitted by the front-end graph.nodes form and
|
||||
// the dsl/testdata templates) is normalized into Param.Prompt, mirroring
|
||||
// Python agent/component/llm.py:119-120 which coerces a string prompts
|
||||
// into [{"role":"user","content":prompts}]. Without this normalization
|
||||
// the string form is silently dropped (the .([]any) assertion fails).
|
||||
func TestNewExtractorComponent_PromptsString(t *testing.T) {
|
||||
withStubChatInvoker(t, stubResponse{Content: "ok"})
|
||||
comp, err := NewExtractorComponent(map[string]any{
|
||||
"field_name": "out",
|
||||
"prompts": "Content: {TitleChunker:FlatMiceFix@chunks}",
|
||||
})
|
||||
if err != nil {
|
||||
t.Fatalf("NewExtractorComponent: %v", err)
|
||||
}
|
||||
ec := comp.(*ExtractorComponent)
|
||||
want := "Content: {TitleChunker:FlatMiceFix@chunks}"
|
||||
if ec.Param.Prompt != want {
|
||||
t.Errorf("Prompt = %q, want %q (string prompts should be normalized)", ec.Param.Prompt, want)
|
||||
}
|
||||
}
|
||||
|
||||
// TestNewExtractorComponent_PromptsString_PromptWins verifies that
|
||||
// "prompt" (string) still takes priority over a string-form "prompts"
|
||||
// when both are present, matching the prompt>prompts precedence of
|
||||
// the list-form path (TestNewExtractorComponent_PromptsArray_PromptWins).
|
||||
func TestNewExtractorComponent_PromptsString_PromptWins(t *testing.T) {
|
||||
withStubChatInvoker(t, stubResponse{Content: "ok"})
|
||||
comp, err := NewExtractorComponent(map[string]any{
|
||||
"field_name": "out",
|
||||
"prompt": "Direct prompt wins.",
|
||||
"prompts": "Should be ignored.",
|
||||
})
|
||||
if err != nil {
|
||||
t.Fatalf("NewExtractorComponent: %v", err)
|
||||
}
|
||||
ec := comp.(*ExtractorComponent)
|
||||
if ec.Param.Prompt != "Direct prompt wins." {
|
||||
t.Errorf("Prompt = %q, want %q", ec.Param.Prompt, "Direct prompt wins.")
|
||||
}
|
||||
}
|
||||
|
||||
// TestNewExtractorComponent_SystemPromptWinsOverSysPrompt verifies
|
||||
// that "system_prompt" takes priority over "sys_prompt".
|
||||
func TestNewExtractorComponent_SystemPromptWinsOverSysPrompt(t *testing.T) {
|
||||
|
||||
@@ -78,7 +78,10 @@ func maybeDispatchMarkdownVision(
|
||||
var images []imgItem
|
||||
for i, item := range dispatched.JSON {
|
||||
kd, _ := item["doc_type_kwd"].(string)
|
||||
if kd != "image" {
|
||||
// Diff 2.5: Python enhances both image and table items
|
||||
// (parser/utils.py:181 checks {"image","table"}); only items
|
||||
// carrying an image are sent to the VLM (utils.py:183).
|
||||
if kd != "image" && kd != "table" {
|
||||
continue
|
||||
}
|
||||
img, _ := item["image"].(string)
|
||||
|
||||
@@ -315,6 +315,21 @@ func maybeDispatchAudio(
|
||||
if outputFormat == "" {
|
||||
outputFormat = "text"
|
||||
}
|
||||
// Diff 2.11: when output_format is "json" the transcription must be
|
||||
// carried as a JSON item. Returning it only in Text made the Invoke
|
||||
// switch silently drop it (the switch has no "json" branch and the
|
||||
// JSON slice was empty). Mirror the JSON-item shape used by the
|
||||
// other parser branches.
|
||||
if outputFormat == "json" {
|
||||
return parserDispatchResult{
|
||||
OutputFormat: "json",
|
||||
DocType: "audio",
|
||||
JSON: []map[string]any{{
|
||||
"text": transcription,
|
||||
"doc_type_kwd": "audio",
|
||||
}},
|
||||
}, true, nil
|
||||
}
|
||||
return parserDispatchResult{
|
||||
OutputFormat: outputFormat,
|
||||
DocType: "audio",
|
||||
|
||||
@@ -116,3 +116,154 @@ func TestMaybeDispatchImage_UsesSystemPrompt(t *testing.T) {
|
||||
t.Fatalf("VLM user text = %q, want %q (image branch must read system_prompt)", got, "自定义视觉提示")
|
||||
}
|
||||
}
|
||||
|
||||
// audioTranscribeDriver is a mock ModelDriver whose TranscribeAudio returns a
|
||||
// fixed transcription, so maybeDispatchAudio can be exercised without a real
|
||||
// ASR provider.
|
||||
type audioTranscribeDriver struct {
|
||||
modelModule.ModelDriver
|
||||
transcription string
|
||||
}
|
||||
|
||||
func (d *audioTranscribeDriver) TranscribeAudio(ctx context.Context, _ *string, _ *string, _ *modelModule.APIConfig, _ *modelModule.ASRConfig, _ *common.ModelUsage) (*modelModule.ASRResponse, error) {
|
||||
return &modelModule.ASRResponse{Text: d.transcription}, nil
|
||||
}
|
||||
|
||||
// TestMaybeDispatchAudio_JSONCarriesTranscription pins diff 2.11: when the
|
||||
// audio family's output_format is "json", the ASR transcription must be
|
||||
// carried in the JSON items (not only in the Text field). Before the fix the
|
||||
// branch returned Text only with an empty JSON slice, and the Invoke switch
|
||||
// silently dropped the transcription because it has no "json" branch.
|
||||
func TestMaybeDispatchAudio_JSONCarriesTranscription(t *testing.T) {
|
||||
origResolver := resolveTenantModelByType
|
||||
defer func() { resolveTenantModelByType = origResolver }()
|
||||
|
||||
const want = "hello world"
|
||||
drv := &audioTranscribeDriver{transcription: want}
|
||||
resolveTenantModelByType = func(tenantID string, modelType entity.ModelType) (modelModule.ModelDriver, string, *modelModule.APIConfig, int, error) {
|
||||
return drv, "asr-model", &modelModule.APIConfig{}, 0, nil
|
||||
}
|
||||
|
||||
setups := defaultSetups()
|
||||
setups["audio"]["output_format"] = "json"
|
||||
|
||||
res, dispatched, err := maybeDispatchAudio(
|
||||
context.Background(),
|
||||
utility.FileTypeAURAL,
|
||||
"test.mp3",
|
||||
[]byte("fake-audio"),
|
||||
map[string]any{"tenant_id": "t1"},
|
||||
setups,
|
||||
)
|
||||
if err != nil {
|
||||
t.Fatalf("maybeDispatchAudio: %v", err)
|
||||
}
|
||||
if !dispatched {
|
||||
t.Fatalf("expected dispatched=true for AURAL file")
|
||||
}
|
||||
if res.OutputFormat != "json" {
|
||||
t.Fatalf("OutputFormat = %q, want json", res.OutputFormat)
|
||||
}
|
||||
if len(res.JSON) != 1 {
|
||||
t.Fatalf("JSON len = %d, want 1 (transcription must be carried as a JSON item)", len(res.JSON))
|
||||
}
|
||||
if got, _ := res.JSON[0]["text"].(string); got != want {
|
||||
t.Fatalf("JSON[0].text = %q, want %q", got, want)
|
||||
}
|
||||
if got, _ := res.JSON[0]["doc_type_kwd"].(string); got != "audio" {
|
||||
t.Fatalf("JSON[0].doc_type_kwd = %q, want audio", got)
|
||||
}
|
||||
}
|
||||
|
||||
// TestMaybeDispatchAudio_TextCarriesTranscription guards the text path: with
|
||||
// output_format "text" the transcription stays in the Text field and JSON is
|
||||
// empty (current default after aligning with Python parser.py:232).
|
||||
func TestMaybeDispatchAudio_TextCarriesTranscription(t *testing.T) {
|
||||
origResolver := resolveTenantModelByType
|
||||
defer func() { resolveTenantModelByType = origResolver }()
|
||||
|
||||
const want = "hello world"
|
||||
drv := &audioTranscribeDriver{transcription: want}
|
||||
resolveTenantModelByType = func(tenantID string, modelType entity.ModelType) (modelModule.ModelDriver, string, *modelModule.APIConfig, int, error) {
|
||||
return drv, "asr-model", &modelModule.APIConfig{}, 0, nil
|
||||
}
|
||||
|
||||
setups := defaultSetups()
|
||||
setups["audio"]["output_format"] = "text"
|
||||
|
||||
res, dispatched, err := maybeDispatchAudio(
|
||||
context.Background(),
|
||||
utility.FileTypeAURAL,
|
||||
"test.mp3",
|
||||
[]byte("fake-audio"),
|
||||
map[string]any{"tenant_id": "t1"},
|
||||
setups,
|
||||
)
|
||||
if err != nil {
|
||||
t.Fatalf("maybeDispatchAudio: %v", err)
|
||||
}
|
||||
if !dispatched {
|
||||
t.Fatalf("expected dispatched=true for AURAL file")
|
||||
}
|
||||
if res.OutputFormat != "text" {
|
||||
t.Fatalf("OutputFormat = %q, want text", res.OutputFormat)
|
||||
}
|
||||
if res.Text != want {
|
||||
t.Fatalf("Text = %q, want %q", res.Text, want)
|
||||
}
|
||||
if len(res.JSON) != 0 {
|
||||
t.Fatalf("JSON len = %d, want 0 for text output", len(res.JSON))
|
||||
}
|
||||
}
|
||||
|
||||
// TestMaybeDispatchMarkdownVision_EnhancesTables pins diff 2.5: markdown
|
||||
// vision enhancement must also process items whose doc_type_kwd is "table"
|
||||
// (Python checks {"image","table"} in parser/utils.py:181), not only "image".
|
||||
// Before the fix the table item was skipped and never sent to the VLM.
|
||||
func TestMaybeDispatchMarkdownVision_EnhancesTables(t *testing.T) {
|
||||
origResolver := resolveTenantModelByType
|
||||
defer func() { resolveTenantModelByType = origResolver }()
|
||||
|
||||
drv := &imagePromptCaptureDriver{}
|
||||
resolveTenantModelByType = func(tenantID string, modelType entity.ModelType) (modelModule.ModelDriver, string, *modelModule.APIConfig, int, error) {
|
||||
return drv, "img-model", &modelModule.APIConfig{}, 0, nil
|
||||
}
|
||||
|
||||
dispatched := parserDispatchResult{
|
||||
OutputFormat: "json",
|
||||
JSON: []map[string]any{
|
||||
{"doc_type_kwd": "table", "image": "base64table", "text": ""},
|
||||
},
|
||||
}
|
||||
|
||||
res, handled, err := maybeDispatchMarkdownVision(
|
||||
context.Background(),
|
||||
utility.FileTypeMarkdown,
|
||||
dispatched,
|
||||
map[string]any{"tenant_id": "t1"},
|
||||
)
|
||||
if err != nil {
|
||||
t.Fatalf("maybeDispatchMarkdownVision: %v", err)
|
||||
}
|
||||
if !handled {
|
||||
t.Fatalf("expected handled=true for markdown with a table image")
|
||||
}
|
||||
if len(res.JSON) != 1 {
|
||||
t.Fatalf("JSON len = %d, want 1", len(res.JSON))
|
||||
}
|
||||
// The table item must have been sent to the VLM and its description appended.
|
||||
if got, _ := res.JSON[0]["text"].(string); got != "captured" {
|
||||
t.Fatalf("table item text = %q, want %q (table items must be vision-enhanced)", got, "captured")
|
||||
}
|
||||
}
|
||||
|
||||
// TestDefaultEmailOutputFormatIsJSON pins diff 2.2: the email family default
|
||||
// output_format must be "json" (matching Python parser.py:212), not "text".
|
||||
// With "text" the structured email fields (from/to/subject/attachments/...) are
|
||||
// flattened into a blob and lost downstream.
|
||||
func TestDefaultEmailOutputFormatIsJSON(t *testing.T) {
|
||||
got, _ := defaultSetups()["email"]["output_format"].(string)
|
||||
if got != "json" {
|
||||
t.Fatalf("email default output_format = %q, want json", got)
|
||||
}
|
||||
}
|
||||
|
||||
@@ -302,7 +302,7 @@ func defaultSetups() map[string]schema.ParserSetup {
|
||||
"from", "to", "cc", "bcc", "date", "subject",
|
||||
"body", "attachments", "metadata",
|
||||
},
|
||||
"output_format": "text",
|
||||
"output_format": "json",
|
||||
},
|
||||
"audio": {
|
||||
"suffix": []string{
|
||||
@@ -310,7 +310,7 @@ func defaultSetups() map[string]schema.ParserSetup {
|
||||
"aiff", "au", "midi", "wma", "realaudio", "vqf",
|
||||
"oggvorbis", "ape",
|
||||
},
|
||||
"output_format": "json",
|
||||
"output_format": "text",
|
||||
},
|
||||
"video": {
|
||||
"suffix": []string{"mp4", "avi", "mkv"},
|
||||
@@ -478,9 +478,10 @@ func (c *ParserComponent) Invoke(ctx context.Context, inputs map[string]any) (ma
|
||||
dispatched = dispatchParse(ctx, fileTypeExt, filename, binary, c.Setups)
|
||||
dispatched = hydrateEmptyDispatchPayload(dispatched, binary)
|
||||
|
||||
// DOCX vision figure enhancement: enrich the markdown
|
||||
// with LLM-generated descriptions of embedded images.
|
||||
// Mirrors Python's vision_figure_parser_docx_wrapper_naive.
|
||||
// DOCX vision figure enhancement: on the JSON output path,
|
||||
// append vision-model descriptions to embedded image items
|
||||
// (doc_type_kwd "image"). Mirrors Python's
|
||||
// enhance_media_sections_with_vision in parser.py:_doc.
|
||||
dispatched, _, _ = maybeDispatchDOCXVision(ctx, fileTypeExt, dispatched, inputs, c.Setups)
|
||||
|
||||
// Markdown vision figure enhancement: enrich parsed
|
||||
|
||||
@@ -358,16 +358,32 @@ func resolvePDFVisionModelID(setup schema.ParserSetup) (string, bool) {
|
||||
return "", false
|
||||
}
|
||||
|
||||
// isNamedPDFParseMethod reports whether raw is a recognized named PDF
|
||||
// parse method (as opposed to a CustomVLM model name). Its membership set
|
||||
// MUST stay aligned with the PDF whitelist enforced by
|
||||
// (*ParserComponent).Check() (parser.go:200-203):
|
||||
//
|
||||
// deepdoc, plain_text, mineru, docling,
|
||||
// opendataloader, tcadp parser, paddleocr, somark
|
||||
//
|
||||
// A parse_method that Check() rejects must not be treated as a named method
|
||||
// here, otherwise it silently falls through to the CustomVLM vision path
|
||||
// instead of failing fast at construction.
|
||||
//
|
||||
// The "@"-suffixed spellings cover the layout_recognizer convention used to
|
||||
// select a specific backend (e.g. "foo@mineru"); "@mineru" mirrors the
|
||||
// MinerU branch in maybeDispatchPDFVision (pdf_vision_dispatch.go:64-68).
|
||||
func isNamedPDFParseMethod(raw string) bool {
|
||||
method := strings.ToLower(strings.TrimSpace(raw))
|
||||
switch {
|
||||
case strings.HasSuffix(method, "@paddleocr"),
|
||||
case strings.HasSuffix(method, "@mineru"),
|
||||
strings.HasSuffix(method, "@paddleocr"),
|
||||
strings.HasSuffix(method, "@somark"),
|
||||
strings.HasSuffix(method, "@opendataloader"):
|
||||
return true
|
||||
}
|
||||
switch method {
|
||||
case "deepdoc", "mineru", "plain_text", "plain text", "plaintext", "paddleocr", "docling", "opendataloader", "somark", "tcadp", "tcadp parser":
|
||||
case "deepdoc", "plain_text", "mineru", "docling", "opendataloader", "tcadp parser", "paddleocr", "somark":
|
||||
return true
|
||||
}
|
||||
return false
|
||||
|
||||
85
internal/ingestion/component/pdf_vision_dispatch_test.go
Normal file
85
internal/ingestion/component/pdf_vision_dispatch_test.go
Normal file
@@ -0,0 +1,85 @@
|
||||
//
|
||||
// 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.
|
||||
|
||||
package component
|
||||
|
||||
import "testing"
|
||||
|
||||
// TestIsNamedPDFParseMethodWhitelistAligned verifies that the runtime
|
||||
// "named parse_method" classifier agrees with (*ParserComponent).Check()'s
|
||||
// PDF whitelist (parser.go:200-203):
|
||||
//
|
||||
// deepdoc, plain_text, mineru, docling,
|
||||
// opendataloader, tcadp parser, paddleocr, somark
|
||||
//
|
||||
// Diff 2.10: a parse_method that Check() rejects must NOT be treated as a
|
||||
// recognized named method by isNamedPDFParseMethod — otherwise it silently
|
||||
// falls through to the CustomVLM vision path instead of failing fast at
|
||||
// construction (and Python would have rejected it outright).
|
||||
func TestIsNamedPDFParseMethodWhitelistAligned(t *testing.T) {
|
||||
// Values that MUST be recognized (subset of the Check() whitelist,
|
||||
// case-insensitive).
|
||||
named := []string{
|
||||
"deepdoc", "plain_text", "mineru", "docling",
|
||||
"opendataloader", "tcadp parser", "paddleocr", "somark",
|
||||
"DeepDoc", "PLAIN_TEXT", "MinerU", "DocLing",
|
||||
"OpenDataLoader", "TCADP Parser", "PaddleOCR", "SoMark",
|
||||
}
|
||||
for _, v := range named {
|
||||
if !isNamedPDFParseMethod(v) {
|
||||
t.Errorf("isNamedPDFParseMethod(%q) = false, want true (in Check() whitelist)", v)
|
||||
}
|
||||
}
|
||||
|
||||
// Values that MUST NOT be recognized. These either duplicate the
|
||||
// whitelist with non-canonical spelling ("plain text"/"plaintext")
|
||||
// or are bare-family abbreviations ("tcadp") that Check() does not
|
||||
// accept, so they should be funneled to the CustomVLM path (or fail
|
||||
// construction) rather than masquerading as a named method.
|
||||
notNamed := []string{
|
||||
"plain text", "plaintext", "tcadp",
|
||||
"CustomVLM", "some_vlm", "gpt-4o",
|
||||
"", " ",
|
||||
}
|
||||
for _, v := range notNamed {
|
||||
if isNamedPDFParseMethod(v) {
|
||||
t.Errorf("isNamedPDFParseMethod(%q) = true, want false (not in Check() whitelist)", v)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// TestIsNamedPDFParseMethodLayoutSuffixes verifies the "@"-suffixed
|
||||
// layout_recognizer spelling is recognized per backend. "@mineru" must be
|
||||
// included to mirror the MinerU dispatch branch in maybeDispatchPDFVision
|
||||
// (pdf_vision_dispatch.go:64-68); without it a "foo@mineru" parse_method
|
||||
// would be misrouted to the generic vision path instead of MinerU.
|
||||
func TestIsNamedPDFParseMethodLayoutSuffixes(t *testing.T) {
|
||||
suffixed := []string{
|
||||
"foo@mineru", "@mineru",
|
||||
"foo@paddleocr", "@paddleocr",
|
||||
"foo@somark", "@somark",
|
||||
"foo@opendataloader", "@opendataloader",
|
||||
}
|
||||
for _, v := range suffixed {
|
||||
if !isNamedPDFParseMethod(v) {
|
||||
t.Errorf("isNamedPDFParseMethod(%q) = false, want true (recognized layout suffix)", v)
|
||||
}
|
||||
}
|
||||
|
||||
// A non-recognized suffix is a CustomVLM name, not a named method.
|
||||
if isNamedPDFParseMethod("foo@unknown") {
|
||||
t.Errorf("isNamedPDFParseMethod(%q) = true, want false", "foo@unknown")
|
||||
}
|
||||
}
|
||||
@@ -434,7 +434,11 @@ func (c *TokenizerComponent) embedChunks(ctx context.Context, tenantID, kbID, em
|
||||
if trimmedName == "" {
|
||||
log.Printf("Tokenizer: empty name provided from upstream, embedding will skip title weighting")
|
||||
} else {
|
||||
titleResults, err := encodeWithTimeout(ctx, embedder, []string{trimmedName})
|
||||
// Encode the raw name (no TrimSpace) to mirror Python
|
||||
// tokenizer.py:95 which passes name verbatim to embedding. The
|
||||
// empty-name guard above still uses TrimSpace, matching Python's
|
||||
// `.strip()==""` check at tokenizer.py:200.
|
||||
titleResults, err := encodeWithTimeout(ctx, embedder, []string{name})
|
||||
if err != nil {
|
||||
return nil, 0, fmt.Errorf("Tokenizer: encode title: %w", err)
|
||||
}
|
||||
|
||||
@@ -591,6 +591,32 @@ func TestTokenizerComponent_Embedding_EmptyNameWarnsAndUsesContentVector(t *test
|
||||
}
|
||||
}
|
||||
|
||||
// Python tokenizer.py:95 passes the raw name to embedding without .strip();
|
||||
// Go must match — the title embedding must receive the original name, not a
|
||||
// TrimSpace'd copy. The empty-name guard still uses TrimSpace (mirroring
|
||||
// Python's `.strip()==""` check at tokenizer.py:200), but the value encoded
|
||||
// is the raw name.
|
||||
func TestTokenizerComponent_Embedding_UsesRawNameNotTrimmed(t *testing.T) {
|
||||
requireTokenizerPool(t)
|
||||
c, stub := withStubEmbedder(t, 2)
|
||||
|
||||
if _, err := c.Invoke(context.Background(), map[string]any{
|
||||
"name": " report.pdf ",
|
||||
"output_format": "chunks",
|
||||
"chunks": []map[string]any{{"text": "alpha"}},
|
||||
}); err != nil {
|
||||
t.Fatalf("Invoke: %v", err)
|
||||
}
|
||||
if len(stub.callInputs) < 1 {
|
||||
t.Fatalf("callInputs len = %d, want >= 1", len(stub.callInputs))
|
||||
}
|
||||
// First call is the title embedding; it must receive the raw name with
|
||||
// surrounding whitespace preserved, matching Python.
|
||||
if got := stub.callInputs[0][0]; got != " report.pdf " {
|
||||
t.Fatalf("title embedding input = %q, want %q (raw, not trimmed)", got, " report.pdf ")
|
||||
}
|
||||
}
|
||||
|
||||
func TestTokenizerComponent_Embedding_TruncatesByMaxTokensMinus10(t *testing.T) {
|
||||
requireTokenizerPool(t)
|
||||
c, stub := withStubEmbedder(t, 2)
|
||||
|
||||
@@ -27,14 +27,15 @@ import (
|
||||
// task_executor.run_dataflow:879 re.split(r"[,,;;、\r\n]+", keywords).
|
||||
var keywordsSplitRE = regexp.MustCompile(`[,,;;、\r\n]+`)
|
||||
|
||||
// nonEmpty drops empty strings from parts and returns nil if none remain. It is
|
||||
// the shared tail of SplitKeywords and SplitQuestions: split by whatever
|
||||
// delimiter, then prune empties and collapse an all-empty result to nil so a
|
||||
// _kwd array is absent (nil) rather than [""].
|
||||
// nonEmpty drops empty and whitespace-only strings from parts and returns nil
|
||||
// if none remain. It is the shared tail of SplitKeywords and SplitQuestions:
|
||||
// split by whatever delimiter, then prune blanks and collapse an all-blank
|
||||
// result to nil so a _kwd array is absent (nil) rather than [""].
|
||||
// Mirrors Python task_executor's `if k.strip()` filter.
|
||||
func nonEmpty(parts []string) []string {
|
||||
out := make([]string, 0, len(parts))
|
||||
for _, p := range parts {
|
||||
if p != "" {
|
||||
if strings.TrimSpace(p) != "" {
|
||||
out = append(out, p)
|
||||
}
|
||||
}
|
||||
|
||||
@@ -16,7 +16,10 @@
|
||||
|
||||
package utility
|
||||
|
||||
import "testing"
|
||||
import (
|
||||
"reflect"
|
||||
"testing"
|
||||
)
|
||||
|
||||
// SplitKeywords - Python task_executor.run_dataflow:879
|
||||
// re.split(r"[,,;;、\r\n]+", keywords) with empty filtering.
|
||||
@@ -49,6 +52,28 @@ func TestSplitKeywords_FiltersEmptyStrings(t *testing.T) {
|
||||
}
|
||||
}
|
||||
|
||||
// Python task_executor.run_dataflow filters split parts with `if k.strip()`,
|
||||
// so whitespace-only segments (a lone space between commas, a tab, etc.) must
|
||||
// be dropped — not kept as " ". Input "a, ,b" must yield ["a","b"], matching
|
||||
// Python, not ["a"," ","b"].
|
||||
func TestSplitKeywords_FiltersWhitespaceOnly(t *testing.T) {
|
||||
cases := []struct {
|
||||
in string
|
||||
want []string
|
||||
}{
|
||||
{"a, ,b", []string{"a", "b"}},
|
||||
{"a,\t,b", []string{"a", "b"}},
|
||||
{" , , ", nil},
|
||||
{"kw1, ,kw2", []string{"kw1", "kw2"}},
|
||||
}
|
||||
for _, c := range cases {
|
||||
got := SplitKeywords(c.in)
|
||||
if !reflect.DeepEqual(got, c.want) {
|
||||
t.Errorf("SplitKeywords(%q) = %v, want %v", c.in, got, c.want)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func TestSplitKeywords_Empty(t *testing.T) {
|
||||
result := SplitKeywords("")
|
||||
if len(result) != 0 {
|
||||
|
||||
Reference in New Issue
Block a user