internal/ingestion/pipeline/pipeline.go

Rama main · réplica de lectura del repositorio original.

//
//  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 pipeline

import (
	"context"
	"errors"
	"fmt"
	"strings"
	"time"

	"ragflow/internal/agent/canvas"
	_ "ragflow/internal/agent/component"
	"ragflow/internal/agent/runtime"
	"ragflow/internal/common"
	kvrocks "ragflow/internal/engine/kvrocks"
	"ragflow/internal/ingestion/component/globals"
	"ragflow/internal/ingestion/component/schema"
	"ragflow/internal/ingestion/mediacache"
	"ragflow/internal/parser/media"
	"ragflow/internal/utility"

	"go.uber.org/zap"
)

var _ mediacache.Backend = (*kvrocks.Client)(nil)

// Pipeline is a compiled ingestion canvas plus task-scoped metadata.
type Pipeline struct {
	taskID     string
	documentID string // owning document; numeric progress is mirrored back to
	// the document table for existing status projections.
	canvas  *canvas.Canvas
	tracker *canvas.RunTracker       // optional injected; nil -> resolve at Run
	kv      *kvrocks.Client          // optional; nil -> no tracker fallback (unit/test runs)
	media   media.Cache              // optional injected run media cache
	factory runtime.ComponentFactory // optional instance-scoped component factory
	sink    ProgressSink             // optional progress sink; nil -> drop events (DB-independent)
}

// PipelineOption mutates a Pipeline before Run. Used to inject test doubles
// (in-memory store / miniredis tracker) or dedicated Redis pools.
type PipelineOption func(*Pipeline)

// WithRunTracker injects a RunTracker for progress mirroring (terminal
// run-status marks). When unset and no Kvrocks client is configured, Run does
// not mirror progress.
func WithRunTracker(t *canvas.RunTracker) PipelineOption {
	return func(p *Pipeline) { p.tracker = t }
}

// WithKV injects the Kvrocks client used to resolve a RunTracker when none is
// injected. It replaces the previous global kvrocks.Get() lookup, so the
// pipeline no longer reads a process-wide singleton (arch_test.go guards that
// rule). A nil client disables the RunTracker fallback.
func WithKV(kv *kvrocks.Client) PipelineOption {
	return func(p *Pipeline) { p.kv = kv }
}

// WithDocumentID binds the pipeline's owning document so numeric progress can
// be mirrored back into the document table.
// Pass the empty string to disable the mirror (e.g. headless/test runs where
// the document row is not materialized).
func WithDocumentID(docID string) PipelineOption {
	return func(p *Pipeline) { p.documentID = docID }
}

// ProgressEvent is a structured component lifecycle event emitted by the
// pipeline to a ProgressSink. The pipeline fills the task/document/component
// identity and phase/status message; the sink caches the denominator (total)
// from OnComponentTotal and needs no canvas knowledge.
type ProgressEvent struct {
	TaskID     string
	DocumentID string
	Component  string
	Message    string
	Phase      int
}

// ProgressSink receives pipeline progress for durable persistence. It is the
// single channel through which the pipeline reports component lifecycle
// events and the component-total denominator; the pipeline itself never
// touches the DAO layer. Implementations live in the orchestration layer
// (internal/ingestion/service). A nil sink is valid: events are dropped and
// the pipeline stays DB-independent (unit tests, headless runs).
type ProgressSink interface {
	OnComponentTotal(ctx context.Context, taskID string, total int)
	OnComponentProgress(ctx context.Context, ev ProgressEvent)
}

// WithProgressSink injects a sink that receives component progress events
// and the component-total denominator. When unset, the pipeline drops
// progress events and stays DB-independent.
func WithProgressSink(s ProgressSink) PipelineOption {
	return func(p *Pipeline) { p.sink = s }
}

// NewPipelineFromDSL compiles the canonical ingestion canvas DSL.
// It accepts either the inner canvas DSL or the template wrapper whose
// top-level `dsl` field carries that canvas.
func NewPipelineFromDSL(dsl []byte, taskID string, opts ...PipelineOption) (*Pipeline, error) {
	taskID = strings.TrimSpace(taskID)
	if taskID == "" {
		return nil, errors.New("pipeline: task id is required")
	}
	// UnwrapCanvasDSL is the single source of truth for stripping the
	// optional {"dsl": {...}} canvas envelope; it also reports a nil/unparseable
	// DSL.
	canvasDSL, err := UnwrapCanvasDSL(dsl)
	if err != nil {
		return nil, err
	}
	cnv, err := canvas.DecodeFromDSL(canvasDSL)
	if err != nil {
		return nil, fmt.Errorf("pipeline: decode canvas DSL: %w", err)
	}
	if err := ValidatePipeline(cnv); err != nil {
		return nil, err
	}
	p := &Pipeline{
		taskID: taskID,
		canvas: cnv,
	}
	for _, o := range opts {
		o(p)
	}
	return p, nil
}

