ragflow/internal/ingestion/component/media_dispatch.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.
//

// Media dispatch runs image OCR, table structure recognition, optional vision
// enhancement, audio transcription, and video dispatch at the component layer.

package component

import (
	"bytes"
	"context"
	"encoding/base64"
	"fmt"
	"image"
	"image/png"
	"os"
	"path/filepath"
	"strings"

	"go.uber.org/zap"

	agentruntime "ragflow/internal/agent/runtime"
	"ragflow/internal/common"
	imageparser "ragflow/internal/deepdoc/parser/image"
	deepdoctype "ragflow/internal/deepdoc/parser/type"
	"ragflow/internal/entity"
	modelModule "ragflow/internal/entity/models"
	"ragflow/internal/ingestion/component/progress"
	"ragflow/internal/ingestion/component/schema"
	"ragflow/internal/parser/media"
	"ragflow/internal/parser/parser"
	"ragflow/internal/utility"

	"gorm.io/gorm"
)

// Video dispatch: IMAGE2TEXT vision chat ---

func maybeDispatchVideo(
	ctx context.Context,
	db *gorm.DB,
	fileType utility.FileType,
	filename string,
	binary []byte,
	inputs map[string]any,
	setups map[string]schema.ParserSetup,
) (parser.ParseResult, bool, error) {
	if fileType != utility.FileTypeVIDEO {
		return parser.ParseResult{}, false, nil
	}
	if _, ok := setups["video"]; !ok {
		return parser.ParseResult{}, false, nil
	}

	// Video parsing is intentionally not implemented yet: the underlying
	// video-analysis capability is pending. The previously-shipped path sent a
	// video_url data URI, but no model driver honors it — OpenAI-compatible
	// drivers only accept image_url (the block is ignored), and Gemini's
	// googleMessageParts only accepts text/image_url and silently drops
	// video_url. Returning an explicit error is safer than silently producing
	// a description from the prompt text alone. The real implementation must
	// be provider-specific (OpenAI-compatible: frame extraction -> image_url;
	// Gemini: raw-bytes inline_data; Qwen: file://).
	// When video analysis is implemented, it emits standard Text JSON items:
	// [{"text": transcript, "doc_type_kwd": "text"}] with output_format "json".
	return parser.ParseResult{}, true,
		fmt.Errorf("Parser: video parsing is not yet supported; underlying video analysis capability is pending")
}

