ragflow/internal/ingestion/task/task_context.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.
//

package task

import (
	"context"
	"fmt"
	"strings"

	"ragflow/internal/common"
	"ragflow/internal/dao"
	"ragflow/internal/entity"
)

// TaskKind discriminates which execution path a queued TaskContext takes.
type TaskKind int

const (
	// TaskKindIngestion is an ingestion document task (IngestionTask set).
	TaskKindIngestion TaskKind = iota
	// TaskKindMemory is an async memory-extraction task (MemoryPayload set,
	// IngestionTask nil). It shares the worker pool with ingestion tasks but
	// runs through executeMemoryTask instead of the ingestion state machine.
	TaskKindMemory
)

// TaskContext holds the execution inputs for an ingestion document task or a
// memory-extraction task. Ingestion tasks populate IngestionTask and the
// document/KB/tenant chain; memory tasks populate MemoryPayload and leave
// IngestionTask nil.
type TaskContext struct {
	Ctx context.Context

	// Kind selects the execution path: TaskKindIngestion or TaskKindMemory.
	Kind TaskKind

	IngestionTask *entity.IngestionTask

	Doc    entity.Document
	KB     entity.Knowledgebase
	Tenant entity.Tenant

	PipelineID string
	File       any

	// MemoryPayload carries the raw task_type="memory" message body for
	// memory tasks (id/memory_id/source_id/message_dict). Only set for
	// TaskKindMemory.
	MemoryPayload map[string]any

	// Handle is the message-queue ack handle for the task message that scheduled
	// this context. The scheduler sets it before queueing; the worker decides
	// the terminal Ack/Nack:
	//   - TaskKindIngestion: ack on a durably-persisted terminal status and
	//     nack otherwise (e.g. shutdown mid-task) so the message is redelivered
	//     and resumed after restart.
	//   - TaskKindMemory: ack on success and on terminal failure (task absent,
	//     already-failed, or progress=-1 persisted by HandleSaveToMemoryTask);
	//     nack on transient failure (task-load DB error before any marker, or
	//     LLM/network error that did not reach progress=-1) so the message is
	//     redelivered. See executeMemoryTask.
	Handle common.TaskHandle
}

// NewMemoryTaskContextForScheduling creates a lightweight TaskContext for a
// memory-extraction task. It only sets the scheduling-related fields, not the
// full ingestion business data.
func NewMemoryTaskContextForScheduling(ctx context.Context, payload map[string]any, handle common.TaskHandle) *TaskContext {
	return &TaskContext{
		Ctx:           ctx,
		Kind:          TaskKindMemory,
		MemoryPayload: payload,
		Handle:        handle,
	}
}

// NewTaskContextForScheduling creates a lightweight TaskContext for queue scheduling.
// This only sets the scheduling-related fields, not the full business data.
func NewTaskContextForScheduling(ctx context.Context, task *entity.IngestionTask) *TaskContext {
	return &TaskContext{
		Ctx:           ctx,
		Kind:          TaskKindIngestion,
		IngestionTask: task,
	}
}

// LoadFromIngestionTask loads the full task context from an IngestionTask.
// It follows the FK chain: ingestion task -> document -> knowledgebase -> tenant.
func LoadFromIngestionTask(ctx context.Context, ingestionTask *entity.IngestionTask) (*TaskContext, error) {
	doc, err := dao.NewDocumentDAO().GetByID(ctx, dao.DB, ingestionTask.DocumentID)
	if err != nil {
		return nil, fmt.Errorf("load document %s: %w", ingestionTask.DocumentID, err)
	}
	if doc == nil {
		return nil, fmt.Errorf("document %s not found", ingestionTask.DocumentID)
	}

	kb, err := dao.NewKnowledgebaseDAO().GetByID(ctx, dao.DB, doc.KbID)
	if err != nil || kb == nil {
		return nil, fmt.Errorf("error when load knowledgebase %s: %w", doc.KbID, err)
	}

	tenant, err := dao.NewTenantDAO().GetByID(ctx, dao.DB, kb.TenantID)
	if err != nil || tenant == nil {
		return nil, fmt.Errorf("error when load tenant %s: %w", kb.TenantID, err)
	}

	pipelineID := resolvePipelineID(doc, kb)

	return &TaskContext{
		Ctx:           ctx,
		IngestionTask: ingestionTask,
		PipelineID:    pipelineID,
		Doc:           *doc,
		KB:            *kb,
		Tenant:        *tenant,
	}, nil
}

func resolvePipelineID(doc *entity.Document, kb *entity.Knowledgebase) string {
	if doc != nil && doc.PipelineID != nil {
		if pipelineID := strings.TrimSpace(*doc.PipelineID); pipelineID != "" {
			return pipelineID
		}
	}
	if kb != nil && kb.PipelineID != nil {
		if pipelineID := strings.TrimSpace(*kb.PipelineID); pipelineID != "" {
			return pipelineID
		}
	}
	return ""
}