// ValidatePipeline enforces ingestion pipeline constraints.
// Specifically, at most one Extractor component is permitted in the graph.
func ValidatePipeline(cnv *canvas.Canvas) error {
	if cnv == nil {
		return nil
	}
	extractorCount := 0
	for id, comp := range cnv.Components {
		if common.IsExtractorComponent(id, comp.Obj.ComponentName) {
			extractorCount++
		}
	}
	if extractorCount > 1 {
		return fmt.Errorf("pipeline validation error: at most 1 Extractor component is allowed, found %d", extractorCount)
	}
	return nil
}

// WithComponentFactory installs an instance-scoped factory override for this
// pipeline. It is used during canvas compilation so one pipeline run can
// construct task-specific component instances without mutating the process-wide
// runtime default factory.
func (p *Pipeline) WithComponentFactory(factory runtime.ComponentFactory) *Pipeline {
	if p != nil {
		p.factory = factory
	}
	return p
}

func mergeInto(dst, src map[string]any) map[string]any {
	if src == nil {
		return dst
	}
	if dst == nil {
		dst = make(map[string]any, len(src))
	}
	for k, v := range src {
		dst[k] = v
	}
	return dst
}

func cloneMapOrEmpty(m map[string]any) map[string]any {
	if m == nil {
		return map[string]any{}
	}
	out := make(map[string]any, len(m))
	for k, v := range m {
		out[k] = v
	}
	return out
}

// defaultTrackerTTL is the expiry applied to the RunTracker run-status hash.
// It matches the model-result cache TTL; media has its own one-day TTL.
var defaultTrackerTTL = 7 * 24 * time.Hour

// Run executes the full ingestion graph described by the canonical DSL.
// There is no pipeline-layer partial resume entry point: execution always
// starts from the graph entry and component-level replay decisions belong to
// the components themselves.
func (p *Pipeline) Run(ctx context.Context, inputs map[string]any, overrideParams map[string]any) (map[string]any, error) {
	if p == nil {
		return nil, fmt.Errorf("pipeline: Run on nil pipeline")
	}
	if p.canvas == nil {
		return nil, fmt.Errorf("pipeline: canvas is nil")
	}
	if runtime.DefaultFactory() == nil {
		runtime.InstallDefaultRegistryFactory()
	}
	if runtime.DefaultFactory() == nil {
		return nil, fmt.Errorf("pipeline: Run: runtime default component factory is not installed")
	}

	compileCtx := ctx
	if p.factory != nil {
		compileCtx = canvas.WithComponentFactory(compileCtx, p.factory)
	}

	// Resolve the run tracker for progress mirroring (terminal
	// MarkSucceeded/MarkFailed/MarkCancelled marks). Resume-from-checkpoint is
	// intentionally not supported: the pipeline always runs the full graph
	// from the entry node, and component-level replay decisions (e.g. serving
	// prior LLM/embedding results from the run cache) belong to the
	// components themselves.
	tracker := p.resolveTracker()

	var compileOpts []canvas.CompileOption
	// Run-level setups (keyed by cpnID) override the DSL-baked component
	// setups at compile time (higher priority; see canvas.WithOverrideParams).
	if overrideParams != nil {
		compileOpts = append(compileOpts, canvas.WithOverrideParams(overrideParams))
	}
	compiled, err := canvas.Compile(compileCtx, p.canvas, compileOpts...)
	if err != nil {
		return nil, fmt.Errorf("pipeline: Run: compile agent: %w", err)
	}

	// Record the component count as the authoritative denominator for
	// progress percentage. Best-effort: a DB failure (or headless run
	// with no DB) must not abort the pipeline — progress is observability.
	if p.sink != nil {
		p.sink.OnComponentTotal(ctx, p.taskID, len(p.canvas.Components))
	}

	runState := canvas.NewCanvasState("", p.taskID)
	runCtx := canvas.WithState(ctx, runState)
	// Framework-level progress fan-out: the canvas framework
	// (realComponentBody) pulls this callback from ctx via
	// runtime.ProgressCallbackFromContext and records every component
	// start/done/fail event as an ingestion_task_log row. The callback
	// is nil when the DB is not initialized (unit tests, headless
	// runs), in which case TrackProgress is a no-op — progress is an
	// observability concern, not a data dependency.
	runCtx = runtime.WithProgressCallback(runCtx, p.componentProgressCallback(ctx))
	runCtx = runtime.WithProgressMessageCallback(runCtx, p.componentProgressMessageCallback(ctx))
	runCtx = runtime.WithProgressFractionCallback(runCtx, p.componentFractionCallback(ctx))

	current := cloneMapOrEmpty(inputs)

	// Seed the workflow-wide Globals bag with the run-level metadata
	// (name, tenant_id, kb_id, model_id, doc_id, ...) once, from the
	// pipeline run inputs. Downstream components read these from ctx
	// instead of relying on every node re-emitting them. The File
	// component re-publishes `name` (and storage refs) as it derives
	// them mid-run.
	globals.SeedIngestionGlobals(runCtx, current)
	// Expose the task id on the run context so every component can register the
	// run-cache keys it writes into the task manifest (chunkcache). Without
	// this the manifest is never populated and PurgeTask on success becomes a
	// no-op, leaving orphaned cache entries until TTL expiry.
	globals.SetTaskID(runCtx, p.taskID)
	mediaCache := p.media
	if mediaCache == nil && p.kv != nil {
		mediaCache = mediacache.New(p.kv, p.taskID)
	}
	if mediaCache != nil {
		runCtx = media.WithCache(runCtx, mediaCache)
	}

	return p.runPlain(runCtx, current, compiled, tracker, runState)
}