// maybeDispatchImage extracts local or selected-model OCR (native text for SVG)
// and emits structured table/text/figure items when the page pipeline finds
// them. Without a table it keeps the single whole-image item; with a table it
// emits text and table/figure crops instead, matching PDF page output.
func maybeDispatchImage(
	ctx context.Context,
	p parser.ParseResultProducer,
	db *gorm.DB,
	fileType utility.FileType,
	filename string,
	binary []byte,
	inputs map[string]any,
	setups map[string]schema.ParserSetup,
	enableVisionEnhancement bool,
	vision visionSettings,
) (parser.ParseResult, bool, error) {
	if fileType != utility.FileTypeVISUAL {
		return parser.ParseResult{}, false, nil
	}
	setup, ok := setups["image"]
	if !ok {
		return parser.ParseResult{}, false, nil
	}
	method := getStringOr(setup, "parse_method", "")
	_, hasOCRFlag := setup["ocr_enabled"].(bool)
	// ocr_enabled is the OCR switch. Legacy setups omit it and encode the same
	// choice in parse_method: "ocr"/empty selects local OCR, any other value is
	// a VLM model reference that also implies OCR is off. The selection itself
	// lives in imageSetupRunsOCR (parser.go), shared with Check().
	useOCR := imageSetupRunsOCR(setup)
	release, err := parser.AcquireImageMedia(ctx)
	if err != nil {
		return parser.ParseResult{}, true, err
	}
	defer release()
	if len(binary) == 0 || len(binary) > media.ImageMaxBytes() {
		return parser.ParseResult{}, true, fmt.Errorf("parser: image payload exceeds size limits")
	}
	// SVG carries its text in the document, so it is read rather than
	// rasterized for OCR, whatever the OCR switch says.
	isSVG := isSVGFilename(filename)
	var text string
	var ocrErr error
	var blocks []imageparser.Block
	if isSVG {
		text, err = extractSVGText(ctx, binary)
		if err != nil {
			return parser.ParseResult{}, true, err
		}
	} else {
		ocrModelID := strings.TrimSpace(getStringOr(setup, "ocr_model_id", ""))
		img, err := decodeDispatchImage(binary, useOCR && ocrModelID == "")
		if err != nil {
			return parser.ParseResult{}, true, err
		}
		if useOCR {
			if ocrModelID == "" {
				var analyzer deepdoctype.DocAnalyzer
				analyzer, ocrErr = parser.GetDocAnalyzer()
				if ocrErr == nil {
					finish := agentruntime.TrackIO(ctx, "DeepDOC", "image OCR", "media=image")
					blocks, ocrErr = imageparser.Parse(ctx, img, analyzer)
					finish()
					if ocrErr == nil {
						text = imageBlockText(blocks)
					}
				}
			} else {
				// Selected OCR uses original bytes, so media admission need not
				// cover database resolution or remote inference.
				release()
				finish := agentruntime.TrackIO(ctx, "Parser", "OCR model", "media=image backend=selected")
				text, ocrErr = extractSelectedImageOCR(ctx, db, getStringOr(inputs, "tenant_id", ""), ocrModelID, binary)
				finish()
			}
		}
	}
	release()
	parsed := dispatchParse(ctx, p, fileType, filename, binary, setups, progress.Window{})
	if parsed.Err != nil {
		return parsed, true, parsed.Err
	}
	if len(parsed.JSON) == 0 {
		return parsed, true, fmt.Errorf("parser: image parser returned no image item")
	}
	imageItem := parsed.JSON[0].ToMap()
	imageData, _ := imageItem["image"].(string)
	hasTables := imageHasTables(blocks)
	var items []parser.DocItem
	if hasTables {
		items = imageBlockItems(blocks)
	} else {
		imageItem["text"] = text
		items = []parser.DocItem{parser.DocItemFromMap(imageItem)}
	}
	if ocrErr != nil {
		parsed.Warnings = append(parsed.Warnings, fmt.Sprintf("image OCR unavailable: %v", ocrErr))
	} else if isSVG && text == "" {
		parsed.Warnings = append(parsed.Warnings, "image SVG contains no text")
	} else if useOCR && strings.TrimSpace(text) == "" && !hasTables {
		parsed.Warnings = append(parsed.Warnings, "image OCR returned no text")
	}
	if !useOCR && !enableVisionEnhancement {
		// Neither text source is enabled: local OCR is off and the vision
		// description is not permitted. The item still carries the image, but the
		// Tokenizer's retrievability filter drops chunks that have no text and no
		// surrounding context, so this document indexes nothing. Say so instead of
		// silently returning an empty result.
		parsed.Warnings = append(parsed.Warnings,
			"image has no text source: OCR is off and vision enhancement is disabled, so the Tokenizer will discard this item")
	}
	if err := ctx.Err(); err != nil {
		return parsed, true, err
	}
	if enableVisionEnhancement && isSVG {
		// Vision providers take raster images; raw SVG is not sent to them.
		parsed.Warnings = append(parsed.Warnings, "image VLM enhancement skipped: SVG is not rasterized for vision models")
	} else if enableVisionEnhancement {
		// A setup with an explicit ocr_enabled switch always uses the global
		// vision settings. Only legacy setups (no switch) carry the VLM model
		// reference in parse_method, which also disabled OCR by construction.
		imageVision := vision
		if !hasOCRFlag && !useOCR {
			imageVision.modelID = method
		}
		description, warnings := describeImage(ctx, db, imageData, getStringOr(inputs, "tenant_id", ""), setup, inputs, imageVision)
		parsed.Warnings = append(parsed.Warnings, warnings...)
		if description != "" {
			if hasTables {
				if imageData != "" {
					items = append(items, parser.ImageItem{Text: description, Image: imageData})
				} else {
					items = append(items, parser.NewTextJSONItem(description))
				}
			} else {
				imageItem["text"] = text
				appendItemText(imageItem, description)
				items[0] = parser.DocItemFromMap(imageItem)
			}
		} else if strings.TrimSpace(text) == "" && !hasTables {
			// Enhancement was permitted and the model produced nothing usable,
			// so the item still has no text at all. Without this the chunk is
			// dropped downstream with no trace, which is what the two other
			// empty-text warnings above exist to prevent.
			parsed.Warnings = append(parsed.Warnings,
				"image has no searchable text: the vision model returned no description, so the Tokenizer will discard this item")
		}
	}
	parsed.JSON = items
	return parsed, true, nil
}

