diff --git a/cmd/ragflow_server.go b/cmd/ragflow_server.go index cf0ec173b2..0f19d8879d 100644 --- a/cmd/ragflow_server.go +++ b/cmd/ragflow_server.go @@ -559,6 +559,10 @@ func runIngestor(ctx context.Context, cancel context.CancelFunc, args *serverArg globalConfig.GetDefaultChatModel().Name, globalConfig.GetDefaultEmbeddingModel().Name, ) + // Memory extraction runs on the Ingestor's shared NATS consumer + worker + // pool (task_type="memory" dispatched by processMessage -> executeMemoryTask), + // so there is no longer a dedicated Redis memory consumer to start. + ingestor.SetMemoryMessageService(service.NewMemoryMessageService(service.NewMemoryService())) // Start returns immediately (it launches the owned consume/compile // goroutines and joins them via Stop); a provisioning failure here is @@ -567,13 +571,6 @@ func runIngestor(ctx context.Context, cancel context.CancelFunc, args *serverArg common.Error("Failed to initialize ingestor", err) } - // Memory extraction consumer: drains task_type="memory" messages - // from the te.0.common Redis stream and runs LLM extraction. - memoryConsumerCtx, stopMemoryConsumer := context.WithCancel(ctx) - defer stopMemoryConsumer() - memoryMessageSvc := service.NewMemoryMessageService(service.NewMemoryService()) - go memoryMessageSvc.StartTaskConsumer(memoryConsumerCtx) - common.Info("\n ____ __ _\n" + " / _/___ ____ ____ _____/ /_(_)___ ____ ________ ______ _____ _____\n" + " / // __ \\/ __ `/ _ \\/ ___/ __/ / __ \\/ __ \\ / ___/ _ \\/ ___/ | / / _ \\/ ___/\n" + diff --git a/internal/common/task.go b/internal/common/task.go index fbca0b55cf..c50bfba9f3 100644 --- a/internal/common/task.go +++ b/internal/common/task.go @@ -16,14 +16,32 @@ package common +import "encoding/json" + const ( + // TaskSubject is the NATS subject on which ingestion and memory tasks are + // published and consumed. Producer and consumer must reference this single + // symbol so the routing contract cannot diverge (mirrors the RAGFLOW_TASKS + // JetStream subject in internal/engine/nats). + TaskSubject = "tasks.RAGFLOW" + TaskTypeIngestionTask = "ingestion_task" TaskTypeIngestionTest = "ingestion_test" + // TaskTypeMemory is the async memory-extraction task type. Memory tasks + // share the tasks.RAGFLOW subject and the Ingestor's consumer + worker + // pool with ingestion tasks; processMessage dispatches them by TaskType. + // The memory-specific payload (message_dict/memory_id/source_id) is + // carried in TaskMessage.Payload. + TaskTypeMemory = "memory" ) type TaskMessage struct { TaskID string `json:"task_id" binding:"required"` TaskType string `json:"task_type" binding:"required"` + // Payload carries the task-specific body for non-ingestion task types + // (e.g. the memory extraction payload). It is left empty for ingestion + // tasks and old messages, so existing consumers are unaffected. + Payload json.RawMessage `json:"payload,omitempty"` } type TaskHandle interface { diff --git a/internal/dao/compilation_template_seed.go b/internal/dao/compilation_template_seed.go index 38c8d690d6..68c3bdab08 100644 --- a/internal/dao/compilation_template_seed.go +++ b/internal/dao/compilation_template_seed.go @@ -51,6 +51,20 @@ var builtinCompilationTemplateKinds = []struct { func strptr(s string) *string { return &s } +// builtinTemplateID derives a deterministic, <=32-byte id for a built-in +// compilation template from "-". A plain concatenation +// overflows the varchar(32) id column for long kinds (e.g. knowledge_graph), +// so we take the last 32 bytes. Each kind ends up inside those last bytes (the +// "-" suffix is preserved and kinds are distinct), so the truncated ids +// stay unique and readable, and the seed stays idempotent (OnConflict on id). +func builtinTemplateID(kind string) string { + full := BuiltinCompilationTemplateGroupID + "-" + kind + if len(full) <= 32 { + return full + } + return full[len(full)-32:] +} + // SeedBuiltinCompilationTemplates provisions the built-in compilation template // group (and its child templates) as a global, tenant-agnostic catalogue, so // the Go compiler (and the frontend "create from template" catalogue) can @@ -89,7 +103,7 @@ func SeedBuiltinCompilationTemplatesForTenant(ctx context.Context, db *gorm.DB, for _, t := range builtinCompilationTemplateKinds { tmpl := &entity.CompilationTemplate{ - ID: BuiltinCompilationTemplateGroupID + "-" + t.Kind, + ID: builtinTemplateID(t.Kind), TenantID: nil, // global built-in catalogue GroupID: strptr(BuiltinCompilationTemplateGroupID), Name: t.Name, diff --git a/internal/dao/compilation_template_seed_test.go b/internal/dao/compilation_template_seed_test.go index dca74cd70c..9b852d7e40 100644 --- a/internal/dao/compilation_template_seed_test.go +++ b/internal/dao/compilation_template_seed_test.go @@ -37,6 +37,26 @@ func TestBuiltinCompilationTemplateKinds_NoDatasetnav(t *testing.T) { } } +// TestBuiltinTemplateIDWithinColumnWidth guards against a regression where the +// built-in template id exceeds the varchar(32) compilation_template.id column, +// which would make the MySQL seed fail with Error 1406 (data too long). +func TestBuiltinTemplateIDWithinColumnWidth(t *testing.T) { + const idColumnSize = 32 + seen := make(map[string]string, len(builtinCompilationTemplateKinds)) + for _, k := range builtinCompilationTemplateKinds { + id := builtinTemplateID(k.Kind) + if len(id) > idColumnSize { + t.Fatalf("built-in template id %q is %d bytes > id column size %d", id, len(id), idColumnSize) + } + // Truncated ids must stay unique so OnConflict(id) does not clobber two + // kinds into one row. + if prev, dup := seen[id]; dup { + t.Fatalf("built-in template id %q collides for kinds %q and %q", id, prev, k.Kind) + } + seen[id] = k.Kind + } +} + func TestSeedBuiltinCompilationTemplatesForTenant(t *testing.T) { db, err := gorm.Open(sqlite.Open(":memory:"), &gorm.Config{}) if err != nil { diff --git a/internal/ingestion/service/ingestion_service.go b/internal/ingestion/service/ingestion_service.go index dcb1c1a16b..34d31c3f99 100644 --- a/internal/ingestion/service/ingestion_service.go +++ b/internal/ingestion/service/ingestion_service.go @@ -18,9 +18,9 @@ package service import ( "context" + "encoding/json" "errors" "fmt" - "ragflow/internal/utility" "runtime" "strings" "sync" @@ -37,6 +37,7 @@ import ( taskpkg "ragflow/internal/ingestion/task" servicepkg "ragflow/internal/service" documentpkg "ragflow/internal/service/document" + "ragflow/internal/utility" "github.com/cenkalti/backoff/v5" ) @@ -72,6 +73,10 @@ type Ingestor struct { ingestionTaskSvc *servicepkg.IngestionTaskService docState *docStateUpdater + // memorySvc runs async memory-extraction tasks (TaskKindMemory) that share + // the worker pool with ingestion tasks. nil disables memory extraction + // (e.g. tests that don't exercise it). + memorySvc *servicepkg.MemoryMessageService // knowledgeCompile is the dataset-level post-processing consumer (ยง11, // Option E) owned by this ingestor. It is driven by kcConcurrency owned @@ -154,7 +159,7 @@ func (e *Ingestor) Start() error { // error is retained and returned to every later caller. func (e *Ingestor) start() error { msgQueueEngine := engine.GetMessageQueueEngine() - if err := msgQueueEngine.InitConsumer("tasks.RAGFLOW"); err != nil { + if err := msgQueueEngine.InitConsumer(common.TaskSubject); err != nil { return err } @@ -211,6 +216,13 @@ func (e *Ingestor) consumeLoop() { } } +// SetMemoryMessageService installs the memory-extraction service used by +// TaskKindMemory tasks that share the worker pool. Call it before Start; a nil +// value disables memory extraction (received memory tasks are ack-skipped). +func (e *Ingestor) SetMemoryMessageService(memorySvc *servicepkg.MemoryMessageService) { + e.memorySvc = memorySvc +} + // SetKnowledgeCompileModelConfig supplies the default LLM/embedding model ids // used by the dataset-level compile consumer's deduper. Call it before Start. func (e *Ingestor) SetKnowledgeCompileModelConfig(llmID, embedding string) { @@ -298,6 +310,40 @@ func (e *Ingestor) processMessage(handle common.TaskHandle) { } }() + // Memory-extraction tasks share the tasks.RAGFLOW consumer and the worker + // pool with ingestion tasks. They do NOT use the ingestion state machine + // (no ingestion_task row): the message body is dispatched straight to a + // worker via TaskContext.Kind==TaskKindMemory, which runs the memory + // extractor and acks/nacks on its own. + if taskMessage.TaskType == common.TaskTypeMemory { + if e.memorySvc == nil { + common.Warn(fmt.Sprintf("memory task %s received but memory extractor is disabled, ack", taskMessage.TaskID)) + if err := handle.Ack(); err != nil { + common.Error(fmt.Sprintf("error ack memory task %s", taskMessage.TaskID), err) + } + return + } + var payload map[string]any + if len(taskMessage.Payload) == 0 || json.Unmarshal(taskMessage.Payload, &payload) != nil { + common.Warn(fmt.Sprintf("memory task %s has no parseable payload, ack", taskMessage.TaskID)) + if err := handle.Ack(); err != nil { + common.Error(fmt.Sprintf("error ack memory task %s", taskMessage.TaskID), err) + } + return + } + taskCtx := taskpkg.NewMemoryTaskContextForScheduling(e.ctx, payload, handle) + select { + case e.taskChan <- taskCtx: + common.Info(fmt.Sprintf("Memory task %s queued (channel: %d/%d)", taskMessage.TaskID, len(e.taskChan), cap(e.taskChan))) + default: + common.Info(fmt.Sprintf("No available slot for memory task %s, nack", taskMessage.TaskID)) + if nackErr := handle.Nack(); nackErr != nil { + common.Error(fmt.Sprintf("error nack memory task %s", taskMessage.TaskID), nackErr) + } + } + return + } + if taskMessage.TaskType != common.TaskTypeIngestionTask { common.Info(fmt.Sprintf("task %s is not an ingestion task", taskMessage.TaskID)) if err := handle.Ack(); err != nil { @@ -402,12 +448,68 @@ func (e *Ingestor) workerLoop(id int32) { case <-e.ctx.Done(): return case taskCtx := <-e.taskChan: + if taskCtx.Kind == taskpkg.TaskKindMemory { + e.executeMemoryTask(e.ctx, taskCtx) + continue + } common.Info("task context:" + taskCtx.IngestionTask.ID) e.executeTask(e.ctx, taskCtx) } } } +// executeMemoryTask runs one async memory-extraction task (TaskKindMemory) on +// a worker of the shared pool. Unlike ingestion tasks, memory tasks have no +// ingestion_task row / state machine: HandleSaveToMemoryTask persists the +// extracted messages and settles task progress on the way out. +// +// Settlement is error-category aware: +// - Terminal failure (task row absent, already-failed, or progress=-1 already +// persisted) is Acked so an already-consumed message is never redelivered +// into an infinite nack loop. +// - Transient failure (a task-load DB error before any durable marker, or an +// LLM/network failure that did not reach progress=-1) is Nacked so the +// message is redelivered and retried instead of being silently dropped. +func (e *Ingestor) executeMemoryTask(ctx context.Context, taskCtx *taskpkg.TaskContext) { + taskID, _ := taskCtx.MemoryPayload["id"].(string) + if taskID == "" { + taskID, _ = taskCtx.MemoryPayload["task_id"].(string) + } + common.Info(fmt.Sprintf("Starting memory task %s", taskID)) + if taskCtx.Handle == nil { + common.Warn("memory task handle is nil, skip") + return + } + if e.memorySvc == nil { + common.Warn(fmt.Sprintf("memory task %s: memory extractor disabled, ack", taskID)) + if err := taskCtx.Handle.Ack(); err != nil { + common.Error(fmt.Sprintf("ack memory task %s", taskID), err) + } + return + } + if err := e.memorySvc.HandleSaveToMemoryTask(ctx, taskCtx.MemoryPayload); err != nil { + // HandleSaveToMemoryTask wraps terminal outcomes in ErrMemoryTaskTerminal + // (durable progress=-1 written, or no row to retry). Everything else is + // transient and must be redelivered rather than dropped. + if errors.Is(err, servicepkg.ErrMemoryTaskTerminal) { + common.Error(fmt.Sprintf("memory task %s failed terminally, ack", taskID), err) + if ackErr := taskCtx.Handle.Ack(); ackErr != nil { + common.Error(fmt.Sprintf("ack failed memory task %s", taskID), ackErr) + } + return + } + common.Error(fmt.Sprintf("memory task %s failed transiently, nack for redelivery", taskID), err) + if nackErr := taskCtx.Handle.Nack(); nackErr != nil { + common.Error(fmt.Sprintf("nack memory task %s", taskID), nackErr) + } + return + } + common.Info(fmt.Sprintf("Memory task %s completed", taskID)) + if err := taskCtx.Handle.Ack(); err != nil { + common.Error(fmt.Sprintf("ack memory task %s", taskID), err) + } +} + func (e *Ingestor) executeTask(ctx context.Context, taskCtx *taskpkg.TaskContext) { task := taskCtx.IngestionTask common.Info(fmt.Sprintf("Starting task %s", task.ID)) diff --git a/internal/ingestion/service/process_message_test.go b/internal/ingestion/service/process_message_test.go index 7bb93fc2c0..2a565c76de 100644 --- a/internal/ingestion/service/process_message_test.go +++ b/internal/ingestion/service/process_message_test.go @@ -1,18 +1,174 @@ package service import ( + "context" + "encoding/json" + "errors" "testing" "ragflow/internal/common" + "ragflow/internal/dao" "ragflow/internal/entity" taskpkg "ragflow/internal/ingestion/task" "ragflow/internal/ingestion/testutil" + "ragflow/internal/service" ) func newFakeHandle(taskID, taskType string) *fakeTaskHandle { return &fakeTaskHandle{msg: common.TaskMessage{TaskID: taskID, TaskType: taskType}} } +// TestProcessMessage_MemoryTaskDispatches verifies that a task_type="memory" +// message (with a NATS payload) is dispatched to the shared worker pool as a +// TaskKindMemory context rather than being acked-skipped as a non-ingestion +// task. It must NOT touch the ingestion state machine (no StartRunning). +func TestProcessMessage_MemoryTaskDispatches(t *testing.T) { + db := testutil.SetupTestDB(t) + cleanup := testutil.ReplaceDBForTest(t, db) + defer cleanup() + + ingestor := NewIngestor("test", 1, []string{"pdf"}) + // Memory extractor must be enabled for the memory branch to enqueue. + ingestor.SetMemoryMessageService(service.NewMemoryMessageService(nil)) + + payload, err := json.Marshal(map[string]any{ + "id": "mem-task-1", + "task_type": "memory", + "memory_id": "mem-1", + "source_id": 42, + "message_dict": map[string]any{ + "user_id": "u1", + "agent_id": "a1", + "session_id": "s1", + "user_input": "hi", + "agent_response": "hello", + }, + }) + if err != nil { + t.Fatalf("marshal payload: %v", err) + } + handle := &fakeTaskHandle{msg: common.TaskMessage{TaskID: "mem-task-1", TaskType: common.TaskTypeMemory, Payload: payload}} + + ingestor.processMessage(handle) + if handle.acks.Load() != 0 || handle.nacks.Load() != 0 { + t.Fatalf("memory dispatch: expected 0 Ack/0 Nack (settled by worker), got acks=%d nacks=%d", handle.acks.Load(), handle.nacks.Load()) + } + if len(ingestor.taskChan) != 1 { + t.Fatalf("expected 1 memory task enqueued, got %d", len(ingestor.taskChan)) + } + taskCtx := <-ingestor.taskChan + if taskCtx.Kind != taskpkg.TaskKindMemory { + t.Fatalf("task kind = %v, want TaskKindMemory", taskCtx.Kind) + } + if taskCtx.IngestionTask != nil { + t.Fatalf("memory task must not carry an IngestionTask (no ingestion state machine), got %+v", taskCtx.IngestionTask) + } + if taskCtx.MemoryPayload == nil || taskCtx.MemoryPayload["memory_id"] != "mem-1" { + t.Fatalf("memory payload not carried correctly: %+v", taskCtx.MemoryPayload) + } +} + +// TestProcessMessage_MemoryTaskDisabledAcks verifies that when the memory +// extractor is not installed (nil), a memory task is acked and skipped so it +// does not loop forever on the worker pool. +func TestProcessMessage_MemoryTaskDisabledAcks(t *testing.T) { + db := testutil.SetupTestDB(t) + cleanup := testutil.ReplaceDBForTest(t, db) + defer cleanup() + + ingestor := NewIngestor("test", 1, []string{"pdf"}) // memorySvc nil by default + handle := newFakeHandle("mem-task-2", common.TaskTypeMemory) + + ingestor.processMessage(handle) + if handle.acks.Load() != 1 || handle.nacks.Load() != 0 { + t.Fatalf("memory disabled: expected 1 Ack/0 Nack, got acks=%d nacks=%d", handle.acks.Load(), handle.nacks.Load()) + } + if len(ingestor.taskChan) != 0 { + t.Fatal("expected no memory task enqueued when extractor disabled") + } +} + +// TestExecuteMemoryTaskAlreadyFailedAcks verifies the Ack-on-failure contract in +// the case that matters: the task row is already persisted with progress=-1 (the +// durable "failed" marker written by HandleSaveToMemoryTask). In that state +// HandleSaveToMemoryTask returns "already failed", and executeMemoryTask must +// Ack โ€” a Nack would redeliver the message into an infinite retry loop against +// an already-failed row. Using a real taskDAO + MemoryService (not a nil-guard) +// so the test exercises the "task already failed" branch, not the nil-guard. +func TestExecuteMemoryTaskAlreadyFailedAcks(t *testing.T) { + db := testutil.SetupTestDB(t) + cleanup := testutil.ReplaceDBForTest(t, db) + defer cleanup() + + // Seed a task row already marked failed (progress=-1), mirroring what + // HandleSaveToMemoryTask persists before returning an error. + taskID := "mem-task-3" + if err := dao.DB.Create(&entity.Task{ID: taskID, DocID: "doc-mem-3", Progress: -1}).Error; err != nil { + t.Fatalf("insert already-failed task: %v", err) + } + + ingestor := NewIngestor("test", 1, []string{"pdf"}) + // Real memory service (non-nil memories) so HandleSaveToMemoryTask gets past + // the nil-guard and reaches the progress==-1 "already failed" branch. + ingestor.SetMemoryMessageService(service.NewMemoryMessageService(service.NewMemoryService())) + + handle := &fakeTaskHandle{msg: common.TaskMessage{TaskID: taskID, TaskType: common.TaskTypeMemory}} + taskCtx := taskpkg.NewMemoryTaskContextForScheduling(context.Background(), map[string]any{ + "id": taskID, "task_type": "memory", "memory_id": "mem-3", "source_id": 7, + "message_dict": map[string]any{"user_id": "u", "agent_id": "a", "session_id": "s"}, + }, handle) + + // Prove the failure precondition: the seeded progress=-1 row must make + // HandleSaveToMemoryTask return the "already failed" error. Otherwise the + // Ack below would only reflect the success path and prove nothing about + // the Ack-on-failure contract. + if err := ingestor.memorySvc.HandleSaveToMemoryTask(context.Background(), taskCtx.MemoryPayload); err == nil { + t.Fatal("expected HandleSaveToMemoryTask to fail on an already-failed (progress=-1) task, got nil") + } + + ingestor.executeMemoryTask(context.Background(), taskCtx) + if handle.acks.Load() != 1 || handle.nacks.Load() != 0 { + t.Fatalf("already-failed memory task: expected 1 Ack/0 Nack, got acks=%d nacks=%d", handle.acks.Load(), handle.nacks.Load()) + } +} + +// TestExecuteMemoryTaskTransientFailureNacks verifies that a transient (non-terminal) +// failure from HandleSaveToMemoryTask โ€” e.g. a task-load DB error before any durable +// progress=-1 marker is written โ€” is Nacked so the message is redelivered, instead of +// being Acked and permanently dropped. The tasks table is intentionally NOT migrated +// so taskDAO.GetByID fails with a "no such table" error rather than gorm.ErrRecordNotFound. +func TestExecuteMemoryTaskTransientFailureNacks(t *testing.T) { + // Migrate only a table unrelated to tasks so GetByID returns a transient + // "no such table: tasks" error (not gorm.ErrRecordNotFound). + db := testutil.SetupTestDB(t, &entity.IngestionTask{}) + cleanup := testutil.ReplaceDBForTest(t, db) + defer cleanup() + + ingestor := NewIngestor("test", 1, []string{"pdf"}) + ingestor.SetMemoryMessageService(service.NewMemoryMessageService(service.NewMemoryService())) + + handle := &fakeTaskHandle{msg: common.TaskMessage{TaskID: "mem-task-x", TaskType: common.TaskTypeMemory}} + taskCtx := taskpkg.NewMemoryTaskContextForScheduling(context.Background(), map[string]any{ + "id": "mem-task-x", "task_type": "memory", "memory_id": "mem-x", "source_id": 1, + "message_dict": map[string]any{"user_id": "u", "agent_id": "a", "session_id": "s"}, + }, handle) + + // Precondition: with no tasks table, HandleSaveToMemoryTask must fail with a + // transient error that is NOT wrapped in ErrMemoryTaskTerminal. + err := ingestor.memorySvc.HandleSaveToMemoryTask(context.Background(), taskCtx.MemoryPayload) + if err == nil { + t.Fatal("expected HandleSaveToMemoryTask to fail with missing tasks table, got nil") + } + if errors.Is(err, service.ErrMemoryTaskTerminal) { + t.Fatalf("expected a transient error, got terminal: %v", err) + } + + ingestor.executeMemoryTask(context.Background(), taskCtx) + if handle.nacks.Load() != 1 || handle.acks.Load() != 0 { + t.Fatalf("transient memory task failure: expected 0 Ack/1 Nack, got acks=%d nacks=%d", handle.acks.Load(), handle.nacks.Load()) + } +} + // TestProcessMessage_NonIngestionTaskAcks: a non-ingestion task is acked and // skipped without touching the task DB or enqueuing. func TestProcessMessage_NonIngestionTaskAcks(t *testing.T) { diff --git a/internal/ingestion/task/task_context.go b/internal/ingestion/task/task_context.go index bcb800f678..efd1af8391 100644 --- a/internal/ingestion/task/task_context.go +++ b/internal/ingestion/task/task_context.go @@ -26,10 +26,28 @@ import ( "ragflow/internal/entity" ) -// TaskContext holds the execution inputs for an ingestion document task. +// 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 @@ -39,18 +57,43 @@ type TaskContext struct { 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 acks on a - // durably-persisted terminal status and nacks otherwise (e.g. shutdown - // mid-task) so the message is redelivered and resumed after restart. + // 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, } } diff --git a/internal/service/memory_extractor.go b/internal/service/memory_extractor.go index ed8489c393..d2afc0b85d 100644 --- a/internal/service/memory_extractor.go +++ b/internal/service/memory_extractor.go @@ -22,13 +22,16 @@ // api/db/joint_services/memory_message_service.py: // handle_save_to_memory_task / save_extracted_to_memory_only / extract_by_llm // -// QueueSaveToMemoryTask persists the raw message and enqueues a -// task_type="memory" message on the Redis stream te..common. -// StartTaskConsumer drains that stream (consumer group -// "rag_flow_svr_task_broker", same as the Python executor), runs LLM -// extraction for the non-raw memory types configured on the memory, and -// persists the extracted messages with source_id pointing at the raw -// message so listMemoryMessages can aggregate them under `extract`. +// QueueSaveToMemoryTask persists the raw message and publishes a +// task_type="memory" TaskMessage on the NATS tasks.RAGFLOW subject. The +// Ingestor's shared consumer + worker pool dispatches it by TaskType to +// HandleSaveToMemoryTask (see internal/ingestion/service/processMessage and +// executeMemoryTask), which runs LLM extraction for the non-raw memory types +// configured on the memory and persists the extracted messages with source_id +// pointing at the raw message so listMemoryMessages can aggregate them under +// `extract`. Publishing over NATS (instead of the Python te.*.common Redis +// stream) keeps Go out of the Python executor's queue and removes the +// cross-consumer contention that previously stole Python dataflow tasks. package service import ( @@ -41,22 +44,24 @@ import ( "ragflow/internal/common" "ragflow/internal/dao" - redisengine "ragflow/internal/engine/redis" "ragflow/internal/entity" models "ragflow/internal/entity/models" - "ragflow/internal/utility" "go.uber.org/zap" + "gorm.io/gorm" ) -// memoryTaskConsumerGroup matches Python common.constants.SVR_CONSUMER_GROUP_NAME -// so Go and Python executors never double-process the same stream entry. -const memoryTaskConsumerGroup = "rag_flow_svr_task_broker" - // memoryTimeLayout is the storage format for valid_at / invalid_at, // matching timestamp_to_date on the Python side. const memoryTimeLayout = "2006-01-02 15:04:05" +// ErrMemoryTaskTerminal marks a memory-task failure that already has a durable +// terminal outcome (task row absent, or progress already persisted as failed), +// so the caller must Ack rather than Nack/redeliver. Transient failures (DB +// read hiccup, LLM/network errors before any durable marker) return plain +// errors so executeMemoryTask can Nack and let the message be redelivered. +var ErrMemoryTaskTerminal = errors.New("memory: terminal task failure, do not redeliver") + // extractedMemory is one LLM-extracted memory item ready for persistence. type extractedMemory struct { MessageType string @@ -65,54 +70,19 @@ type extractedMemory struct { InvalidAt string // empty means still valid } -// StartTaskConsumer is the long-running loop that drains memory -// extraction tasks from the Redis stream. It returns when ctx is -// cancelled. Per-message failures are logged and acked so one bad -// message cannot stall the stream. -func (s *MemoryMessageService) StartTaskConsumer(ctx context.Context) { - redisClient := redisengine.Get() - if redisClient == nil { - common.Error("memory task consumer: Redis is not available", nil) - return - } - consumerName := fmt.Sprintf("go_memory_extractor_%s", utility.GenerateUUID()) - queueName := memoryTaskQueueName(0) - common.Info(fmt.Sprintf("Memory task consumer %s started on queue %s", consumerName, queueName)) - - for { - if ctx.Err() != nil { - return - } - msg, err := redisClient.QueueConsumer(ctx, queueName, memoryTaskConsumerGroup, consumerName, ">") - if err != nil { - common.Error("memory task consumer: consume error", err) - select { - case <-time.After(time.Second): - case <-ctx.Done(): - return - } - continue - } - if msg == nil { - continue - } - payload := msg.GetMessage() - if taskType, _ := payload["task_type"].(string); taskType != "memory" { - common.Warn(fmt.Sprintf("memory task consumer: skip task_type %q", taskType)) - msg.Ack(ctx) - continue - } - if err := s.HandleSaveToMemoryTask(ctx, payload); err != nil { - common.Error("memory task consumer: handle task failed", err) - } - msg.Ack(ctx) - } -} - // HandleSaveToMemoryTask processes one queued memory task. Mirrors // Python handle_save_to_memory_task: validate the task row, then // extract + persist, settling task progress on the way out. +// +// The returned error is wrapped in ErrMemoryTaskTerminal when the failure has +// already produced a durable terminal outcome (dependency/config error, task +// row absent, task already failed, or extraction failed after progress=-1 was +// persisted). Transient failures (a task-load DB error before any marker was +// written) return an unwrapped error so the caller can Nack and redeliver. func (s *MemoryMessageService) HandleSaveToMemoryTask(ctx context.Context, payload map[string]any) error { + if s == nil || s.taskDAO == nil || s.memories == nil { + return fmt.Errorf("%w: memory: nil MemoryMessageService or memory dependency", ErrMemoryTaskTerminal) + } taskID, _ := payload["id"].(string) if taskID == "" { taskID, _ = payload["task_id"].(string) @@ -130,15 +100,21 @@ func (s *MemoryMessageService) HandleSaveToMemoryTask(ctx context.Context, paylo task, err := s.taskDAO.GetByID(ctx, dao.DB, taskID) if err != nil { - return fmt.Errorf("memory: task %s is not found", taskID) + // Record-not-found is terminal: no row to retry against. Any other + // task-load error is transient and must be redelivered (no progress=-1 + // marker was written). + if errors.Is(err, gorm.ErrRecordNotFound) { + return fmt.Errorf("%w: memory: task %s is not found", ErrMemoryTaskTerminal, taskID) + } + return fmt.Errorf("memory: load task %s: %w", taskID, err) } if task.Progress == -1 { - return fmt.Errorf("memory: task %s is already failed", taskID) + return fmt.Errorf("%w: memory: task %s is already failed", ErrMemoryTaskTerminal, taskID) } if err := s.saveExtractedToMemory(ctx, memoryID, msg, sourceID, taskID); err != nil { s.updateTaskProgress(taskID, -1, err.Error()) - return err + return fmt.Errorf("%w: %v", ErrMemoryTaskTerminal, err) } return nil } diff --git a/internal/service/memory_message_service.go b/internal/service/memory_message_service.go index 707aa7b73e..dd3eb6b450 100644 --- a/internal/service/memory_message_service.go +++ b/internal/service/memory_message_service.go @@ -54,15 +54,18 @@ package service import ( "context" + "encoding/json" "errors" "fmt" - "ragflow/internal/utility" "time" + "ragflow/internal/common" "ragflow/internal/dao" + "ragflow/internal/engine" redisengine "ragflow/internal/engine/redis" "ragflow/internal/entity" models "ragflow/internal/entity/models" + "ragflow/internal/utility" ) // MemoryMessage is the wire shape for QueueSaveToMemoryTask. It @@ -352,16 +355,34 @@ func queueMemoryTask(ctx context.Context, memoryID, tenantID string, rawMessageI "agent_response": msg.AgentResponse, }, } - if redisClient := redisengine.Get(); redisClient == nil || !redisClient.QueueProduct(ctx, memoryTaskQueueName(0), message) { - return errors.New("Can't access Redis.") + // Publish the memory-extraction task to NATS (tasks.RAGFLOW) so it is + // consumed by the Ingestor's shared consumer + worker pool, dispatched by + // TaskType=="memory" in processMessage. This keeps Go out of the Python + // te.*.common Redis stream entirely, removing the cross-consumer + // contention that previously stole Python dataflow tasks. + mq := engine.GetMessageQueueEngine() + if mq == nil { + return errors.New("can't access message queue engine") + } + payload, err := json.Marshal(message) + if err != nil { + return fmt.Errorf("marshal memory task payload: %w", err) + } + taskMessage := common.TaskMessage{ + TaskID: taskID, + TaskType: common.TaskTypeMemory, + Payload: payload, + } + tmPayload, err := json.Marshal(taskMessage) + if err != nil { + return fmt.Errorf("marshal memory task message: %w", err) + } + if err := mq.PublishTask(common.TaskSubject, tmPayload); err != nil { + return fmt.Errorf("publish memory task %s: %w", taskID, err) } return nil } -func memoryTaskQueueName(priority int) string { - return fmt.Sprintf("te.%d.common", priority) -} - func mapStringAny(in map[string]any) map[string]interface{} { out := make(map[string]interface{}, len(in)) for k, v := range in { diff --git a/rag/flow/parser/parser.py b/rag/flow/parser/parser.py index 7306232575..cf9048d02a 100644 --- a/rag/flow/parser/parser.py +++ b/rag/flow/parser/parser.py @@ -20,7 +20,8 @@ import random import re from functools import partial -from litellm import logging +import logging + import numpy as np from PIL import Image