// resolveTracker resolves the RunTracker used for progress mirroring. When an
// instance-scoped tracker was injected it wins; otherwise one is created on
// the global Kvrocks client when available.
func (p *Pipeline) resolveTracker() *canvas.RunTracker {
	if p.tracker != nil {
		return p.tracker
	}
	if p.kv != nil {
		return canvas.NewRunTrackerWithClient(p.kv.GetClient(), defaultTrackerTTL)
	}
	return nil
}

// runPlain executes the full ingestion graph with a single Invoke and records
// the terminal run status via the sink/RunTracker when available.
func (p *Pipeline) runPlain(runCtx context.Context, current map[string]any, compiled *canvas.CompiledCanvas, tracker *canvas.RunTracker, runState *canvas.CanvasState) (map[string]any, error) {
	// Terminal tracker writes must survive a run ctx that gets cancelled
	// (the common failure/cancel path). Derive a fresh detached ctx per
	// terminal branch, bounded so a hung Redis cannot stall — mirrors
	// markStopped/markFailed in ingestion_service.go.
	detached := func() (context.Context, context.CancelFunc) {
		return context.WithTimeout(context.WithoutCancel(runCtx), 5*time.Second)
	}

	out, err := compiled.Workflow.Invoke(runCtx, current)
	if err != nil {
		if errors.Is(runCtx.Err(), context.Canceled) || errors.Is(runCtx.Err(), context.DeadlineExceeded) {
			if tracker != nil {
				stateCtx, cancel := detached()
				utility.BestEffort(fmt.Sprintf("MarkCancelled for %s", p.taskID), func() error { return tracker.MarkCancelled(stateCtx, p.taskID) })
				cancel()
			}
			return current, fmt.Errorf("pipeline: run cancelled: %w", runCtx.Err())
		}
		if tracker != nil {
			stateCtx, cancel := detached()
			utility.BestEffort(fmt.Sprintf("MarkFailed for %s", p.taskID), func() error { return tracker.MarkFailed(stateCtx, p.taskID, err.Error()) })
			cancel()
		}
		return current, fmt.Errorf("pipeline: run canvas workflow: %w", err)
	}
	if tracker != nil {
		stateCtx, cancel := detached()
		utility.BestEffort(fmt.Sprintf("MarkSucceeded for %s", p.taskID), func() error { return tracker.MarkSucceeded(stateCtx, p.taskID) })
		cancel()
	}
	return finalizeResult(current, out, runState), nil
}

// finalizeResult merges the graph output into the input map and attaches the
// canvas state snapshot as the success payload.
func finalizeResult(current, out map[string]any, runState *canvas.CanvasState) map[string]any {
	if out == nil {
		current["state"] = runState.Snapshot()
		return current
	}
	merged := mergeInto(current, out)
	merged["state"] = runState.Snapshot()
	return merged
}

