From 01d667296d9be047ce99e36706913446b3a55a47 Mon Sep 17 00:00:00 2001 From: Zhichang Yu Date: Sun, 2 Aug 2026 17:06:29 +0800 Subject: [PATCH] refactor(knowledge_compile): global compile pool, token-budget batching, and DocEngine-only deletion (#17679) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit ## Summary This PR refactors the Go knowledge-compilation ingestion pipeline (`internal/ingestion/knowledge_compile` + `internal/ingestion/component/knowledge_compiler`) with three related changes: - **Token-budget batching for LLM merge decisions.** `LLMMergeDecider.DecideBatch` previously stuffed every `(existing, candidate)` pair into a single LLM call, risking `max_token` overflow. It now splits pairs into token-bounded sub-batches (budget = `llmMaxTokens * 0.85`) via `tokenizer.NumTokensFromString`, runs them concurrently while preserving the global pair index, and never reindexes. - **Process-level global compile pool.** Introduces a single vCPU-sized goroutine pool (`pool.go`, env `KC_COMPILE_CONCURRENCY`) dedicated to *all* knowledge-compilation stages. KNN search loop, `DecideBatch` sub-batches, `WriteMerged`/`DeleteMerged` internals, and the component-level (structure/mindmap) per-call pools are all unified into it via an injected submitter. No more per-job short-lived goroutines in `runCompilerJobs` (futures are collected then awaited on the caller). Fan-out stays bounded by the pool worker count; these stages are docengine-bounded / LLM-bounded, not CPU-bounded. - **DocEngine-only deletion.** `Consumer.processBatch` deletion no longer loads the deleted docs' products into memory. Two sequential DocEngine calls replace the old in-memory surgery: - `DeleteDocLevelForDocs` — one `DeleteChunks` over `doc_id IN deletedDocIDs` (merged rows carry `doc_id == kb`, so only per-doc products match). - `StripMergedSources` — one `Search` of `kc_merged=1` rows filtered by `source_doc_ids IN deletedDocIDs` (intersection pushed down to the engine), `UpdateChunks` the source array of survivors, and `DeleteChunks` the rows whose array became empty. ## Changes - `internal/ingestion/knowledge_compile/pool.go` (new): global `compilerPool` + `runCompilerJobs`/`SubmitCompilerJob`/`SubmitCompilerJobs`. - `internal/ingestion/knowledge_compile/consumer.go`: deletion rewritten to the two DocEngine calls; `mergedBase`/`toDelete`/`stripDeletedSources` removed. - `internal/ingestion/knowledge_compile/writer.go`: `DeleteDocLevelForDocs` + `StripMergedSources` replace `DeleteMergedForDoc`/`DeleteMerged`. - `internal/ingestion/knowledge_compile/reader.go`: drop `LoadMergedBySourceDoc` + `containsString` (keep `LoadDocProducts` for the completion branch). - `internal/ingestion/knowledge_compile/dedup.go`: `NewLLMDeduper` takes `llmMaxTokens`; wires `SetMaxBatchTokens`/`SetSubmitter`. - `internal/ingestion/knowledge_compiler/{structure,merge}.go`, `mindmap/mindmap.go`, `pool_wiring.go`: token-budget split + submitter injection. - Tests: `structure_test.go` (token-budget split), `dedup_test.go`, `consumer_test.go` (tombstone + DocEngine deletion assertions) updated. ## Validation `bash build.sh --test -race ./internal/ingestion/knowledge_compile/... ./internal/ingestion/component/knowledge_compiler/...` passes (unit tier, no external services). 🤖 Generated with [CodeBuddy](https://www.codebuddy.ai) --------- Co-authored-by: yuzhichang --- internal/dao/database.go | 2 +- internal/entity/knowledge_compile_doc.go | 6 +- .../knowledge_compiler/common/memstore.go | 175 ++++++-- .../knowledge_compiler/common/types.go | 5 + .../knowledge_compiler/mindmap/mindmap.go | 59 ++- .../knowledge_compiler/pool_wiring.go | 38 ++ .../knowledge_compiler/structure/merge.go | 242 +++++++++- .../knowledge_compiler/structure/structure.go | 59 ++- .../structure/structure_test.go | 162 +++++++ .../ingestion/knowledge_compile/consumer.go | 362 ++++++++++----- .../knowledge_compile/consumer_test.go | 138 ++++-- internal/ingestion/knowledge_compile/dedup.go | 136 +++++- .../ingestion/knowledge_compile/dedup_test.go | 122 +++++ internal/ingestion/knowledge_compile/event.go | 45 +- internal/ingestion/knowledge_compile/pool.go | 147 ++++++ .../ingestion/knowledge_compile/reader.go | 201 +++++++-- .../ingestion/knowledge_compile/scheduler.go | 417 ++++++++++-------- .../ingestion/knowledge_compile/service.go | 15 +- .../ingestion/knowledge_compile/writer.go | 175 ++++++-- .../ingestion/service/ingestion_service.go | 2 +- 20 files changed, 2009 insertions(+), 499 deletions(-) create mode 100644 internal/ingestion/component/knowledge_compiler/pool_wiring.go create mode 100644 internal/ingestion/knowledge_compile/dedup_test.go create mode 100644 internal/ingestion/knowledge_compile/pool.go diff --git a/internal/dao/database.go b/internal/dao/database.go index 303eb70b8c..818b0097d4 100644 --- a/internal/dao/database.go +++ b/internal/dao/database.go @@ -157,7 +157,7 @@ func InitDB(ctx context.Context, migrateDB bool) error { &entity.IngestionTaskLog{}, &entity.FileCommit{}, &entity.FileCommitItem{}, - &entity.KnowledgeCompileDoc{}, + &entity.KnowledgeCompileDataset{}, // Knowledge-compile compilation templates and their groups. The Go // KnowledgeCompilerComponent resolves a compilation_template (or group) // from these tables at runtime, so the Go side must guarantee they exist. diff --git a/internal/entity/knowledge_compile_doc.go b/internal/entity/knowledge_compile_doc.go index 5ac1c205e6..84f144b0c6 100644 --- a/internal/entity/knowledge_compile_doc.go +++ b/internal/entity/knowledge_compile_doc.go @@ -17,7 +17,7 @@ package entity import "time" -// KnowledgeCompileDoc is the MySQL scheduling row for the dataset-level +// KnowledgeCompileDataset is the MySQL scheduling row for the dataset-level // post-processing consumer (knowledge_compile_design.md §11.4, Option E). It is // the scheduling system of record: backlog_doc_ids holds the not-yet-processed // doc entries for the KB, inflight_doc_ids the ones a worker has claimed (the @@ -29,7 +29,7 @@ import "time" // (doc_id + event_type + seq) as TEXT so the consumer can re-apply the same // out-of-order / tombstone guards as the broker-based design without re-reading // the queue. -type KnowledgeCompileDoc struct { +type KnowledgeCompileDataset struct { DatasetID string `gorm:"primaryKey;column:dataset_id;size:64" json:"dataset_id"` TenantID string `gorm:"column:tenant_id;size:64;not null;default:''" json:"tenant_id"` // The *_doc_ids columns store a JSON array as TEXT. No DDL default is set: @@ -47,4 +47,4 @@ type KnowledgeCompileDoc struct { } // TableName pins the scheduling table name. -func (KnowledgeCompileDoc) TableName() string { return "knowledge_compile_docs" } +func (KnowledgeCompileDataset) TableName() string { return "knowledge_compile_docs" } diff --git a/internal/ingestion/component/knowledge_compiler/common/memstore.go b/internal/ingestion/component/knowledge_compiler/common/memstore.go index 220f246534..983d33e55a 100644 --- a/internal/ingestion/component/knowledge_compiler/common/memstore.go +++ b/internal/ingestion/component/knowledge_compiler/common/memstore.go @@ -4,6 +4,8 @@ import ( "math" "sort" "sync" + + "gonum.org/v1/gonum/mat" ) // Hit is one TopK match returned by MemStore. @@ -15,14 +17,24 @@ type Hit struct { // MemStore is an in-memory product store with exact-cosine TopK retrieval. // It is the single source of truth for in-run dedup (replacing Python's ES -// KNN over the current run's products). Vectors are kept alongside precomputed -// L2 norms so TopK is a single matVec pass. +// KNN over the current run's products). +// +// Vectors are stored as one contiguous float64 row-major matrix (matDense) so +// the per-query dot-product pass — the hot path for document-level structure +// dedup and dataset-level cross-document dedup — is a single level-2 BLAS +// matrix-vector product (mat.VecDense.MulVec) instead of N separate scalar +// loops. gonum only supports float64, so each incoming float32 vector is +// widened when written. All vectors in one store are assumed to share the same +// dimension (embeddings from one model); addLocked pads/truncates a stray +// vector to cols defensively. Precomputed L2 norms complete the cosine. type MemStore struct { - mu sync.RWMutex - items []Product - vectors [][]float32 - norms []float64 - byID map[string]int + mu sync.RWMutex + items []Product + cols int // uniform embedding dimension; 0 while empty + matData []float64 // row-major vectors, len == cols*len(items) + matDense *mat.Dense // view over matData, rebuilt when the row count changes + norms []float64 + byID map[string]int } // NewMemStore constructs an empty MemStore. @@ -66,23 +78,7 @@ func (m *MemStore) Add(p Product) { // applied under a write lock using the previously resolved index. func (m *MemStore) DedupeAdd(row Product, threshold float64, cb DedupCallback) (KeepAction, error) { m.mu.RLock() - qn := l2Norm(row.Vector) - if qn == 0 { - qn = 1 - } - bestIdx, bestScore := -1, 0.0 - for i, v := range m.vectors { - vn := m.norms[i] - if vn == 0 { - vn = 1 - } - score := dotProduct(row.Vector, v) / (qn * vn) - if threshold <= 0 || score >= threshold { - if bestIdx == -1 || score > bestScore { - bestIdx, bestScore = i, score - } - } - } + bestIdx, bestScore := m.bestMatchLocked(row.Vector, threshold) var best Product if bestIdx >= 0 && bestIdx < len(m.items) { best = m.items[bestIdx] @@ -118,7 +114,7 @@ func (m *MemStore) DedupeAdd(row Product, threshold float64, cb DedupCallback) ( // Merges preserve the existing entry's identity (Python preserve_id). replacement.ID = m.items[resolvedIdx].ID m.items[resolvedIdx] = replacement - m.vectors[resolvedIdx] = replacement.Vector + m.replaceMatLocked(resolvedIdx, replacement.Vector) m.norms[resolvedIdx] = l2Norm(replacement.Vector) return KeepMerge, nil default: @@ -129,7 +125,7 @@ func (m *MemStore) DedupeAdd(row Product, threshold float64, cb DedupCallback) ( func (m *MemStore) addLocked(p Product) { m.items = append(m.items, p) - m.vectors = append(m.vectors, p.Vector) + m.appendMatLocked(p.Vector) m.norms = append(m.norms, l2Norm(p.Vector)) if p.ID != "" { m.byID[p.ID] = len(m.items) - 1 @@ -142,14 +138,11 @@ func (m *MemStore) Upsert(p Product) { defer m.mu.Unlock() if idx, ok := m.byID[p.ID]; ok { m.items[idx] = p - m.vectors[idx] = p.Vector + m.replaceMatLocked(idx, p.Vector) m.norms[idx] = l2Norm(p.Vector) return } - m.items = append(m.items, p) - m.vectors = append(m.vectors, p.Vector) - m.norms = append(m.norms, l2Norm(p.Vector)) - m.byID[p.ID] = len(m.items) - 1 + m.addLocked(p) } // Delete removes a product by ID. @@ -161,7 +154,7 @@ func (m *MemStore) Delete(id string) { return } m.items = append(m.items[:idx], m.items[idx+1:]...) - m.vectors = append(m.vectors[:idx], m.vectors[idx+1:]...) + m.removeMatLocked(idx) m.norms = append(m.norms[:idx], m.norms[idx+1:]...) delete(m.byID, id) m.reindexLocked() @@ -189,6 +182,10 @@ func (m *MemStore) Len() int { func (m *MemStore) TopK(vec []float32, k int, threshold float64) []Hit { m.mu.RLock() defer m.mu.RUnlock() + if m.matDense == nil { + return nil + } + dots := m.matDotLocked(vec) qn := l2Norm(vec) if qn == 0 { qn = 1 @@ -198,12 +195,12 @@ func (m *MemStore) TopK(vec []float32, k int, threshold float64) []Hit { score float64 } var cands []cand - for i, v := range m.vectors { + for i := range m.items { vn := m.norms[i] if vn == 0 { vn = 1 } - score := dotProduct(vec, v) / (qn * vn) + score := dots[i] / (qn * vn) if threshold <= 0 || score >= threshold { cands = append(cands, cand{i, score}) } @@ -261,16 +258,110 @@ func (m *MemStore) MergeSourceChunkIDs(id string, chunkIDs []string) { m.items[idx] = p } -func dotProduct(a, b []float32) float64 { - n := len(a) - if len(b) < n { - n = len(b) +// bestMatchLocked returns the index and cosine similarity of the stored vector +// most similar to q among those meeting threshold (threshold<=0 keeps every +// candidate). The dot-product pass is a single BLAS gemv over the row-major +// matrix. Caller must hold at least a read lock. +func (m *MemStore) bestMatchLocked(q []float32, threshold float64) (int, float64) { + if m.matDense == nil { + return -1, 0 } - var s float64 - for i := 0; i < n; i++ { - s += float64(a[i]) * float64(b[i]) + dots := m.matDotLocked(q) + qn := l2Norm(q) + if qn == 0 { + qn = 1 } - return s + bestIdx, bestScore := -1, 0.0 + for i := range m.items { + vn := m.norms[i] + if vn == 0 { + vn = 1 + } + score := dots[i] / (qn * vn) + if threshold <= 0 || score >= threshold { + if bestIdx == -1 || score > bestScore { + bestIdx, bestScore = i, score + } + } + } + return bestIdx, bestScore +} + +// matDotLocked returns matDense * q — the raw dot products of q against every +// stored vector — computed in one level-2 BLAS operation. Caller must hold at +// least a read lock. +func (m *MemStore) matDotLocked(q []float32) []float64 { + qv := make([]float64, m.cols) + for i := 0; i < m.cols && i < len(q); i++ { + qv[i] = float64(q[i]) + } + var out mat.VecDense + out.MulVec(m.matDense, mat.NewVecDense(m.cols, qv)) + return out.RawVector().Data +} + +// appendMatLocked widens vec to cols and appends it as a new matrix row. Caller +// must hold a write lock. +func (m *MemStore) appendMatLocked(v []float32) { + if m.cols == 0 { + m.cols = len(v) + } + m.matData = append(m.matData, float32ToF64(v, m.cols)...) + // gonum panics on non-positive dimensions, so skip the view until a real + // dimension is known (e.g. a product arrived with an empty vector). The + // TopK/bestMatch paths treat a nil matDense as "no usable vectors". + if m.cols > 0 { + m.matDense = mat.NewDense(len(m.items), m.cols, m.matData) + } else { + m.matDense = nil + } +} + +// replaceMatLocked overwrites matrix row idx with vec (widened to cols, missing +// trailing dims zero-filled, extra dims truncated). Caller must hold a write +// lock. +func (m *MemStore) replaceMatLocked(idx int, v []float32) { + if m.cols == 0 { + m.cols = len(v) + if m.cols > 0 { + m.matDense = mat.NewDense(len(m.items), m.cols, m.matData) + } + } + base := idx * m.cols + for j := 0; j < m.cols; j++ { + if j < len(v) { + m.matData[base+j] = float64(v[j]) + } else { + m.matData[base+j] = 0 + } + } +} + +// removeMatLocked removes matrix row idx and rebuilds the view. Caller must +// hold a write lock. +func (m *MemStore) removeMatLocked(idx int) { + copy(m.matData[idx*m.cols:], m.matData[(idx+1)*m.cols:]) + m.matData = m.matData[:len(m.matData)-m.cols] + if len(m.matData) == 0 { + m.cols = 0 + m.matDense = nil + return + } + if m.cols > 0 { + m.matDense = mat.NewDense(len(m.items), m.cols, m.matData) + } else { + m.matDense = nil + } +} + +// float32ToF64 widens v to a row of cols float64 values, zero-filling missing +// trailing dims and truncating excess ones. +func float32ToF64(v []float32, cols int) []float64 { + out := make([]float64, cols) + for i := 0; i < cols && i < len(v); i++ { + out[i] = float64(v[i]) + } + return out } func l2Norm(v []float32) float64 { diff --git a/internal/ingestion/component/knowledge_compiler/common/types.go b/internal/ingestion/component/knowledge_compiler/common/types.go index 397516ba48..97375efcb0 100644 --- a/internal/ingestion/component/knowledge_compiler/common/types.go +++ b/internal/ingestion/component/knowledge_compiler/common/types.go @@ -122,6 +122,11 @@ type Product struct { Vector []float32 ParentID string Meta map[string]any + // Merged marks rows that already went through dataset-level dedup + // (kc_merged=1, doc_id=kb). The consumer distinguishes these from the + // per-document compiled rows (doc_id=, no kc_merged) so it can + // KNN against only the merged set instead of re-deduping the whole KB. + Merged bool } // Outputs is the result of a variant Run. All compiled products are buffered diff --git a/internal/ingestion/component/knowledge_compiler/mindmap/mindmap.go b/internal/ingestion/component/knowledge_compiler/mindmap/mindmap.go index 52b8848d4b..eea837dd5d 100644 --- a/internal/ingestion/component/knowledge_compiler/mindmap/mindmap.go +++ b/internal/ingestion/component/knowledge_compiler/mindmap/mindmap.go @@ -20,6 +20,36 @@ import ( "ragflow/internal/utility" ) +// batchSubmitter fans out the batch extraction jobs on the process-wide +// knowledge-compilation pool. It is injected by the knowledge_compiler wiring +// (component.go) so every stage shares one vCPU-sized concurrency bound; when +// nil the batches run sequentially (the historic default). +var batchSubmitter func(ctx context.Context, jobs []func() error) error + +// SetBatchSubmitter installs the shared-pool fan-out used by Run's extraction +// stage. Pass nil to revert to serial execution. +func SetBatchSubmitter(submit func(ctx context.Context, jobs []func() error) error) { + batchSubmitter = submit +} + +// runBatches mirrors structure.runBatches: concurrent under the wired global +// compiler pool, or serial when no submitter is set. The first error is +// returned after all jobs settle; the global pool is never StopWait'd. +func runBatches(ctx context.Context, jobs []func() error) error { + if len(jobs) == 0 { + return nil + } + if batchSubmitter != nil { + return batchSubmitter(ctx, jobs) + } + for _, j := range jobs { + if err := j(); err != nil { + return err + } + } + return nil +} + // Run executes the mindmap variant. func Run(ctx context.Context, deps common.Deps, param common.Param, inputs common.Inputs) (common.Outputs, error) { docID := firstNonEmpty(inputs.DocID, deps.DatasetID) @@ -38,19 +68,10 @@ func Run(ctx context.Context, deps common.Deps, param common.Param, inputs commo // One LLM task per token-budget batch (mirrors __call__'s task fan-out). batches := packSections(sections, deps.Tokenizer) results := make([]utility.OMap, len(batches)) - // Bounded concurrency. Python fans out with unbounded asyncio.gather plus a - // global chat_limiter semaphore; here we cap with a worker pool sized by - // MaxWorkers (defaulting to 1, i.e. serial). - n := param.MaxWorkers - if n <= 0 { - n = 1 - } - wp := utility.NewWorkerPool[func() error, struct{}](n, n, - func(_ context.Context, fn func() error) (struct{}, error) { return struct{}{}, fn() }) - var futs []utility.WorkerPoolFuture[func() error, struct{}] + jobs := make([]func() error, 0, len(batches)) for i, text := range batches { i, text := i, text - f, err := wp.Submit(ctx, func() error { + jobs = append(jobs, func() error { resp, err := deps.Chat.Chat(ctx, common.ChatRequest{ LLMID: llmID, SystemPrompt: renderPrompt(text), @@ -59,20 +80,16 @@ func Run(ctx context.Context, deps common.Deps, param common.Param, inputs commo if err != nil { return err } + // Distinct slice index per batch → no cross-goroutine contention. results[i] = utility.Todict(utility.Dictify(utility.StripFences(resp.Content))) return nil }) - if err != nil { - wp.StopWait() - return common.Outputs{}, err - } - futs = append(futs, f) } - wp.StopWait() - for _, f := range futs { - if res, _ := f.Wait(ctx); res.Err != nil { - return common.Outputs{}, res.Err - } + // The extraction batches are LLM-bounded, not CPU-bounded: run them on the + // shared global compiler pool (vCPU-sized) when a submitter is wired in, + // otherwise fall back to serial execution (historic default). + if err := runBatches(ctx, jobs); err != nil { + return common.Outputs{}, err } // Merge batch dicts in batch order (mirrors reduce(self._merge, res)) and diff --git a/internal/ingestion/component/knowledge_compiler/pool_wiring.go b/internal/ingestion/component/knowledge_compiler/pool_wiring.go new file mode 100644 index 0000000000..6eb4ec017a --- /dev/null +++ b/internal/ingestion/component/knowledge_compiler/pool_wiring.go @@ -0,0 +1,38 @@ +// +// 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 knowledge_compiler + +import ( + "context" + + "ragflow/internal/ingestion/component/knowledge_compiler/mindmap" + "ragflow/internal/ingestion/component/knowledge_compiler/structure" + "ragflow/internal/ingestion/knowledge_compile" +) + +// init wires every knowledge-compiler variant's batch fan-out to the single +// process-wide compiler pool (vCPU-sized, defined in the knowledge_compile +// package). This unifies the component-level MAP/extraction stages with the +// consumer-level KNN / LLM-merge / write stages under one shared concurrency +// bound, so docengine-bounded and LLM-bounded work never exceed the host's +// core count regardless of which stage is running. +func init() { + submit := func(ctx context.Context, jobs []func() error) error { + return knowledge_compile.SubmitCompilerJobs(ctx, jobs) + } + structure.SetBatchSubmitter(submit) + mindmap.SetBatchSubmitter(submit) +} diff --git a/internal/ingestion/component/knowledge_compiler/structure/merge.go b/internal/ingestion/component/knowledge_compiler/structure/merge.go index f375875079..5faa0333e8 100644 --- a/internal/ingestion/component/knowledge_compiler/structure/merge.go +++ b/internal/ingestion/component/knowledge_compiler/structure/merge.go @@ -7,8 +7,14 @@ import ( "sync" "ragflow/internal/ingestion/component/knowledge_compiler/common" + "ragflow/internal/tokenizer" ) +// maxBatchTokenReserve is the share of the model's token budget kept in reserve +// for the system prompt, the batch instruction, and the model's JSON reply, so +// a single DecideBatch sub-call never overflows max_token. +const maxBatchTokenReserve = 0.15 + // Merge prompts, verbatim from Python structure.py. The merge is driven by // the user-supplied prompts; the decision instruction lets us branch on the // LLM's verdict via GenJSON. @@ -33,6 +39,19 @@ Return ONLY a JSON object with this exact structure (no markdown fences, no comm "merged": }` +// batchMergeDecisionInstruction mirrors mergeDecisionInstruction but for many +// pairs at once: the model judges every pair in a single round-trip and returns +// an array (wrapped under "pairs") so the consumer can fold an entire batch of +// (existing, incoming) groups without one LLM call per pair. +const batchMergeDecisionInstruction = `You are judging several (Item A, Item B) pairs at once. For each pair decide whether they refer to the same logical entity (for entities) or the same logical relation (for relations), and merge them only if they are the same. + +Return ONLY a JSON object with exactly one key "pairs", whose value is a JSON array with one element per pair, in the same order as listed. Each element MUST have this exact structure (no markdown fences, no commentary): +{ + "index": , + "duplicated": , + "merged": +}` + // mergeJudgeTemperature mirrors Python's gen_conf for merge judging // (_struct_merge_pair uses temperature 0.0). var mergeJudgeTemperature = 0.0 @@ -80,6 +99,19 @@ type LLMMergeDecider struct { Embed common.Embedder Threshold float64 + // maxBatchTokens caps the estimated prompt size of a single DecideBatch + // sub-call. When 0 (the default) the whole batch is sent in one call, + // preserving the historic behavior. When positive the batch is split into + // token-bounded sub-calls so a large candidate set can never overflow the + // model's max_token window. + maxBatchTokens int + + // submit runs one LLM sub-batch off the shared knowledge-compilation pool. + // When nil each sub-batch runs inline (sequentially). It is injected by the + // knowledge_compile package so every stage shares one process-wide, + // vCPU-sized concurrency bound. + submit func(ctx context.Context, fn func() error) error + mu sync.Mutex aliases map[string]string } @@ -89,6 +121,31 @@ func NewLLMMergeDecider(chat common.ChatInvoker, llmID string, embed common.Embe return &LLMMergeDecider{Chat: chat, LLMID: llmID, Embed: embed, Threshold: threshold, aliases: map[string]string{}} } +// SetMaxBatchTokens caps the estimated prompt size of a single DecideBatch +// sub-call. Non-positive values disable per-call batching. The effective budget +// is modelMaxTokens*(1-maxBatchTokenReserve), keeping headroom for the system +// prompt and the JSON reply. +func (d *LLMMergeDecider) SetMaxBatchTokens(modelMaxTokens int) { + if modelMaxTokens <= 0 { + d.maxBatchTokens = 0 + return + } + budget := int(float64(modelMaxTokens) * (1 - maxBatchTokenReserve)) + if budget <= 0 { + d.maxBatchTokens = 0 + return + } + d.maxBatchTokens = budget +} + +// SetSubmitter injects the shared knowledge-compilation pool so DecideBatch can +// run its token-bounded sub-batches concurrently. A nil submitter (the default) +// falls back to running sub-batches sequentially. The submitter must run fn on +// a bounded pool and return its error. +func (d *LLMMergeDecider) SetSubmitter(submit func(ctx context.Context, fn func() error) error) { + d.submit = submit +} + // Decide implements MergeDecider. func (d *LLMMergeDecider) Decide(ctx context.Context, existing, incoming common.Product, bestScore float64) (MergeDecision, common.Product, error) { if bestScore < d.Threshold { @@ -101,7 +158,19 @@ func (d *LLMMergeDecider) Decide(ctx context.Context, existing, incoming common. if merged == nil { return DecisionKeepBoth, common.Product{}, nil } + replacement, err := d.BuildReplacement(ctx, existing, incoming, merged) + if err != nil { + return DecisionKeepBoth, common.Product{}, err + } + return DecisionMerge, replacement, nil +} +// buildReplacement folds an LLM-merged payload back into a replacement Product +// for the existing row: aliases are recorded (entities), merge invariants are +// applied (relations), provenance is unioned, and the payload is re-embedded. +// Shared by Decide (single pair) and the batched judge so both paths produce +// identical merged rows. +func (d *LLMMergeDecider) BuildReplacement(ctx context.Context, existing, incoming common.Product, merged map[string]any) (common.Product, error) { kind, _ := existing.Meta["kind"].(string) if kind == "entity" { oldName := entityNameValue(existing) @@ -122,10 +191,10 @@ func (d *LLMMergeDecider) Decide(ctx context.Context, existing, incoming common. texts := []string{payloadDescription(merged)} vecs, err := d.Embed.Encode(ctx, texts) if err != nil { - return DecisionKeepBoth, common.Product{}, err + return common.Product{}, err } if len(vecs) == 0 { - return DecisionKeepBoth, common.Product{}, fmt.Errorf("knowledge_compiler: re-embed of merged payload returned no vector") + return common.Product{}, fmt.Errorf("knowledge_compiler: re-embed of merged payload returned no vector") } meta := map[string]any{} @@ -134,7 +203,7 @@ func (d *LLMMergeDecider) Decide(ctx context.Context, existing, incoming common. } meta["source_chunk_ids"] = chunkIDs refreshMetaFromPayload(meta, kind, merged) - replacement := common.Product{ + return common.Product{ ID: existing.ID, DocID: existing.DocID, TenantID: existing.TenantID, @@ -142,8 +211,173 @@ func (d *LLMMergeDecider) Decide(ctx context.Context, existing, incoming common. Content: payloadJSON(merged), Vector: vecs[0], Meta: meta, + }, nil +} + +// MergePairInput is one (existing, incoming) pair fed to the batched judge. +type MergePairInput struct { + Index int + Existing string + Incoming string +} + +// BatchMergeResult is the judge's verdict for one pair. +type BatchMergeResult struct { + Index int + Duplicated bool + Merged map[string]any +} + +// DecideBatch judges every pair and returns the verdicts in input order. It is +// the batched counterpart of Decide's single-pair judge. When a token budget is +// configured it splits the pairs into token-bounded sub-batches (each sub-batch +// still judged in a single LLM call, preserving pair order) so a large candidate +// set can never overflow the model's max_token window. +func (d *LLMMergeDecider) DecideBatch(ctx context.Context, pairs []MergePairInput) ([]BatchMergeResult, error) { + if len(pairs) == 0 { + return nil, nil } - return DecisionMerge, replacement, nil + if d.maxBatchTokens <= 0 { + return mergePairsBatch(ctx, d.Chat, d.LLMID, pairs) + } + chunks := splitByTokens(pairs, d.maxBatchTokens) + // The merge decisions are LLM-bounded, not CPU-bounded: run the token-bounded + // sub-batches concurrently on the injected shared pool (vCPU-sized) when more + // than one chunk exists. Order is preserved by writing each chunk's verdicts + // into its own slot (the model echoes the global pair index, so callers still + // key by input index regardless of execution order). + if len(chunks) == 1 || d.submit == nil { + out := make([]BatchMergeResult, 0, len(pairs)) + for _, chunk := range chunks { + sub, err := mergePairsBatch(ctx, d.Chat, d.LLMID, chunk) + if err != nil { + return nil, err + } + out = append(out, sub...) + } + return out, nil + } + results := make([][]BatchMergeResult, len(chunks)) + var ( + wg sync.WaitGroup + errOnce sync.Once + firstErr error + ) + for i, chunk := range chunks { + i, chunk := i, chunk + wg.Add(1) + err := d.submit(ctx, func() error { + defer wg.Done() + sub, err := mergePairsBatch(ctx, d.Chat, d.LLMID, chunk) + if err != nil { + errOnce.Do(func() { firstErr = err }) + return err + } + results[i] = sub + return nil + }) + if err != nil { + // The job was never enqueued, so the closure's wg.Done() will never + // run — decrement here and record the failure so we don't deadlock + // on wg.Wait() and don't silently drop the submit error. + wg.Done() + errOnce.Do(func() { firstErr = err }) + } + } + wg.Wait() + if firstErr != nil { + return nil, firstErr + } + out := make([]BatchMergeResult, 0, len(pairs)) + for _, sub := range results { + out = append(out, sub...) + } + return out, nil +} + +// splitByTokens groups pairs into contiguous chunks whose estimated prompt +// size stays within budget. The original (global) Index of every pair is +// preserved — the model echoes it back — so callers can still key verdicts by +// the input index. A single pair that alone exceeds budget is sent on its own +// (the model may still truncate, but we never silently drop it). +func splitByTokens(pairs []MergePairInput, budget int) [][]MergePairInput { + var chunks [][]MergePairInput + cur := make([]MergePairInput, 0, len(pairs)) + used := 0 + for _, p := range pairs { + // The pair contents dominate the prompt size; the "Pair N:" wrappers + // are a negligible constant overhead per pair, ignored here. + est := tokenizer.NumTokensFromString(p.Existing) + tokenizer.NumTokensFromString(p.Incoming) + if len(cur) > 0 && used+est > budget { + chunks = append(chunks, cur) + // Fresh backing array: cur[:0] shares the buffer with the chunk we + // just appended, so the next iteration would overwrite it. + cur = make([]MergePairInput, 0, len(pairs)) + used = 0 + } + cur = append(cur, p) + used += est + } + if len(cur) > 0 { + chunks = append(chunks, cur) + } + return chunks +} + +// mergePairsBatch mirrors mergePair but judges every pair in one LLM call and +// returns the verdicts in input order. Pairs whose inputs are unparseable or +// whose verdict is missing are reported as not-duplicated (mirrors the +// per-pair skip behavior: Python logs and keeps both). +func mergePairsBatch(ctx context.Context, chat common.ChatInvoker, llmID string, pairs []MergePairInput) ([]BatchMergeResult, error) { + if len(pairs) == 0 { + return nil, nil + } + var b strings.Builder + for _, p := range pairs { + b.WriteString(fmt.Sprintf("Pair %d:\nItem A (existing):\n%s\n\nItem B (incoming):\n%s\n\n", p.Index, p.Existing, p.Incoming)) + } + res, err := common.GenJSON(ctx, chat, common.ChatRequest{ + LLMID: llmID, + SystemPrompt: mergeSystemPrompt + "\n\n" + batchMergeDecisionInstruction, + UserPrompt: b.String(), + Temperature: &mergeJudgeTemperature, + }) + if err != nil { + return nil, err + } + arr, ok := res["pairs"].([]any) + if !ok { + // Model returned no array: treat every pair as not duplicated rather + // than failing the whole batch. + out := make([]BatchMergeResult, len(pairs)) + for i, p := range pairs { + out[i] = BatchMergeResult{Index: p.Index} + } + return out, nil + } + byIndex := make(map[int]BatchMergeResult, len(arr)) + for _, el := range arr { + obj, ok := el.(map[string]any) + if !ok { + continue + } + idx, _ := obj["index"].(float64) + dup, _ := obj["duplicated"].(bool) + var merged map[string]any + if m, ok := obj["merged"].(map[string]any); ok { + merged = m + } + byIndex[int(idx)] = BatchMergeResult{Index: int(idx), Duplicated: dup, Merged: merged} + } + out := make([]BatchMergeResult, len(pairs)) + for i, p := range pairs { + if r, ok := byIndex[p.Index]; ok { + out[i] = r + } else { + out[i] = BatchMergeResult{Index: p.Index} + } + } + return out, nil } // recordAlias adds one alias→canonical mapping (thread-safe). diff --git a/internal/ingestion/component/knowledge_compiler/structure/structure.go b/internal/ingestion/component/knowledge_compiler/structure/structure.go index 9908e6bc1e..34d659eab0 100644 --- a/internal/ingestion/component/knowledge_compiler/structure/structure.go +++ b/internal/ingestion/component/knowledge_compiler/structure/structure.go @@ -11,9 +11,40 @@ import ( "fmt" "ragflow/internal/ingestion/component/knowledge_compiler/common" - "ragflow/internal/utility" ) +// batchSubmitter fans out the MAP-stage extraction jobs on the process-wide +// knowledge-compilation pool. It is injected by the knowledge_compiler wiring +// (component.go) so every stage shares one vCPU-sized concurrency bound; when +// nil the batches run sequentially (the historic default). +var batchSubmitter func(ctx context.Context, jobs []func() error) error + +// SetBatchSubmitter installs the shared-pool fan-out used by Run's MAP stage. +// Pass nil to revert to serial execution. +func SetBatchSubmitter(submit func(ctx context.Context, jobs []func() error) error) { + batchSubmitter = submit +} + +// runBatches executes the MAP-stage jobs. When a shared-pool submitter is +// wired in, the jobs run concurrently under the single process-wide, vCPU-sized +// compiler-pool concurrency bound; otherwise they run sequentially. On any +// error the first non-nil error is returned after all jobs settle — the global +// pool is never StopWait'd, so an error here does not disrupt other stages. +func runBatches(ctx context.Context, jobs []func() error) error { + if len(jobs) == 0 { + return nil + } + if batchSubmitter != nil { + return batchSubmitter(ctx, jobs) + } + for _, j := range jobs { + if err := j(); err != nil { + return err + } + } + return nil +} + // structureBatchTokenBudget caps one extraction batch's packed chunk tokens. // Python derives the budget from chat_mdl.max_length minus the prompt // overhead; the Go ChatInvoker seam does not expose the model window, so we @@ -54,16 +85,10 @@ func Run(ctx context.Context, deps common.Deps, param common.Param, inputs commo // ---- MAP ---- batches := common.PackBatches(inputs.Chunks, structureBatchTokenBudget, deps.Tokenizer) perBatch := make([][]common.Product, len(batches)) - n := param.MaxWorkers - if n <= 0 { - n = 1 - } - wp := utility.NewWorkerPool[func() error, struct{}](n, n, - func(_ context.Context, fn func() error) (struct{}, error) { return struct{}{}, fn() }) - var futs []utility.WorkerPoolFuture[func() error, struct{}] + jobs := make([]func() error, 0, len(batches)) for i, batch := range batches { i, batch := i, batch - f, err := wp.Submit(ctx, func() error { + jobs = append(jobs, func() error { packed, batchIDs := PackBatch(batch) if len(batchIDs) == 0 { return nil @@ -76,20 +101,16 @@ func Run(ctx context.Context, deps common.Deps, param common.Param, inputs commo if err != nil { return err } + // Distinct slice index per batch → no cross-goroutine contention. perBatch[i] = rows return nil }) - if err != nil { - wp.StopWait() - return common.Outputs{}, err - } - futs = append(futs, f) } - wp.StopWait() - for _, f := range futs { - if res, _ := f.Wait(ctx); res.Err != nil { - return common.Outputs{}, res.Err - } + // The extraction batches are LLM-bounded, not CPU-bounded: run them on the + // shared global compiler pool (vCPU-sized) when a submitter is wired in, + // otherwise fall back to serial execution (historic default). + if err := runBatches(ctx, jobs); err != nil { + return common.Outputs{}, err } // ---- DEDUP ---- diff --git a/internal/ingestion/component/knowledge_compiler/structure/structure_test.go b/internal/ingestion/component/knowledge_compiler/structure/structure_test.go index 2815f98ff0..9103860ba5 100644 --- a/internal/ingestion/component/knowledge_compiler/structure/structure_test.go +++ b/internal/ingestion/component/knowledge_compiler/structure/structure_test.go @@ -96,6 +96,19 @@ func chunkIDsFromPrompt(prompt string) []string { func (m *graphChat) Chat(_ context.Context, req common.ChatRequest) (*common.ChatResponse, error) { ids := chunkIDsFromPrompt(req.UserPrompt) switch { + case strings.HasPrefix(req.UserPrompt, "Pair "): + // Batched merge judge: one call covers every (Item A, Item B) pair. + m.mergeCalls++ + pairs := parseBatchPairs(req.UserPrompt) + var results []string + for _, p := range pairs { + if sameLogicalItem(p.a, p.b) { + results = append(results, fmt.Sprintf(`{"index":%d,"duplicated":true,"merged":%s}`, p.index, payloadJSON(p.a))) + } else { + results = append(results, fmt.Sprintf(`{"index":%d,"duplicated":false,"merged":null}`, p.index)) + } + } + return &common.ChatResponse{Content: fmt.Sprintf(`{"pairs":[%s]}`, strings.Join(results, ","))}, nil case strings.HasPrefix(req.UserPrompt, "Item A (existing):"): m.mergeCalls++ a, b := parseMergeItems(req.UserPrompt) @@ -153,6 +166,40 @@ func parseMergeItems(prompt string) (map[string]any, map[string]any) { // sameLogicalItem mirrors what a reasonable judge would decide on the // fixtures: entities are duplicates when their names match (or the known // Al≡Alpha alias pair); relations when source+target+type all match. +func parseBatchPairs(prompt string) []struct { + index int + a, b map[string]any +} { + sections := strings.Split(prompt, "\nPair ") + var pairs []struct { + index int + a, b map[string]any + } + for i, s := range sections { + if i == 0 { + // First section may start with "Pair 0:" without a leading "\n". + if !strings.HasPrefix(s, "Pair ") { + continue + } + } + // Drop the leading "N:\n". + rest := s + if idx := strings.Index(rest, ":\n"); idx >= 0 { + rest = rest[idx+2:] + } + aBody, bBody, ok := strings.Cut(rest, "\n\nItem B (incoming):\n") + if !ok { + continue + } + aBody = strings.TrimPrefix(aBody, "Item A (existing):\n") + pairs = append(pairs, struct { + index int + a, b map[string]any + }{index: i, a: parsePayload(aBody), b: parsePayload(bBody)}) + } + return pairs +} + func sameLogicalItem(a, b map[string]any) bool { if a == nil || b == nil { return false @@ -561,6 +608,121 @@ func TestLLMMergeDeciderContracts(t *testing.T) { } } +// TestLLMMergeDeciderDecideBatch locks the batched judge: every pair is +// judged in a single LLM call (one mergeCalls increment), and the verdict +// array is returned in input order with duplicated/merged fields set. +func TestLLMMergeDeciderDecideBatch(t *testing.T) { + chat := &graphChat{} + d := NewLLMMergeDecider(chat, "llm1", hashEmbedder{dim: 8}, 0.99) + + alpha := common.Product{ + ID: "row-alpha", + DocID: "kb1", + Content: payloadJSON(map[string]any{"type": "letter", "name": "Alpha", "description": "the letter Alpha"}), + Meta: map[string]any{"kind": "entity", "name": "Alpha", "source_chunk_ids": []string{"c1"}}, + } + al := common.Product{ + Content: payloadJSON(map[string]any{"type": "letter", "name": "Al", "description": "short for Alpha"}), + Meta: map[string]any{"kind": "entity", "name": "Al", "source_chunk_ids": []string{"c2"}}, + } + beta := common.Product{ + Content: payloadJSON(map[string]any{"type": "letter", "name": "Beta", "description": "the letter Beta"}), + Meta: map[string]any{"kind": "entity", "name": "Beta", "source_chunk_ids": []string{"c3"}}, + } + + pairs := []MergePairInput{ + {Index: 0, Existing: alpha.Content, Incoming: al.Content}, // duplicated (Al≡Alpha) + {Index: 1, Existing: alpha.Content, Incoming: beta.Content}, // distinct + } + before := chat.mergeCalls + results, err := d.DecideBatch(context.Background(), pairs) + if err != nil { + t.Fatalf("DecideBatch: %v", err) + } + if chat.mergeCalls != before+1 { + t.Errorf("batched judge made %d LLM calls, want 1", chat.mergeCalls-before) + } + if len(results) != 2 { + t.Fatalf("want 2 results, got %d", len(results)) + } + if !results[0].Duplicated || results[0].Merged == nil { + t.Errorf("pair 0 (Al→Alpha) should be duplicated") + } + if results[1].Duplicated || results[1].Merged != nil { + t.Errorf("pair 1 (Beta) should be distinct") + } + if results[0].Index != 0 || results[1].Index != 1 { + t.Errorf("results must preserve input index order") + } +} + +// TestLLMMergeDeciderDecideBatchEmpty locks that an empty pair slice is a +// no-op and never invokes the LLM. +func TestLLMMergeDeciderDecideBatchEmpty(t *testing.T) { + chat := &graphChat{} + d := NewLLMMergeDecider(chat, "llm1", hashEmbedder{dim: 8}, 0.99) + before := chat.mergeCalls + out, err := d.DecideBatch(context.Background(), nil) + if err != nil || out != nil { + t.Fatalf("empty DecideBatch: err=%v out=%v", err, out) + } + if chat.mergeCalls != before { + t.Errorf("empty DecideBatch must not call the LLM") + } +} + +// TestLLMMergeDeciderDecideBatchSplitsByTokenBudget locks that a tight token +// budget forces DecideBatch to issue several LLM calls (sub-batches) while +// still returning every verdict keyed by its original global index. This is +// the guard against overflowing the model's max_token on large candidate sets. +func TestLLMMergeDeciderDecideBatchSplitsByTokenBudget(t *testing.T) { + chat := &graphChat{} + d := NewLLMMergeDecider(chat, "llm1", hashEmbedder{dim: 8}, 0.99) + // Tiny model budget → one pair per sub-batch (exercises the split path). + // 20 * (1-0.15) = 17 token budget, well under each ~40-token pair. + d.SetMaxBatchTokens(20) + + alpha := common.Product{ + ID: "row-alpha", + DocID: "kb1", + Content: payloadJSON(map[string]any{"type": "letter", "name": "Alpha", "description": "the letter Alpha"}), + Meta: map[string]any{"kind": "entity", "name": "Alpha", "source_chunk_ids": []string{"c1"}}, + } + al := common.Product{ + Content: payloadJSON(map[string]any{"type": "letter", "name": "Al", "description": "short for Alpha"}), + Meta: map[string]any{"kind": "entity", "name": "Al", "source_chunk_ids": []string{"c2"}}, + } + beta := common.Product{ + Content: payloadJSON(map[string]any{"type": "letter", "name": "Beta", "description": "the letter Beta"}), + Meta: map[string]any{"kind": "entity", "name": "Beta", "source_chunk_ids": []string{"c3"}}, + } + + pairs := []MergePairInput{ + {Index: 0, Existing: alpha.Content, Incoming: al.Content}, // duplicated (Al≡Alpha) + {Index: 1, Existing: alpha.Content, Incoming: beta.Content}, // distinct + {Index: 2, Existing: beta.Content, Incoming: alpha.Content}, // distinct + } + before := chat.mergeCalls + results, err := d.DecideBatch(context.Background(), pairs) + if err != nil { + t.Fatalf("DecideBatch: %v", err) + } + // Budget=1 → one LLM call per pair. + if chat.mergeCalls != before+3 { + t.Errorf("token-split DecideBatch made %d LLM calls, want 3", chat.mergeCalls-before) + } + if len(results) != 3 { + t.Fatalf("want 3 results, got %d", len(results)) + } + // Verdicts must stay keyed by the original global index, not re-indexed. + wantDup := map[int]bool{0: true, 1: false, 2: false} + for _, r := range results { + if r.Duplicated != wantDup[r.Index] { + t.Errorf("pair %d: duplicated=%v, want %v", r.Index, r.Duplicated, wantDup[r.Index]) + } + } +} + func TestApplyMergeInvariants(t *testing.T) { existing := common.Product{ Content: payloadJSON(map[string]any{"type": "linked", "source": "Alpha", "target": "Beta"}), diff --git a/internal/ingestion/knowledge_compile/consumer.go b/internal/ingestion/knowledge_compile/consumer.go index e7b2a1b9aa..1d8bb6c3df 100644 --- a/internal/ingestion/knowledge_compile/consumer.go +++ b/internal/ingestion/knowledge_compile/consumer.go @@ -17,8 +17,12 @@ package knowledge_compile import ( "context" + "sort" "sync" "time" + + "ragflow/internal/engine" + kccommon "ragflow/internal/ingestion/component/knowledge_compiler/common" ) // Consumer is the dataset-level post-processing worker (§11.5). Multiple @@ -28,37 +32,37 @@ import ( // broker — is the scheduling system of record and the source of same-KB // serialization. type Consumer struct { - scheduler Scheduler + scheduler Claimer reader Reader writer Writer factory DeduperFactory - batchSize int - ttl time.Duration - heartbeat time.Duration - pollInterval time.Duration - sweepInterval time.Duration + ttl time.Duration + heartbeat time.Duration + pollInterval time.Duration + sweepInterval time.Duration + mergeThreshold float64 // KNN similarity threshold for "existing merged row is a duplicate" mu sync.Mutex seqs map[string]map[string]uint64 // dataset -> docID -> last applied seq (per-doc out-of-order guard) tombs map[string]map[string]uint64 // dataset -> docID -> delete event seq (tombstone) } -// NewConsumer constructs a Consumer driven by the given Scheduler. Tests pass a +// NewConsumer constructs a Consumer driven by the given Claimer. Tests pass a // FakeScheduler and override the Reader/Writer/Deduper via options. -func NewConsumer(scheduler Scheduler, opts ...Option) *Consumer { +func NewConsumer(scheduler Claimer, opts ...Option) *Consumer { c := &Consumer{ - scheduler: scheduler, - reader: infinityReader{}, - writer: infinityWriter{}, - factory: defaultDeduperFactory, - batchSize: 32, - ttl: 2 * time.Minute, - heartbeat: 20 * time.Second, - pollInterval: 2 * time.Second, - sweepInterval: 30 * time.Second, - seqs: map[string]map[string]uint64{}, - tombs: map[string]map[string]uint64{}, + scheduler: scheduler, + reader: engineReader{eng: engine.Get()}, + writer: engineWriter{eng: engine.Get()}, + factory: defaultDeduperFactory, + ttl: 2 * time.Minute, + heartbeat: 20 * time.Second, + pollInterval: 2 * time.Second, + sweepInterval: 30 * time.Second, + mergeThreshold: 0.99, + seqs: map[string]map[string]uint64{}, + tombs: map[string]map[string]uint64{}, } for _, o := range opts { o(c) @@ -82,46 +86,52 @@ func (c *Consumer) Run(ctx context.Context) { select { case <-ctx.Done(): return - case ds := <-notifyCh: - if ds != "" { - c.processDataset(ctx, ds) + case datasetID, ok := <-notifyCh: + if !ok { + notifyCh = nil + continue } + // The notify carries the dataset that just received backlog, so + // claim that specific dataset directly instead of probing for an + // arbitrary claimable one. + c.claimAndProcess(ctx, datasetID) case <-poll.C: - c.claimOne(ctx) + c.tryClaimAndProcess(ctx) case <-sweep.C: - // Recover inflight left by crashed workers (crash recovery, §11.5). - if _, err := c.scheduler.ReclaimExpired(ctx, time.Now()); err != nil { - // best-effort; next tick retries - _ = err - } + // Crash recovery: TryClaim reclaims expired inflight leases + // (§11.5) before claiming any ready batch. + c.tryClaimAndProcess(ctx) } } } -// claimOne finds a claimable KB and processes it. FindClaimable returns at most -// one dataset so the worker handles it before looking for more. -func (c *Consumer) claimOne(ctx context.Context) { - ids, err := c.scheduler.FindClaimable(ctx, 1) - if err != nil || len(ids) == 0 { +// claimAndProcess claims the given dataset (a notify pointed at it) and processes +// the closed batch. ok=false means the dataset has no ready batch or a live lease +// already holds it (the race was lost). +func (c *Consumer) claimAndProcess(ctx context.Context, datasetID string) { + cr, ok, err := c.scheduler.Claim(ctx, datasetID) + if err != nil || !ok || len(cr.Entries) == 0 { return } - for _, ds := range ids { - c.processDataset(ctx, ds) - } + c.processClaim(ctx, cr) } -// processDataset claims the closed batch for datasetID, processes it, and acks. -// It is the Option E replacement for the old processOnce (lease + drain + merge): -// the claim transaction returns a frozen batch boundary, so there is no moving -// target and no Nak-churn routing. -func (c *Consumer) processDataset(ctx context.Context, datasetID string) { - cr, ok, err := c.scheduler.Claim(ctx, datasetID, c.batchSize) - if err != nil { - return - } - if !ok || len(cr.Entries) == 0 { - return // race lost or nothing to claim +// tryClaimAndProcess claims one closed batch (ready or reclaimed) and processes +// it. ok=false means there was nothing to do this tick. +func (c *Consumer) tryClaimAndProcess(ctx context.Context) { + cr, ok, err := c.scheduler.TryClaim(ctx) + if err != nil || !ok || len(cr.Entries) == 0 { + return // nothing to claim, or the race was lost } + c.processClaim(ctx, cr) +} + +// processClaim processes an already-claimed batch (cr) and acks on success. It +// is the Option E replacement for the old processOnce (lease + drain + merge): +// the claim returned a frozen batch boundary, so there is no moving target and +// no Nak-churn routing. +func (c *Consumer) processClaim(ctx context.Context, cr ClaimResult) { + datasetID := cr.DatasetID // Heartbeat refreshes the claim TTL while we process; a failed touch means // the lease was taken over (or reclaimed) and we must abort without acking. @@ -200,37 +210,69 @@ func (c *Consumer) processBatch(ctx context.Context, tenant, kb string, entries } var completed []BacklogEntry var deleted []string + // Per-document seq advances and tombstone clears are staged locally and + // committed only after the write/delete paths below return without error, + // so a failed batch leaves c.seqs/c.tombs untouched (except tombstones + // recorded for the deletion events already present in this batch) and can + // be retried on reclaim. + var pendingSeq []BacklogEntry + var pendingTombClear []string + // Reconcile each document to its highest-sequence event before classifying + // it as deleted or completed. A completion and a deletion for the same doc + // can land in the same batch, and completions/deletions can also arrive + // out of order across batches. The last event (by per-doc seq) wins, so we + // must not delete a doc that was re-ingested after its deletion, nor accept + // a stale completion shadowed by a later deletion. + type docState struct { + winEvent EventType + winSeq uint64 + } + byDoc := make(map[string]*docState, len(entries)) for _, e := range entries { - switch EventType(e.EventType) { + et := EventType(e.EventType) + st := byDoc[e.DocID] + if st == nil { + byDoc[e.DocID] = &docState{winEvent: et, winSeq: e.Seq} + continue + } + if e.Seq >= st.winSeq { + st.winEvent = et + st.winSeq = e.Seq + } + } + for docID, st := range byDoc { + switch st.winEvent { case EventTypeDeleted: + // Record the tombstone for this batch's deletion so a stale + // completion (seq <= delete seq) is skipped. The tombstone is + // committed immediately (deletions are not rolled back on a + // failed batch because re-running the delete is idempotent). if tomb == nil { tomb = map[string]uint64{} c.tombs[kb] = tomb } - tomb[e.DocID] = e.Seq - deleted = append(deleted, e.DocID) + tomb[docID] = st.winSeq + deleted = append(deleted, docID) case EventTypeCompleted: - // A tombstone means the doc was deleted; a completion with a seq - // <= the delete seq is the stale original completion (skip it). - // A completion with a higher seq is a genuine re-ingest after - // deletion: clear the tombstone so the doc is processed again - // (otherwise the tombstone would skip it forever and grow - // unbounded across the consumer's lifetime). - if delSeq, ok := tomb[e.DocID]; ok { - if e.Seq <= delSeq { - continue // deleted (or a stale completion) before it completed - } - delete(tomb, e.DocID) + // A re-ingest after an earlier deletion: clear the prior + // tombstone (deferred until the batch succeeds). + if _, hadTomb := tomb[docID]; hadTomb { + pendingTombClear = append(pendingTombClear, docID) } - // Seq is per-document, so the stale/duplicate check must be scoped - // to the document, not the whole dataset (C4). - if prev, ok := docSeqs[e.DocID]; ok && e.Seq <= prev { + // Seq is per-document, so the stale/duplicate check must be + // scoped to the document, not the whole dataset (C4). + if prev, ok := docSeqs[docID]; ok && st.winSeq <= prev { continue // stale / duplicate completion for this doc } - docSeqs[e.DocID] = e.Seq - completed = append(completed, e) + // The seq advance is deferred to after the batch succeeds (see + // below), so a transient reader/deduper/writer failure does not + // permanently drop the completion on the next reclaim+retry. + pendingSeq = append(pendingSeq, BacklogEntry{DocID: docID, EventType: string(EventTypeCompleted), Seq: st.winSeq}) + completed = append(completed, BacklogEntry{DocID: docID, EventType: string(EventTypeCompleted), Seq: st.winSeq}) } } + // Sort for deterministic iteration in the delete/load passes below. + sort.Strings(deleted) c.mu.Unlock() if len(deleted) == 0 && len(completed) == 0 { @@ -242,44 +284,158 @@ func (c *Consumer) processBatch(ctx context.Context, tenant, kb string, entries deduper = NewNoopDeduper() } - products, err := c.reader.LoadCompiledProducts(ctx, tenant, kb) - if err != nil { - return err - } - if len(products) == 0 { - // Nothing to merge. Still drop fully-orphaned merged products for - // deleted docs, then treat as success (nothing to do). - for _, d := range deleted { - _ = c.writer.DeleteMergedForDoc(ctx, tenant, kb, d) - } - return nil - } - - // With deletions present, recompute the whole KB; otherwise scope to the - // contributing documents of this batch (efficiency). - if len(deleted) == 0 { - docSet := make(map[string]bool, len(completed)) - for _, e := range completed { - docSet[e.DocID] = true - } - scoped := products[:0] - for _, p := range products { - if docSet[p.DocID] { - scoped = append(scoped, p) - } - } - products = scoped - } - - merged, err := deduper.Dedup(ctx, products) - if err != nil { - return err - } - if err := c.writer.WriteMerged(ctx, tenant, kb, merged); err != nil { - return err - } + deletedSet := make(map[string]bool, len(deleted)) for _, d := range deleted { - _ = c.writer.DeleteMergedForDoc(ctx, tenant, kb, d) + deletedSet[d] = true } + + // --- Deletion (two sequential DocEngine calls, no in-memory load) --- + // Deleted wins regardless of batch order, so we process deletions first. + // The deleted docs' products are never loaded into memory: the DocEngine + // does all the work in two calls: + // 1. DeleteDocLevelForDocs drops every per-document (doc-level) product of + // the deleted docs in a single engine call. + // 2. StripMergedSources removes the deleted doc ids from the source_doc_ids + // array of every dataset-level merged product and deletes any product + // whose array became empty. + // A merged product referencing several deleted docs is pruned in one pass, + // and an emptied product is removed exactly once. + if len(deleted) > 0 { + delIDs := make([]string, 0, len(deletedSet)) + for d := range deletedSet { + delIDs = append(delIDs, d) + } + if err := c.writer.DeleteDocLevelForDocs(ctx, tenant, kb, delIDs); err != nil { + return err + } + if err := c.writer.StripMergedSources(ctx, tenant, kb, delIDs); err != nil { + return err + } + } + + // --- Completion merge --- + // Load only the per-document products of the completed (and not deleted) + // docs — bounded by this batch, never the whole KB. A doc that is both + // completed and deleted is a stale tombstone: the deletion wins, so we skip + // its completion. + var incoming []kccommon.Product + for _, e := range completed { + if deletedSet[e.DocID] { + continue + } + docProducts, err := c.reader.LoadDocProducts(ctx, tenant, kb, e.DocID) + if err != nil { + return err + } + incoming = append(incoming, docProducts...) + } + + // In-memory dedup among the completed batch first. + candidates, err := deduper.Dedup(ctx, incoming) + if err != nil { + return err + } + + // Then dedup each candidate against the DocEngine via KNN top1 + LLM judge, + // mirroring Python _struct_doc_storage_dedup_batch. Candidates that KNN-hit + // the same existing merged row are grouped so the row is merged/updated once + // instead of being rewritten per candidate. + type mergeGroup struct { + existing kccommon.Product + candidates []kccommon.Product + score float64 + } + groupsByID := make(map[string]*mergeGroup, len(candidates)) + var ( + unmatchedMu sync.Mutex + unmatched []kccommon.Product + groupsMu sync.Mutex + ) + // The KNN pass is docengine-bounded (vector search), not CPU-bounded, so we + // fan it out across the shared global compilerPool (vCPU-sized). Output order + // is irrelevant: merged rows are upserted by their idempotent dataset-level + // id, and each candidate lands in exactly one group / the unmatched set. + jobs := make([]compilerJob, 0, len(candidates)) + for _, cand := range candidates { + cand := cand + jobs = append(jobs, func() error { + var vec64 []float64 + if len(cand.Vector) > 0 { + vec64 = make([]float64, len(cand.Vector)) + for i, v := range cand.Vector { + vec64[i] = float64(v) + } + } + hit, score, err := c.reader.SearchSimilar(ctx, tenant, kb, cand.Variant, vec64, 1, c.mergeThreshold) + if err != nil { + return err + } + if hit.ID == "" { + // No sufficiently-similar merged row: insert the candidate as a new + // merged row. + cand.Merged = true + cand.DocID = kb + unmatchedMu.Lock() + unmatched = append(unmatched, cand) + unmatchedMu.Unlock() + return nil + } + groupsMu.Lock() + g := groupsByID[hit.ID] + if g == nil { + g = &mergeGroup{existing: hit, score: score} + groupsByID[hit.ID] = g + } + g.candidates = append(g.candidates, cand) + groupsMu.Unlock() + return nil + }) + } + if err := runCompilerJobs(ctx, jobs); err != nil { + return err + } + + // Fold every KNN group into the LLM in a single batch round-trip (one + // DecideBatch call instead of one Decide per pair), then collect the + // updated existing rows and the candidates judged distinct (new rows). + var newMerged []kccommon.Product + if len(groupsByID) > 0 { + batched := make([]MergeGroup, 0, len(groupsByID)) + for _, g := range groupsByID { + batched = append(batched, MergeGroup{ + Existing: g.existing, + Candidates: g.candidates, + Score: g.score, + }) + } + batched, err = deduper.DecideBatch(ctx, batched) + if err != nil { + return err + } + for _, g := range batched { + newMerged = append(newMerged, g.Merged) + unmatched = append(unmatched, g.Distinct...) + } + } + + // Write the surviving merged set (updated existing + new distinct rows). + mergedFinal := make([]kccommon.Product, 0, len(newMerged)+len(unmatched)) + mergedFinal = append(mergedFinal, newMerged...) + mergedFinal = append(mergedFinal, unmatched...) + if err := c.writer.WriteMerged(ctx, tenant, kb, mergedFinal); err != nil { + return err + } + + // All merge and delete paths succeeded: commit the staged per-document seq + // advances and tombstone clears. These are only now persisted so a failed + // batch leaves c.seqs/c.tombs untouched and a later reclaim can retry. + c.mu.Lock() + for _, e := range pendingSeq { + c.seqs[kb][e.DocID] = e.Seq + } + for _, docID := range pendingTombClear { + delete(c.tombs[kb], docID) + } + c.mu.Unlock() return nil } diff --git a/internal/ingestion/knowledge_compile/consumer_test.go b/internal/ingestion/knowledge_compile/consumer_test.go index fb2c1ee988..2c3c9d9035 100644 --- a/internal/ingestion/knowledge_compile/consumer_test.go +++ b/internal/ingestion/knowledge_compile/consumer_test.go @@ -24,30 +24,42 @@ import ( kccommon "ragflow/internal/ingestion/component/knowledge_compiler/common" ) -// fakeReader returns a fixed product set. +// fakeReader returns a fixed per-document product set, keyed by docID. type fakeReader struct { mu sync.Mutex products []kccommon.Product calls int } -func (r *fakeReader) LoadCompiledProducts(_ context.Context, _, _ string) ([]kccommon.Product, error) { +func (r *fakeReader) LoadDocProducts(_ context.Context, _, _, docID string) ([]kccommon.Product, error) { r.mu.Lock() defer r.mu.Unlock() r.calls++ - out := make([]kccommon.Product, len(r.products)) - copy(out, r.products) + var out []kccommon.Product + for _, p := range r.products { + if p.DocID == docID { + out = append(out, p) + } + } return out, nil } +func (r *fakeReader) SearchSimilar(_ context.Context, _, _ string, _ kccommon.Variant, _ []float64, _ int, _ float64) (kccommon.Product, float64, error) { + return kccommon.Product{}, 0, nil +} + // fakeWriter captures written merged products. type fakeWriter struct { - mu sync.Mutex - written [][]kccommon.Product - deleted []string + mu sync.Mutex + written [][]kccommon.Product + deletedDocLevel []string + strippedSources []string } func (w *fakeWriter) WriteMerged(_ context.Context, _, _ string, products []kccommon.Product) error { + if len(products) == 0 { + return nil + } w.mu.Lock() defer w.mu.Unlock() cp := make([]kccommon.Product, len(products)) @@ -56,10 +68,17 @@ func (w *fakeWriter) WriteMerged(_ context.Context, _, _ string, products []kcco return nil } -func (w *fakeWriter) DeleteMergedForDoc(_ context.Context, _, _, docID string) error { +func (w *fakeWriter) DeleteDocLevelForDocs(_ context.Context, _, _ string, docIDs []string) error { w.mu.Lock() defer w.mu.Unlock() - w.deleted = append(w.deleted, docID) + w.deletedDocLevel = append(w.deletedDocLevel, docIDs...) + return nil +} + +func (w *fakeWriter) StripMergedSources(_ context.Context, _, _ string, docIDs []string) error { + w.mu.Lock() + defer w.mu.Unlock() + w.strippedSources = append(w.strippedSources, docIDs...) return nil } @@ -77,7 +96,6 @@ func newTestConsumer(sch *FakeScheduler, r *fakeReader, w *fakeWriter, factory D WithReader(r), WithWriter(w), WithDeduperFactory(factory), - WithBatchSize(32), ) } @@ -87,10 +105,10 @@ func TestConsumerCompletedWritesMerged(t *testing.T) { w := &fakeWriter{} c := newTestConsumer(sch, r, w, func(string) (Deduper, error) { return NewNoopDeduper(), nil }) - if err := sch.AppendBacklog(context.Background(), "t1", "kb1", "d1", string(EventTypeCompleted), 1); err != nil { + if err := sch.Publish(context.Background(), "t1", "kb1", "d1", string(EventTypeCompleted), 1); err != nil { t.Fatalf("append: %v", err) } - c.processDataset(context.Background(), "kb1") + c.tryClaimAndProcess(context.Background()) w.mu.Lock() defer w.mu.Unlock() @@ -101,8 +119,8 @@ func TestConsumerCompletedWritesMerged(t *testing.T) { t.Fatalf("expected 2 merged products, got %d", len(w.written[0])) } // After ack, the claim row must be cleared (no live lease left behind). - if got, _ := sch.FindClaimable(context.Background(), 1); len(got) != 0 { - t.Fatalf("expected no claimable row after ack, got %v", got) + if _, ok, _ := sch.TryClaim(context.Background()); ok { + t.Fatalf("expected no claimable row after ack") } } @@ -112,19 +130,65 @@ func TestConsumerTombstoneSkipsCompletedBeforeDeleted(t *testing.T) { w := &fakeWriter{} c := newTestConsumer(sch, r, w, func(string) (Deduper, error) { return NewNoopDeduper(), nil }) - // deleted before completed -> completed must be skipped. - if err := sch.AppendBacklog(context.Background(), "t1", "kb1", "d1", string(EventTypeDeleted), 1); err != nil { - t.Fatalf("append deleted: %v", err) - } - if err := sch.AppendBacklog(context.Background(), "t1", "kb1", "d1", string(EventTypeCompleted), 2); err != nil { + // completed (seq 1) before deleted (seq 2): the completion is the stale + // original, so it must be skipped and d1's per-doc products orphaned. + if err := sch.Publish(context.Background(), "t1", "kb1", "d1", string(EventTypeCompleted), 1); err != nil { t.Fatalf("append completed: %v", err) } - c.processDataset(context.Background(), "kb1") + if err := sch.Publish(context.Background(), "t1", "kb1", "d1", string(EventTypeDeleted), 2); err != nil { + t.Fatalf("append deleted: %v", err) + } + c.tryClaimAndProcess(context.Background()) w.mu.Lock() defer w.mu.Unlock() - if len(w.deleted) != 1 { - t.Fatalf("expected 1 orphan-delete call for d1, got %d", len(w.deleted)) + // No merged write (completed skipped). The deletion is handled entirely on + // the DocEngine: d1's per-doc products are dropped in one call and d1 is + // stripped from every dataset-level product in one call. No products are + // loaded into memory. + if len(w.written) != 0 { + t.Fatalf("expected no merged write, got %d", len(w.written)) + } + if len(w.deletedDocLevel) != 1 || w.deletedDocLevel[0] != "d1" { + t.Fatalf("expected DeleteDocLevelForDocs([d1]), got %v", w.deletedDocLevel) + } + if len(w.strippedSources) != 1 || w.strippedSources[0] != "d1" { + t.Fatalf("expected StripMergedSources([d1]), got %v", w.strippedSources) + } +} + +func TestConsumerReingestAfterDeletionWins(t *testing.T) { + sch := NewFakeScheduler() + r := &fakeReader{products: sampleProducts()} + w := &fakeWriter{} + c := newTestConsumer(sch, r, w, func(string) (Deduper, error) { return NewNoopDeduper(), nil }) + + // deleted (seq 1) then completed (seq 2): the completion has the higher + // sequence, so it wins — the doc is re-ingested, NOT deleted. The deletion + // must not drop its per-doc products, and the completion must be merged. + if err := sch.Publish(context.Background(), "t1", "kb1", "d1", string(EventTypeDeleted), 1); err != nil { + t.Fatalf("append deleted: %v", err) + } + if err := sch.Publish(context.Background(), "t1", "kb1", "d1", string(EventTypeCompleted), 2); err != nil { + t.Fatalf("append completed: %v", err) + } + c.tryClaimAndProcess(context.Background()) + + w.mu.Lock() + defer w.mu.Unlock() + // No deletion calls: the higher-seq completion overrides the deletion. + if len(w.deletedDocLevel) != 0 { + t.Fatalf("expected no DeleteDocLevelForDocs, got %v", w.deletedDocLevel) + } + if len(w.strippedSources) != 0 { + t.Fatalf("expected no StripMergedSources, got %v", w.strippedSources) + } + // The completion is merged into the dataset-level products. + if len(w.written) != 1 { + t.Fatalf("expected 1 WriteMerged call, got %d", len(w.written)) + } + if len(w.written[0]) != 2 { + t.Fatalf("expected 2 merged products, got %d", len(w.written[0])) } } @@ -132,12 +196,12 @@ func TestSchedulerClaimClosedBatch(t *testing.T) { sch := NewFakeScheduler() for i := 0; i < 40; i++ { docID := "d" + string(rune('a'+i%26)) + string(rune('0'+i/26)) - if err := sch.AppendBacklog(context.Background(), "t1", "kb1", docID, string(EventTypeCompleted), uint64(i)); err != nil { + if err := sch.Publish(context.Background(), "t1", "kb1", docID, string(EventTypeCompleted), uint64(i)); err != nil { t.Fatalf("append: %v", err) } } - // First claim returns the bounded prefix (batchSize=32), not all 40. - cr1, ok, err := sch.Claim(context.Background(), "kb1", 32) + // First claim returns the bounded prefix (default batch=32), not all 40. + cr1, ok, err := sch.Claim(context.Background(), "kb1") if err != nil || !ok { t.Fatalf("claim1: ok=%v err=%v", ok, err) } @@ -146,7 +210,7 @@ func TestSchedulerClaimClosedBatch(t *testing.T) { } // A second claim by the same holder (still live lease) must not re-claim // the same dataset until the first batch is acked. - _, ok2, _ := sch.Claim(context.Background(), "kb1", 32) + _, ok2, _ := sch.Claim(context.Background(), "kb1") if ok2 { t.Fatalf("second claim should have lost the race (live lease)") } @@ -154,7 +218,7 @@ func TestSchedulerClaimClosedBatch(t *testing.T) { if _, err := sch.Ack(context.Background(), "kb1", cr1.Token, cr1.Entries); err != nil { t.Fatalf("ack: %v", err) } - cr3, ok3, err := sch.Claim(context.Background(), "kb1", 32) + cr3, ok3, err := sch.Claim(context.Background(), "kb1") if err != nil || !ok3 { t.Fatalf("claim3: ok=%v err=%v", ok3, err) } @@ -165,23 +229,19 @@ func TestSchedulerClaimClosedBatch(t *testing.T) { func TestSchedulerReclaimExpired(t *testing.T) { sch := NewFakeScheduler() - if err := sch.AppendBacklog(context.Background(), "t1", "kb1", "d1", string(EventTypeCompleted), 1); err != nil { + if err := sch.Publish(context.Background(), "t1", "kb1", "d1", string(EventTypeCompleted), 1); err != nil { t.Fatalf("append: %v", err) } - _, ok, err := sch.Claim(context.Background(), "kb1", 32) + _, ok, err := sch.Claim(context.Background(), "kb1") if err != nil || !ok { t.Fatalf("claim: ok=%v err=%v", ok, err) } - // Simulate a crash: the inflight is never acked. Reclaim moves it back. - n, err := sch.ReclaimExpired(context.Background(), time.Now().Add(10*time.Minute)) - if err != nil { - t.Fatalf("reclaim: %v", err) - } - if n != 1 { - t.Fatalf("expected 1 reclaimed entry, got %d", n) - } - // The entry is claimable again. - cr2, ok2, err := sch.Claim(context.Background(), "kb1", 32) + // Simulate a crash: the inflight is never acked and the lease has expired. + past := time.Now().Add(-time.Hour) + sch.rows["kb1"].expires = &past + // TryClaim reclaims the expired lease back into backlog and immediately + // claims it again. + cr2, ok2, err := sch.TryClaim(context.Background()) if err != nil || !ok2 { t.Fatalf("reclaim claim: ok=%v err=%v", ok2, err) } diff --git a/internal/ingestion/knowledge_compile/dedup.go b/internal/ingestion/knowledge_compile/dedup.go index fabc7412db..6c0cec3993 100644 --- a/internal/ingestion/knowledge_compile/dedup.go +++ b/internal/ingestion/knowledge_compile/dedup.go @@ -25,8 +25,33 @@ import ( // Deduper folds a set of per-document compiled Products into the dataset-level // merged set. The LLM-backed implementation reuses the component's // GroupedDeduper + LLMMergeDecider (§11.6). +// +// Dedup folds the completed batch's products among themselves (in-memory, before +// any engine round-trip). Decide judges whether an incoming product duplicates +// an existing merged row found by KNN and, when so, returns the merged row +// (mirrors Python _struct_doc_storage_dedup_batch: KNN top1 + LLM merge). type Deduper interface { Dedup(ctx context.Context, rows []kccommon.Product) ([]kccommon.Product, error) + Decide(ctx context.Context, existing, incoming kccommon.Product, bestScore float64) (kccommon.Product, bool, error) + // DecideBatch judges every (existing, candidates) group in a single LLM + // round-trip, folding each group's candidates into its existing row and + // reporting the merged row plus any candidates judged distinct (new rows). + // It replaces the per-pair Decide loop at the batch hot path. + DecideBatch(ctx context.Context, groups []MergeGroup) ([]MergeGroup, error) +} + +// MergeGroup is one KNN-found existing merged row plus the batch of incoming +// products that all KNN-hit it. DecideBatch folds the candidates into Existing +// and fills Merged (the updated row), Duplicate (whether anything was merged), +// and Distinct (candidates judged not-duplicates, i.e. new merged rows). +type MergeGroup struct { + Existing kccommon.Product + Candidates []kccommon.Product + Score float64 + + Merged kccommon.Product + Duplicate bool + Distinct []kccommon.Product } // DeduperFactory builds a per-tenant Deduper. It is invoked once per batch so @@ -42,8 +67,20 @@ type llmDeduper struct { } // NewLLMDeduper builds a KB-scoped deduper from the runtime chat/embed deps. -func NewLLMDeduper(chat kccommon.ChatInvoker, embed kccommon.Embedder, llmID string, threshold float64) Deduper { +// llmMaxTokens is the chat model's token budget (0 disables per-batch token +// splitting in DecideBatch). +func NewLLMDeduper(chat kccommon.ChatInvoker, embed kccommon.Embedder, llmID string, threshold float64, llmMaxTokens int) Deduper { decider := structure.NewLLMMergeDecider(chat, llmID, embed, threshold) + decider.SetMaxBatchTokens(llmMaxTokens) + // Share the process-wide, vCPU-sized compiler pool so DecideBatch's + // token-bounded sub-batches run concurrently with the rest of the pipeline + // (LLM-bounded), all under one concurrency limit. SubmitCompilerJobs enqueues + // every sub-batch then waits on their futures on the caller goroutine — it + // never blocks on a single job, so a stopped pool returns an error instead of + // hanging DecideBatch. + decider.SetSubmitter(func(ctx context.Context, fn func() error) error { + return SubmitCompilerJobs(ctx, []compilerJob{fn}) + }) return &llmDeduper{group: structure.NewGroupedDeduper(decider), decider: decider, embed: embed} } @@ -61,6 +98,79 @@ func (x *llmDeduper) Dedup(ctx context.Context, rows []kccommon.Product) ([]kcco return x.group.Rows(), nil } +// Decide delegates the per-pair duplicate judgment to the LLM merge decider, +// which re-embeds the merged payload and unions provenance on a duplicate verdict. +func (x *llmDeduper) Decide(ctx context.Context, existing, incoming kccommon.Product, bestScore float64) (kccommon.Product, bool, error) { + decision, merged, err := x.decider.Decide(ctx, existing, incoming, bestScore) + if err != nil { + return kccommon.Product{}, false, err + } + if decision == structure.DecisionMerge { + return merged, true, nil + } + return kccommon.Product{}, false, nil +} + +// DecideBatch folds every group with a single LLM call. All candidate pairs +// across all groups are judged at once (mergePairsBatch); each group then folds +// its candidates into its existing row in order so a chain of merges within a +// group accumulates correctly. +func (x *llmDeduper) DecideBatch(ctx context.Context, groups []MergeGroup) ([]MergeGroup, error) { + // Assign a flat pair index to every (group, candidate). + var inputs []structure.MergePairInput + pairIndexOf := make([][]int, len(groups)) + for gi := range groups { + pairIndexOf[gi] = make([]int, len(groups[gi].Candidates)) + for ci := range groups[gi].Candidates { + idx := len(inputs) + pairIndexOf[gi][ci] = idx + inputs = append(inputs, structure.MergePairInput{ + Index: idx, + Existing: groups[gi].Existing.Content, + Incoming: groups[gi].Candidates[ci].Content, + }) + } + } + if len(inputs) == 0 { + return groups, nil + } + results, err := x.decider.DecideBatch(ctx, inputs) + if err != nil { + return nil, err + } + byIndex := make(map[int]structure.BatchMergeResult, len(results)) + for _, r := range results { + byIndex[r.Index] = r + } + + for gi := range groups { + existing := groups[gi].Existing + var distinct []kccommon.Product + duplicated := false + for ci, cand := range groups[gi].Candidates { + r := byIndex[pairIndexOf[gi][ci]] + if !r.Duplicated || r.Merged == nil { + // Judged distinct: keep it as its own new merged row. + c := cand + c.Merged = true + c.DocID = existing.DocID + distinct = append(distinct, c) + continue + } + replacement, err := x.decider.BuildReplacement(ctx, existing, cand, r.Merged) + if err != nil { + return nil, err + } + existing = replacement + duplicated = true + } + groups[gi].Merged = existing + groups[gi].Duplicate = duplicated + groups[gi].Distinct = distinct + } + return groups, nil +} + // noopDeduper performs no LLM merge; it returns the input rows unchanged so // the writer still emits dataset-level products (without cross-document merging). // Used as a safe fallback when LLM deps are unavailable. @@ -70,5 +180,29 @@ func (noopDeduper) Dedup(_ context.Context, rows []kccommon.Product) ([]kccommon return rows, nil } +// Decide never merges: without an LLM judge every incoming product is kept as a +// distinct merged row (matches the noop fallback's "no cross-document merge"). +func (noopDeduper) Decide(_ context.Context, _, incoming kccommon.Product, _ float64) (kccommon.Product, bool, error) { + return kccommon.Product{}, false, nil +} + +// DecideBatch never merges: every candidate becomes its own new merged row. +func (noopDeduper) DecideBatch(_ context.Context, groups []MergeGroup) ([]MergeGroup, error) { + for gi := range groups { + existing := groups[gi].Existing + var distinct []kccommon.Product + for _, cand := range groups[gi].Candidates { + c := cand + c.Merged = true + c.DocID = existing.DocID + distinct = append(distinct, c) + } + groups[gi].Merged = existing + groups[gi].Duplicate = false + groups[gi].Distinct = distinct + } + return groups, nil +} + // NewNoopDeduper builds the fallback deduper. func NewNoopDeduper() Deduper { return noopDeduper{} } diff --git a/internal/ingestion/knowledge_compile/dedup_test.go b/internal/ingestion/knowledge_compile/dedup_test.go new file mode 100644 index 0000000000..a4a178356d --- /dev/null +++ b/internal/ingestion/knowledge_compile/dedup_test.go @@ -0,0 +1,122 @@ +// +// 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 knowledge_compile + +import ( + "context" + "testing" + + kccommon "ragflow/internal/ingestion/component/knowledge_compiler/common" +) + +// recordingDeduper is a fake Deduper that records the Decide/DecideBatch calls +// and applies a simple policy: a candidate whose content is "dup" is merged +// into the existing row, otherwise it stays distinct. +type recordingDeduper struct { + noopDeduper + decideBatchCalls int + lastGroups []MergeGroup +} + +func (r *recordingDeduper) DecideBatch(ctx context.Context, groups []MergeGroup) ([]MergeGroup, error) { + r.decideBatchCalls++ + r.lastGroups = groups + for gi := range groups { + existing := groups[gi].Existing + var distinct []kccommon.Product + dup := false + for _, cand := range groups[gi].Candidates { + if stringOf(cand.Content) == "dup" { + merged := existing + merged.Content = "merged-" + existing.Content + existing = merged + dup = true + continue + } + c := cand + c.Merged = true + c.DocID = existing.DocID + distinct = append(distinct, c) + } + groups[gi].Merged = existing + groups[gi].Duplicate = dup + groups[gi].Distinct = distinct + } + return groups, nil +} + +func stringOf(s any) string { + if v, ok := s.(string); ok { + return v + } + return "" +} + +func TestDeduperDecideBatchFoldsGroups(t *testing.T) { + d := &recordingDeduper{} + existing := kccommon.Product{ID: "row-1", DocID: "kb1", Content: "base"} + groups := []MergeGroup{ + { + Existing: existing, + Candidates: []kccommon.Product{{ID: "a", Content: "dup"}, {ID: "b", Content: "fresh"}}, + Score: 1.0, + }, + } + out, err := d.DecideBatch(context.Background(), groups) + if err != nil { + t.Fatalf("DecideBatch: %v", err) + } + if d.decideBatchCalls != 1 { + t.Errorf("expected exactly one DecideBatch call, got %d", d.decideBatchCalls) + } + g := out[0] + if !g.Duplicate { + t.Errorf("group with a dup candidate should report Duplicate=true") + } + if g.Merged.Content != "merged-base" { + t.Errorf("merged content = %q, want merged-base", g.Merged.Content) + } + if len(g.Distinct) != 1 || stringOf(g.Distinct[0].Content) != "fresh" { + t.Errorf("distinct = %v, want one fresh candidate", g.Distinct) + } + if g.Distinct[0].DocID != "kb1" || !g.Distinct[0].Merged { + t.Errorf("distinct candidate should be a new merged row under the KB") + } +} + +func TestNoopDeduperDecideBatchNoMerge(t *testing.T) { + d := NewNoopDeduper() + existing := kccommon.Product{ID: "row-1", DocID: "kb1", Content: "base"} + groups := []MergeGroup{ + {Existing: existing, Candidates: []kccommon.Product{{ID: "a", Content: "x"}, {ID: "b", Content: "y"}}}, + } + out, err := d.DecideBatch(context.Background(), groups) + if err != nil { + t.Fatalf("DecideBatch: %v", err) + } + g := out[0] + if g.Duplicate { + t.Errorf("noop deduper must never merge") + } + if len(g.Distinct) != 2 { + t.Fatalf("noop deduper should report both candidates distinct, got %d", len(g.Distinct)) + } + for _, c := range g.Distinct { + if c.DocID != "kb1" || !c.Merged { + t.Errorf("distinct candidate should be a new merged row: %+v", c) + } + } +} diff --git a/internal/ingestion/knowledge_compile/event.go b/internal/ingestion/knowledge_compile/event.go index 25e1f3248c..c0225ea752 100644 --- a/internal/ingestion/knowledge_compile/event.go +++ b/internal/ingestion/knowledge_compile/event.go @@ -78,38 +78,47 @@ func ParseEvent(data []byte) (KCCompileEvent, error) { return e, nil } -// defaultScheduler is the package-level Scheduler used by the publishing path +// defaultPublisher is the package-level Publisher used by the publishing path // (PublishCompleted / PublishDeleted). It is installed by Provision (called // once by the owning Ingestor at startup). Until then, publishing is a no-op. -var defaultScheduler Scheduler +var defaultPublisher Publisher -// SetScheduler installs the package-level Scheduler used for publishing. -func SetScheduler(s Scheduler) { defaultScheduler = s } +// defaultClaimer is the package-level Claimer handed to the consumer workers. +// It is the same underlying instance as defaultPublisher (a *mysqlScheduler +// satisfies both roles), set together by SetScheduler. +var defaultClaimer Claimer -// DefaultScheduler returns the package-level Scheduler (nil until Provision). -func DefaultScheduler() Scheduler { return defaultScheduler } +// SetScheduler installs the package-level Publisher and Claimer used for +// publishing and consuming. The same *mysqlScheduler satisfies both interfaces. +func SetScheduler(s Publisher) { + defaultPublisher = s + if c, ok := s.(Claimer); ok { + defaultClaimer = c + } +} + +// DefaultPublisher returns the package-level Publisher used by the producer path. +func DefaultPublisher() Publisher { return defaultPublisher } + +// DefaultClaimer returns the package-level Claimer used by consumer workers. +func DefaultClaimer() Claimer { return defaultClaimer } // PublishCompleted records a doc_completed event: it appends the doc to the // KB's durable MySQL backlog and wakes idle workers over NATS. It is a no-op -// when no Scheduler has been installed (e.g. DB unavailable). A failure is +// when no Publisher has been installed (e.g. DB unavailable). A failure is // returned so callers can log but never fail the pipeline on it. func PublishCompleted(ctx context.Context, tenantID, datasetID, docID string, seq uint64) error { - if defaultScheduler == nil { + if defaultPublisher == nil { return nil } - if err := defaultScheduler.AppendBacklog(ctx, tenantID, datasetID, docID, string(EventTypeCompleted), seq); err != nil { - return err - } - return defaultScheduler.Notify(ctx, datasetID) + return defaultPublisher.Publish(ctx, tenantID, datasetID, docID, string(EventTypeCompleted), seq) } -// PublishDeleted records a doc_deleted event the same way (append + notify). +// PublishDeleted records a doc_deleted event the same way (Publish handles the +// append + notify pairing). func PublishDeleted(ctx context.Context, tenantID, datasetID, docID string, seq uint64) error { - if defaultScheduler == nil { + if defaultPublisher == nil { return nil } - if err := defaultScheduler.AppendBacklog(ctx, tenantID, datasetID, docID, string(EventTypeDeleted), seq); err != nil { - return err - } - return defaultScheduler.Notify(ctx, datasetID) + return defaultPublisher.Publish(ctx, tenantID, datasetID, docID, string(EventTypeDeleted), seq) } diff --git a/internal/ingestion/knowledge_compile/pool.go b/internal/ingestion/knowledge_compile/pool.go new file mode 100644 index 0000000000..b6f5dd31a7 --- /dev/null +++ b/internal/ingestion/knowledge_compile/pool.go @@ -0,0 +1,147 @@ +// +// 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 knowledge_compile + +import ( + "context" + "os" + "runtime" + "strconv" + + "ragflow/internal/utility" +) + +// compilerJob is one unit of knowledge-compilation work (an I/O- or +// LLM-bounded task) executed on the shared global pool. It is a type alias for +// func() error so callers can pass plain []func() error slices without a cast. +type compilerJob = func() error + +// compilerPool is the process-wide bounded worker pool that drives cross-doc +// concurrency for every knowledge-compilation stage: the DocEngine KNN pass in +// processBatch, the LLM merge-decision batches inside DecideBatch, and the +// merged-product writes/deletes. It mirrors internal/ingestion/component/ +// extractor.go's extractorPool: held globally so every Consumer invocation +// shares one rate limiter instead of spinning up a pool per batch. The pool +// only bounds concurrency (it is never StopWait'd), so concurrent processBatch +// calls do not disturb each other — each call tracks completion with its own +// WaitGroup + first-error collection. +// +// The fixed size is the host vCPU count: the stages are docengine-bounded +// (KNN / write / delete) or LLM-bounded (merge decisions) rather than +// CPU-bounded, so the degree of useful parallelism is capped by the number of +// available cores rather than by a hand-tuned constant. +var compilerPool = utility.NewWorkerPool[compilerJob, struct{}]( + compilerConcurrency(), + compilerConcurrency()*4, + func(_ context.Context, j compilerJob) (struct{}, error) { return struct{}{}, j() }, +) + +// compilerConcurrency resolves the global pool size. It defaults to the host +// vCPU count, overridable via KC_COMPILE_CONCURRENCY (mirroring the extractor +// pool's MAX_CONCURRENT_CHATS tuning knob). +func compilerConcurrency() int { + if v := os.Getenv("KC_COMPILE_CONCURRENCY"); v != "" { + if n, err := strconv.Atoi(v); err == nil && n > 0 { + return n + } + } + n := runtime.NumCPU() + if n <= 0 { + return 1 + } + return n +} + +// SetCompilerConcurrency overrides the global pool size at runtime (e.g. from +// service init or tests). Mirrors SetExtractorConcurrency. +func SetCompilerConcurrency(n int) { + if n > 0 { + compilerPool.Resize(n) + } +} + +// runCompilerJobs submits every job to the global pool and waits for all to +// finish, returning the first non-nil error (if any). ctx cancellation aborts +// outstanding jobs. +// +// No per-job goroutines are spun up: Submit is non-blocking until the pool's +// input buffer fills (vCPU*4 deep), so we first collect one future per job and +// then Wait on each in a second pass on the calling goroutine. This keeps the +// fan-out bounded by the shared pool's worker count while avoiding len(jobs) +// short-lived goroutines. +func runCompilerJobs(ctx context.Context, jobs []compilerJob) error { + if len(jobs) == 0 { + return nil + } + futures := make([]utility.WorkerPoolFuture[compilerJob, struct{}], 0, len(jobs)) + var firstErr error + for _, j := range jobs { + f, err := compilerPool.Submit(ctx, j) + if err != nil { + // Pool stopped / ctx done before we could enqueue the rest: + // remember it and stop submitting; we still await what is queued. + if firstErr == nil { + firstErr = err + } + break + } + futures = append(futures, f) + } + for _, f := range futures { + res, werr := f.Wait(ctx) + if werr != nil { + // Wait returns the context error (not a result error) when ctx wins + // the select; surface it so callers don't see a clean nil while jobs + // are incomplete. + if firstErr == nil { + firstErr = werr + } + continue + } + if res.Err != nil && firstErr == nil { + firstErr = res.Err + } + } + return firstErr +} + +// SubmitCompilerJob runs a single job on the global pool and waits for it, +// returning its error. Used to inject bounded parallelism into lower-level +// packages (e.g. structure.LLMMergeDecider) without creating an import cycle. +func SubmitCompilerJob(ctx context.Context, fn compilerJob) error { + f, err := compilerPool.Submit(ctx, fn) + if err != nil { + return err + } + res, werr := f.Wait(ctx) + if werr != nil { + return werr + } + return res.Err +} + +// CompilerBatchSubmitter is the fan-out contract injected into lower-level +// knowledge_compiler variant packages (structure/mindmap) so every stage shares +// the one process-wide compiler pool. Implementations must submit every job to +// the shared pool, wait for all to finish, and return the first non-nil error +// (without StopWait-ing the global pool). +type CompilerBatchSubmitter func(ctx context.Context, jobs []compilerJob) error + +// SubmitCompilerJobs fans out a batch of jobs on the global pool and returns the +// first error. This is the CompilerBatchSubmitter handed to variant packages. +func SubmitCompilerJobs(ctx context.Context, jobs []compilerJob) error { + return runCompilerJobs(ctx, jobs) +} diff --git a/internal/ingestion/knowledge_compile/reader.go b/internal/ingestion/knowledge_compile/reader.go index 93c0a54104..3b0332a85e 100644 --- a/internal/ingestion/knowledge_compile/reader.go +++ b/internal/ingestion/knowledge_compile/reader.go @@ -17,45 +17,84 @@ package knowledge_compile import ( "context" + "encoding/json" "fmt" + "strconv" "ragflow/internal/engine" "ragflow/internal/engine/types" kccommon "ragflow/internal/ingestion/component/knowledge_compiler/common" ) -// Reader loads the per-document compiled products (available_int=0) of a KB so -// the consumer can recompute the dataset-level merged set (§11.6 step 1, §11.7 -// incremental re-dedup). +// Reader finds the compiled products needed for incremental dedup without +// loading the whole KB into memory (§11.6 step 1, §11.7 incremental re-dedup). +// +// The dedup between an incoming per-document product and the already-merged +// rows lives in external storage (DocEngine): the consumer only keeps the +// in-flight batch in memory and asks the engine for nearest matches via KNN. +// It never scans every compiled chunk of the KB, which would OOM on a large +// knowledge base — this mirrors Python's _struct_doc_storage_dedup_batch, which +// takes only the just-compiled docs and KNN-queries the store per doc. type Reader interface { - LoadCompiledProducts(ctx context.Context, tenant, kb string) ([]kccommon.Product, error) + // LoadDocProducts returns the per-document compiled rows for a single + // document (doc_id == source_doc). Bounded by one document, never the whole + // KB. + LoadDocProducts(ctx context.Context, tenant, kb, docID string) ([]kccommon.Product, error) + + // SearchSimilar runs a dense (KNN) search over the existing merged rows of + // the given variant and returns the single most-similar row whose score is + // at least minScore, plus that score. It returns a zero Product when nothing + // clears the threshold. This mirrors Python's _struct_doc_storage_knn_candidate + // (topn=1, similarity_threshold): find the dot product above the threshold + // and maximum, then decide duplication with the LLM. + SearchSimilar(ctx context.Context, tenant, kb string, variant kccommon.Variant, vector []float64, topN int, minScore float64) (kccommon.Product, float64, error) } -type infinityReader struct{} +// engineReader loads the per-document compiled products through the global +// DocEngine (§11.6 step 1, §11.7 incremental re-dedup). It depends on the +// process-wide DocEngine obtained via engine.Get(); the engine abstraction owns +// the storage schema, so this reader is not backend-specific. +type engineReader struct { + eng engine.DocEngine +} -func (infinityReader) LoadCompiledProducts(ctx context.Context, tenant, kb string) ([]kccommon.Product, error) { - eng := engine.Get() +// compiledSelectFields are the columns needed to reconstruct a Product from a +// stored compiled chunk document. +var compiledSelectFields = []string{ + "id", "doc_id", "tenant_id", "compile_kwd", + "content_with_weight", "kc_payload", + "source_chunk_ids", "source_doc_ids", + "name_kwd", "entity_type_kwd", "from_entity_kwd", "to_entity_kwd", + "slug_kwd", "type", +} + +// loadDocProductsLimit is the per-page size used when scrolling a single +// document's compiled rows. A document can compile more than this many rows, so +// LoadDocProducts pages until the engine returns fewer than a full page. +const loadDocProductsLimit = 5000 + +// LoadDocProducts returns the per-document compiled rows for a single document. +// It is bounded to one document, so the consumer never loads the whole KB. The +// results are paged so a document with more than loadDocProductsLimit rows is +// not silently truncated. +func (r engineReader) LoadDocProducts(ctx context.Context, tenant, kb, docID string) ([]kccommon.Product, error) { + eng := r.eng + if eng == nil { + eng = engine.Get() + } if eng == nil { return nil, nil } - baseName := fmt.Sprintf("ragflow_%s", tenant) - const batchSize = 5000 var out []kccommon.Product offset := 0 for { res, err := eng.Search(ctx, &types.SearchRequest{ - IndexNames: []string{baseName}, - KbIDs: []string{kb}, - Filter: map[string]interface{}{"available_int": 0}, - SelectFields: []string{ - "id", "doc_id", "tenant_id", "compile_kwd", - "content_with_weight", "kc_payload", - "source_chunk_ids", "source_doc_ids", - "name_kwd", "entity_type_kwd", "from_entity_kwd", "to_entity_kwd", - "slug_kwd", "type", - }, - Limit: batchSize, - Offset: offset, + IndexNames: []string{fmt.Sprintf("ragflow_%s", tenant)}, + KbIDs: []string{kb}, + Filter: map[string]interface{}{"doc_id": docID}, + SelectFields: compiledSelectFields, + Limit: loadDocProductsLimit, + Offset: offset, }) if err != nil { return nil, err @@ -69,13 +108,10 @@ func (infinityReader) LoadCompiledProducts(ctx context.Context, tenant, kb strin out = append(out, p) } } - // The KB-wide scan is paginated: keep fetching until a page returns - // fewer than batchSize rows, so a KB larger than the cap merges against - // the full compiled set instead of a truncated slice. - if len(res.Chunks) < batchSize { + if len(res.Chunks) < loadDocProductsLimit { break } - offset += batchSize + offset += loadDocProductsLimit } return out, nil } @@ -94,6 +130,7 @@ func productFromChunkMap(c map[string]interface{}, tenant string) (kccommon.Prod id, _ := c["id"].(string) docID, _ := c["doc_id"].(string) variant, _ := c["compile_kwd"].(string) + merged := isMerged(c["kc_merged"]) meta := map[string]any{} if v, ok := c["name_kwd"].(string); ok && v != "" { @@ -137,5 +174,117 @@ func productFromChunkMap(c map[string]interface{}, tenant string) (kccommon.Prod Content: content, Vector: vec, Meta: meta, + Merged: merged, }, true } + +// SearchSimilar runs a dense KNN over the existing merged rows (kc_merged=1, +// compile_kwd=variant) of the KB and returns the closest hit above minScore. +func (r engineReader) SearchSimilar(ctx context.Context, tenant, kb string, variant kccommon.Variant, vector []float64, topN int, minScore float64) (kccommon.Product, float64, error) { + eng := r.eng + if eng == nil { + eng = engine.Get() + } + if eng == nil { + return kccommon.Product{}, 0, nil + } + if topN <= 0 { + topN = 1 + } + dim := len(vector) + req := &types.SearchRequest{ + IndexNames: []string{fmt.Sprintf("ragflow_%s", tenant)}, + KbIDs: []string{kb}, + Limit: topN, + SelectFields: []string{"id", "doc_id", "kb_id", "content_with_weight", "kc_payload", + "name_kwd", "entity_type_kwd", "from_entity_kwd", "to_entity_kwd", "slug_kwd", + "type", "source_chunk_ids", "source_doc_ids", "kc_merged", "compile_kwd"}, + Filter: map[string]interface{}{ + "kc_merged": 1, + "compile_kwd": string(variant), + }, + MatchExprs: []interface{}{ + &types.MatchDenseExpr{ + VectorColumnName: fmt.Sprintf("q_%d_vec", dim), + EmbeddingData: vector, + DistanceType: "cosine", + TopN: topN, + ExtraOptions: map[string]interface{}{"min_score": minScore}, + }, + }, + } + res, err := eng.Search(ctx, req) + if err != nil { + return kccommon.Product{}, 0, err + } + for _, c := range res.Chunks { + p, ok := productFromChunkMap(c, tenant) + if !ok || !p.Merged { + continue + } + score := toFloat64(c["score"]) + return p, score, nil + } + return kccommon.Product{}, 0, nil +} + +// isMerged normalizes the boxed kc_merged field returned by the DocEngine, +// which may be stored as a string ("1"/"0"/"true"), a bool, or a numeric, +// depending on backend and mapping. Returns true for any positive/true form. +func isMerged(v interface{}) bool { + switch t := v.(type) { + case nil: + return false + case bool: + return t + case string: + switch t { + case "1", "true", "True", "TRUE": + return true + } + if f, err := strconv.ParseFloat(t, 64); err == nil { + return f > 0 + } + return false + case int: + return t > 0 + case int64: + return t > 0 + case float64: + return t > 0 + case float32: + return t > 0 + case json.Number: + if f, err := t.Float64(); err == nil { + return f > 0 + } + } + return false +} + +// toFloat64 normalizes the boxed score field returned by the DocEngine into a +// float64, accepting float32, float64, numeric strings, and json.Number. It +// returns 0 when the value is missing or not numeric. +func toFloat64(v interface{}) float64 { + switch t := v.(type) { + case nil: + return 0 + case float64: + return t + case float32: + return float64(t) + case int: + return float64(t) + case int64: + return float64(t) + case string: + if f, err := strconv.ParseFloat(t, 64); err == nil { + return f + } + case json.Number: + if f, err := t.Float64(); err == nil { + return f + } + } + return 0 +} diff --git a/internal/ingestion/knowledge_compile/scheduler.go b/internal/ingestion/knowledge_compile/scheduler.go index 43abe90edc..323311877c 100644 --- a/internal/ingestion/knowledge_compile/scheduler.go +++ b/internal/ingestion/knowledge_compile/scheduler.go @@ -54,53 +54,61 @@ type BacklogEntry struct { // ClaimResult is returned by Scheduler.Claim: the KB's tenant plus the closed // batch of entries moved into inflight, and the token identifying this claim. type ClaimResult struct { - TenantID string - Entries []BacklogEntry - Token string + DatasetID string + TenantID string + Entries []BacklogEntry + Token string } -// Scheduler owns the claim/ack lifecycle against the durable scheduling store -// (Option E §11.4/§11.5). MySQL is the system of record; NATS notify is only a -// wake-up. The claim is a move (backlog -> inflight), never a copy, so a doc id -// is visible to at most one worker at a time; same-KB serialization follows -// from the inflight set, not from the broker. -type Scheduler interface { +// Publisher is the producer-side role (Option E §11.4). It appends a document +// event to the dataset's durable backlog and wakes idle workers. MySQL is the +// system of record; NATS notify is only a best-effort wake-up. Publish is the +// single producer entry point so a publisher never forgets to Notify after +// enqueue (the two are now one atomic call site). +type Publisher interface { // Provision ensures the backing store (table + notify subject) exists. Provision(ctx context.Context) error - // AppendBacklog adds one doc event to the KB's backlog. The append is - // transactional and preserves any existing backlog (concurrent publishers - // do not clobber each other). - AppendBacklog(ctx context.Context, tenantID, datasetID, docID, eventType string, seq uint64) error - // Claim moves a bounded prefix of backlog -> inflight for datasetID and - // returns the closed batch. acquired=false when the row is held by another - // live lease (the race was lost) or the backlog is empty. - Claim(ctx context.Context, datasetID string, batchSize int) (ClaimResult, bool, error) - // Ack removes the claimed batch from inflight; clears the claim metadata - // only when backlog is also empty. Returns the remaining backlog size. - Ack(ctx context.Context, datasetID, token string, batch []BacklogEntry) (backlogRemaining int, err error) + + // Publish records one doc event in the dataset's durable backlog and wakes + // idle workers. It is transactional so concurrent publishers do not clobber + // each other's backlog, and it always pairs the append with a notify. + Publish(ctx context.Context, tenantID, datasetID, docID, eventType string, seq uint64) error +} + +// Claimer is the consumer-side role (Option E §11.5). A cluster of competing +// workers claims closed batches (backlog -> inflight, a move not a copy) per +// dataset, keeps the lease alive while processing, and acks when done. The claim +// is per-KB, so the same dataset is serialized by its single live lease. +type Claimer interface { + // TryClaim finds a dataset with ready backlog and no live lease, or reclaims + // an expired lease, claims it atomically, and returns the closed batch. + // ok is false when there is nothing to claim or reclaim. + TryClaim(ctx context.Context) (ClaimResult, bool, error) + // Claim claims datasetID directly (no batch-size argument; the implementation + // decides the batch boundary). acquired=false when a live lease already holds + // the row (the race was lost) or the backlog is empty. + Claim(ctx context.Context, datasetID string) (ClaimResult, bool, error) // TouchClaim extends the lease while processing (heartbeat). Returns false // when the lease is gone (taken over / reclaimed) so the worker must abort. TouchClaim(ctx context.Context, datasetID, token string, ttl time.Duration) (bool, error) - // FindClaimable returns up to limit dataset ids with non-empty backlog and - // no live lease (a free token to claim). - FindClaimable(ctx context.Context, limit int) ([]string, error) - // Notify wakes idle workers about a dataset that just got backlog. - Notify(ctx context.Context, datasetID string) error - // SubscribeNotify returns a channel of dataset ids pushed by Notify, or nil + // Ack removes the claimed batch from inflight; clears the claim metadata + // only when backlog is also empty. + Ack(ctx context.Context, datasetID, token string, batch []BacklogEntry) (backlogRemaining int, err error) + + // SubscribeNotify returns a channel of dataset ids pushed by Publish, or nil // when the implementation has no push wake-up (callers fall back to polling). SubscribeNotify(ctx context.Context) (<-chan string, error) - // ReclaimExpired moves expired inflight entries back to backlog (sweeper). - ReclaimExpired(ctx context.Context, now time.Time) (int, error) } // newScheduler is the production constructor: a MySQL-backed scheduler with an // optional NATS wake-up. holder identifies this ingestor instance; ttl is the -// claim lease duration. -func newScheduler(db *gorm.DB, mq engine.MessageQueue, holder string, ttl time.Duration) Scheduler { +// claim lease duration. The returned *mysqlScheduler satisfies both Publisher +// and Claimer. +func newScheduler(db *gorm.DB, mq engine.MessageQueue, holder string, ttl time.Duration) *mysqlScheduler { if ttl <= 0 { ttl = 2 * time.Minute } - return &mysqlScheduler{db: db, mq: mq, holder: holder, leaseTTL: ttl} + return &mysqlScheduler{db: db, mq: mq, holder: holder, leaseTTL: ttl, claimBatchSize: defaultClaimBatch} } // ---- JSON helpers (the *_doc_ids columns are TEXT holding []BacklogEntry) ---- @@ -134,25 +142,34 @@ func minInt(a, b int) int { // ---- MySQL-backed scheduler ---- type mysqlScheduler struct { - db *gorm.DB - mq engine.MessageQueue - holder string - leaseTTL time.Duration + db *gorm.DB + mq engine.MessageQueue + holder string + leaseTTL time.Duration + claimBatchSize int } +// defaultClaimBatch is the closed-batch boundary when Claim is invoked without +// an explicit size (the previous batchSize argument). +const defaultClaimBatch = 32 + func (s *mysqlScheduler) Provision(ctx context.Context) error { if s.db == nil { return nil } - return s.db.WithContext(ctx).AutoMigrate(&entity.KnowledgeCompileDoc{}) + return s.db.WithContext(ctx).AutoMigrate(&entity.KnowledgeCompileDataset{}) } -func (s *mysqlScheduler) AppendBacklog(ctx context.Context, tenantID, datasetID, docID, eventType string, seq uint64) error { +// Publish appends one doc event to the dataset's durable backlog and wakes idle +// workers. The append is transactional (concurrent publishers do not clobber +// each other's backlog), and the notify is always paired with the append so a +// producer never needs a separate Notify call. +func (s *mysqlScheduler) Publish(ctx context.Context, tenantID, datasetID, docID, eventType string, seq uint64) error { if s.db == nil { return nil } entry := BacklogEntry{DocID: docID, EventType: eventType, Seq: seq} - // KnowledgeCompileDoc.dataset_id is the PRIMARY KEY, so the SELECT ... FOR + // KnowledgeCompileDataset.dataset_id is the PRIMARY KEY, so the SELECT ... FOR // UPDATE below also takes an InnoDB gap lock for a not-yet-existing dataset; // concurrent publishers for the same dataset serialize on that lock and the // loser observes the row already present. FirstOrCreate then reliably finds @@ -163,7 +180,7 @@ func (s *mysqlScheduler) AppendBacklog(ctx context.Context, tenantID, datasetID, // Where) so the inserted row is fully populated and later queries by // dataset_id find it (M19: a raw-string Where alone leaves DatasetID // blank on the created row). - row := entity.KnowledgeCompileDoc{ + row := entity.KnowledgeCompileDataset{ DatasetID: datasetID, TenantID: tenantID, BacklogDocIDs: "[]", @@ -183,50 +200,62 @@ func (s *mysqlScheduler) AppendBacklog(ctx context.Context, tenantID, datasetID, return tx.Save(&row).Error }) if err != nil { - return fmt.Errorf("knowledge_compile: append backlog %s: %w", datasetID, err) + return fmt.Errorf("knowledge_compile: publish backlog %s: %w", datasetID, err) } - return nil + return s.notify(ctx, datasetID) } -func (s *mysqlScheduler) Claim(ctx context.Context, datasetID string, batchSize int) (ClaimResult, bool, error) { +// claimRow atomically claims the closed batch from the row identified by +// datasetID: it takes a FOR UPDATE row lock, refuses a live lease, moves up to +// claimBatchSize entries from backlog to inflight, and stamps the lease. The +// callers (Claim and TryClaim) differ only in how they pick the datasetID; the +// claim itself is identical so both share this helper and stay race-free. +func (s *mysqlScheduler) claimRow(ctx context.Context, tx *gorm.DB, datasetID string) (ClaimResult, bool, error) { + var row entity.KnowledgeCompileDataset + if err := tx.Clauses(clause.Locking{Strength: "UPDATE"}). + Where("dataset_id = ?", datasetID).First(&row).Error; err != nil { + if errors.Is(err, gorm.ErrRecordNotFound) { + return ClaimResult{}, false, nil // nothing to claim + } + return ClaimResult{}, false, err + } + now := time.Now() + liveLease := row.ClaimOwner != "" && row.ClaimExpiresAt != nil && row.ClaimExpiresAt.After(now) + if liveLease { + return ClaimResult{}, false, nil // a live lease means some worker is already processing this KB + } + backlog := parseEntries(row.BacklogDocIDs) + if len(backlog) == 0 { + return ClaimResult{}, false, nil + } + n := minInt(s.claimBatchSize, len(backlog)) + batch := backlog[:n] + inflight := parseEntries(row.InflightDocIDs) + inflight = append(inflight, batch...) + row.BacklogDocIDs = marshalEntries(backlog[n:]) + row.InflightDocIDs = marshalEntries(inflight) + row.ClaimOwner = s.holder + row.ClaimToken = generateHolder() + exp := now.Add(s.leaseTTL) + row.ClaimExpiresAt = &exp + if err := tx.Save(&row).Error; err != nil { + return ClaimResult{}, false, err + } + return ClaimResult{DatasetID: row.DatasetID, TenantID: row.TenantID, Entries: batch, Token: row.ClaimToken}, true, nil +} + +func (s *mysqlScheduler) Claim(ctx context.Context, datasetID string) (ClaimResult, bool, error) { if s.db == nil { return ClaimResult{}, false, nil } var res ClaimResult var acquired bool err := s.db.WithContext(ctx).Transaction(func(tx *gorm.DB) error { - var row entity.KnowledgeCompileDoc - if err := tx.Clauses(clause.Locking{Strength: "UPDATE"}). - Where("dataset_id = ?", datasetID).First(&row).Error; err != nil { - if errors.Is(err, gorm.ErrRecordNotFound) { - return nil // nothing to claim - } - return err + cr, ok, e := s.claimRow(ctx, tx, datasetID) + if e != nil { + return e } - now := time.Now() - liveLease := row.ClaimOwner != "" && row.ClaimExpiresAt != nil && row.ClaimExpiresAt.After(now) - if liveLease { - return nil // a live lease means some worker is already processing this KB - } - backlog := parseEntries(row.BacklogDocIDs) - if len(backlog) == 0 { - return nil - } - n := minInt(batchSize, len(backlog)) - batch := backlog[:n] - inflight := parseEntries(row.InflightDocIDs) - inflight = append(inflight, batch...) - row.BacklogDocIDs = marshalEntries(backlog[n:]) - row.InflightDocIDs = marshalEntries(inflight) - row.ClaimOwner = s.holder - row.ClaimToken = generateHolder() - exp := now.Add(s.leaseTTL) - row.ClaimExpiresAt = &exp - if err := tx.Save(&row).Error; err != nil { - return err - } - res = ClaimResult{TenantID: row.TenantID, Entries: batch, Token: row.ClaimToken} - acquired = true + res, acquired = cr, ok return nil }) if err != nil { @@ -241,7 +270,7 @@ func (s *mysqlScheduler) Ack(ctx context.Context, datasetID, token string, batch } var remaining int err := s.db.WithContext(ctx).Transaction(func(tx *gorm.DB) error { - var row entity.KnowledgeCompileDoc + var row entity.KnowledgeCompileDataset if err := tx.Clauses(clause.Locking{Strength: "UPDATE"}). Where("dataset_id = ?", datasetID).First(&row).Error; err != nil { return err @@ -273,7 +302,7 @@ func (s *mysqlScheduler) TouchClaim(ctx context.Context, datasetID, token string if s.db == nil { return false, nil } - res := s.db.WithContext(ctx).Model(&entity.KnowledgeCompileDoc{}). + res := s.db.WithContext(ctx).Model(&entity.KnowledgeCompileDataset{}). Where("dataset_id = ? AND claim_token = ?", datasetID, token). Update("claim_expires_at", time.Now().Add(ttl)) if res.Error != nil { @@ -282,23 +311,111 @@ func (s *mysqlScheduler) TouchClaim(ctx context.Context, datasetID, token string return res.RowsAffected > 0, nil } -func (s *mysqlScheduler) FindClaimable(ctx context.Context, limit int) ([]string, error) { - if s.db == nil || limit <= 0 { - return nil, nil +// TryClaim finds a dataset with ready backlog and no live lease, or reclaims an +// expired lease, claims it atomically within a single transaction, and returns +// the closed batch. ok is false when there is nothing to claim or reclaim. +// +// The find and the claim run inside one locked transaction so there is no race +// window between picking a dataset and taking its lease: a concurrent worker +// cannot claim the same row out from under us. Reclaim is inlined here so the +// sweeper and the poller share one entry point (no separate ReclaimExpired). +func (s *mysqlScheduler) TryClaim(ctx context.Context) (ClaimResult, bool, error) { + if s.db == nil { + return ClaimResult{}, false, nil } - var ids []string - err := s.db.WithContext(ctx).Model(&entity.KnowledgeCompileDoc{}). - Where("(claim_expires_at IS NULL OR claim_expires_at <= ?) AND backlog_doc_ids <> '[]' AND backlog_doc_ids <> '' AND backlog_doc_ids IS NOT NULL", time.Now()). - Order("priority DESC, updated_at ASC"). - Limit(limit). - Pluck("dataset_id", &ids).Error + var res ClaimResult + var acquired bool + err := s.db.WithContext(ctx).Transaction(func(tx *gorm.DB) error { + now := time.Now() + // 1) a ready dataset (backlog, no live lease) — claim it directly. + if id, ok := s.findClaimableID(ctx, tx, now); ok { + cr, ok, e := s.claimRow(ctx, tx, id) + if e != nil { + return e + } + res, acquired = cr, ok + return nil + } + // 2) reclaim an expired lease back into backlog, then claim it. + if id, ok, e := s.reclaimOne(ctx, tx, now); e != nil { + return e + } else if ok { + cr, ok, e := s.claimRow(ctx, tx, id) + if e != nil { + return e + } + res, acquired = cr, ok + return nil + } + return nil + }) if err != nil { - return nil, err + return ClaimResult{}, false, err } - return ids, nil + return res, acquired, nil } -func (s *mysqlScheduler) Notify(ctx context.Context, datasetID string) error { +// findClaimableID returns one dataset id with non-empty backlog and no live +// lease, or ok=false when none exists. It runs inside caller's transaction so +// the returned id is still locked when the caller claims it. +func (s *mysqlScheduler) findClaimableID(ctx context.Context, tx *gorm.DB, now time.Time) (string, bool) { + var id string + err := tx.Model(&entity.KnowledgeCompileDataset{}). + Where("(claim_expires_at IS NULL OR claim_expires_at <= ?) AND backlog_doc_ids <> '[]' AND backlog_doc_ids <> '' AND backlog_doc_ids IS NOT NULL", now). + Order("priority DESC, updated_at ASC"). + Limit(1). + Pluck("dataset_id", &id).Error + if err != nil || id == "" { + return "", false + } + return id, true +} + +// reclaimOne moves one expired inflight batch back to backlog (crash recovery) +// and returns that dataset id, or ok=false when nothing is expired. It runs +// inside caller's transaction so the reclaimed row remains locked when the +// caller claims it. +func (s *mysqlScheduler) reclaimOne(ctx context.Context, tx *gorm.DB, now time.Time) (string, bool, error) { + var rows []entity.KnowledgeCompileDataset + if err := tx.Model(&entity.KnowledgeCompileDataset{}). + Where("claim_expires_at IS NOT NULL AND claim_expires_at <= ? AND inflight_doc_ids <> '[]' AND inflight_doc_ids <> ''", now). + Find(&rows).Error; err != nil { + return "", false, err + } + for _, row := range rows { + var cur entity.KnowledgeCompileDataset + if err := tx.Clauses(clause.Locking{Strength: "UPDATE"}). + Where("dataset_id = ?", row.DatasetID).First(&cur).Error; err != nil { + if errors.Is(err, gorm.ErrRecordNotFound) { + continue + } + return "", false, err + } + if cur.ClaimExpiresAt != nil && cur.ClaimExpiresAt.After(now) { + continue + } + inflight := parseEntries(cur.InflightDocIDs) + if len(inflight) == 0 { + continue + } + backlog := parseEntries(cur.BacklogDocIDs) + backlog = append(backlog, inflight...) + cur.BacklogDocIDs = marshalEntries(backlog) + cur.InflightDocIDs = "[]" + cur.ClaimOwner = "" + cur.ClaimToken = "" + cur.ClaimExpiresAt = nil + if err := tx.Save(&cur).Error; err != nil { + return "", false, err + } + return cur.DatasetID, true, nil + } + return "", false, nil +} + +// notify is the best-effort wake-up published after a successful Publish. It is +// a no-op without a configured NATS subject; callers fall back to polling. +func (s *mysqlScheduler) notify(ctx context.Context, datasetID string) error { if s.mq == nil { return nil } @@ -313,50 +430,6 @@ func (s *mysqlScheduler) SubscribeNotify(ctx context.Context) (<-chan string, er return s.mq.SubscribeNotify(ctx) } -func (s *mysqlScheduler) ReclaimExpired(ctx context.Context, now time.Time) (int, error) { - if s.db == nil { - return 0, nil - } - var rows []entity.KnowledgeCompileDoc - if err := s.db.WithContext(ctx).Model(&entity.KnowledgeCompileDoc{}). - Where("claim_expires_at IS NOT NULL AND claim_expires_at <= ? AND inflight_doc_ids <> '[]' AND inflight_doc_ids <> ''", now). - Find(&rows).Error; err != nil { - return 0, err - } - total := 0 - for _, row := range rows { - inflight := parseEntries(row.InflightDocIDs) - if len(inflight) == 0 { - continue - } - total += len(inflight) - err := s.db.WithContext(ctx).Transaction(func(tx *gorm.DB) error { - var cur entity.KnowledgeCompileDoc - if err := tx.Clauses(clause.Locking{Strength: "UPDATE"}). - Where("dataset_id = ?", row.DatasetID).First(&cur).Error; err != nil { - return err - } - // Re-check liveness under the lock so we don't clobber a lease that - // was refreshed between the scan and now. - if cur.ClaimExpiresAt != nil && cur.ClaimExpiresAt.After(now) { - return nil - } - backlog := parseEntries(cur.BacklogDocIDs) - backlog = append(backlog, parseEntries(cur.InflightDocIDs)...) - cur.BacklogDocIDs = marshalEntries(backlog) - cur.InflightDocIDs = "[]" - cur.ClaimOwner = "" - cur.ClaimToken = "" - cur.ClaimExpiresAt = nil - return tx.Save(&cur).Error - }) - if err != nil { - return total, err - } - } - return total, nil -} - // removeEntries drops every entry in batch from the inflight slice (match by // doc_id + event_type + seq). The batch is the exact set the worker claimed, so // this is a precise removal, not a "clear all" — any inflight added by a @@ -390,8 +463,8 @@ type fakeRow struct { expires *time.Time } -// FakeScheduler is an in-memory Scheduler used by tests. It mirrors the -// MySQL semantics (move-not-copy claim, token-checked ack, lease takeover). +// FakeScheduler is an in-memory Publisher + Claimer used by tests. It mirrors +// the MySQL semantics (move-not-copy claim, token-checked ack, lease takeover). type FakeScheduler struct { mu sync.Mutex rows map[string]*fakeRow @@ -412,7 +485,8 @@ func NewFakeScheduler() *FakeScheduler { func (f *FakeScheduler) Provision(_ context.Context) error { return nil } -func (f *FakeScheduler) AppendBacklog(_ context.Context, tenantID, datasetID, docID, eventType string, seq uint64) error { +// Publish appends one doc event and pushes a notify (same as the MySQL path). +func (f *FakeScheduler) Publish(_ context.Context, tenantID, datasetID, docID, eventType string, seq uint64) error { f.mu.Lock() defer f.mu.Unlock() r, ok := f.rows[datasetID] @@ -424,10 +498,14 @@ func (f *FakeScheduler) AppendBacklog(_ context.Context, tenantID, datasetID, do r.tenant = tenantID } r.backlog = append(r.backlog, BacklogEntry{DocID: docID, EventType: eventType, Seq: seq}) + select { + case f.notifyCh <- datasetID: + default: + } return nil } -func (f *FakeScheduler) Claim(_ context.Context, datasetID string, batchSize int) (ClaimResult, bool, error) { +func (f *FakeScheduler) Claim(_ context.Context, datasetID string) (ClaimResult, bool, error) { f.mu.Lock() defer f.mu.Unlock() r, ok := f.rows[datasetID] @@ -442,7 +520,7 @@ func (f *FakeScheduler) Claim(_ context.Context, datasetID string, batchSize int if len(r.backlog) == 0 { return ClaimResult{}, false, nil } - n := minInt(batchSize, len(r.backlog)) + n := minInt(defaultClaimBatch, len(r.backlog)) batch := append([]BacklogEntry{}, r.backlog[:n]...) r.inflight = append(r.inflight, batch...) r.backlog = append([]BacklogEntry{}, r.backlog[n:]...) @@ -450,7 +528,7 @@ func (f *FakeScheduler) Claim(_ context.Context, datasetID string, batchSize int r.token = generateHolder() exp := now.Add(f.leaseTTL) r.expires = &exp - return ClaimResult{TenantID: r.tenant, Entries: batch, Token: r.token}, true, nil + return ClaimResult{DatasetID: datasetID, TenantID: r.tenant, Entries: batch, Token: r.token}, true, nil } func (f *FakeScheduler) Ack(_ context.Context, datasetID, token string, batch []BacklogEntry) (int, error) { @@ -482,53 +560,42 @@ func (f *FakeScheduler) TouchClaim(_ context.Context, datasetID, token string, _ return true, nil } -func (f *FakeScheduler) FindClaimable(_ context.Context, limit int) ([]string, error) { +// TryClaim mirrors the production flow: claim a ready dataset, otherwise reclaim +// an expired lease and claim it. +func (f *FakeScheduler) TryClaim(ctx context.Context) (ClaimResult, bool, error) { f.mu.Lock() - defer f.mu.Unlock() - var ids []string now := time.Now() + var readyID, expiredID string for id, r := range f.rows { - if len(r.backlog) == 0 { - continue - } - if r.owner != "" && r.expires != nil && r.expires.After(now) { - continue - } - ids = append(ids, id) - if limit > 0 && len(ids) >= limit { + if len(r.backlog) > 0 && (r.owner == "" || r.expires == nil || !r.expires.After(now)) { + readyID = id break } } - return ids, nil -} - -func (f *FakeScheduler) Notify(_ context.Context, datasetID string) error { - select { - case f.notifyCh <- datasetID: - default: + if readyID == "" { + for id, r := range f.rows { + if r.owner != "" && r.expires != nil && r.expires.After(now) { + continue + } + if len(r.inflight) > 0 { + r.backlog = append(r.backlog, r.inflight...) + r.inflight = nil + r.owner, r.token, r.expires = "", "", nil + expiredID = id + break + } + } } - return nil + f.mu.Unlock() + if readyID != "" { + return f.Claim(ctx, readyID) + } + if expiredID != "" { + return f.Claim(ctx, expiredID) + } + return ClaimResult{}, false, nil } func (f *FakeScheduler) SubscribeNotify(_ context.Context) (<-chan string, error) { return f.notifyCh, nil } - -func (f *FakeScheduler) ReclaimExpired(_ context.Context, now time.Time) (int, error) { - f.mu.Lock() - defer f.mu.Unlock() - total := 0 - for _, r := range f.rows { - if r.owner != "" && r.expires != nil && r.expires.After(now) { - continue - } - if len(r.inflight) == 0 { - continue - } - total += len(r.inflight) - r.backlog = append(r.backlog, r.inflight...) - r.inflight = nil - r.owner, r.token, r.expires = "", "", nil - } - return total, nil -} diff --git a/internal/ingestion/knowledge_compile/service.go b/internal/ingestion/knowledge_compile/service.go index e7401047fb..7faedcd0f4 100644 --- a/internal/ingestion/knowledge_compile/service.go +++ b/internal/ingestion/knowledge_compile/service.go @@ -30,15 +30,6 @@ import ( // Option configures a Consumer. type Option func(*Consumer) -// WithBatchSize sets the per-KB batch size trigger (closed-batch boundary). -func WithBatchSize(n int) Option { - return func(c *Consumer) { - if n > 0 { - c.batchSize = n - } - } -} - // WithTTL sets the per-KB claim lease TTL. func WithTTL(d time.Duration) Option { return func(c *Consumer) { @@ -67,8 +58,8 @@ func WithPollInterval(d time.Duration) Option { } } -// WithSweepInterval sets how often the worker runs ReclaimExpired to recover -// inflight left by crashed workers. +// WithSweepInterval sets how often the worker tick calls TryClaim, which +// reclaims inflight left by crashed workers before claiming any ready batch. func WithSweepInterval(d time.Duration) Option { return func(c *Consumer) { if d > 0 { @@ -110,7 +101,7 @@ func defaultDeduperFactory(tenant string) (Deduper, error) { if err != nil { return nil, err } - return NewLLMDeduper(deps.Chat, deps.Embed, defaultLLMID, 0.99), nil + return NewLLMDeduper(deps.Chat, deps.Embed, defaultLLMID, 0.99, deps.LLMMaxLength), nil } func generateHolder() string { diff --git a/internal/ingestion/knowledge_compile/writer.go b/internal/ingestion/knowledge_compile/writer.go index ec4631b146..126e13d71c 100644 --- a/internal/ingestion/knowledge_compile/writer.go +++ b/internal/ingestion/knowledge_compile/writer.go @@ -31,29 +31,62 @@ import ( type Writer interface { // WriteMerged upserts the dataset-level merged products (available_int=1). WriteMerged(ctx context.Context, tenant, kb string, products []kccommon.Product) error - // DeleteMergedForDoc removes merged products that became fully orphaned by - // the deletion of docID. Multi-doc merged products are recomputed by the - // caller (via the Reader + Deduper), not here. - DeleteMergedForDoc(ctx context.Context, tenant, kb, docID string) error + // DeleteDocLevelForDocs drops every per-document (doc-level, kc_merged != 1) + // product of the deleted docs in a single DocEngine call. Dataset-level + // merged rows are not targeted because their doc_id equals the kb, never a + // deleted source doc id. + DeleteDocLevelForDocs(ctx context.Context, tenant, kb string, deletedDocIDs []string) error + // StripMergedSources removes deletedDocIDs from the source_doc_ids array of + // every dataset-level (kc_merged=1) product for the dataset. It searches the + // merged set once, rewrites the source array of every non-empty survivor in a + // single update pass, and deletes (in one call) any product whose array + // became empty. + StripMergedSources(ctx context.Context, tenant, kb string, deletedDocIDs []string) error } -type infinityWriter struct{} +// engineWriter persists dataset-level merged products through the global +// DocEngine (§11.7). Like engineReader, it depends on the process-wide DocEngine +// obtained via engine.Get(); the storage schema lives behind the engine +// abstraction rather than in this package. +type engineWriter struct { + eng engine.DocEngine +} -func (infinityWriter) WriteMerged(ctx context.Context, tenant, kb string, products []kccommon.Product) error { +// writeMergedBatchSize bounds how many rows each parallel InsertChunks call +// carries, so the DocEngine write fan-out stays granular under the shared pool. +const writeMergedBatchSize = 200 + +func (w engineWriter) WriteMerged(ctx context.Context, tenant, kb string, products []kccommon.Product) error { if len(products) == 0 { return nil } - eng := engine.Get() + eng := w.eng + if eng == nil { + eng = engine.Get() + } if eng == nil { return nil } baseName := fmt.Sprintf("ragflow_%s", tenant) - chunks := make([]map[string]interface{}, 0, len(products)) - for _, p := range products { - chunks = append(chunks, mergedChunkMap(tenant, kb, p)) + // Shard the rows and drive the inserts through the shared global pool + // (docengine-bounded) instead of one monolithic InsertChunks call. + jobs := make([]compilerJob, 0, (len(products)+writeMergedBatchSize-1)/writeMergedBatchSize) + for start := 0; start < len(products); start += writeMergedBatchSize { + end := start + writeMergedBatchSize + if end > len(products) { + end = len(products) + } + batch := products[start:end] + jobs = append(jobs, func() error { + chunks := make([]map[string]interface{}, 0, len(batch)) + for _, p := range batch { + chunks = append(chunks, mergedChunkMap(tenant, kb, p)) + } + _, err := eng.InsertChunks(ctx, chunks, baseName, kb) + return err + }) } - _, err := eng.InsertChunks(ctx, chunks, baseName, kb) - return err + return runCompilerJobs(ctx, jobs) } // mergedChunkMap builds the chunk-index document for a dataset-level merged @@ -77,33 +110,79 @@ func mergedChunkMap(tenant, kb string, p kccommon.Product) map[string]interface{ } // Persist the merged product's embedding under the dimension-suffixed column // used elsewhere in the index, so dataset-level rows remain vector-searchable - // and Reader.LoadCompiledProducts can reconstruct them (otherwise the vector - // is silently dropped and search returns nothing for merged rows). + // and the Reader can reconstruct them (otherwise the vector is silently + // dropped and KNN search returns nothing for merged rows). if dim := len(p.Vector); dim > 0 { m[fmt.Sprintf("q_%d_vec", dim)] = p.Vector } return m } -func (infinityWriter) DeleteMergedForDoc(ctx context.Context, tenant, kb, docID string) error { - eng := engine.Get() +// DeleteDocLevelForDocs removes the per-document (doc-level) products of every +// deleted doc in a single DocEngine call. The table is scoped to the dataset +// (kb), and merged rows carry doc_id == kb, so filtering on doc_id IN +// deletedDocIDs can only match the per-document products of the deleted docs. +func (w engineWriter) DeleteDocLevelForDocs(ctx context.Context, tenant, kb string, deletedDocIDs []string) error { + if len(deletedDocIDs) == 0 { + return nil + } + eng := w.eng + if eng == nil { + eng = engine.Get() + } if eng == nil { return nil } baseName := fmt.Sprintf("ragflow_%s", tenant) + _, err := eng.DeleteChunks(ctx, map[string]interface{}{ + "doc_id": deletedDocIDs, + }, baseName, kb) + return err +} + +// StripMergedSources removes deletedDocIDs from the source_doc_ids array of +// every dataset-level (kc_merged=1) product for the dataset. The query filters +// on source_doc_ids IN deletedDocIDs so the engine only returns rows that +// actually reference a deleted doc (intersection pushed down); the survivors' +// source arrays are rewritten in a single update pass driven by the shared +// pool, and any product whose array became empty is deleted in one call. The +// deleted docs' products themselves are never loaded into memory. +func (w engineWriter) StripMergedSources(ctx context.Context, tenant, kb string, deletedDocIDs []string) error { + if len(deletedDocIDs) == 0 { + return nil + } + eng := w.eng + if eng == nil { + eng = engine.Get() + } + if eng == nil { + return nil + } + baseName := fmt.Sprintf("ragflow_%s", tenant) + delSet := make(map[string]bool, len(deletedDocIDs)) + for _, d := range deletedDocIDs { + delSet[d] = true + } + const batchSize = 2000 - // Re-query from the start each iteration (deletions shift the result set), - // deleting every fully-orphaned merged row we find, until a page returns - // fewer than batchSize candidates. Without pagination, orphaned merged rows - // beyond the 2000 cap would never be cleaned up. + var toDeleteIDs []string + var jobs []compilerJob + offset := 0 for { res, err := eng.Search(ctx, &types.SearchRequest{ - IndexNames: []string{baseName}, - KbIDs: []string{kb}, - Filter: map[string]interface{}{"available_int": 1, "kc_merged": 1}, + IndexNames: []string{baseName}, + KbIDs: []string{kb}, + // kc_merged=1 isolates dataset-level rows; source_doc_ids IN + // deletedDocIDs pushes the intersection test into the engine so only + // rows that actually reference a deleted doc are returned (Infinity + // array IN means "contains at least one of"). + Filter: map[string]interface{}{ + "kc_merged": 1, + "source_doc_ids": deletedDocIDs, + }, SelectFields: []string{"id", "source_doc_ids"}, Limit: batchSize, - Offset: 0, + Offset: offset, }) if err != nil { return err @@ -112,21 +191,49 @@ func (infinityWriter) DeleteMergedForDoc(ctx context.Context, tenant, kb, docID break } for _, c := range res.Chunks { - ids := metaStringSlice(c, "source_doc_ids") - // Fully orphaned: the only contributor is the deleted document. - if len(ids) == 1 && ids[0] == docID { - if _, err := eng.DeleteChunks(ctx, map[string]interface{}{ - "id": c["id"], - "kb_id": kb, - "available_int": 1, - }, baseName, kb); err != nil { - return err - } + id, _ := c["id"].(string) + if id == "" { + continue } + src := metaStringSlice(c, "source_doc_ids") + kept := make([]string, 0, len(src)) + changed := false + for _, d := range src { + if delSet[d] { + changed = true + continue + } + kept = append(kept, d) + } + if !changed { + continue + } + if len(kept) == 0 { + toDeleteIDs = append(toDeleteIDs, id) + continue + } + keptCopy := append([]string(nil), kept...) + idCopy := id + jobs = append(jobs, func() error { + return eng.UpdateChunks(ctx, map[string]interface{}{"id": idCopy}, + map[string]interface{}{"source_doc_ids": keptCopy}, baseName, kb) + }) } if len(res.Chunks) < batchSize { break } + offset += batchSize + } + if err := runCompilerJobs(ctx, jobs); err != nil { + return err + } + if len(toDeleteIDs) > 0 { + if _, err := eng.DeleteChunks(ctx, map[string]interface{}{ + "id": toDeleteIDs, + "kb_id": kb, + }, baseName, kb); err != nil { + return err + } } return nil } diff --git a/internal/ingestion/service/ingestion_service.go b/internal/ingestion/service/ingestion_service.go index 877c97186b..a73aec65ef 100644 --- a/internal/ingestion/service/ingestion_service.go +++ b/internal/ingestion/service/ingestion_service.go @@ -247,7 +247,7 @@ func (e *Ingestor) startDatasetKnowledgeCompile() { common.Warn(fmt.Sprintf("dataset-level compile consumer unavailable; compiled chunks will not be merged: %v", err)) return } - e.knowledgeCompile = knowledge_compile.NewConsumer(knowledge_compile.DefaultScheduler()) + e.knowledgeCompile = knowledge_compile.NewConsumer(knowledge_compile.DefaultClaimer()) n := e.kcConcurrency if n <= 0 { n = int32(runtime.NumCPU())