mirror of
https://github.com/infiniflow/ragflow.git
synced 2026-07-27 02:42:00 +08:00
Fix parser params
This commit is contained in:
@@ -16,13 +16,14 @@
|
|||||||
|
|
||||||
// DOCX vision figure dispatch: enriches the parse result with
|
// DOCX vision figure dispatch: enriches the parse result with
|
||||||
// LLM-generated descriptions of embedded images, mirroring
|
// LLM-generated descriptions of embedded images, mirroring
|
||||||
// Python's vision_figure_parser_docx_wrapper_naive in
|
// Python's enhance_media_sections_with_vision in
|
||||||
// deepdoc/parser/figure_parser.py.
|
// rag/flow/parser/utils.py (invoked from parser.py:_doc's JSON branch).
|
||||||
//
|
//
|
||||||
// Unlike the PDF vision path (which replaces dispatchParse
|
// Unlike the PDF vision path (which replaces dispatchParse entirely),
|
||||||
// entirely), DOCX vision is a post-processing step: it takes
|
// DOCX vision is a post-processing step. It mirrors Python exactly:
|
||||||
// the already-parsed markdown + extracted figures and augments
|
// vision enrichment happens ONLY on the JSON output path, where each
|
||||||
// the markdown text with vision model descriptions.
|
// 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
|
package component
|
||||||
|
|
||||||
@@ -31,7 +32,6 @@ import (
|
|||||||
"fmt"
|
"fmt"
|
||||||
"os"
|
"os"
|
||||||
"path/filepath"
|
"path/filepath"
|
||||||
"sort"
|
|
||||||
"strings"
|
"strings"
|
||||||
"sync"
|
"sync"
|
||||||
|
|
||||||
@@ -59,18 +59,20 @@ var (
|
|||||||
docxVisionPromptMu sync.RWMutex
|
docxVisionPromptMu sync.RWMutex
|
||||||
)
|
)
|
||||||
|
|
||||||
// maybeDispatchDOCXVision checks whether the dispatch result for a
|
// maybeDispatchDOCXVision enriches a DOCX parse result with vision-model
|
||||||
// DOCX file contains embedded image figures and, when a vision
|
// descriptions of embedded images. It mirrors Python's
|
||||||
// model is available, enriches the markdown with AI-generated
|
// enhance_media_sections_with_vision (rag/flow/parser/utils.py:162), which
|
||||||
// figure descriptions. It mirrors the Python flow:
|
// runs only in the JSON output branch of parser.py:_doc.
|
||||||
//
|
//
|
||||||
// 1. naive_merge_docx → chunks (text + images + context)
|
// For each JSON item whose doc_type_kwd is "image" or "table" AND that
|
||||||
// 2. vision_figure_parser_docx_wrapper_naive → LLM descriptions
|
// 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
|
// The markdown output path receives no vision enrichment — Python's DOCX
|
||||||
// path produces figures in dispatched.File["figures"].
|
// markdown branch only concatenates text and never calls the vision model.
|
||||||
// It returns (result, handled, error). handled is true when the
|
|
||||||
// dispatched result was modified.
|
|
||||||
func maybeDispatchDOCXVision(
|
func maybeDispatchDOCXVision(
|
||||||
ctx context.Context,
|
ctx context.Context,
|
||||||
fileType utility.FileType,
|
fileType utility.FileType,
|
||||||
@@ -81,11 +83,9 @@ func maybeDispatchDOCXVision(
|
|||||||
if fileType != utility.FileTypeDOCX {
|
if fileType != utility.FileTypeDOCX {
|
||||||
return dispatched, false, nil
|
return dispatched, false, nil
|
||||||
}
|
}
|
||||||
if dispatched.Err != nil || dispatched.OutputFormat != "markdown" {
|
// Python triggers vision enrichment only on the JSON path
|
||||||
return dispatched, false, nil
|
// (parser.py:_doc → enhance_media_sections_with_vision).
|
||||||
}
|
if dispatched.Err != nil || dispatched.OutputFormat != "json" || len(dispatched.JSON) == 0 {
|
||||||
figs, hasFigures := extractDOCXFiguresFromDispatch(dispatched)
|
|
||||||
if !hasFigures {
|
|
||||||
return dispatched, false, nil
|
return dispatched, false, nil
|
||||||
}
|
}
|
||||||
|
|
||||||
@@ -102,105 +102,74 @@ func maybeDispatchDOCXVision(
|
|||||||
return dispatched, false, nil
|
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
|
var wg sync.WaitGroup
|
||||||
sem := make(chan struct{}, docxVisionConcurrency)
|
sem := make(chan struct{}, docxVisionConcurrency)
|
||||||
|
|
||||||
for i, fig := range figs {
|
for slot, tg := range targets {
|
||||||
wg.Add(1)
|
wg.Add(1)
|
||||||
go func(idx int, f map[string]any) {
|
go func(slot int, itemIdx int) {
|
||||||
defer wg.Done()
|
defer wg.Done()
|
||||||
sem <- struct{}{}
|
sem <- struct{}{}
|
||||||
defer func() { <-sem }()
|
defer func() { <-sem }()
|
||||||
|
|
||||||
imageB64, _ := f["image"].(string)
|
img, _ := dispatched.JSON[itemIdx]["image"].(string)
|
||||||
ctxAbove, _ := f["context_above"].(string)
|
if img == "" {
|
||||||
ctxBelow, _ := f["context_below"].(string)
|
|
||||||
|
|
||||||
if imageB64 == "" {
|
|
||||||
return
|
return
|
||||||
}
|
}
|
||||||
|
// DOCX JSON items have no surrounding context (unlike the
|
||||||
prompt, err := docxVisionPromptBuilder(ctxAbove, ctxBelow)
|
// former markdown path), so use the bare figure prompt —
|
||||||
if err != nil {
|
// matching Python's VisionFigureParser(context_size=0).
|
||||||
|
prompt, perr := docxVisionPromptBuilder("", "")
|
||||||
|
if perr != nil {
|
||||||
return
|
return
|
||||||
}
|
}
|
||||||
|
messages := buildVisionMessages(prompt, img)
|
||||||
messages := buildVisionMessages(prompt, imageB64)
|
resp, ierr := visionChatInvoker(ctx, driver, modelName, messages, apiConfig)
|
||||||
resp, err := visionChatInvoker(ctx, driver, modelName, messages, apiConfig)
|
if ierr != nil {
|
||||||
if err != nil {
|
|
||||||
return
|
return
|
||||||
}
|
}
|
||||||
descriptions[idx] = extractDOCXVisionAnswer(resp)
|
descriptions[slot] = extractDOCXVisionAnswer(resp)
|
||||||
}(i, fig)
|
}(slot, tg.idx)
|
||||||
}
|
}
|
||||||
wg.Wait()
|
wg.Wait()
|
||||||
|
|
||||||
// Insert each description at the figure's position in the markdown,
|
modified := false
|
||||||
// matching Python's `chunks[idx]["text"] += description`.
|
for slot, tg := range targets {
|
||||||
// Figures carry a "marker" (text immediately before the image) to
|
desc := strings.TrimSpace(descriptions[slot])
|
||||||
// locate the insertion point. Process in reverse order so earlier
|
if desc == "" {
|
||||||
// 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 == "" {
|
|
||||||
continue
|
continue
|
||||||
}
|
}
|
||||||
if i >= len(figs) {
|
existing, _ := dispatched.JSON[tg.idx]["text"].(string)
|
||||||
continue
|
if existing != "" {
|
||||||
|
dispatched.JSON[tg.idx]["text"] = existing + "\n" + desc
|
||||||
|
} else {
|
||||||
|
dispatched.JSON[tg.idx]["text"] = desc
|
||||||
}
|
}
|
||||||
marker, _ := figs[i]["marker"].(string)
|
modified = true
|
||||||
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})
|
|
||||||
}
|
|
||||||
// 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
|
return dispatched, modified, 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
|
|
||||||
}
|
}
|
||||||
|
|
||||||
// buildDOCXVisionPrompt loads the figure-describe prompt template
|
// 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 {
|
if v, ok := params["prompt"].(string); ok {
|
||||||
p.Prompt = v
|
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 {
|
} else if promptsRaw, ok := params["prompts"].([]any); ok && len(promptsRaw) > 0 {
|
||||||
if first, ok := promptsRaw[0].(map[string]any); ok {
|
if first, ok := promptsRaw[0].(map[string]any); ok {
|
||||||
if content, ok := first["content"].(string); 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
|
// TestNewExtractorComponent_SystemPromptWinsOverSysPrompt verifies
|
||||||
// that "system_prompt" takes priority over "sys_prompt".
|
// that "system_prompt" takes priority over "sys_prompt".
|
||||||
func TestNewExtractorComponent_SystemPromptWinsOverSysPrompt(t *testing.T) {
|
func TestNewExtractorComponent_SystemPromptWinsOverSysPrompt(t *testing.T) {
|
||||||
|
|||||||
@@ -78,7 +78,10 @@ func maybeDispatchMarkdownVision(
|
|||||||
var images []imgItem
|
var images []imgItem
|
||||||
for i, item := range dispatched.JSON {
|
for i, item := range dispatched.JSON {
|
||||||
kd, _ := item["doc_type_kwd"].(string)
|
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
|
continue
|
||||||
}
|
}
|
||||||
img, _ := item["image"].(string)
|
img, _ := item["image"].(string)
|
||||||
|
|||||||
@@ -315,6 +315,21 @@ func maybeDispatchAudio(
|
|||||||
if outputFormat == "" {
|
if outputFormat == "" {
|
||||||
outputFormat = "text"
|
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{
|
return parserDispatchResult{
|
||||||
OutputFormat: outputFormat,
|
OutputFormat: outputFormat,
|
||||||
DocType: "audio",
|
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, "自定义视觉提示")
|
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",
|
"from", "to", "cc", "bcc", "date", "subject",
|
||||||
"body", "attachments", "metadata",
|
"body", "attachments", "metadata",
|
||||||
},
|
},
|
||||||
"output_format": "text",
|
"output_format": "json",
|
||||||
},
|
},
|
||||||
"audio": {
|
"audio": {
|
||||||
"suffix": []string{
|
"suffix": []string{
|
||||||
@@ -310,7 +310,7 @@ func defaultSetups() map[string]schema.ParserSetup {
|
|||||||
"aiff", "au", "midi", "wma", "realaudio", "vqf",
|
"aiff", "au", "midi", "wma", "realaudio", "vqf",
|
||||||
"oggvorbis", "ape",
|
"oggvorbis", "ape",
|
||||||
},
|
},
|
||||||
"output_format": "json",
|
"output_format": "text",
|
||||||
},
|
},
|
||||||
"video": {
|
"video": {
|
||||||
"suffix": []string{"mp4", "avi", "mkv"},
|
"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 = dispatchParse(ctx, fileTypeExt, filename, binary, c.Setups)
|
||||||
dispatched = hydrateEmptyDispatchPayload(dispatched, binary)
|
dispatched = hydrateEmptyDispatchPayload(dispatched, binary)
|
||||||
|
|
||||||
// DOCX vision figure enhancement: enrich the markdown
|
// DOCX vision figure enhancement: on the JSON output path,
|
||||||
// with LLM-generated descriptions of embedded images.
|
// append vision-model descriptions to embedded image items
|
||||||
// Mirrors Python's vision_figure_parser_docx_wrapper_naive.
|
// (doc_type_kwd "image"). Mirrors Python's
|
||||||
|
// enhance_media_sections_with_vision in parser.py:_doc.
|
||||||
dispatched, _, _ = maybeDispatchDOCXVision(ctx, fileTypeExt, dispatched, inputs, c.Setups)
|
dispatched, _, _ = maybeDispatchDOCXVision(ctx, fileTypeExt, dispatched, inputs, c.Setups)
|
||||||
|
|
||||||
// Markdown vision figure enhancement: enrich parsed
|
// Markdown vision figure enhancement: enrich parsed
|
||||||
|
|||||||
@@ -358,16 +358,32 @@ func resolvePDFVisionModelID(setup schema.ParserSetup) (string, bool) {
|
|||||||
return "", false
|
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 {
|
func isNamedPDFParseMethod(raw string) bool {
|
||||||
method := strings.ToLower(strings.TrimSpace(raw))
|
method := strings.ToLower(strings.TrimSpace(raw))
|
||||||
switch {
|
switch {
|
||||||
case strings.HasSuffix(method, "@paddleocr"),
|
case strings.HasSuffix(method, "@mineru"),
|
||||||
|
strings.HasSuffix(method, "@paddleocr"),
|
||||||
strings.HasSuffix(method, "@somark"),
|
strings.HasSuffix(method, "@somark"),
|
||||||
strings.HasSuffix(method, "@opendataloader"):
|
strings.HasSuffix(method, "@opendataloader"):
|
||||||
return true
|
return true
|
||||||
}
|
}
|
||||||
switch method {
|
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 true
|
||||||
}
|
}
|
||||||
return false
|
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 == "" {
|
if trimmedName == "" {
|
||||||
log.Printf("Tokenizer: empty name provided from upstream, embedding will skip title weighting")
|
log.Printf("Tokenizer: empty name provided from upstream, embedding will skip title weighting")
|
||||||
} else {
|
} 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 {
|
if err != nil {
|
||||||
return nil, 0, fmt.Errorf("Tokenizer: encode title: %w", err)
|
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) {
|
func TestTokenizerComponent_Embedding_TruncatesByMaxTokensMinus10(t *testing.T) {
|
||||||
requireTokenizerPool(t)
|
requireTokenizerPool(t)
|
||||||
c, stub := withStubEmbedder(t, 2)
|
c, stub := withStubEmbedder(t, 2)
|
||||||
|
|||||||
@@ -27,14 +27,15 @@ import (
|
|||||||
// task_executor.run_dataflow:879 re.split(r"[,,;;、\r\n]+", keywords).
|
// task_executor.run_dataflow:879 re.split(r"[,,;;、\r\n]+", keywords).
|
||||||
var keywordsSplitRE = regexp.MustCompile(`[,,;;、\r\n]+`)
|
var keywordsSplitRE = regexp.MustCompile(`[,,;;、\r\n]+`)
|
||||||
|
|
||||||
// nonEmpty drops empty strings from parts and returns nil if none remain. It is
|
// nonEmpty drops empty and whitespace-only strings from parts and returns nil
|
||||||
// the shared tail of SplitKeywords and SplitQuestions: split by whatever
|
// if none remain. It is the shared tail of SplitKeywords and SplitQuestions:
|
||||||
// delimiter, then prune empties and collapse an all-empty result to nil so a
|
// split by whatever delimiter, then prune blanks and collapse an all-blank
|
||||||
// _kwd array is absent (nil) rather than [""].
|
// 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 {
|
func nonEmpty(parts []string) []string {
|
||||||
out := make([]string, 0, len(parts))
|
out := make([]string, 0, len(parts))
|
||||||
for _, p := range parts {
|
for _, p := range parts {
|
||||||
if p != "" {
|
if strings.TrimSpace(p) != "" {
|
||||||
out = append(out, p)
|
out = append(out, p)
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|||||||
@@ -16,7 +16,10 @@
|
|||||||
|
|
||||||
package utility
|
package utility
|
||||||
|
|
||||||
import "testing"
|
import (
|
||||||
|
"reflect"
|
||||||
|
"testing"
|
||||||
|
)
|
||||||
|
|
||||||
// SplitKeywords - Python task_executor.run_dataflow:879
|
// SplitKeywords - Python task_executor.run_dataflow:879
|
||||||
// re.split(r"[,,;;、\r\n]+", keywords) with empty filtering.
|
// 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) {
|
func TestSplitKeywords_Empty(t *testing.T) {
|
||||||
result := SplitKeywords("")
|
result := SplitKeywords("")
|
||||||
if len(result) != 0 {
|
if len(result) != 0 {
|
||||||
|
|||||||
Reference in New Issue
Block a user