// componentProgressCallback returns a runtime.ProgressCallback that forwards
// every component lifecycle event (start/done/fail) to the pipeline's
// ProgressSink. The sink owns all persistence; this callback only shapes the
// event - deriving the message string the frontend expects - so the pipeline
// never touches the DAO layer. Returns nil when no sink is attached, leaving
// TrackProgress a no-op and the pipeline DB-independent (unit tests, headless
// runs).
func (p *Pipeline) componentProgressCallback(ctx context.Context) runtime.ProgressCallback {
	if p.sink == nil {
		return nil
	}
	return func(ev runtime.ProgressEvent) {
		componentName := ev.Component
		if comp, ok := p.canvas.Components[ev.Component]; ok && comp.Obj.ComponentName != "" {
			componentName = comp.Obj.ComponentName
		}
		var msg string
		switch ev.Phase {
		case runtime.PhaseEnter:
			msg = componentName + " Started"
		case runtime.PhaseExit:
			msg = componentName + " Done"
		case runtime.PhaseError:
			msg = formatComponentFailure(componentName, ev.Err)
		}
		// Surface every component lifecycle event as a structured log line so
		// a component failure (e.g. an LLM/client error) is captured in
		// ingestor_server.log even if the wrapped error never reaches the
		// higher-level "Task ... failed" branch.
		switch ev.Phase {
		case runtime.PhaseError:
			if errors.Is(ev.Err, context.Canceled) {
				// A user cancel is a normal control path; the authoritative
				// "Task ... cancelled" line is logged by the service layer, so
				// keep this at debug to avoid duplicate noise.
				common.Debug("component progress: canceled",
					zap.String("component", ev.Component),
					zap.String("task_id", p.taskID),
					zap.String("document_id", p.documentID))
			} else if ev.Err != nil {
				common.Error("component progress: error", ev.Err,
					zap.String("component", ev.Component),
					zap.String("task_id", p.taskID),
					zap.String("document_id", p.documentID))
			} else {
				common.Info("component progress: error",
					zap.String("component", ev.Component),
					zap.String("task_id", p.taskID),
					zap.String("document_id", p.documentID))
			}
		default:
			// Keep the message constant: msg may carry component names or
			// error-derived text, and a newline in it could forge a log record
			// (CWE-117). Pass msg and the phase as structured fields instead.
			common.Info("component progress",
				zap.String("message", msg),
				zap.Int("phase", int(ev.Phase)),
				zap.String("component", ev.Component),
				zap.String("task_id", p.taskID),
				zap.String("document_id", p.documentID))
		}
		sinkCtx, cancel := progressSinkContext(ctx)
		defer cancel()
		p.sink.OnComponentProgress(sinkCtx, ProgressEvent{
			TaskID:     p.taskID,
			DocumentID: p.documentID,
			Component:  ev.Component,
			Message:    msg,
			Phase:      int(ev.Phase),
		})
	}
}

func formatComponentFailure(componentName string, err error) string {
	if err == nil {
		return componentName + " Error"
	}
	if errors.Is(err, context.DeadlineExceeded) || strings.Contains(strings.ToLower(err.Error()), context.DeadlineExceeded.Error()) {
		switch componentName {
		case schema.ComponentNameExtractor:
			return "Extractor: Document extraction timed out while waiting for the configured LLM. Check that the model service is reachable, then retry."
		case schema.ComponentNameCompiler:
			return "Compiler: Document compilation timed out. Check that the configured model service is reachable and responsive, then retry."
		}
	}

	message := err.Error()
	prefix := componentName + ":"
	for componentName != "" && strings.HasPrefix(strings.ToLower(message), strings.ToLower(prefix)) {
		message = strings.TrimSpace(message[len(prefix):])
	}
	if componentName == "" {
		return message
	}
	return componentName + ": " + message
}

type detailedProgressSink interface {
	OnComponentMessage(ctx context.Context, taskID, documentID, component, message string)
}

// fractionProgressSink is the optional interface through which the pipeline
// forwards in-flight component fractions (pages parsed, chunks embedded) to
// the sink's progress tracker. Mirrors detailedProgressSink: a sink that does
// not implement it simply receives no fraction channel.
type fractionProgressSink interface {
	OnComponentFraction(ctx context.Context, component string, fraction float64)
}

// componentFractionCallback forwards fraction reports to the sink. The sink
// only mutates in-memory state here (its flusher owns persistence), so the
// run context is passed through without the WithoutCancel fallback that the
// I/O-bound callbacks need.
func (p *Pipeline) componentFractionCallback(ctx context.Context) runtime.ProgressFractionCallback {
	sink, ok := p.sink.(fractionProgressSink)
	if !ok {
		return nil
	}
	return func(component string, fraction float64) {
		sink.OnComponentFraction(ctx, component, fraction)
	}
}

func (p *Pipeline) componentProgressMessageCallback(ctx context.Context) runtime.ProgressMessageCallback {
	sink, ok := p.sink.(detailedProgressSink)
	if !ok {
		return nil
	}
	return func(component, message string) {
		common.Info("component progress detail",
			zap.String("component", component),
			zap.String("task_id", p.taskID),
			zap.String("document_id", p.documentID),
			zap.String("message", message))
		sinkCtx, cancel := progressSinkContext(ctx)
		defer cancel()
		sink.OnComponentMessage(sinkCtx, p.taskID, p.documentID, component, message)
	}
}

func progressSinkContext(ctx context.Context) (context.Context, context.CancelFunc) {
	if ctx.Err() == nil {
		return context.WithCancel(ctx)
	}
	return context.WithTimeout(context.WithoutCancel(ctx), 5*time.Second)
}