mirror of
https://github.com/infiniflow/ragflow.git
synced 2026-07-16 20:57:21 +08:00
301 lines
8.7 KiB
Go
301 lines
8.7 KiB
Go
//
|
|
// 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.
|
|
//
|
|
|
|
// 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.
|
|
//
|
|
// 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.
|
|
|
|
package component
|
|
|
|
import (
|
|
"fmt"
|
|
"os"
|
|
"path/filepath"
|
|
"sort"
|
|
"strings"
|
|
"sync"
|
|
|
|
"ragflow/internal/entity"
|
|
modelModule "ragflow/internal/entity/models"
|
|
"ragflow/internal/ingestion/component/schema"
|
|
"ragflow/internal/utility"
|
|
)
|
|
|
|
var (
|
|
docxVisionPromptBuilder = buildDOCXVisionPrompt
|
|
visionChatInvoker = defaultVisionChatInvoker
|
|
docxVisionConcurrency uint = 10
|
|
)
|
|
|
|
const (
|
|
docxVisionPromptFile = "vision_llm_figure_describe_prompt.md"
|
|
docxVisionPromptWithContextFile = "vision_llm_figure_describe_prompt_with_context.md"
|
|
)
|
|
|
|
var (
|
|
docxVisionPromptsBase string
|
|
docxVisionPromptsOnce sync.Once
|
|
docxVisionPromptCache = make(map[string]string)
|
|
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:
|
|
//
|
|
// 1. naive_merge_docx → chunks (text + images + context)
|
|
// 2. vision_figure_parser_docx_wrapper_naive → LLM descriptions
|
|
//
|
|
// 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.
|
|
func maybeDispatchDOCXVision(
|
|
fileType utility.FileType,
|
|
dispatched parserDispatchResult,
|
|
inputs map[string]any,
|
|
setups map[string]schema.ParserSetup,
|
|
) (parserDispatchResult, bool, error) {
|
|
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 {
|
|
return dispatched, false, nil
|
|
}
|
|
|
|
tenantID := getStringOr(inputs, "tenant_id", "")
|
|
if tenantID == "" {
|
|
return dispatched, false, nil
|
|
}
|
|
|
|
// Resolve the tenant's IMAGE2TEXT model.
|
|
driver, modelName, apiConfig, _, err := resolveTenantModelByType(tenantID, entity.ModelTypeImage2Text)
|
|
if err != nil {
|
|
// Model not available — skip vision enhancement silently,
|
|
// matching Python's try/except pass behaviour.
|
|
return dispatched, false, nil
|
|
}
|
|
|
|
descriptions := make([]string, len(figs))
|
|
var wg sync.WaitGroup
|
|
sem := make(chan struct{}, docxVisionConcurrency)
|
|
|
|
for i, fig := range figs {
|
|
wg.Add(1)
|
|
go func(idx int, f map[string]any) {
|
|
defer wg.Done()
|
|
sem <- struct{}{}
|
|
defer func() { <-sem }()
|
|
|
|
imageB64, _ := f["image"].(string)
|
|
ctxAbove, _ := f["context_above"].(string)
|
|
ctxBelow, _ := f["context_below"].(string)
|
|
|
|
if imageB64 == "" {
|
|
return
|
|
}
|
|
|
|
prompt, err := docxVisionPromptBuilder(ctxAbove, ctxBelow)
|
|
if err != nil {
|
|
return
|
|
}
|
|
|
|
messages := buildVisionMessages(prompt, imageB64)
|
|
resp, err := visionChatInvoker(driver, modelName, messages, apiConfig)
|
|
if err != nil {
|
|
return
|
|
}
|
|
descriptions[idx] = extractDOCXVisionAnswer(resp)
|
|
}(i, fig)
|
|
}
|
|
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 == "" {
|
|
continue
|
|
}
|
|
if i >= len(figs) {
|
|
continue
|
|
}
|
|
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})
|
|
}
|
|
// 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
|
|
}
|
|
|
|
// buildDOCXVisionPrompt loads the figure-describe prompt template
|
|
// and, when context text is available, renders it with the
|
|
// with-context variant. Mirrors Python:
|
|
//
|
|
// if context_above or context_below:
|
|
// prompt = vision_llm_figure_describe_prompt_with_context(context_above, context_below)
|
|
// else:
|
|
// prompt = vision_llm_figure_describe_prompt()
|
|
func buildDOCXVisionPrompt(contextAbove, contextBelow string) (string, error) {
|
|
hasContext := strings.TrimSpace(contextAbove) != "" || strings.TrimSpace(contextBelow) != ""
|
|
|
|
var templateName string
|
|
if hasContext {
|
|
templateName = docxVisionPromptWithContextFile
|
|
} else {
|
|
templateName = docxVisionPromptFile
|
|
}
|
|
|
|
template, err := loadDOCXVisionPromptFile(templateName)
|
|
if err != nil {
|
|
return "", err
|
|
}
|
|
|
|
if hasContext {
|
|
template = strings.ReplaceAll(template, "{{ context_above }}", contextAbove)
|
|
template = strings.ReplaceAll(template, "{{ context_below }}", contextBelow)
|
|
}
|
|
return template, nil
|
|
}
|
|
|
|
func loadDOCXVisionPromptFile(filename string) (string, error) {
|
|
docxVisionPromptMu.RLock()
|
|
if cached, ok := docxVisionPromptCache[filename]; ok {
|
|
docxVisionPromptMu.RUnlock()
|
|
return cached, nil
|
|
}
|
|
docxVisionPromptMu.RUnlock()
|
|
|
|
baseDir, err := docxVisionPromptsBaseDir()
|
|
if err != nil {
|
|
return "", err
|
|
}
|
|
promptPath := filepath.Join(baseDir, "rag", "prompts", filename)
|
|
content, err := os.ReadFile(promptPath)
|
|
if err != nil {
|
|
return "", fmt.Errorf("docx vision prompt %q: %w", filename, err)
|
|
}
|
|
cached := strings.TrimSpace(string(content))
|
|
docxVisionPromptMu.Lock()
|
|
docxVisionPromptCache[filename] = cached
|
|
docxVisionPromptMu.Unlock()
|
|
return cached, nil
|
|
}
|
|
|
|
func docxVisionPromptsBaseDir() (string, error) {
|
|
var initErr error
|
|
docxVisionPromptsOnce.Do(func() {
|
|
root := utility.GetProjectRoot()
|
|
if _, statErr := os.Stat(filepath.Join(root, "rag", "prompts")); statErr == nil {
|
|
docxVisionPromptsBase = root
|
|
return
|
|
}
|
|
initErr = fmt.Errorf("rag/prompts not found under project root %q", root)
|
|
})
|
|
if initErr != nil {
|
|
return "", initErr
|
|
}
|
|
return docxVisionPromptsBase, nil
|
|
}
|
|
|
|
func buildVisionMessages(prompt, imageBase64 string) []modelModule.Message {
|
|
dataURI := "data:image/png;base64," + imageBase64
|
|
return []modelModule.Message{{
|
|
Role: "user",
|
|
Content: []interface{}{
|
|
map[string]any{"type": "text", "text": prompt},
|
|
map[string]any{"type": "image_url", "image_url": map[string]any{"url": dataURI}},
|
|
},
|
|
}}
|
|
}
|
|
|
|
func extractDOCXVisionAnswer(resp *modelModule.ChatResponse) string {
|
|
if resp == nil || resp.Answer == nil {
|
|
return ""
|
|
}
|
|
return strings.TrimSpace(*resp.Answer)
|
|
}
|
|
|
|
func defaultVisionChatInvoker(
|
|
driver modelModule.ModelDriver,
|
|
modelName string,
|
|
messages []modelModule.Message,
|
|
apiConfig *modelModule.APIConfig,
|
|
) (*modelModule.ChatResponse, error) {
|
|
vision := true
|
|
return driver.ChatWithMessages(modelName, messages, apiConfig, &modelModule.ChatConfig{Vision: &vision})
|
|
}
|