// decodeDispatchImage validates dimensions and raster limits for all image
// sources, decoding pixels only when local DeepDoc OCR needs them.
func decodeDispatchImage(data []byte, decodeRaster bool) (image.Image, error) {
	config, _, err := image.DecodeConfig(bytes.NewReader(data))
	if err != nil {
		return nil, fmt.Errorf("parser: decode image: %w", err)
	}
	if config.Width <= 0 || config.Height <= 0 {
		return nil, fmt.Errorf("parser: image has invalid dimensions %dx%d", config.Width, config.Height)
	}
	// Length/width are unbounded (mirrors Python's picture.py); the decoded
	// raster byte ceiling is opt-in via RAGFLOW_IMAGE_RASTER_MAX_BYTES (0 =
	// unlimited).
	if err := parser.CheckImageRasterLimit(config.Width, config.Height); err != nil {
		return nil, fmt.Errorf("parser: %w", err)
	}
	// VLM and PaddleOCR consume the original bytes; only DeepDoc needs a raster.
	if !decodeRaster {
		return nil, nil
	}
	img, _, err := image.Decode(bytes.NewReader(data))
	if err != nil {
		return nil, fmt.Errorf("parser: decode image: %w", err)
	}
	return img, nil
}

// imageBlockText is the flat OCR text used by the single-image fallback and
// the vision-description path. Table content stays in the table items.
func imageBlockText(blocks []imageparser.Block) string {
	texts := make([]string, 0, len(blocks))
	for _, block := range blocks {
		if block.Kind != imageparser.BlockText && block.Kind != imageparser.BlockImage {
			continue
		}
		if text := strings.TrimSpace(block.Text); text != "" {
			texts = append(texts, text)
		}
	}
	return strings.Join(texts, "\n")
}

func imageHasTables(blocks []imageparser.Block) bool {
	for _, block := range blocks {
		if block.Kind == imageparser.BlockTable && block.Table != nil {
			return true
		}
	}
	return false
}

func imageBlockItems(blocks []imageparser.Block) []parser.DocItem {
	items := make([]parser.DocItem, 0, len(blocks))
	for _, block := range blocks {
		switch block.Kind {
		case imageparser.BlockText:
			if strings.TrimSpace(block.Text) == "" {
				continue
			}
			items = append(items, parser.TextItem{Text: block.Text})
		case imageparser.BlockTable:
			if block.Table == nil {
				continue
			}
			extra := map[string]any{"layout": "table"}
			if uri := encodeMediaImage(block.Image); uri != "" {
				extra["image"] = uri
			}
			items = append(items, parser.TableItem{Table: block.Table, Extra: extra})
		case imageparser.BlockImage:
			if uri := encodeMediaImage(block.Image); uri != "" {
				items = append(items, parser.ImageItem{
					Text:  block.Text,
					Image: uri,
					Extra: map[string]any{"layout": "figure"},
				})
			}
		}
	}
	return items
}

