From c524f450e0c3f7208f7bb0f949b0fb19597901c7 Mon Sep 17 00:00:00 2001 From: xugangqiang Date: Fri, 24 Jul 2026 14:51:07 +0800 Subject: [PATCH] Fix parser params --- .../component/docx_vision_dispatch.go | 167 +++++++--------- .../component/docx_vision_dispatch_test.go | 181 ++++++++++++++++++ internal/ingestion/component/extractor.go | 7 + .../ingestion/component/extractor_test.go | 42 ++++ .../component/markdown_vision_dispatch.go | 5 +- .../ingestion/component/media_dispatch.go | 15 ++ .../component/media_dispatch_test.go | 151 +++++++++++++++ internal/ingestion/component/parser.go | 11 +- .../component/pdf_vision_dispatch.go | 20 +- .../component/pdf_vision_dispatch_test.go | 85 ++++++++ internal/ingestion/component/tokenizer.go | 6 +- .../ingestion/component/tokenizer_test.go | 26 +++ internal/utility/split.go | 11 +- internal/utility/split_test.go | 27 ++- 14 files changed, 640 insertions(+), 114 deletions(-) create mode 100644 internal/ingestion/component/docx_vision_dispatch_test.go create mode 100644 internal/ingestion/component/pdf_vision_dispatch_test.go diff --git a/internal/ingestion/component/docx_vision_dispatch.go b/internal/ingestion/component/docx_vision_dispatch.go index 622bb968e5..5c0ac2ad70 100644 --- a/internal/ingestion/component/docx_vision_dispatch.go +++ b/internal/ingestion/component/docx_vision_dispatch.go @@ -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 diff --git a/internal/ingestion/component/docx_vision_dispatch_test.go b/internal/ingestion/component/docx_vision_dispatch_test.go new file mode 100644 index 0000000000..7554a2adac --- /dev/null +++ b/internal/ingestion/component/docx_vision_dispatch_test.go @@ -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": "
", "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 != "
" { + 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: "![Image](data:image/png;base64,abc)", + 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 != "![Image](data:image/png;base64,abc)" { + t.Errorf("markdown mutated: %q", res.Markdown) + } +} diff --git a/internal/ingestion/component/extractor.go b/internal/ingestion/component/extractor.go index 309830fc57..ac4ae73e5d 100644 --- a/internal/ingestion/component/extractor.go +++ b/internal/ingestion/component/extractor.go @@ -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 { diff --git a/internal/ingestion/component/extractor_test.go b/internal/ingestion/component/extractor_test.go index 79e828e0d5..79ccb1df0e 100644 --- a/internal/ingestion/component/extractor_test.go +++ b/internal/ingestion/component/extractor_test.go @@ -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) { diff --git a/internal/ingestion/component/markdown_vision_dispatch.go b/internal/ingestion/component/markdown_vision_dispatch.go index 59f3793d6b..069900b7df 100644 --- a/internal/ingestion/component/markdown_vision_dispatch.go +++ b/internal/ingestion/component/markdown_vision_dispatch.go @@ -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) diff --git a/internal/ingestion/component/media_dispatch.go b/internal/ingestion/component/media_dispatch.go index e7a14f6721..444730d037 100644 --- a/internal/ingestion/component/media_dispatch.go +++ b/internal/ingestion/component/media_dispatch.go @@ -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", diff --git a/internal/ingestion/component/media_dispatch_test.go b/internal/ingestion/component/media_dispatch_test.go index 0b12589fad..b2bde8e5dd 100644 --- a/internal/ingestion/component/media_dispatch_test.go +++ b/internal/ingestion/component/media_dispatch_test.go @@ -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) + } +} diff --git a/internal/ingestion/component/parser.go b/internal/ingestion/component/parser.go index 6288455c3e..fd9f6157bb 100644 --- a/internal/ingestion/component/parser.go +++ b/internal/ingestion/component/parser.go @@ -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 diff --git a/internal/ingestion/component/pdf_vision_dispatch.go b/internal/ingestion/component/pdf_vision_dispatch.go index 8354a833c9..6700d620c1 100644 --- a/internal/ingestion/component/pdf_vision_dispatch.go +++ b/internal/ingestion/component/pdf_vision_dispatch.go @@ -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 diff --git a/internal/ingestion/component/pdf_vision_dispatch_test.go b/internal/ingestion/component/pdf_vision_dispatch_test.go new file mode 100644 index 0000000000..5af4b97b6d --- /dev/null +++ b/internal/ingestion/component/pdf_vision_dispatch_test.go @@ -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") + } +} diff --git a/internal/ingestion/component/tokenizer.go b/internal/ingestion/component/tokenizer.go index dcae6b5ea9..5029772032 100644 --- a/internal/ingestion/component/tokenizer.go +++ b/internal/ingestion/component/tokenizer.go @@ -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) } diff --git a/internal/ingestion/component/tokenizer_test.go b/internal/ingestion/component/tokenizer_test.go index b02b6bf9de..4883504aab 100644 --- a/internal/ingestion/component/tokenizer_test.go +++ b/internal/ingestion/component/tokenizer_test.go @@ -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) diff --git a/internal/utility/split.go b/internal/utility/split.go index d3163260e0..0334fa09d3 100644 --- a/internal/utility/split.go +++ b/internal/utility/split.go @@ -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) } } diff --git a/internal/utility/split_test.go b/internal/utility/split_test.go index 6225edf666..3fb4acd29c 100644 --- a/internal/utility/split_test.go +++ b/internal/utility/split_test.go @@ -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 {