func encodeMediaImage(img image.Image) string {
	if img == nil {
		return ""
	}
	var buf bytes.Buffer
	if err := png.Encode(&buf, img); err != nil {
		return ""
	}
	return "data:image/png;base64," + base64.StdEncoding.EncodeToString(buf.Bytes())
}

func describeImage(
	ctx context.Context,
	db *gorm.DB,
	dataURI string,
	tenantID string,
	setup schema.ParserSetup,
	inputs map[string]any,
	vision visionSettings,
) (string, []string) {
	// --- Optional VLM description ---
	// Captions follow the knowledge base language, which also drives
	// tokenization, so indexed text and analysis can never disagree. The family
	// setup's lang is not consulted here: it is an OCR engine input (see
	// pdf_vision_dispatch.go) and its old implicit default was the subject of
	// issue #20727.
	lang := resolveVisionLanguage(inputs)
	if tenantID == "" {
		return "", []string{"image VLM enhancement skipped: tenant ID is missing"}
	}

	// Use the selected description model or the tenant default.
	modelRef := vision.modelID
	var driver modelModule.ModelDriver
	var modelName string
	var apiConfig *modelModule.APIConfig
	var err error
	if modelRef != "" {
		driver, modelName, apiConfig, _, err = resolveModelConfig(ctx, db, tenantID, entity.ModelTypeImage2Text, modelRef)
		if err != nil {
			common.Warn("media dispatch image: per-call VLM resolve failed, falling back to tenant default",
				zap.String("modelRef", modelRef), zap.String("tenant", tenantID), zap.Error(err))
			driver, modelName, apiConfig, _, err = resolveTenantModelByType(ctx, db, tenantID, entity.ModelTypeImage2Text)
		}
	} else {
		driver, modelName, apiConfig, _, err = resolveTenantModelByType(ctx, db, tenantID, entity.ModelTypeImage2Text)
	}
	if err == nil && driver == nil {
		err = fmt.Errorf("no usable vision model")
	}
	if err != nil {
		return "", []string{fmt.Sprintf("image VLM enhancement skipped: model unavailable: %v", err)}
	}

	prompt := defaultImageVisionPrompt(lang)
	// The global enhancement prompt wins; the image family's legacy system_prompt
	// (the key the Python parser's image branch used) stays readable for canvases
	// saved before the move. Do NOT read setup["prompt"] here — that key is for
	// the video family, not image.
	if v := common.FirstNonEmpty(vision.systemPrompt, getStringOr(setup, "system_prompt", "")); v != "" {
		prompt = v
	}
	messages := []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}},
		},
	}}
	visionFlag := true
	chatModel := modelModule.NewChatModel(driver, &modelName, apiConfig)
	finish := agentruntime.TrackIO(ctx, "Parser", "VLM", "media=image")
	resp, err := chatModel.ChatWithMessages(ctx, messages, &modelModule.ChatConfig{Vision: &visionFlag}, nil)
	finish()
	if err != nil {
		return "", []string{fmt.Sprintf("image VLM enhancement failed: %v", err)}
	}
	vlmText := ""
	if resp != nil && resp.Answer != nil {
		vlmText = strings.TrimSpace(*resp.Answer)
	}

	return vlmText, nil
}

// Audio dispatch: SPEECH2TEXT transcription ---
// Mirrors Python's rag/app/audio.py:chunk():
//   - Writes the audio binary to a temp file (extension-preserving)
//   - Calls the tenant's SPEECH2TEXT model via TranscribeAudio()
//   - Returns the transcription as text

func maybeDispatchAudio(
	ctx context.Context,
	db *gorm.DB,
	fileType utility.FileType,
	filename string,
	binary []byte,
	inputs map[string]any,
	setups map[string]schema.ParserSetup,
) (parser.ParseResult, bool, error) {
	if fileType != utility.FileTypeAURAL {
		return parser.ParseResult{}, false, nil
	}
	setup, ok := setups["audio"]
	if !ok {
		return parser.ParseResult{}, false, nil
	}
	tenantID := getStringOr(inputs, "tenant_id", "")
	if tenantID == "" {
		return parser.ParseResult{}, true,
			fmt.Errorf("parser: audio requires tenant_id")
	}

	modelRef := configuredAudioModelID(setup)
	var driver modelModule.ModelDriver
	var modelName string
	var apiConfig *modelModule.APIConfig
	var err error
	if modelRef != "" {
		driver, modelName, apiConfig, _, err = resolveModelConfig(ctx, db, tenantID, entity.ModelTypeSpeech2Text, modelRef)
		if err != nil {
			common.Warn("media dispatch audio: per-call VLM resolve failed, falling back to tenant default",
				zap.String("modelRef", modelRef), zap.String("tenant", tenantID), zap.Error(err))
			driver, modelName, apiConfig, _, err = resolveTenantModelByType(ctx, db, tenantID, entity.ModelTypeSpeech2Text)
		}
	} else {
		driver, modelName, apiConfig, _, err = resolveTenantModelByType(ctx, db, tenantID, entity.ModelTypeSpeech2Text)
	}
	if err != nil {
		return parser.ParseResult{}, true,
			fmt.Errorf("parser: audio speech2text model: %w", err)
	}

	tmpFile, err := writeTempAudioFile(filename, binary)
	if err != nil {
		return parser.ParseResult{}, true,
			fmt.Errorf("parser: audio temp file: %w", err)
	}
	defer os.Remove(tmpFile)

	asrModel := modelModule.NewASRModel(driver, &modelName, apiConfig)
	finish := agentruntime.TrackIO(ctx, "Parser", "ASR", "media=audio")
	resp, err := asrModel.Transcribe(ctx, &tmpFile, nil, nil)
	finish()
	if err != nil {
		return parser.ParseResult{}, true,
			fmt.Errorf("Parser: audio transcription: %w", err)
	}

	transcription := ""
	if resp != nil {
		transcription = resp.Text
	}

	return parser.ParseResult{
		JSON: []parser.DocItem{parser.NewTextJSONItem(transcription)},
	}, true, nil
}

// writeTempAudioFile writes binary to a temp file preserving the
// original extension so the ASR provider can detect the format.
func writeTempAudioFile(filename string, binary []byte) (string, error) {
	ext := filepath.Ext(filename)
	tmp, err := os.CreateTemp("", "ragflow_audio_*"+ext)
	if err != nil {
		return "", err
	}
	defer tmp.Close()
	if _, err := tmp.Write(binary); err != nil {
		os.Remove(tmp.Name())
		return "", err
	}
	return tmp.Name(), nil
}

// videoMIME maps common video filename extensions to MIME types
// for constructing base64 data URIs. Retained as a reference for the
// future real video-parsing implementation (provider-specific frame
// extraction / inline_data / file://); not currently used because
// maybeDispatchVideo returns an explicit unsupported error.
func videoMIME(filename string) string {
	dot := strings.LastIndex(filename, ".")
	if dot == -1 {
		return "video/mp4"
	}
	switch strings.ToLower(filename[dot+1:]) {
	case "mp4":
		return "video/mp4"
	case "avi":
		return "video/x-msvideo"
	case "mkv":
		return "video/x-matroska"
	case "mov":
		return "video/quicktime"
	case "wmv":
		return "video/x-ms-wmv"
	case "flv":
		return "video/x-flv"
	case "webm":
		return "video/webm"
	case "mpeg", "mpg":
		return "video/mpeg"
	case "3gp":
		return "video/3gpp"
	default:
		return "video/mp4"
	}
}