diff --git a/agent/templates/compiler.json b/agent/templates/compiler.json index 3895f903eb..ce0324648c 100644 --- a/agent/templates/compiler.json +++ b/agent/templates/compiler.json @@ -7,6 +7,7 @@ "params": { "compilation_template_group_id": "", "llm_id": "", + "plan": false, "outputs": { "chunks": { "type": "Array", @@ -411,6 +412,7 @@ "form": { "compilation_template_group_id": "", "llm_id": "", + "plan": false, "outputs": { "chunks": { "type": "Array", diff --git a/internal/ingestion/component/knowledge_compiler/common/types.go b/internal/ingestion/component/knowledge_compiler/common/types.go index 04d68a4849..31445d54e6 100644 --- a/internal/ingestion/component/knowledge_compiler/common/types.go +++ b/internal/ingestion/component/knowledge_compiler/common/types.go @@ -49,6 +49,11 @@ type Param struct { SimilarityThreshold float64 MaxWorkers int EnableHistoricalDedup bool + // Plan selects the wiki compilation mode. nil means the DSL did not specify a + // mode, so template config may decide; the effective default is false (Mode A). + // true selects B-mode (LLM planning + reconcile), while false selects A-mode + // (one entity/concept = one deterministic flat page). + Plan *bool // Extra carries arbitrary caller-provided overrides merged into the // resolved template config. Extra map[string]any @@ -168,12 +173,23 @@ func ParseParam(m map[string]any) (Param, error) { if v, ok := m["enable_historical_dedup"].(bool); ok { p.EnableHistoricalDedup = v } + // Keep presence separate from the boolean value: nil means template config may + // supply the mode, while an explicit false still overrides template plan:true. + if v, ok := m["plan"].(bool); ok { + p.Plan = &v + } if raw, ok := m["extra"].(map[string]any); ok { p.Extra = raw } return p, nil } +// PlanEnabled reports whether wiki B-mode is selected. An unset value defaults +// to Mode A. +func (p Param) PlanEnabled() bool { + return p.Plan != nil && *p.Plan +} + // KindToVariant maps a compilation_template.kind to the Go compiler Variant. // // The Python model uses richer kind values (mind_map, page_index, diff --git a/internal/ingestion/component/knowledge_compiler/component.go b/internal/ingestion/component/knowledge_compiler/component.go index fa44d109cb..a52135dea8 100644 --- a/internal/ingestion/component/knowledge_compiler/component.go +++ b/internal/ingestion/component/knowledge_compiler/component.go @@ -440,6 +440,14 @@ func overlayTemplateConfig(param *common.Param, cfg map[string]any) { if v, ok := cfg["enable_historical_dedup"].(bool); ok { param.EnableHistoricalDedup = v } + if param.Plan == nil { + if v, ok := cfg["no_plan"].(bool); ok && v { + disabled := false + param.Plan = &disabled + } else if v, ok := cfg["plan"].(bool); ok { + param.Plan = &v + } + } // llm_id / embedding_model are optional per-call overrides documented on // Invoke. The template config supplies defaults, so only apply them when // the caller has not already provided an explicit value (the caller wins). diff --git a/internal/ingestion/component/knowledge_compiler/component_test.go b/internal/ingestion/component/knowledge_compiler/component_test.go index a9aaa5418a..d9f796e81a 100644 --- a/internal/ingestion/component/knowledge_compiler/component_test.go +++ b/internal/ingestion/component/knowledge_compiler/component_test.go @@ -730,7 +730,7 @@ func TestKnowledgeCompiler_Wiki_HistoricalDedupDropsDuplicates(t *testing.T) { installVariantTemplateResolver(t, "wiki") c, err := NewKnowledgeCompilerComponent("Compiler", map[string]any{ - "compilation_template_id": "tpl-wiki", "llm_id": "llm1", "embedding_model": "emb1", + "compilation_template_id": "tpl-wiki", "llm_id": "llm1", "embedding_model": "emb1", "plan": true, }) if err != nil { t.Fatalf("NewKnowledgeCompilerComponent: %v", err) @@ -789,7 +789,7 @@ func TestKnowledgeCompiler_Wiki_UpdateMergesExistingPage(t *testing.T) { installVariantTemplateResolver(t, "wiki") c, err := NewKnowledgeCompilerComponent("Compiler", map[string]any{ - "compilation_template_id": "tpl-wiki", "llm_id": "llm1", "embedding_model": "emb1", + "compilation_template_id": "tpl-wiki", "llm_id": "llm1", "embedding_model": "emb1", "plan": true, }) if err != nil { t.Fatalf("NewKnowledgeCompilerComponent: %v", err) @@ -1296,6 +1296,53 @@ func TestKnowledgeCompiler_TenantFromGlobals(t *testing.T) { } } +func TestOverlayTemplateConfigPlanPrecedence(t *testing.T) { + boolPtr := func(value bool) *bool { return &value } + + cases := []struct { + name string + initial *bool + cfg map[string]any + want bool + }{ + { + name: "unset_uses_template_plan", + cfg: map[string]any{"plan": true}, + want: true, + }, + { + name: "no_plan_disables_plan", + cfg: map[string]any{"no_plan": true, "plan": true}, + want: false, + }, + { + name: "explicit_dsl_false_overrides_template", + initial: boolPtr(false), + cfg: map[string]any{"plan": true}, + want: false, + }, + { + name: "explicit_dsl_true_overrides_no_plan", + initial: boolPtr(true), + cfg: map[string]any{"no_plan": true}, + want: true, + }, + } + + for _, tc := range cases { + t.Run(tc.name, func(t *testing.T) { + param := common.Param{Plan: tc.initial} + overlayTemplateConfig(¶m, tc.cfg) + if param.Plan == nil { + t.Fatal("Plan = nil after template config") + } + if *param.Plan != tc.want { + t.Errorf("Plan = %t, want %t", *param.Plan, tc.want) + } + }) + } +} + // TestKnowledgeCompiler_BuildInputsAcceptsMapSliceChunks locks the chunk-carrier // contract: the upstream chunker hands chunks over as []map[string]any (see the // chunk map shape {"id","text","ck_type","doc_type_kwd","tk_nums"} observed from diff --git a/internal/ingestion/component/knowledge_compiler/wiki/prompt.go b/internal/ingestion/component/knowledge_compiler/wiki/prompt.go index 033f745824..25ddba76c7 100644 --- a/internal/ingestion/component/knowledge_compiler/wiki/prompt.go +++ b/internal/ingestion/component/knowledge_compiler/wiki/prompt.go @@ -227,6 +227,9 @@ const wikiRefineWriterUserTemplate = `## Task ## Available pages (ONLY use these slugs for [[wikilinks]]) {all_plan_slugs} +## Related KB pages (cross-link only those that are also in the available pages list above) +{related_kb_pages} + {existing_section} ## Source document text diff --git a/internal/ingestion/component/knowledge_compiler/wiki/wiki.go b/internal/ingestion/component/knowledge_compiler/wiki/wiki.go index 1a2097ddd6..d5caebdf07 100644 --- a/internal/ingestion/component/knowledge_compiler/wiki/wiki.go +++ b/internal/ingestion/component/knowledge_compiler/wiki/wiki.go @@ -365,9 +365,23 @@ func (p *wikiPipeline) run() error { zap.String("doc_id", p.runKey()), zap.Int("entities", len(p.reduced.Entities)), zap.Int("claims", len(p.reduced.Claims))) - plan, err := p.runPlan() - if err != nil { - return err + // PLAN (B-mode only): LLM-based page plan + reconcile against existing pages. + // Mode A (param.PlanEnabled() == false, the default) skips the planner entirely — every + // extracted entity/concept becomes its own flat page (1 identity = 1 page), so + // the wiki is a flat encyclopedia without PLAN-grouped pages. See buildModeAPlan. + var plan wikiPlan + var err error + if p.param.PlanEnabled() { + plan, err = p.runPlan() + if err != nil { + return err + } + appcommon.Info("wiki: PLAN (B-mode) done", zap.String("dataset_id", p.datasetID), zap.String("doc_id", p.runKey()), zap.Int("plan_pages", len(plan.Pages))) + } else { + // wiki_incremental port (T1 + M1 Mode A): deterministic flat plan from the + // reduced graph. No LLM, no reconcile. + plan = p.buildModeAPlan() + appcommon.Info("wiki: PLAN (A-mode flat) done", zap.String("dataset_id", p.datasetID), zap.String("doc_id", p.runKey()), zap.Int("plan_pages", len(plan.Pages))) } p.plan = plan appcommon.Info("wiki: PLAN done", @@ -762,6 +776,7 @@ func (p *wikiPipeline) runRefinePage( "title": firstNonEmpty(planItem.Title, planItem.Slug), "page_type": firstNonEmpty(planItem.PageType, "concept"), "all_plan_slugs": strings.Join(available, "\n"), + "related_kb_pages": formatWikiRelatedKB(planItem.RelatedKB, pageTitles), "existing_section": existingSection, "source_context": sourceContext, "evidence_count": fmt.Sprintf("%d", len(evidence)), @@ -794,6 +809,14 @@ func (p *wikiPipeline) runRefinePage( } } contentRendered, outlinks := transformWikiLinks(contentRaw, firstNonEmpty(p.datasetID, p.docID), pageTitles, slugToPageType) + // wiki_incremental port (O2): deterministically guarantee a "See also" + // cross-link section from RelatedKB. transformWikiLinks/resolveSlug drop + // links whose target cannot be resolved to a known slug, so an LLM-authored + // page could silently lose every RelatedKB edge. We append the related pages + // AFTER transformWikiLinks and do NOT re-run the resolver, so dead links are + // bypassed and the missing target does not drop the edge — each RelatedKB + // full-slug is linked once, only if not already present in Outlinks. + contentRendered, outlinks = appendWikiSeeAlso(contentRendered, outlinks, planItem.RelatedKB, p.datasetID, pageTitles, slugToPageType) sourceDocIDs := collectWikiSourceDocIDs(p.inputs.Chunks, sourceChunkIDs, p.docID) summary := firstParagraph(contentRendered) if summary == "" { @@ -1545,6 +1568,81 @@ func wikiTitleKey(pageType, title string) string { return strings.TrimSpace(pageType) + "\x00" + normKey(title) } +// buildModeAPlan synthesizes a FLAT wiki plan (Mode A) directly from the reduced +// extract graph — no LLM planner, no reconcile. Every extracted entity and +// concept becomes its own canonical page (1 identity = 1 page), slugged as +// "entity/" / "concept/". Cross-links (RelatedKB) are derived purely +// from the reduced relations: each relation end is mapped to its full-slug page +// and the counterpart is added to the page's RelatedKB. This is the deterministic +// counterpart of Python's no_plan wiki mode. +func (p *wikiPipeline) buildModeAPlan() wikiPlan { + reduced := p.reduced + // Build the full-slug index for every entity/concept so relation endpoints + // (which are bare names) resolve to canonical pages. + fullSlugFor := func(name, pageType string) string { + return pageType + "/" + normalizeWikiSlugHyphens(slugify(name)) + } + slugToIndex := map[string]int{} + var pages []wikiPlanPage + addPage := func(slug, title, pageType string, entityNames []string) { + if _, ok := slugToIndex[slug]; ok { + return + } + page := wikiPlanPage{ + Action: "CREATE", + Slug: slug, + Title: title, + PageType: pageType, + Topic: title, + EntityNames: entityNames, + Priority: len(pages) + 1, + } + slugToIndex[slug] = len(pages) + pages = append(pages, page) + } + for _, e := range reduced.Entities { + name := strings.TrimSpace(e.Name) + if name == "" { + continue + } + slug := fullSlugFor(name, "entity") + addPage(slug, name, "entity", []string{name}) + } + for _, c := range reduced.Concepts { + term := strings.TrimSpace(c.Term) + if term == "" { + continue + } + slug := fullSlugFor(term, "concept") + addPage(slug, term, "concept", []string{term}) + } + // Resolve relations into RelatedKB (full-slug cross-links) on both endpoints. + for _, rel := range reduced.Relations { + from := strings.TrimSpace(rel.From) + to := strings.TrimSpace(rel.To) + if from == "" || to == "" { + continue + } + fromSlug := fullSlugFor(from, "entity") + toSlug := fullSlugFor(to, "entity") + if fromSlug == toSlug { + continue + } + if i, ok := slugToIndex[fromSlug]; ok { + pages[i].RelatedKB = append(pages[i].RelatedKB, toSlug) + } + if i, ok := slugToIndex[toSlug]; ok { + pages[i].RelatedKB = append(pages[i].RelatedKB, fromSlug) + } + } + for i := range pages { + pages[i].RelatedKB = uniqueStrings(pages[i].RelatedKB) + } + plan := wikiPlan{Pages: pages} + plan.Pages = normalizeWikiPlanPages(plan.Pages, reduced) + return plan +} + func buildWikiFallbackPages(reduced wikiExtract) []wikiPlanPage { var out []wikiPlanPage seen := map[string]bool{} @@ -2159,47 +2257,64 @@ func pageTypeOf(slug string, slugToPageType map[string]string) string { return "page" } -func transformWikiLinks(content, kbID string, pageTitles, slugToPageType map[string]string) (string, []string) { - kbID = strings.TrimSpace(kbID) - // Resolve a wikitext slug (which may be a bare "name" or a full - // "/") to the canonical full slug used as the page - // identifier (slug_kwd). This mirrors Python's _wiki_resolve_dead_slug: - // plain names / titles are reverse-mapped to the full pid, and slugs that - // cannot be resolved are dropped (dead links). Without this, bare slugs - // written by the LLM never match the full-slug page index, so wiki - // relations (edges) are silently lost. - bareToSlug := map[string]string{} - titleToSlug := map[string]string{} - // Plan slugs are canonicalized to the hyphen style upstream - // (normalizeWikiPlanPage -> normalizeWikiSlugHyphens), so slugToPageType / - // pageTitles keys are already hyphen full-slugs. LLM-authored wikitext - // links may still carry underscores (e.g. "dong_zhuo"); normalize both - // sides to hyphens so the reverse lookup matches, and always emit the - // resolved outlink in the canonical hyphen full-slug form. +// wikiSlugResolver resolves a wikitext slug (a bare "name" or a full +// "/") to the canonical hyphen full-slug used as the page +// identifier (slug_kwd), mirroring Python's _wiki_resolve_dead_slug: plain +// names / titles are reverse-mapped to the full pid. An empty result means the +// target is not a known page and must not be emitted as a link or graph edge. +type wikiSlugResolver struct { + slugToPageType map[string]string + bareToSlug map[string]string + titleToSlug map[string]string +} + +// newWikiSlugResolver builds the reverse lookup maps from the available-page +// index. Plan slugs are canonicalized to the hyphen style upstream +// (normalizeWikiPlanPage -> normalizeWikiSlugHyphens), so slugToPageType / +// pageTitles keys are already hyphen full-slugs. LLM-authored wikitext links may +// still carry underscores (e.g. "dong_zhuo"); normalize both sides to hyphens so +// the reverse lookup matches. +func newWikiSlugResolver(pageTitles, slugToPageType map[string]string) *wikiSlugResolver { + r := &wikiSlugResolver{ + slugToPageType: slugToPageType, + bareToSlug: map[string]string{}, + titleToSlug: map[string]string{}, + } for fullSlug := range slugToPageType { - bareToSlug[normalizeWikiSlugHyphens(lastPathSlug(fullSlug))] = fullSlug + r.bareToSlug[normalizeWikiSlugHyphens(lastPathSlug(fullSlug))] = fullSlug } for fullSlug, title := range pageTitles { if t := strings.TrimSpace(title); t != "" { - titleToSlug[t] = fullSlug + r.titleToSlug[t] = fullSlug } } - resolveSlug := func(slug string) string { - slug = strings.TrimSpace(slug) - if slug == "" { - return "" - } - if _, ok := slugToPageType[slug]; ok { - return slug - } - if full, ok := bareToSlug[normalizeWikiSlugHyphens(lastPathSlug(slug))]; ok { - return full - } - if full, ok := titleToSlug[slug]; ok { - return full - } + return r +} + +// resolve returns the canonical hyphen full-slug for target, or "" when the +// target is not a known page (dead link). The result is always emitted in the +// canonical hyphen form so outlinks_kwd and slug_kwd agree in format. +func (r *wikiSlugResolver) resolve(target string) string { + slug := strings.TrimSpace(target) + if slug == "" { return "" } + if _, ok := r.slugToPageType[slug]; ok { + return normalizeWikiSlugHyphens(slug) + } + if full, ok := r.bareToSlug[normalizeWikiSlugHyphens(lastPathSlug(slug))]; ok { + return normalizeWikiSlugHyphens(full) + } + if full, ok := r.titleToSlug[slug]; ok { + return normalizeWikiSlugHyphens(full) + } + return "" +} + +func transformWikiLinks(content, kbID string, pageTitles, slugToPageType map[string]string) (string, []string) { + kbID = strings.TrimSpace(kbID) + resolver := newWikiSlugResolver(pageTitles, slugToPageType) + resolveSlug := resolver.resolve seen := map[string]bool{} var outlinks []string track := func(slug string) { @@ -2317,6 +2432,91 @@ func transformWikiLinks(content, kbID string, pageTitles, slugToPageType map[str return out, outlinks } +// formatWikiRelatedKB renders the RelatedKB page list for the refine-writer +// prompt (O1): one "- [[]]" bullet per related full-slug, with the page +// title as display text when known. It is a prompt hint only — the authoritative +// See-also section is appended deterministically by appendWikiSeeAlso after the +// LLM output is transformed. +func formatWikiRelatedKB(related []string, pageTitles map[string]string) string { + if len(related) == 0 { + return "(none)" + } + var b strings.Builder + for _, slug := range related { + slug = strings.TrimSpace(slug) + if slug == "" { + continue + } + canon := normalizeWikiSlugHyphens(slug) + if title := strings.TrimSpace(pageTitles[canon]); title != "" { + b.WriteString("- [[" + canon + "|" + title + "]]\n") + } else { + b.WriteString("- [[" + canon + "]]\n") + } + } + if b.Len() == 0 { + return "(none)" + } + return strings.TrimRight(b.String(), "\n") +} + +// appendWikiSeeAlso guarantees a cross-link section from RelatedKB (O2). It runs +// AFTER transformWikiLinks, which drops dead links (targets that cannot be +// resolved). To avoid losing every RelatedKB edge when a target is missing, we do +// NOT re-run the resolver here: each RelatedKB full-slug is linked once as +// artifact///, only if it is not already present in +// outlinks. This bypasses the dead-link drop and keeps the graph edge. +func appendWikiSeeAlso(content string, outlinks, related []string, kbID string, pageTitles, slugToPageType map[string]string) (string, []string) { + if len(related) == 0 { + return content, outlinks + } + resolver := newWikiSlugResolver(pageTitles, slugToPageType) + have := make(map[string]bool, len(outlinks)) + for _, o := range outlinks { + have[normalizeWikiSlugHyphens(o)] = true + } + var bullets []string + for _, slug := range related { + // Resolve the RelatedKB target against the current page index. An + // unresolved target is not a known page: omit it entirely (no broken + // artifact link) and do NOT append it to outlinks, so no phantom graph + // edge is persisted for a page that was never compiled. + canon := resolver.resolve(slug) + if canon == "" || have[canon] { + continue + } + have[canon] = true + bareSlug := canon + pageType := pageTypeOf(canon, slugToPageType) + if idx := strings.Index(canon, "/"); idx >= 0 && canon[:idx] == pageType { + bareSlug = canon[idx+1:] + } + if pageType == "" { + pageType = "page" + } + label := pageTitles[canon] + if label == "" { + label = strings.Title(strings.ReplaceAll(lastPathSlug(canon), "-", " ")) + } + bullets = append(bullets, "- ["+label+"](artifact/"+kbID+"/"+pageType+"/"+bareSlug+")") + outlinks = append(outlinks, canon) + } + if len(bullets) == 0 { + return content, outlinks + } + // Append a "See also" section if one is not already present; otherwise add + // the bullets under the existing "See also" heading. + body := strings.TrimRight(content, "\n") + if seeAlsoRe.MatchString(body) { + return body + "\n" + strings.Join(bullets, "\n") + "\n", outlinks + } + return body + "\n\n## See also\n\n" + strings.Join(bullets, "\n") + "\n", outlinks +} + +// seeAlsoRe matches a Markdown "See also" heading (case-insensitive) so we can +// append related links under an existing section instead of duplicating it. +var seeAlsoRe = regexp.MustCompile(`(?m)^#{1,6}\s*see\s*also\s*$`) + func firstN(s []string, n int) []string { if len(s) <= n { return s diff --git a/internal/ingestion/knowledge_compile/consumer.go b/internal/ingestion/knowledge_compile/consumer.go index 8f7a61b38f..753b017201 100644 --- a/internal/ingestion/knowledge_compile/consumer.go +++ b/internal/ingestion/knowledge_compile/consumer.go @@ -17,18 +17,26 @@ package knowledge_compile import ( "context" + "errors" "fmt" "sort" "sync" "time" "ragflow/internal/common" + "ragflow/internal/dao" "ragflow/internal/engine" kccommon "ragflow/internal/ingestion/component/knowledge_compiler/common" "go.uber.org/zap" + "gorm.io/gorm" ) +// kcDB is the package-level MySQL handle installed by Provision. It backs the +// doc enumeration used by RebuildDataset; it is nil (and docLister returns no +// docs) when the scheduler was provisioned without a DB. +var kcDB *gorm.DB + // Consumer is the dataset-level post-processing worker (§11.5). Multiple // instances compete on the MySQL scheduling rows; each KB is processed by at // most one instance at a time via the per-KB claim (so the same KB is handled @@ -49,6 +57,20 @@ type Consumer struct { mu sync.Mutex tombs map[string]map[string]uint64 // dataset -> docID -> delete marker (tombstone) + + // rebuildPause suppresses local claims while this consumer rebuilds a dataset. + rebuildPause map[string]bool + + // docLister enumerates every doc id in a KB so a rewrite can republish them. + docLister func(ctx context.Context, tenant, kb string) ([]string, error) +} + +// rewriteScheduler is the subset of the scheduler API a dataset rewrite needs, +// satisfied by *mysqlScheduler (and FakeScheduler in tests). Using a narrow +// interface keeps the Claimer surface unchanged. +type rewriteScheduler interface { + Publish(ctx context.Context, tenantID, datasetID, docID, eventType string) error + CancelInflight(ctx context.Context, datasetID, token string) error } // NewConsumer constructs a Consumer driven by the given Claimer. Tests pass a @@ -65,6 +87,8 @@ func NewConsumer(scheduler Claimer, opts ...Option) *Consumer { sweepInterval: 30 * time.Second, mergeThreshold: 0.85, tombs: map[string]map[string]uint64{}, + rebuildPause: map[string]bool{}, + docLister: defaultDocLister, } for _, o := range opts { o(c) @@ -135,6 +159,13 @@ func (c *Consumer) tryClaimAndProcess(ctx context.Context) { func (c *Consumer) processClaim(ctx context.Context, cr ClaimResult) { datasetID := cr.DatasetID + c.mu.Lock() + paused := c.rebuildPause[datasetID] + c.mu.Unlock() + if paused { + return + } + // 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. stopHb := make(chan struct{}) @@ -165,7 +196,7 @@ func (c *Consumer) processClaim(ctx context.Context, cr ClaimResult) { var batchErr error go func() { defer close(done) - batchErr = c.processBatch(ctx, cr.TenantID, datasetID, cr.Entries) + batchErr = c.processBatch(ctx, cr.TenantID, datasetID, cr.Token, cr.Entries) }() select { @@ -185,6 +216,9 @@ func (c *Consumer) processClaim(ctx context.Context, cr ClaimResult) { // must leave the claimed batch in the backlog for reclamation/retry rather // than silently dropping it (C5: never ack what we failed to merge). if batchErr != nil { + if errors.Is(batchErr, errClaimSuperseded) { + return + } common.Error("knowledge_compile: batch processing failed, leaving batch for retry", batchErr, zap.String("dataset_id", datasetID), @@ -201,11 +235,142 @@ func (c *Consumer) processClaim(ctx context.Context, cr ClaimResult) { } } +var errClaimSuperseded = errors.New("knowledge_compile: claim superseded by rewrite") + +// withWriteLock runs a destructive side effect fn under the scheduler's +// per-dataset write/rebuild lock, verifying the claim token inside the lock +// immediately before fn. This closes the TOCTOU gap between cancellation and a +// writer side effect: a cancelled worker cannot write after the rebuild clears +// storage, and a rebuild cannot interleave while fn runs. +func (c *Consumer) withWriteLock(ctx context.Context, kb, token string, fn func() error) error { + return c.scheduler.WithDatasetLock(ctx, kb, func(currentToken string) error { + if currentToken != token { + return errClaimSuperseded + } + return fn() + }) +} + +// RebuildDataset performs a full incremental rewrite of a KB's dataset-level +// merged products (W1/W2/W5/W6). The order is fixed to close both the +// enumerate-then-clear race (M19/C-race) AND the cross-process stale-write +// window: +// 1. set the rewrite pause (in-process only; clears via defer on every path); +// 2. CancelInflight invalidates any in-flight claim before clearing storage. +// Each worker compares its claim token under the write lock and therefore +// cannot write after cancellation; +// 3. DeleteMerged + DropWikiGraph clear the old merged + graph state INSIDE the +// same per-dataset row lock (WithDatasetLock). This is what actually closes +// the TOCTOU gap: the generation check + write of every worker run under this +// lock, so the rebuild's clear either runs after any in-flight worker write +// has finished (it is later removed by the clear) or after a stale worker has +// self-dropped — a worker can never repopulate cleared state with old results +// because it cannot hold the lock concurrently with the clear; +// 4. enumerate every doc and republish; +// 5. clear the local pause. +// +// mode is "incremental" or "rewrite"; today both take the same clean-and-rebuild +// path, with mode retained for future differential strategies. +func (c *Consumer) RebuildDataset(ctx context.Context, tenant, kb, mode string) error { + rs, ok := c.scheduler.(rewriteScheduler) + if !ok { + return fmt.Errorf("knowledge_compile: scheduler %T does not support rewrite", c.scheduler) + } + + // 1. pause the rewrite window. The pause is cleared by a deferred cleanup so a + // failure mid-rebuild (DeleteMerged/DropWikiGraph/CancelInflight/Bump/Publish) + // can never leave the dataset permanently paused (which would drop every + // future claim). Partially published docs are retried on the next claim. + c.mu.Lock() + c.rebuildPause[kb] = true + c.mu.Unlock() + defer func() { + c.mu.Lock() + c.rebuildPause[kb] = false + c.mu.Unlock() + }() + + // 2. Cancel any in-flight claim before clearing storage. A worker that was + // already running observes its revoked claim token inside withWriteLock and + // returns without writing. + if err := rs.CancelInflight(ctx, kb, ""); err != nil { + return fmt.Errorf("knowledge_compile: rebuild cancel inflight: %w", err) + } + + // 3. clear old merged + graph state (structural filter is the source of + // truth; no per-doc rows are touched) under the per-dataset row lock. This + // guarantees mutual exclusion with every worker's destructive write (which + // runs under the same lock via withWriteLock): the clear either waits for an + // in-flight worker write to finish (then removes its old-generation result) + // or runs after stale workers have self-dropped. A worker cannot hold the + // lock while the clear runs, so it can never repopulate cleared state with + // old results. + if err := c.scheduler.WithDatasetLock(ctx, kb, func(_ string) error { + if derr := c.writer.DeleteMerged(ctx, tenant, kb); derr != nil { + return derr + } + return c.writer.DropWikiGraph(ctx, tenant, kb) + }); err != nil { + return fmt.Errorf("knowledge_compile: rebuild clear merged+graph: %w", err) + } + + // 4. enumerate every doc and republish. + docs, err := c.docLister(ctx, tenant, kb) + if err != nil { + return fmt.Errorf("knowledge_compile: rebuild list docs: %w", err) + } + for _, docID := range docs { + if perr := rs.Publish(ctx, tenant, kb, docID, string(EventTypeCompleted)); perr != nil { + return fmt.Errorf("knowledge_compile: rebuild republish %s: %w", docID, perr) + } + } + + // 6. resweep: the deferred cleanup clears the pause so any claim arriving + // post-enumeration is processed (and stale ones self-drop in processClaim). + common.Info("knowledge_compile: dataset rebuild complete", + zap.String("dataset_id", kb), + zap.String("mode", mode), + zap.Int("docs", len(docs))) + return nil +} + +// defaultDocLister enumerates every doc id in a KB via the Document DAO. When +// no DB was provisioned it returns no docs (a rewrite becomes a no-op clean). +func defaultDocLister(ctx context.Context, tenant, kb string) ([]string, error) { + if kcDB == nil { + return nil, nil + } + return dao.NewDocumentDAO().ListIDsByKBIDWithOptions(ctx, kcDB, dao.DocumentListOptions{KbID: kb}) +} + +// filterWikiPageCandidates returns only the candidates the dataset-level merge +// is allowed to fold: for the wiki variant, that is strictly page-kind products +// (Meta.kind == "page"). Sections are a doc-level concern and must never enter +// the dataset-level page bucket, and the deduper must not carry an implicit +// section-filter contract. Non-wiki variants pass through unchanged. Legacy wiki +// rows whose kind is empty are derived in the Reader; only truly page-kind wiki +// products proceed. +func filterWikiPageCandidates(candidates []kccommon.Product) []kccommon.Product { + // Allocate a fresh slice: reusing candidates[:0] would overwrite the caller's + // backing array in place, corrupting any other reference (e.g. the vector + // audit that still reads `incoming`). + filtered := make([]kccommon.Product, 0, len(candidates)) + for _, cand := range candidates { + if cand.Variant == kccommon.VariantWiki { + if metaString(cand.Meta, "kind") != "page" { + continue + } + } + filtered = append(filtered, cand) + } + return filtered +} + // processBatch applies out-of-order / tombstone handling, then recomputes and // writes the dataset-level merged products for the claimed closed batch. It // returns an error if any reader/dedup/writer step fails so the caller can // leave the batch for reclamation instead of acking dropped work. -func (c *Consumer) processBatch(ctx context.Context, tenant, kb string, entries []BacklogEntry) error { +func (c *Consumer) processBatch(ctx context.Context, tenant, kb, token string, entries []BacklogEntry) error { common.Info("knowledge_compile: processing claimed batch", zap.String("dataset_id", kb), zap.String("tenant_id", tenant), @@ -326,10 +491,31 @@ func (c *Consumer) processBatch(ctx context.Context, tenant, kb string, entries for d := range deletedSet { delIDs = append(delIDs, d) } - if err := c.writer.DeleteDocLevelForDocs(ctx, tenant, kb, delIDs); err != nil { + // Each destructive write runs under the scheduler's per-dataset + // write/rebuild lock with the generation check performed INSIDE the lock, + // so a rewrite can neither land between the check and the write nor + // interleave with the write itself (it must wait for this lock). This + // closes the TOCTOU window where a worker past its fence could repopulate + // storage the rebuild just cleared. + if err := c.withWriteLock(ctx, kb, token, func() error { + return c.writer.DeleteDocLevelForDocs(ctx, tenant, kb, delIDs) + }); err != nil { + if errors.Is(err, errClaimSuperseded) { + common.Info("knowledge_compile: batch stale before delete, aborting (rewrite barrier)", + zap.String("dataset_id", kb)) + } return err } - if err := c.writer.StripMergedSources(ctx, tenant, kb, delIDs); err != nil { + // The second destructive call is its own locked section: a rewrite can + // land between the two, in which case the stale batch must not apply its + // second side effect under the new generation. + if err := c.withWriteLock(ctx, kb, token, func() error { + return c.writer.StripMergedSources(ctx, tenant, kb, delIDs) + }); err != nil { + if errors.Is(err, errClaimSuperseded) { + common.Info("knowledge_compile: batch stale before strip, aborting (rewrite barrier)", + zap.String("dataset_id", kb)) + } return err } } @@ -351,10 +537,19 @@ func (c *Consumer) processBatch(ctx context.Context, tenant, kb string, entries incoming = append(incoming, docProducts...) } + // wiki_incremental port (M1): the dataset-level merge only processes wiki + // PAGES. A wiki doc yields both page and section products (Meta.kind + // "page"/"section"); sections are a doc-level concern and must never be folded + // into the dataset-level page bucket. Filter BEFORE the in-memory dedup so + // sections never enter the deduper (and never trigger LLM/embedding/alias + // processing) — the decider must not carry an implicit section-filter + // contract. Legacy rows whose kind is empty are derived in the Reader; only + // truly page-kind wiki products proceed. + candidates := filterWikiPageCandidates(incoming) // In-memory dedup among the completed batch first. - candidates, err := deduper.Dedup(ctx, incoming) - if err != nil { - return err + candidates, dedupErr := deduper.Dedup(ctx, candidates) + if dedupErr != nil { + return dedupErr } // Diagnostics: after batch dedup, verify the per-doc products still carry // their embedding before the KNN/merge path consumes cand.Vector. If @@ -433,7 +628,12 @@ func (c *Consumer) processBatch(ctx context.Context, tenant, kb string, entries vec64[i] = float64(v) } } - hit, score, err := c.reader.SearchSimilar(ctx, tenant, kb, cand.Variant, vec64, 1, c.mergeThreshold) + // wiki_incremental port (B2): request topN >= 2 so SearchSimilar can + // apply its score-descending skip rule on a dirty top-1 (the reader + // drops rows whose compile_kwd does not map to the searched variant; + // a re-query with topN=2 lets the next clean candidate surface instead + // of falling through to "no hit"). + hit, score, err := c.reader.SearchSimilar(ctx, tenant, kb, cand.Variant, vec64, 2, c.mergeThreshold) if err != nil { return err } @@ -518,11 +718,21 @@ func (c *Consumer) processBatch(ctx context.Context, tenant, kb string, entries } } - // Write the surviving merged set (updated existing + new distinct rows). + // Write the surviving merged set (updated existing + new distinct rows). The + // generation check + WriteMerged run atomically under the per-dataset + // write/rebuild lock, so a rewrite that lands during the (now long) KNN + LLM + // merge cannot leak into the freshly rewritten index, and a rewrite that is + // clearing storage cannot interleave with the write. 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 { + if err := c.withWriteLock(ctx, kb, token, func() error { + return c.writer.WriteMerged(ctx, tenant, kb, mergedFinal) + }); err != nil { + if errors.Is(err, errClaimSuperseded) { + common.Info("knowledge_compile: batch stale before write, aborting (rewrite barrier)", + zap.String("dataset_id", kb)) + } return err } @@ -532,8 +742,15 @@ func (c *Consumer) processBatch(ctx context.Context, tenant, kb string, entries // when the current batch only prunes (deletes) wiki pages, and even for a // non-wiki batch where the dataset already has no wiki pages (ProjectWikiGraph // then just drops any stale graph rows). The graph is reconstructible, so the - // cost of an occasional no-op reprojection is accepted. - if err := c.writer.ProjectWikiGraph(ctx, tenant, kb); err != nil { + // cost of an occasional no-op reprojection is accepted. It is also a locked + // destructive side effect for the same TOCTOU reasons as WriteMerged. + if err := c.withWriteLock(ctx, kb, token, func() error { + return c.writer.ProjectWikiGraph(ctx, tenant, kb) + }); err != nil { + if errors.Is(err, errClaimSuperseded) { + common.Info("knowledge_compile: batch stale before graph projection, aborting (rewrite barrier)", + zap.String("dataset_id", kb)) + } return err } diff --git a/internal/ingestion/knowledge_compile/consumer_test.go b/internal/ingestion/knowledge_compile/consumer_test.go deleted file mode 100644 index 5104ad0dae..0000000000 --- a/internal/ingestion/knowledge_compile/consumer_test.go +++ /dev/null @@ -1,424 +0,0 @@ -// -// 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" - "sync" - "testing" - "time" - - kccommon "ragflow/internal/ingestion/component/knowledge_compiler/common" -) - -// 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) LoadDocProducts(_ context.Context, _, _, docID string) ([]kccommon.Product, error) { - r.mu.Lock() - defer r.mu.Unlock() - r.calls++ - 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 - deletedDocLevel []string - strippedSources []string - projected []string - dropped []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)) - copy(cp, products) - w.written = append(w.written, cp) - return nil -} - -func (w *fakeWriter) DeleteDocLevelForDocs(_ context.Context, _, _ string, docIDs []string) error { - w.mu.Lock() - defer w.mu.Unlock() - 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 -} - -func (w *fakeWriter) ProjectWikiGraph(_ context.Context, _, kb string) error { - w.mu.Lock() - defer w.mu.Unlock() - w.projected = append(w.projected, kb) - return nil -} - -func (w *fakeWriter) DropWikiGraph(_ context.Context, _, kb string) error { - w.mu.Lock() - defer w.mu.Unlock() - w.dropped = append(w.dropped, kb) - return nil -} - -func sampleProducts() []kccommon.Product { - return []kccommon.Product{ - {ID: "p1", DocID: "d1", TenantID: "t1", Variant: kccommon.Variant("structure"), - Content: `{"name":"X"}`, Meta: map[string]any{"name": "X", "kind": "entity"}}, - {ID: "p2", DocID: "d1", TenantID: "t1", Variant: kccommon.Variant("structure"), - Content: `{"name":"Y"}`, Meta: map[string]any{"name": "Y", "kind": "entity"}}, - } -} - -func newTestConsumer(sch *FakeScheduler, r *fakeReader, w *fakeWriter, factory DeduperFactory) *Consumer { - return NewConsumer(sch, - WithReader(r), - WithWriter(w), - WithDeduperFactory(factory), - ) -} - -func TestConsumerCompletedWritesMerged(t *testing.T) { - sch := NewFakeScheduler() - r := &fakeReader{products: sampleProducts()} - w := &fakeWriter{} - c := newTestConsumer(sch, r, w, func(string) (Deduper, error) { return NewNoopDeduper(), nil }) - - if err := sch.Publish(context.Background(), "t1", "kb1", "d1", string(EventTypeCompleted)); err != nil { - t.Fatalf("append: %v", err) - } - c.tryClaimAndProcess(context.Background()) - - w.mu.Lock() - defer w.mu.Unlock() - 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])) - } - if len(w.projected) != 1 || w.projected[0] != "kb1" { - t.Fatalf("expected 1 ProjectWikiGraph(kb1) call, got %v", w.projected) - } - // After ack, the claim row must be cleared (no live lease left behind). - if _, ok, _ := sch.TryClaim(context.Background()); ok { - t.Fatalf("expected no claimable row after ack") - } -} - -func TestConsumerTombstoneSkipsCompletedBeforeDeleted(t *testing.T) { - sch := NewFakeScheduler() - r := &fakeReader{products: sampleProducts()} - w := &fakeWriter{} - c := newTestConsumer(sch, r, w, func(string) (Deduper, error) { return NewNoopDeduper(), nil }) - - // completed before deleted: the deletion is the last event for d1, so it - // wins and d1's per-doc products are orphaned (completion skipped). - if err := sch.Publish(context.Background(), "t1", "kb1", "d1", string(EventTypeCompleted)); err != nil { - t.Fatalf("append completed: %v", err) - } - if err := sch.Publish(context.Background(), "t1", "kb1", "d1", string(EventTypeDeleted)); err != nil { - t.Fatalf("append deleted: %v", err) - } - c.tryClaimAndProcess(context.Background()) - - w.mu.Lock() - defer w.mu.Unlock() - // No merged write (completed skipped, since deletion is the last event). - // 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 then completed: the completion is the last event for d1, 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)); err != nil { - t.Fatalf("append deleted: %v", err) - } - if err := sch.Publish(context.Background(), "t1", "kb1", "d1", string(EventTypeCompleted)); err != nil { - t.Fatalf("append completed: %v", err) - } - c.tryClaimAndProcess(context.Background()) - - w.mu.Lock() - defer w.mu.Unlock() - // No deletion calls: the last (completion) event 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])) - } - if len(w.projected) != 1 || w.projected[0] != "kb1" { - t.Fatalf("expected 1 ProjectWikiGraph(kb1) call, got %v", w.projected) - } -} - -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.Publish(context.Background(), "t1", "kb1", docID, string(EventTypeCompleted)); err != nil { - t.Fatalf("append: %v", err) - } - } - // 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) - } - if len(cr1.Entries) != 32 { - t.Fatalf("expected 32 entries in first claim, got %d", len(cr1.Entries)) - } - // 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") - if ok2 { - t.Fatalf("second claim should have lost the race (live lease)") - } - // Ack the first batch, then the remaining 8 become claimable. - 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") - if err != nil || !ok3 { - t.Fatalf("claim3: ok=%v err=%v", ok3, err) - } - if len(cr3.Entries) != 8 { - t.Fatalf("expected 8 remaining entries, got %d", len(cr3.Entries)) - } -} - -func TestSchedulerReclaimExpired(t *testing.T) { - sch := NewFakeScheduler() - if err := sch.Publish(context.Background(), "t1", "kb1", "d1", string(EventTypeCompleted)); err != nil { - t.Fatalf("append: %v", err) - } - _, 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 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) - } - if len(cr2.Entries) != 1 { - t.Fatalf("expected 1 entry after reclaim, got %d", len(cr2.Entries)) - } -} - -// rowCounts snapshots a fake scheduling row's inflight/backlog entry counts. -type rowCounts struct { - inflight int - backlog int -} - -func fakeRowCounts(sch *FakeScheduler, datasetID string) (rowCounts, string) { - sch.mu.Lock() - defer sch.mu.Unlock() - r := sch.rows[datasetID] - if r == nil { - return rowCounts{}, "" - } - return rowCounts{inflight: len(r.inflight), backlog: len(r.backlog)}, r.state -} - -// TestSchedulerStateMachineLocksStateAndCounts locks the full lifecycle state -// machine (plan v4.1 §9.2) on the FakeScheduler, which mirrors the MySQL -// scheduler's transitions: -// -// Publish -> pending; Claim -> running (+error cleared); -// SetError (failed batch left for retry) keeps running + records error; -// lease expiry -> reclaimOne -> pending (inflight moved back to backlog); -// Ack with backlog drained -> completed. -func TestSchedulerStateMachineLocksStateAndCounts(t *testing.T) { - sch := NewFakeScheduler() - - // Publish two docs: state=pending, backlog=2, inflight=0. - if err := sch.Publish(context.Background(), "t1", "kb1", "d1", string(EventTypeCompleted)); err != nil { - t.Fatalf("publish d1: %v", err) - } - if err := sch.Publish(context.Background(), "t1", "kb1", "d2", string(EventTypeCompleted)); err != nil { - t.Fatalf("publish d2: %v", err) - } - if c, s := fakeRowCounts(sch, "kb1"); s != DatasetStatePending || c.backlog != 2 || c.inflight != 0 { - t.Fatalf("after publish: want state=pending backlog=2 inflight=0, got state=%s %+v", s, c) - } - - // Claim: state=running, backlog moves to inflight (batch=2), error cleared. - cr, ok, err := sch.Claim(context.Background(), "kb1") - if err != nil || !ok { - t.Fatalf("claim: ok=%v err=%v", ok, err) - } - if len(cr.Entries) != 2 { - t.Fatalf("expected 2-entry claim, got %d", len(cr.Entries)) - } - if c, s := fakeRowCounts(sch, "kb1"); s != DatasetStateRunning || c.backlog != 0 || c.inflight != 2 { - t.Fatalf("after claim: want state=running backlog=0 inflight=2, got state=%s %+v", s, c) - } - - // Failed batch left for retry: SetError records a diagnostic, state stays running. - if err := sch.SetError(context.Background(), "kb1", cr.Token, "boom"); err != nil { - t.Fatalf("set error: %v", err) - } - sch.mu.Lock() - gotErr := sch.rows["kb1"].errorMsg - sch.mu.Unlock() - if gotErr != "boom" { - t.Fatalf("expected errorMsg=boom, got %q", gotErr) - } - if _, s := fakeRowCounts(sch, "kb1"); s != DatasetStateRunning { - t.Fatalf("failed batch must stay running (left for retry), got state=%s", s) - } - - // Lease expires: reclaimOne moves inflight back to backlog, clears lease -> pending. - past := time.Now().Add(-time.Hour) - sch.mu.Lock() - sch.rows["kb1"].expires = &past - sch.mu.Unlock() - // Lock the reclaim transition in isolation (before any re-claim): the fake's - // reclaim helper is the same code path TryClaim uses, mirroring reclaimOne. - sch.mu.Lock() - if id := sch.fakeReclaimExpired(time.Now()); id != "kb1" { - sch.mu.Unlock() - t.Fatalf("expected kb1 to be reclaimed, got %q", id) - } - sch.mu.Unlock() - if c, s := fakeRowCounts(sch, "kb1"); s != DatasetStatePending || c.backlog != 2 || c.inflight != 0 { - t.Fatalf("after reclaim: want state=pending backlog=2 inflight=0, got state=%s %+v", s, c) - } - - // Claim again then Ack the drained batch -> completed, counts zeroed. - cr2, ok3, err := sch.Claim(context.Background(), "kb1") - if err != nil || !ok3 { - t.Fatalf("re-claim: ok=%v err=%v", ok3, err) - } - if _, err := sch.Ack(context.Background(), "kb1", cr2.Token, cr2.Entries); err != nil { - t.Fatalf("ack: %v", err) - } - if c, s := fakeRowCounts(sch, "kb1"); s != DatasetStateCompleted || c.backlog != 0 || c.inflight != 0 { - t.Fatalf("after ack drain: want state=completed backlog=0 inflight=0, got state=%s %+v", s, c) - } -} - -// TestSchedulerSetErrorScopedToClaimToken locks the concurrency guard on -// SetError: a failed batch is diagnosed only while its own claim token is still -// live. If worker A's lease expires and is reclaimed, worker B re-claims and -// completes; a late SetError from A (stale token) must NOT overwrite the row's -// diagnostic, otherwise a completed state would be misread as failed. -func TestSchedulerSetErrorScopedToClaimToken(t *testing.T) { - sch := NewFakeScheduler() - if err := sch.Publish(context.Background(), "t1", "kb1", "d1", string(EventTypeCompleted)); err != nil { - t.Fatalf("publish: %v", err) - } - - // A claims and begins processing (running). - crA, ok, err := sch.Claim(context.Background(), "kb1") - if err != nil || !ok { - t.Fatalf("claim A: ok=%v err=%v", ok, err) - } - - // A's lease expires before it finishes; the sweeper reclaims the inflight - // batch back to backlog (pending) and B takes over. - past := time.Now().Add(-time.Hour) - sch.mu.Lock() - sch.rows["kb1"].expires = &past - sch.mu.Unlock() - if id := sch.fakeReclaimExpired(time.Now()); id != "kb1" { - t.Fatalf("expected kb1 reclaimed, got %q", id) - } - crB, okB, err := sch.Claim(context.Background(), "kb1") - if err != nil || !okB { - t.Fatalf("claim B: ok=%v err=%v", okB, err) - } - if crA.Token == crB.Token { - t.Fatalf("expected distinct claim tokens, got %q", crA.Token) - } - - // B succeeds and drains the backlog to completed. - if _, err := sch.Ack(context.Background(), "kb1", crB.Token, crB.Entries); err != nil { - t.Fatalf("ack B: %v", err) - } - if c, s := fakeRowCounts(sch, "kb1"); s != DatasetStateCompleted || c.backlog != 0 || c.inflight != 0 { - t.Fatalf("after B ack: want completed empty, got state=%s %+v", s, c) - } - - // A's late failure arrives with its now-stale token: it must be ignored. - if err := sch.SetError(context.Background(), "kb1", crA.Token, "stale failure from A"); err != nil { - t.Fatalf("stale set error: %v", err) - } - sch.mu.Lock() - gotErr := sch.rows["kb1"].errorMsg - sch.mu.Unlock() - if gotErr != "" { - t.Fatalf("stale SetError overwrote diagnostic: got errorMsg=%q want empty", gotErr) - } - if _, s := fakeRowCounts(sch, "kb1"); s != DatasetStateCompleted { - t.Fatalf("stale SetError changed state: got %q want completed", s) - } -} diff --git a/internal/ingestion/knowledge_compile/dedup.go b/internal/ingestion/knowledge_compile/dedup.go index 3018b4b157..37785f6481 100644 --- a/internal/ingestion/knowledge_compile/dedup.go +++ b/internal/ingestion/knowledge_compile/dedup.go @@ -111,27 +111,57 @@ func (x *llmDeduper) Decide(ctx context.Context, existing, incoming kccommon.Pro 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)) +// DecideBatch folds every group into its dataset-level merged row. Wiki groups +// use a REPLACE-ONLY strategy (wikiDecideBatch) that never invokes the LLM JSON +// merge — Markdown wiki pages must not be concatenated/merged by the generic +// structure decider. Structure (and other non-wiki) groups are judged by the LLM +// merge decider as before. +// splitWikiGroups partitions the batch into wiki groups (replace-only, no LLM) +// and structure groups (LLM-merged). structIdx[i] is the ORIGINAL position in +// `groups` of structGroups[i]; it is what the fold uses to write results back to +// the right slice element. Recording the position inside structGroups (i.e. +// len(structGroups)) instead would always equal i and misroute structure results +// to the wrong groups whenever a wiki group appears earlier in the batch. +func splitWikiGroups(groups []MergeGroup) (wikiIdx, structIdx []int, structGroups []MergeGroup) { for gi := range groups { - pairIndexOf[gi] = make([]int, len(groups[gi].Candidates)) - for ci := range groups[gi].Candidates { + if isWikiGroup(groups[gi]) { + wikiIdx = append(wikiIdx, gi) + } else { + structIdx = append(structIdx, gi) + structGroups = append(structGroups, groups[gi]) + } + } + return wikiIdx, structIdx, structGroups +} + +func (x *llmDeduper) DecideBatch(ctx context.Context, groups []MergeGroup) ([]MergeGroup, error) { + // Split into wiki (replace-only, no LLM) and structure (LLM-merged) groups. + wikiIdx, structIdx, structGroups := splitWikiGroups(groups) + // Wiki groups: replace-only, in place. + for _, gi := range wikiIdx { + groups[gi] = wikiDecideBatch(ctx, []MergeGroup{groups[gi]})[0] + } + if len(structGroups) == 0 { + return groups, nil + } + + // Assign a flat pair index to every (group, candidate) of the structure groups. + var inputs []structure.MergePairInput + pairIndexOf := make([][]int, len(structGroups)) + for gi := range structGroups { + pairIndexOf[gi] = make([]int, len(structGroups[gi].Candidates)) + for ci := range structGroups[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, + Existing: structGroups[gi].Existing.Content, + Incoming: structGroups[gi].Candidates[ci].Content, }) } } if len(inputs) == 0 { + // Only wiki groups had candidates; copy them back (already folded above). return groups, nil } results, err := x.decider.DecideBatch(ctx, inputs) @@ -143,12 +173,12 @@ func (x *llmDeduper) DecideBatch(ctx context.Context, groups []MergeGroup) ([]Me byIndex[r.Index] = r } - for gi := range groups { - existing := groups[gi].Existing + for si, gi := range structIdx { + existing := structGroups[si].Existing var distinct []kccommon.Product duplicated := false - for ci, cand := range groups[gi].Candidates { - r := byIndex[pairIndexOf[gi][ci]] + for ci, cand := range structGroups[si].Candidates { + r := byIndex[pairIndexOf[si][ci]] if !r.Duplicated || r.Merged == nil { // Judged distinct: keep it as its own new merged row. c := cand diff --git a/internal/ingestion/knowledge_compile/dedup_test.go b/internal/ingestion/knowledge_compile/dedup_test.go index a4a178356d..c90193401d 100644 --- a/internal/ingestion/knowledge_compile/dedup_test.go +++ b/internal/ingestion/knowledge_compile/dedup_test.go @@ -97,6 +97,47 @@ func TestDeduperDecideBatchFoldsGroups(t *testing.T) { } } +// TestSplitWikiGroupsMapsStructureIndexes locks the Critical fix in +// llmDeduper.DecideBatch: structIdx[i] must hold the ORIGINAL position in the +// batch, not the position inside the structure-only slice. When a wiki group +// appears before a structure group, a buggy len(structGroups) mapping would make +// structIdx[0]==0 and fold structure results back into the wiki group instead of +// the structure group at index 1. +func TestSplitWikiGroupsMapsStructureIndexes(t *testing.T) { + groups := []MergeGroup{ + {Existing: kccommon.Product{ID: "wiki", Variant: kccommon.VariantWiki, Content: "w"}}, + {Existing: kccommon.Product{ID: "struct-0", Variant: kccommon.Variant("structure"), Content: "s0"}}, + {Existing: kccommon.Product{ID: "wiki-2", Variant: kccommon.VariantWiki, Content: "w2"}}, + {Existing: kccommon.Product{ID: "struct-1", Variant: kccommon.Variant("structure"), Content: "s1"}}, + } + wikiIdx, structIdx, structGroups := splitWikiGroups(groups) + + if len(wikiIdx) != 2 || wikiIdx[0] != 0 || wikiIdx[1] != 2 { + t.Fatalf("wikiIdx = %v, want [0 2]", wikiIdx) + } + if len(structIdx) != 2 { + t.Fatalf("structIdx = %v, want two structure entries", structIdx) + } + // The structure groups appear at original batch indices 1 and 3. + wantStructIdx := []int{1, 3} + for i, gi := range structIdx { + if gi != wantStructIdx[i] { + t.Fatalf("structIdx[%d] = %d, want %d (original batch index)", i, gi, wantStructIdx[i]) + } + // groups[structIdx[i]] must be the exact element copied into structGroups[i], + // so folding back via groups[structIdx[i]] lands on the correct group. + if groups[gi].Existing.ID != structGroups[i].Existing.ID { + t.Fatalf("groups[%d] (%s) is not the same element as structGroups[%d] (%s)", + gi, groups[gi].Existing.ID, i, structGroups[i].Existing.ID) + } + } + // The buggy mapping would have produced structIdx == [0 1], pointing at the + // wiki groups; assert we did NOT regress to that. + if structIdx[0] == 0 || structIdx[1] == 1 { + t.Fatalf("structIdx mapped to the wiki groups' positions: %v", structIdx) + } +} + func TestNoopDeduperDecideBatchNoMerge(t *testing.T) { d := NewNoopDeduper() existing := kccommon.Product{ID: "row-1", DocID: "kb1", Content: "base"} diff --git a/internal/ingestion/knowledge_compile/fake_engine_test.go b/internal/ingestion/knowledge_compile/fake_engine_test.go new file mode 100644 index 0000000000..48fe8e98ef --- /dev/null +++ b/internal/ingestion/knowledge_compile/fake_engine_test.go @@ -0,0 +1,114 @@ +// +// 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" + + "ragflow/internal/engine/types" + + "gorm.io/gorm" +) + +// fakeEngine is a minimal engine.DocEngine for unit tests. It records the +// Search filter/SelectFields and the DeleteChunks condition so tests can assert +// the structural scope of a reader/writer call, and returns a canned chunk set. +type fakeEngine struct { + // searchChunks is the result returned by Search. + searchChunks []map[string]interface{} + // lastSearchReq captures the most recent SearchRequest for assertions. + lastSearchReq *types.SearchRequest + // lastDeleteCond captures the most recent DeleteChunks condition. + lastDeleteCond map[string]interface{} + // deleteCount is the number returned by DeleteChunks. + deleteCount int64 +} + +func (f *fakeEngine) Search(_ context.Context, req *types.SearchRequest) (*types.SearchResult, error) { + f.lastSearchReq = req + return &types.SearchResult{Chunks: f.searchChunks}, nil +} + +func (f *fakeEngine) DeleteChunks(_ context.Context, condition map[string]interface{}, _, _ string) (int64, error) { + f.lastDeleteCond = condition + return f.deleteCount, nil +} + +// --- stubs: the reader/writer path under test never exercises these. --- + +func (f *fakeEngine) CreateChunkStore(context.Context, string, string, int, string) error { + return nil +} +func (f *fakeEngine) InsertChunks(context.Context, []map[string]interface{}, string, string) ([]string, error) { + return nil, nil +} +func (f *fakeEngine) UpdateChunks(context.Context, map[string]interface{}, map[string]interface{}, string, string) error { + return nil +} +func (f *fakeEngine) GetChunk(context.Context, string, string, []string) (interface{}, error) { + return nil, nil +} +func (f *fakeEngine) DropChunkStore(context.Context, string, string) error { return nil } +func (f *fakeEngine) ChunkStoreExists(context.Context, string, string) (bool, error) { + return true, nil +} +func (f *fakeEngine) CreateMetadataStore(context.Context, string) error { return nil } +func (f *fakeEngine) InsertMetadata(context.Context, []map[string]interface{}, string) ([]string, error) { + return nil, nil +} +func (f *fakeEngine) UpdateMetadata(context.Context, string, string, map[string]interface{}, string) error { + return nil +} +func (f *fakeEngine) DeleteMetadata(context.Context, map[string]interface{}, string) (int64, error) { + return 0, nil +} +func (f *fakeEngine) DeleteMetadataKeys(context.Context, string, string, []string, string) error { + return nil +} +func (f *fakeEngine) DropMetadataStore(context.Context, string) error { return nil } +func (f *fakeEngine) MetadataStoreExists(context.Context, string) (bool, error) { return true, nil } +func (f *fakeEngine) SearchMetadata(context.Context, *types.SearchMetadataRequest) (*types.SearchMetadataResult, error) { + return nil, nil +} +func (f *fakeEngine) IndexDocument(context.Context, string, string, interface{}) error { return nil } +func (f *fakeEngine) DeleteDocument(context.Context, string, string) error { return nil } +func (f *fakeEngine) BulkIndex(context.Context, string, []interface{}) (interface{}, error) { + return nil, nil +} +func (f *fakeEngine) GetFields([]map[string]interface{}, []string) map[string]map[string]interface{} { + return nil +} +func (f *fakeEngine) GetAggregation([]map[string]interface{}, string) []map[string]interface{} { + return nil +} +func (f *fakeEngine) GetHighlight([]map[string]interface{}, []string, string) map[string]string { + return nil +} +func (f *fakeEngine) RunSQL(context.Context, string, string, []string, string) ([]map[string]interface{}, error) { + return nil, nil +} +func (f *fakeEngine) GetChunkIDs([]map[string]interface{}) []string { return nil } +func (f *fakeEngine) KNNScores(context.Context, []map[string]interface{}, []float64, int) (map[string]interface{}, error) { + return nil, nil +} +func (f *fakeEngine) GetScores(map[string]interface{}) map[string]float64 { return nil } +func (f *fakeEngine) Ping(context.Context) error { return nil } +func (f *fakeEngine) Close() error { return nil } +func (f *fakeEngine) GetType() string { return "fake" } +func (f *fakeEngine) SupportsPageRank() bool { return false } +func (f *fakeEngine) FilterDocIdsByMetaPushdown(context.Context, *gorm.DB, []string, []map[string]interface{}, string) []string { + return nil +} diff --git a/internal/ingestion/knowledge_compile/reader.go b/internal/ingestion/knowledge_compile/reader.go index 7141fe6c57..809605b86f 100644 --- a/internal/ingestion/knowledge_compile/reader.go +++ b/internal/ingestion/knowledge_compile/reader.go @@ -63,6 +63,14 @@ type engineReader struct { // compiledSelectFields are the columns needed to reconstruct a Product from a // stored compiled chunk document. +// +// wiki_incremental port: the list also selects `kc_kind` and +// `create_timestamp_flt` (+`create_time`) so the reader can round-trip the +// wiki product kind (page/section) and the original creation timestamp without +// re-deriving them (see productFromChunkMap). Without these in the SELECT list, +// the stored values are invisible to the reader and every merged row would fall +// back to the compile_kwd-derived kind / a fresh now() timestamp — which both +// breaks the page/section filter and re-stamps creation time on every rebuild. var compiledSelectFields = []string{ "id", "doc_id", "tenant_id", "compile_kwd", "available_int", @@ -70,6 +78,7 @@ var compiledSelectFields = []string{ "source_chunk_ids", "source_doc_ids", "name_kwd", "entity_type_kwd", "from_entity_kwd", "to_entity_kwd", "slug_kwd", "type", + "kc_kind", "create_timestamp_flt", "create_time", } // wikiSelectFields are the additional columns a wiki page carries (beyond @@ -126,7 +135,13 @@ func (r engineReader) LoadDocProducts(ctx context.Context, tenant, kb, docID str if _, ok := c["compile_kwd"]; !ok { continue } - if p, ok := productFromChunkMap(c, tenant); ok { + // Reverse-map the row's compile_kwd; reject dirty/unknown kinds + // (kwdToVariant error) so a malformed row never loads as a product. + rowVariant, verr := kwdToVariant(asString(c["compile_kwd"])) + if verr != nil { + continue + } + if p, ok := productFromChunkMap(c, tenant, rowVariant); ok { out = append(out, p) } } @@ -159,7 +174,14 @@ func (r engineReader) LoadDocProducts(ctx context.Context, tenant, kb, docID str // productFromChunkMap reconstructs a kccommon.Product from a stored compiled // chunk document. It reads the payload from kc_payload (falling back to // content_with_weight) and the embedding from the q__vec column. -func productFromChunkMap(c map[string]interface{}, tenant string) (kccommon.Product, bool) { +// +// expect is the variant the caller is querying for. The stored compile_kwd is +// reverse-mapped via kwdToVariant and compared against expect; a mismatch (or +// an unknown/dirty compile_kwd that does not map to any known variant) causes +// the row to be skipped. This is the canonical dirty-row contract promised by +// the plan: we never rely on the raw string equality alone, so unknown kinds +// are rejected consistently rather than leaking into the wrong bucket. +func productFromChunkMap(c map[string]interface{}, tenant string, expect kccommon.Variant) (kccommon.Product, bool) { content, _ := c["kc_payload"].(string) if content == "" { content, _ = c["content_with_weight"].(string) @@ -169,7 +191,17 @@ func productFromChunkMap(c map[string]interface{}, tenant string) (kccommon.Prod } id, _ := c["id"].(string) docID, _ := c["doc_id"].(string) - variant, _ := c["compile_kwd"].(string) + // Normalize through the shared asString helper (matching LoadDocProducts) so + // a non-string scalar or a list-wrapped keyword column from the engine is + // reverse-mapped consistently instead of being rejected as a dirty row. + variant := asString(c["compile_kwd"]) + // Dirty-row contract: the raw compile_kwd must reverse-map to the expected + // variant. An unknown/dirty kwd (kwdToVariant error) or a mapped variant + // that differs from expect is rejected. + mapped, err := kwdToVariant(variant) + if err != nil || mapped != expect { + return kccommon.Product{}, false + } merged := isAvailable(c["available_int"]) meta := map[string]any{} @@ -226,6 +258,30 @@ func productFromChunkMap(c map[string]interface{}, tenant string) (kccommon.Prod meta["kind"] = "entity" } } + // wiki_incremental port: round-trip the wiki product kind so the + // dataset-level merge can reliably distinguish pages from sections. The + // merged writer stores kc_kind; when present it is authoritative. Without + // it (legacy rows), derive from compile_kwd: wiki_page -> "page", + // wiki_section -> "section". This fix is what stops the processBatch + // "Meta.kind==page" filter from deleting every wiki page (previously kind + // was empty for wiki pages that had no entity/relation endpoint). + if v, ok := c["kc_kind"].(string); ok && v != "" { + meta["kind"] = v + } else if variant == compileKwdWikiPage { + meta["kind"] = "page" + } else if variant == compileKwdWikiSection { + meta["kind"] = "section" + } + // wiki_incremental port: restore the original creation timestamp so a + // replace-only merge (wikiMerge.Replace) preserves it instead of stamping + // a fresh now(). existing rows carry create_timestamp_flt (and optionally + // a human-readable create_time string). + if v, ok := metaFloat(c, "create_timestamp_flt"); ok { + meta["created_at_unix"] = v + } + if v, ok := c["create_time"].(string); ok && v != "" { + meta["created_at"] = v + } if v := metaStringSlice(c, "source_chunk_ids"); len(v) > 0 { meta["source_chunk_ids"] = v } @@ -238,7 +294,7 @@ func productFromChunkMap(c map[string]interface{}, tenant string) (kccommon.Prod ID: id, DocID: docID, TenantID: tenant, - Variant: kccommon.Variant(variant), + Variant: expect, Content: content, Vector: vec, Meta: meta, @@ -266,11 +322,12 @@ func (r engineReader) SearchSimilar(ctx context.Context, tenant, kb string, vari Limit: topN, SelectFields: append([]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", "available_int", "compile_kwd"}, + "type", "source_chunk_ids", "source_doc_ids", "available_int", "compile_kwd", + "kc_kind", "create_timestamp_flt", "create_time"}, wikiSelectFields...), Filter: map[string]interface{}{ "available_int": 1, - "compile_kwd": string(variant), + "compile_kwd": compileKwdForVariant(variant), }, MatchExprs: []interface{}{ &types.MatchDenseExpr{ @@ -286,8 +343,23 @@ func (r engineReader) SearchSimilar(ctx context.Context, tenant, kb string, vari if err != nil { return kccommon.Product{}, 0, err } + // The KNN Filter scopes compile_kwd, but a foreign/legacy row that slipped + // past it must not poison the merge candidate. Reject by the raw compile_kwd + // (wiki_page vs wiki_section are distinct keywords even though both map to + // VariantWiki); the page/section distinction is resolved downstream by + // Meta.kind (see productFromChunkMap). + expectKwd := compileKwdForVariant(variant) for _, c := range res.Chunks { - p, ok := productFromChunkMap(c, tenant) + // The KNN Filter already scopes compile_kwd, but a foreign/legacy row that + // slipped past it must not poison the merge candidate. Reject by the + // reverse-mapped variant (the dirty-row contract): productFromChunkMap + // validates kwdToVariant(c) == variant and drops dirty/unknown kinds. The + // raw-keyword check below is a fast pre-filter before the full product + // reconstruction. + if asString(c["compile_kwd"]) != expectKwd { + continue + } + p, ok := productFromChunkMap(c, tenant, variant) if !ok || !p.Merged { continue } @@ -335,6 +407,36 @@ func isAvailable(v interface{}) bool { return false } +// asString normalizes a boxed chunk-map value into a string (used where the +// backend may return a typed string, a json.Number, or a single-element list for +// a keyword column). A list-wrapped keyword (e.g. []string{"wiki_page"} or +// []any{"wiki_page"}) is unwrapped to its first element so reverse-mapping and +// raw-keyword filters behave consistently across engine backends. +func asString(v interface{}) string { + switch t := v.(type) { + case nil: + return "" + case string: + return t + case json.Number: + return t.String() + case []string: + if len(t) == 1 { + return t[0] + } + return "" + case []interface{}: + if len(t) == 1 { + if s, ok := t[0].(string); ok { + return s + } + return fmt.Sprintf("%v", t[0]) + } + return "" + } + return fmt.Sprintf("%v", v) +} + // 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. diff --git a/internal/ingestion/knowledge_compile/reader_test.go b/internal/ingestion/knowledge_compile/reader_test.go new file mode 100644 index 0000000000..9642e1a89d --- /dev/null +++ b/internal/ingestion/knowledge_compile/reader_test.go @@ -0,0 +1,125 @@ +// +// 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" +) + +// TestSearchSimilarFiltersByVariant asserts the B1/KNN contract: SearchSimilar +// scopes the engine query to available_int=1 AND compile_kwd=variant, and the +// in-memory dirty-row guard drops any row whose compile_kwd does not map back to +// the variant (so a foreign row can never become the merge candidate). +func TestSearchSimilarFiltersByVariant(t *testing.T) { + eng := &fakeEngine{ + searchChunks: []map[string]interface{}{ + // dirty row: wrong variant (must be skipped by the in-memory guard). + { + "id": "dirty", + "doc_id": "kb", + "available_int": 1, + "compile_kwd": "wiki_section", // not wiki_page + "kc_payload": "{\"c\":1}", + "_score": 0.99, + }, + // good row: correct variant + kc_kind round-trip. + { + "id": "page1", + "doc_id": "kb", + "available_int": 1, + "compile_kwd": "wiki_page", + "kc_kind": "page", + "kc_payload": "{\"c\":1}", + "create_timestamp_flt": 1700000000.0, + "create_time": "2023-11-14T22:13:20Z", + "_score": 0.95, + }, + }, + } + r := engineReader{eng: eng} + + p, score, err := r.SearchSimilar(context.Background(), "t1", "kb", kccommon.VariantWiki, []float64{0.1, 0.2, 0.3}, 2, 0.5) + if err != nil { + t.Fatalf("SearchSimilar: %v", err) + } + if p.ID != "page1" { + t.Fatalf("expected the wiki_page row to win, got %q", p.ID) + } + if p.Meta["kind"] != "page" { + t.Fatalf("expected kc_kind round-trip to 'page', got %v", p.Meta["kind"]) + } + if p.Merged != true { + t.Fatalf("expected merged=true for available_int=1 row") + } + if _, ok := p.Meta["created_at_unix"]; !ok { + t.Fatalf("expected created_at_unix round-trip from create_timestamp_flt") + } + if score != 0.95 { + t.Fatalf("expected _score=0.95, got %v", score) + } + + // Structural filter: the engine request must scope to the variant. + if eng.lastSearchReq == nil { + t.Fatal("engine.Search was not called") + } + if eng.lastSearchReq.Filter["available_int"] != 1 { + t.Fatalf("expected available_int=1 filter, got %v", eng.lastSearchReq.Filter["available_int"]) + } + if eng.lastSearchReq.Filter["compile_kwd"] != string(compileKwdWikiPage) { + t.Fatalf("expected compile_kwd=%q filter, got %v", compileKwdWikiPage, eng.lastSearchReq.Filter["compile_kwd"]) + } + // SelectFields must include the round-trip columns added in the 8th review. + for _, want := range []string{"kc_kind", "create_timestamp_flt", "create_time"} { + found := false + for _, f := range eng.lastSearchReq.SelectFields { + if f == want { + found = true + break + } + } + if !found { + t.Fatalf("SearchSimilar SelectFields missing %q", want) + } + } +} + +// TestSearchSimilarSkipsNonMerged asserts a row with available_int=0 (a +// per-document row, not a dataset-level merged row) is never a KNN candidate. +func TestSearchSimilarSkipsNonMerged(t *testing.T) { + eng := &fakeEngine{ + searchChunks: []map[string]interface{}{ + { + "id": "doc1", + "doc_id": "doc", + "available_int": 0, + "compile_kwd": "wiki_page", + "kc_payload": "{\"c\":1}", + "_score": 0.99, + }, + }, + } + r := engineReader{eng: eng} + p, _, err := r.SearchSimilar(context.Background(), "t1", "kb", kccommon.VariantWiki, []float64{0.1}, 1, 0.5) + if err != nil { + t.Fatalf("SearchSimilar: %v", err) + } + if p.ID != "" { + t.Fatalf("expected no candidate from a non-merged row, got %q", p.ID) + } +} diff --git a/internal/ingestion/knowledge_compile/scheduler.go b/internal/ingestion/knowledge_compile/scheduler.go index 479d285279..bf9d7f136b 100644 --- a/internal/ingestion/knowledge_compile/scheduler.go +++ b/internal/ingestion/knowledge_compile/scheduler.go @@ -114,6 +114,26 @@ type Claimer interface { // status of the worker that took over. SetError(ctx context.Context, datasetID, token, errMsg string) error + // CancelInflight drops the live claim for a dataset back into the backlog so + // a rewrite can start from a clean index. The token is intentionally NOT + // checked: this is the rewrite's authoritative cancel. Workers verify their + // claim token inside the dataset write lock before every destructive write, so + // a cancelled worker cannot write after the rebuild clears storage. It is a safe no-op when + // there is no live lease (the dataset is idle or already drained). The dropped + // batch is re-marked pending (not deleted) so it is reclaimed and re-checked + CancelInflight(ctx context.Context, datasetID, token string) error + + // WithDatasetLock executes fn while holding an exclusive per-dataset lock + // that serializes against both writer side effects and dataset rebuilds + // (RebuildDataset). It is the synchronization that closes the TOCTOU gap + // between a worker's claim validation and its destructive write: the caller + // must compare its claim token and perform the write inside fn. fn + // runs inside a transaction that holds the dataset row lock for its whole + // duration (MySQL) or a per-row mutex (Fake), so rebuild and writes are + // mutually exclusive. claimToken is read under the lock; the caller must not + // re-read it on a separate connection. + WithDatasetLock(ctx context.Context, datasetID string, fn func(claimToken string) error) 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) @@ -184,6 +204,23 @@ func (s *mysqlScheduler) Provision(ctx context.Context) error { // 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) error { + return s.publish(ctx, tenantID, datasetID, docID, eventType) +} + +// loadRow fetches the dataset scheduling row (no lock). Returns nil (no error) +// when the row does not exist yet. +func (s *mysqlScheduler) loadRow(ctx context.Context, datasetID string) (*entity.KnowledgeCompileDataset, error) { + var row entity.KnowledgeCompileDataset + if err := s.db.WithContext(ctx).Where("dataset_id = ?", datasetID).First(&row).Error; err != nil { + if errors.Is(err, gorm.ErrRecordNotFound) { + return nil, nil + } + return nil, err + } + return &row, nil +} + +func (s *mysqlScheduler) publish(ctx context.Context, tenantID, datasetID, docID, eventType string) error { if s.db == nil { return nil } @@ -235,6 +272,39 @@ func (s *mysqlScheduler) Publish(ctx context.Context, tenantID, datasetID, docID return s.notify(ctx, datasetID) } +// WithDatasetLock executes fn under an exclusive per-dataset row lock. It locks +// the dataset row (SELECT ... FOR UPDATE) and runs fn inside the same +// transaction, so the row lock is held for the entire duration of fn and any +// concurrent WithDatasetLock on the same dataset blocks until fn returns. This +// is what makes a worker's claim-token check + destructive write atomic against +// a concurrent RebuildDataset. +// +// The claim token passed to fn is read within this transaction, so the check in +// fn does not need (and must not attempt) a separate DB read that would deadlock +// against the held row lock. fn runs inside the transaction; if it returns an +// error the transaction is rolled back +// and the error is propagated unchanged (any DB write performed inside fn is +// rolled back, so no partial side effect persists on an aborted write). +func (s *mysqlScheduler) WithDatasetLock(ctx context.Context, datasetID string, fn func(claimToken string) error) error { + if s.db == nil { + // No DB (test/no-op scheduler): nothing to serialize against. + return fn("") + } + return s.db.WithContext(ctx).Transaction(func(tx *gorm.DB) error { + row := entity.KnowledgeCompileDataset{ + DatasetID: datasetID, + BacklogDocIDs: "[]", + InflightDocIDs: "[]", + } + if err := tx.Clauses(clause.Locking{Strength: "UPDATE"}). + Where("dataset_id = ?", datasetID). + FirstOrCreate(&row).Error; err != nil { + return err + } + return fn(row.ClaimToken) + }) +} + // 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 @@ -381,6 +451,46 @@ func (s *mysqlScheduler) SetError(ctx context.Context, datasetID, token, errMsg return nil } +// CancelInflight drops the live claim for a dataset back into the backlog so a +// rewrite can start from a clean index. It is a no-op when there is no live +// lease. The dropped batch is re-marked pending and is reprocessed after the +// rebuild republishes the dataset documents. +func (s *mysqlScheduler) CancelInflight(ctx context.Context, datasetID, token string) error { + if s.db == nil { + return nil + } + err := s.db.WithContext(ctx).Transaction(func(tx *gorm.DB) 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 nil + } + return err + } + // No live lease (idle or already drained) → nothing to cancel. + now := time.Now() + liveLease := row.ClaimOwner != "" && row.ClaimExpiresAt != nil && row.ClaimExpiresAt.After(now) + if !liveLease { + return nil + } + inflight := parseEntries(row.InflightDocIDs) + backlog := parseEntries(row.BacklogDocIDs) + backlog = append(backlog, inflight...) + row.InflightDocIDs = "[]" + row.ClaimOwner = "" + row.ClaimToken = "" + row.ClaimExpiresAt = nil + row.BacklogDocIDs = marshalEntries(backlog) + row.State = DatasetStatePending + return tx.Save(&row).Error + }) + if err != nil { + return fmt.Errorf("knowledge_compile: cancel inflight %s: %w", datasetID, err) + } + return nil +} + func (s *mysqlScheduler) TouchClaim(ctx context.Context, datasetID, token string, ttl time.Duration) (bool, error) { if s.db == nil { return false, nil @@ -561,6 +671,9 @@ type fakeRow struct { expires *time.Time state string errorMsg string + // writeMu serializes per-dataset writer side effects and rebuilds, mirroring + // the MySQL scheduler's WithDatasetLock row lock. + writeMu sync.Mutex } // FakeScheduler is an in-memory Publisher + Claimer used by tests. It mirrors @@ -585,8 +698,12 @@ func NewFakeScheduler() *FakeScheduler { func (f *FakeScheduler) Provision(_ context.Context) error { return nil } -// Publish appends one doc event and pushes a notify (same as the MySQL path). +// Publish appends one doc event and pushes a notify (same contract as MySQL). func (f *FakeScheduler) Publish(_ context.Context, tenantID, datasetID, docID, eventType string) error { + return f.publish(tenantID, datasetID, docID, eventType) +} + +func (f *FakeScheduler) publish(tenantID, datasetID, docID, eventType string) error { f.mu.Lock() defer f.mu.Unlock() r, ok := f.rows[datasetID] @@ -683,6 +800,47 @@ func (f *FakeScheduler) SetError(_ context.Context, datasetID, token, errMsg str return nil } +// WithDatasetLock executes fn under the per-row write mutex, mirroring the MySQL +// scheduler's row-lock semantics: a rebuild and a writer's destructive side +// effect on the same dataset are mutually exclusive. claimToken is read under +// the lock. +func (f *FakeScheduler) WithDatasetLock(_ context.Context, datasetID string, fn func(claimToken string) error) error { + f.mu.Lock() + r, ok := f.rows[datasetID] + if !ok { + r = &fakeRow{} + f.rows[datasetID] = r + } + f.mu.Unlock() + + r.writeMu.Lock() + defer r.writeMu.Unlock() + f.mu.Lock() + claimToken := r.token + f.mu.Unlock() + return fn(claimToken) +} + +// CancelInflight drops the live claim back into the backlog (mirrors MySQL). +func (f *FakeScheduler) CancelInflight(_ context.Context, datasetID, _ string) error { + f.mu.Lock() + defer f.mu.Unlock() + r, ok := f.rows[datasetID] + if !ok { + return nil + } + now := time.Now() + live := r.owner != "" && r.expires != nil && r.expires.After(now) + if !live { + return nil + } + r.backlog = append(r.backlog, r.inflight...) + r.inflight = nil + r.owner, r.token, r.expires = "", "", nil + r.state = DatasetStatePending + return nil +} + // fakeReclaimExpired finds one dataset with an expired lease and non-empty // inflight, moves the inflight batch back to backlog, clears the lease, and // marks the row pending — mirroring the MySQL scheduler's reclaimOne. It returns diff --git a/internal/ingestion/knowledge_compile/service.go b/internal/ingestion/knowledge_compile/service.go index 4acbb0975a..6377a38269 100644 --- a/internal/ingestion/knowledge_compile/service.go +++ b/internal/ingestion/knowledge_compile/service.go @@ -123,6 +123,7 @@ func Provision(ctx context.Context, mq engine.MessageQueue, db *gorm.DB) error { if db == nil { return nil } + kcDB = db s := newScheduler(db, mq, generateHolder(), 2*time.Minute) // Bound the startup AutoMigrate so a slow/unreachable DB cannot block // startup indefinitely. The caller's ctx is also honoured (cancelled on diff --git a/internal/ingestion/knowledge_compile/wiki_graph_test.go b/internal/ingestion/knowledge_compile/wiki_graph_test.go index 0fc55eacaf..a0e39e3c82 100644 --- a/internal/ingestion/knowledge_compile/wiki_graph_test.go +++ b/internal/ingestion/knowledge_compile/wiki_graph_test.go @@ -156,6 +156,45 @@ func TestProjectWikiGraphRowsWeightEqualsOutlinkCount(t *testing.T) { } } +// TestProjectWikiGraphRowsBareSlugOutlinkCompletesGraph asserts the G1 reader +// parity fix: an outlink given as a bare slug ("beta", not "entity/beta") is +// resolved to the full slug via the page slug map and still yields a relation, +// so the wiki nav graph is complete (matching Python dataset_wiki_generator). +func TestProjectWikiGraphRowsBareSlugOutlinkCompletesGraph(t *testing.T) { + w := engineWriter{} + pages := []wikiPageProjection{ + page("entity", "alpha", "beta"), // bare outlink resolves to entity/beta + page("entity", "beta"), + } + rows, err := w.projectWikiGraphRows(context.Background(), "t1", kbForTest, pages) + if err != nil { + t.Fatalf("project: %v", err) + } + if findRelation(rows, "entity/alpha", "entity/beta") == nil { + t.Fatalf("bare-slug outlink must resolve to a relation entity/alpha->entity/beta") + } +} + +// TestProjectWikiGraphRowsDropsDanglingAndSelfLoop asserts a dangling outlink +// (target page absent) and a self-loop (alpha->alpha) are skipped, so the graph +// never references non-existent vertices. +func TestProjectWikiGraphRowsDropsDanglingAndSelfLoop(t *testing.T) { + w := engineWriter{} + pages := []wikiPageProjection{ + page("entity", "alpha", "alpha", "ghost"), // self-loop + dangling + } + rows, err := w.projectWikiGraphRows(context.Background(), "t1", kbForTest, pages) + if err != nil { + t.Fatalf("project: %v", err) + } + if findRelation(rows, "entity/alpha", "entity/alpha") != nil { + t.Fatalf("self-loop must be skipped") + } + if findRelation(rows, "entity/alpha", "entity/ghost") != nil { + t.Fatalf("dangling outlink must be skipped") + } +} + // TestProjectWikiGraphRowsFieldMapping verifies the reader-facing display // columns are written (aliases_kwd from entity_names_kwd, description_with_weight // from summary_with_weight), and the id is the xxhash colon-namespaced form. diff --git a/internal/ingestion/knowledge_compile/wiki_merge.go b/internal/ingestion/knowledge_compile/wiki_merge.go new file mode 100644 index 0000000000..f86e2b0cc3 --- /dev/null +++ b/internal/ingestion/knowledge_compile/wiki_merge.go @@ -0,0 +1,156 @@ +// +// 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" + + kccommon "ragflow/internal/ingestion/component/knowledge_compiler/common" +) + +// wikiReplace folds a wiki page candidate into its existing dataset-level merged +// row using a REPLACE-ONLY strategy. Wiki pages are Markdown, so the generic +// structure.LLMMergeDecider JSON merge (which concatenates/merges arbitrary JSON +// payloads) would mangle the page body — it must never run on wiki content. +// +// The existing row keeps its identity (id / doc_id / creation time) and provenance +// is unioned (source_doc_ids + source_chunk_ids of existing ∪ candidate). The +// page body is replaced by the candidate's Markdown verbatim. This mirrors +// Python's wiki incremental mode where a page re-compiled from a newer document +// simply supersedes the stored page of the same slug, rather than being merged. +func wikiReplace(existing, candidate kccommon.Product) kccommon.Product { + merged := existing + + // Body: the newer (incoming) Markdown wins verbatim. + merged.Content = candidate.Content + // The embedding must match the replacement Markdown. WriteMerged persists + // p.Vector without re-embedding, so keeping existing.Vector would leave a + // stale embedding for the old page body and break KNN similarity searches. + merged.Vector = candidate.Vector + + // Identity preserved: keep existing id, doc_id (== kb), and the original + // creation timestamp (carried via Meta.created_at_unix by the Reader). + merged.ID = existing.ID + merged.DocID = existing.DocID + + // Provenance union (deduped) so the merged page references every source doc + // and chunk that contributed to it across runs. + merged.Meta = unionWikiProvenance(existing.Meta, candidate.Meta) + + // Re-stamp the resolver/run id from the incoming candidate so the row is + // attributed to the latest batch, but the created_at_unix stays with existing. + if v, ok := candidate.Meta["run_id"]; ok { + merged.Meta["run_id"] = v + } + return merged +} + +// unionWikiProvenance returns a new Meta map based on a (the existing row). The +// candidate b overwrites kind / slug / page_type / title / summary / +// entity_names / related_kb_pages / outlinks. Identity and creation time +// (created_at_unix, created_at) stay from a. Source provenance arrays +// (source_doc_ids, source_chunk_ids) are unioned and deduped. +func unionWikiProvenance(a, b map[string]any) map[string]any { + out := map[string]any{} + for k, v := range a { + out[k] = v + } + // The incoming page replaces current page metadata. Identity and creation + // time remain from the existing row. + for _, key := range []string{"slug", "page_type", "title", "summary", "kind"} { + if v, ok := b[key]; ok { + out[key] = v + } + } + out["source_doc_ids"] = unionStrs(metaStringSliceAny(a, "source_doc_ids"), metaStringSliceAny(b, "source_doc_ids")) + out["source_chunk_ids"] = unionStrs(metaStringSliceAny(a, "source_chunk_ids"), metaStringSliceAny(b, "source_chunk_ids")) + if v, ok := b["entity_names"]; ok { + out["entity_names"] = v + } + if v, ok := b["related_kb_pages"]; ok { + out["related_kb_pages"] = v + } + if v, ok := b["outlinks"]; ok { + out["outlinks"] = v + } + return out +} + +func unionStrs(a, b []string) []string { + if len(a) == 0 { + return b + } + if len(b) == 0 { + return a + } + seen := make(map[string]struct{}, len(a)+len(b)) + out := make([]string, 0, len(a)+len(b)) + for _, s := range append(append([]string(nil), a...), b...) { + if _, ok := seen[s]; ok { + continue + } + seen[s] = struct{}{} + out = append(out, s) + } + return out +} + +// metaStringSliceAny reads a []string from a meta map that may box the value as +// []string or []any (engine/JSON round-trip does not guarantee a single type). +func metaStringSliceAny(m map[string]any, key string) []string { + switch v := m[key].(type) { + case []string: + return v + case []any: + out := make([]string, 0, len(v)) + for _, e := range v { + if s, ok := e.(string); ok { + out = append(out, s) + } + } + return out + } + return nil +} + +// isWikiGroup reports whether a merge group targets the wiki variant and must use +// replace-only (never the JSON-merge decider). +func isWikiGroup(g MergeGroup) bool { + return g.Existing.Variant == kccommon.VariantWiki +} + +// wikiDecideBatch folds every wiki-group candidate into its existing row using +// replace-only semantics, WITHOUT any LLM call. It mutates the passed groups in +// place and returns them. +func wikiDecideBatch(_ context.Context, groups []MergeGroup) []MergeGroup { + for gi := range groups { + existing := groups[gi].Existing + var distinct []kccommon.Product + duplicated := false + for _, cand := range groups[gi].Candidates { + // A wiki page candidate is always a replacement of the existing row + // (same slug, newer content). It is never kept as an additional distinct + // row — distinctness for wiki is decided by KNN (different slug -> different + // existing row -> different group). + existing = wikiReplace(existing, cand) + duplicated = true + } + groups[gi].Merged = existing + groups[gi].Duplicate = duplicated + groups[gi].Distinct = distinct + } + return groups +} diff --git a/internal/ingestion/knowledge_compile/wiki_merge_test.go b/internal/ingestion/knowledge_compile/wiki_merge_test.go new file mode 100644 index 0000000000..75ae2aa08c --- /dev/null +++ b/internal/ingestion/knowledge_compile/wiki_merge_test.go @@ -0,0 +1,89 @@ +// +// 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" +) + +// TestWikiReplaceCopiesCandidateVector locks the wiki replace-only merge contract +// (CodeRabbit Major): wikiReplace swaps the page body for the candidate's Markdown +// AND must copy the candidate's embedding too. WriteMerged persists Product.Vector +// without re-embedding, so keeping the existing row's vector would leave a stale +// embedding for the old page body and break KNN similarity searches. +func TestWikiReplaceCopiesCandidateVector(t *testing.T) { + existing := kccommon.Product{ + ID: "kb/wiki/alpha", + DocID: "kb", + TenantID: "t1", + Variant: kccommon.VariantWiki, + Content: "# old markdown", + Vector: []float32{1, 0, 0, 0}, + Meta: map[string]any{ + "kind": "page", + "slug": "page/alpha", + "created_at": "2024-01-01", + }, + } + candidate := kccommon.Product{ + ID: "kb/wiki/alpha", + DocID: "kb", + TenantID: "t1", + Variant: kccommon.VariantWiki, + Content: "# new markdown", + Vector: []float32{0, 1, 0, 0}, + Meta: map[string]any{ + "kind": "page", + "slug": "page/alpha", + "run_id": "run-9", + }, + } + + groups := wikiDecideBatch(context.Background(), []MergeGroup{{ + Existing: existing, + Candidates: []kccommon.Product{candidate}, + }}) + + merged := groups[0].Merged + if merged.Content != "# new markdown" { + t.Fatalf("content = %q, want candidate markdown", merged.Content) + } + if !groups[0].Duplicate { + t.Fatalf("wiki candidate must be marked duplicate (replacement)") + } + if len(merged.Vector) != len(candidate.Vector) { + t.Fatalf("merged vector must be copied from the candidate, got %v", merged.Vector) + } + for i := range candidate.Vector { + if merged.Vector[i] != candidate.Vector[i] { + t.Fatalf("merged vector[%d] = %v, want candidate %v", i, merged.Vector, candidate.Vector) + } + } + // Identity and creation time come from the existing row; content-bearing and + // page metadata fields come from the candidate. + if merged.ID != existing.ID || merged.DocID != existing.DocID { + t.Fatalf("identity must be preserved: id=%q doc_id=%q", merged.ID, merged.DocID) + } + if merged.Meta["created_at"] != "2024-01-01" { + t.Fatalf("creation time must be preserved from the existing row, got %v", merged.Meta["created_at"]) + } + if merged.Meta["run_id"] != "run-9" { + t.Fatalf("run_id must come from the candidate, got %v", merged.Meta["run_id"]) + } +} diff --git a/internal/ingestion/knowledge_compile/writer.go b/internal/ingestion/knowledge_compile/writer.go index 15a006cbbe..97d2450540 100644 --- a/internal/ingestion/knowledge_compile/writer.go +++ b/internal/ingestion/knowledge_compile/writer.go @@ -60,6 +60,14 @@ type Writer interface { ProjectWikiGraph(ctx context.Context, tenant, kb string) error // DropWikiGraph deletes every wiki_entity / wiki_relation row for the dataset. DropWikiGraph(ctx context.Context, tenant, kb string) error + // DeleteMerged removes the dataset-level merged rows for a KB so an + // incremental build can start from a clean slate. The structural filter + // deletes only rows produced by the dataset-level merge — kb_id == kb AND + // available_int == 1 AND compile_kwd is a wiki variant (wiki_page/ + // wiki_section). Per-document rows (doc_id == doc, available_int == 0) and + // rows for other tenants / variants are untouched. The match_kwd guard is the + // in-memory safety net; the structural filter is the source of truth. + DeleteMerged(ctx context.Context, tenant, kb string) error } // engineWriter persists dataset-level merged products through the global @@ -188,6 +196,26 @@ func mergedChunkMap(tenant, kb, runID, inputHash string, now time.Time, p kccomm "create_time": now.Format("2006-01-02 15:04:05"), "create_timestamp_flt": float64(now.Unix()), } + // wiki_incremental port: persist the product kind so the Reader can round-trip + // page vs section without re-deriving it from compile_kwd. The merged writer + // carries the authoritative kc_kind; legacy rows without it are derived in + // productFromChunkMap (compile_kwd wiki_page -> "page", wiki_section -> + // "section"). Without this, the dataset-level merge could not distinguish a + // wiki page from a section and the processBatch "Meta.kind==page" filter would + // be unreliable. + if kind := metaString(p.Meta, "kind"); kind != "" { + m["kc_kind"] = kind + } + // wiki_incremental port: preserve the original creation timestamp across a + // replace-only merge. If the incoming merged product already carries + // created_at_unix (restored by the Reader from create_timestamp_flt), reuse + // it; otherwise stamp a fresh now() (first creation). This is what stops every + // rebuild from re-stamping the creation time. + if v, ok := metaFloat(p.Meta, "created_at_unix"); ok { + m["create_timestamp_flt"] = v + // Rebuild the human-readable form from the preserved unix time. + m["create_time"] = time.Unix(int64(v), 0).Format("2006-01-02 15:04:05") + } // Carry the wiki page metadata onto the merged row so the dataset-level // products keep the fields the artifact API (ListArtifacts/ListWikiTopics) // and page renderers read. Without this the merged rows lose page_type_kwd / @@ -454,6 +482,42 @@ func metaInt(m map[string]any, key string) (int64, bool) { return 0, false } +// metaFloat extracts a float64 from a map value that may be boxed as float64, +// int64, int, or string — the engine/JSON round-trip does not guarantee a +// single numeric type. Used to recover create_timestamp_flt so the reader can +// preserve the original creation time across a replace-only merge. +func metaFloat(m map[string]any, key string) (float64, bool) { + switch v := m[key].(type) { + case float64: + return v, true + case int64: + return float64(v), true + case int: + return float64(v), true + case string: + var f float64 + if _, err := fmt.Sscanf(v, "%f", &f); err == nil { + return f, true + } + } + return 0, false +} + +// kwdToVariant is the inverse of compileKwdForVariant: it maps a stored +// compile_kwd back to its compiler Variant. Both wiki_page and wiki_section +// map to VariantWiki (same product family); the page/section distinction is +// carried by the kc_kind field, not the variant. Returns an error for an +// unknown kwd so callers can reject dirty/foreign rows. +func kwdToVariant(kwd string) (kccommon.Variant, error) { + switch kwd { + case compileKwdWikiPage, compileKwdWikiSection, compileKwdWikiEntity, compileKwdWikiRelation: + return kccommon.VariantWiki, nil + case string(kccommon.VariantStructure), string(kccommon.VariantTree), string(kccommon.VariantMindmap): + return kccommon.Variant(kwd), nil + } + return "", fmt.Errorf("unknown compile_kwd %q", kwd) +} + // --- Wiki page graph materialization (wiki_entity / wiki_relation) --- // compile_kwd values for the dataset-level products this package writes. The @@ -463,6 +527,7 @@ func metaInt(m map[string]any, key string) (int64, bool) { // dedicated wiki_entity / wiki_relation buckets. const ( compileKwdWikiPage = "wiki_page" + compileKwdWikiSection = "wiki_section" compileKwdWikiEntity = "wiki_entity" compileKwdWikiRelation = "wiki_relation" compileKwdWikiPageGraph = "wiki_page_graph" // legacy Python blob, swept on drop @@ -569,6 +634,34 @@ func (w engineWriter) DropWikiGraph(ctx context.Context, tenant, kb string) erro return w.dropWikiGraph(ctx, tenant, kb) } +// DeleteMerged removes the dataset-level (available_int=1) wiki merged rows for +// a KB so an incremental build can start from a clean slate. The structural +// filter (kb_id + available_int=1 + wiki page/section compile_kwd variants) is +// the source of truth; it never targets per-document rows (available_int=0) nor +// rows of other tenants / variants, so a wrong tenantID / kb cannot cascade. +func (w engineWriter) DeleteMerged(ctx context.Context, tenant, kb string) error { + 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{}{ + "kb_id": kb, + "available_int": 1, + "compile_kwd": []string{ + compileKwdWikiPage, + compileKwdWikiSection, + }, + }, baseName, kb) + if err != nil { + return fmt.Errorf("delete merged: %w", err) + } + return nil +} + // dropWikiGraph is the shared delete path for both ProjectWikiGraph (when the // page set is empty, or before re-inserting) and DropWikiGraph. It deletes by // kb_id + compile_kwd IN (the graph buckets), also sweeping any legacy @@ -773,9 +866,6 @@ func (w engineWriter) projectWikiGraphRows(_ context.Context, tenant, kb string, // bare slug (the latter is what the LLM emits in wikitext links); resolve // both so a bare outlink still produces an edge. for _, tgt := range p.Outlinks { - if tgt == p.Slug { - continue // self-loop: Python skips src == tgt - } tp, ok := bySlug[tgt] if !ok { // Bare-slug outlink: normalize (strip prefix, "_"->"-") and look @@ -793,6 +883,9 @@ func (w engineWriter) projectWikiGraphRows(_ context.Context, tenant, kb string, if !ok { continue // dangling edge: target page not in this projection } + if tgt == p.Slug { + continue // self-loop: Python skips src == tgt (full or bare slug) + } relID := wikiGraphXXHash("wiki_relation", kb, p.Slug+":"+tgt) if seen[relID] { continue diff --git a/internal/ingestion/knowledge_compile/writer_test.go b/internal/ingestion/knowledge_compile/writer_test.go index b45a04f5c1..9d445b2f89 100644 --- a/internal/ingestion/knowledge_compile/writer_test.go +++ b/internal/ingestion/knowledge_compile/writer_test.go @@ -1,6 +1,7 @@ package knowledge_compile import ( + "context" "testing" "time" @@ -87,7 +88,7 @@ func TestProductFromChunkMapRestoresWikiFields(t *testing.T) { "related_kb_pages_kwd": []interface{}{"entity/beta"}, "section_level_int": float64(2), } - p, ok := productFromChunkMap(c, "t1") + p, ok := productFromChunkMap(c, "t1", kccommon.VariantWiki) if !ok { t.Fatalf("productFromChunkMap returned not-ok") } @@ -110,3 +111,73 @@ func TestProductFromChunkMapRestoresWikiFields(t *testing.T) { t.Errorf("per-doc row must not be marked merged") } } + +// TestDeleteMergedScopesToMergedWikiRows locks the W1 contract: DeleteMerged +// must only target dataset-level (available_int=1) wiki merged rows. The +// structural filter (kb_id + available_int + wiki page/section compile_kwd) is +// the source of truth, so a wrong tenant/kb can never cascade to per-document +// rows or to rows of other variants. +func TestDeleteMergedScopesToMergedWikiRows(t *testing.T) { + eng := &fakeEngine{} + w := engineWriter{eng: eng} + if err := w.DeleteMerged(context.Background(), "t1", "kb1"); err != nil { + t.Fatalf("DeleteMerged: %v", err) + } + cond := eng.lastDeleteCond + if cond == nil { + t.Fatal("engine.DeleteChunks was not called") + } + if cond["kb_id"] != "kb1" { + t.Errorf("DeleteMerged kb_id = %v, want kb1", cond["kb_id"]) + } + if cond["available_int"] != 1 { + t.Errorf("DeleteMerged must scope to available_int=1, got %v", cond["available_int"]) + } + variants, ok := cond["compile_kwd"].([]string) + if !ok { + t.Fatalf("DeleteMerged must pass a compile_kwd string slice, got %T", cond["compile_kwd"]) + } + if len(variants) != 2 || variants[0] != compileKwdWikiPage || variants[1] != compileKwdWikiSection { + t.Errorf("DeleteMerged compile_kwd = %v, want [%q %q]", variants, compileKwdWikiPage, compileKwdWikiSection) + } +} + +// TestProductFromChunkMapRejectsDirtyKwd locks the dirty-row contract from the +// wiki_incremental plan (Claim 4): productFromChunkMap must reverse-map the raw +// compile_kwd via kwdToVariant and reject any row whose kwd does not map to the +// expected variant, including unknown / malformed kinds (e.g. "artifact_page", +// "garbage", empty). It must not silently fall back to a raw-string comparison. +func TestProductFromChunkMapRejectsDirtyKwd(t *testing.T) { + base := map[string]interface{}{ + "id": "wiki/1", + "doc_id": "d1", + "content_with_weight": "# Alpha", + } + dirtyKwds := []string{"artifact_page", "garbage", ""} + for _, kwd := range dirtyKwds { + c := map[string]interface{}{} + for k, v := range base { + c[k] = v + } + c["compile_kwd"] = kwd + if p, ok := productFromChunkMap(c, "t1", kccommon.VariantWiki); ok { + t.Errorf("dirty compile_kwd %q should be rejected, got product %+v", kwd, p) + } + } + + // Note: wiki_section maps to VariantWiki (page and section share the wiki + // variant); the page/section distinction lives in Meta.kind and is enforced + // by filterWikiPageCandidates, NOT by the variant dirty-row check. So a + // wiki_section row legitimately satisfies a VariantWiki query here — the + // section is dropped later by the kind filter. + + // A clean wiki_page row must pass for VariantWiki. + good := map[string]interface{}{} + for k, v := range base { + good[k] = v + } + good["compile_kwd"] = compileKwdWikiPage + if _, ok := productFromChunkMap(good, "t1", kccommon.VariantWiki); !ok { + t.Errorf("clean wiki_page row should satisfy VariantWiki query") + } +} diff --git a/internal/ingestion/pipeline/pipeline_knowledge_compiler_dsl_test.go b/internal/ingestion/pipeline/pipeline_knowledge_compiler_dsl_test.go index 7a23c9e57f..b0208af05d 100644 --- a/internal/ingestion/pipeline/pipeline_knowledge_compiler_dsl_test.go +++ b/internal/ingestion/pipeline/pipeline_knowledge_compiler_dsl_test.go @@ -242,6 +242,33 @@ func TestKnowledgeCompilerDSL_ParamBinding(t *testing.T) { if p.EnableHistoricalDedup { t.Errorf("EnableHistoricalDedup default = true, want false") } + if p.Plan != nil { + t.Errorf("Plan = %v, want nil when omitted from DSL", *p.Plan) + } + + for _, tc := range []struct { + name string + plan bool + }{ + {name: "mode_a", plan: false}, + {name: "mode_b", plan: true}, + } { + t.Run(tc.name, func(t *testing.T) { + parsed, err := kc.ParseParam(map[string]any{ + "compilation_template_id": "t1", + "plan": tc.plan, + }) + if err != nil { + t.Fatalf("ParseParam: %v", err) + } + if parsed.Plan == nil || *parsed.Plan != tc.plan { + t.Fatalf("Plan = %v, want explicit %t", parsed.Plan, tc.plan) + } + if parsed.PlanEnabled() != tc.plan { + t.Errorf("PlanEnabled() = %t, want %t", parsed.PlanEnabled(), tc.plan) + } + }) + } } // TestKnowledgeCompilerDSL_KindToVariant locks the compilation_template.kind -> diff --git a/internal/ingestion/pipeline/template/ingestion_pipeline_knowledge_compiler.json b/internal/ingestion/pipeline/template/ingestion_pipeline_knowledge_compiler.json index db0ab19d38..dd2f890d62 100644 --- a/internal/ingestion/pipeline/template/ingestion_pipeline_knowledge_compiler.json +++ b/internal/ingestion/pipeline/template/ingestion_pipeline_knowledge_compiler.json @@ -187,9 +187,9 @@ "obj": { "component_name": "Compiler", "params": { + "compilation_template_group_id": "", "llm_id": "", - "variant": "structure", - "language": "English" + "plan": false } }, "upstream": [ diff --git a/internal/ingestion/task/pipeline_executor_defaults_test.go b/internal/ingestion/task/pipeline_executor_defaults_test.go index fc483af648..0f2268c4b7 100644 --- a/internal/ingestion/task/pipeline_executor_defaults_test.go +++ b/internal/ingestion/task/pipeline_executor_defaults_test.go @@ -51,7 +51,7 @@ var builtinComponentParamsGolden = map[string]string{ "qa": "{\"File\": {}, \"Parser:HipSignsRhyme\": {\"docx\": {\"flatten_media_to_text\": false, \"output_format\": \"json\", \"preprocess\": [\"main_content\"], \"suffix\": [\"docx\"], \"vlm\": {}}, \"markdown\": {\"flatten_media_to_text\": false, \"output_format\": \"json\", \"preprocess\": [\"main_content\"], \"suffix\": [\"md\", \"markdown\", \"mdx\"], \"vlm\": {}}, \"pdf\": {\"flatten_media_to_text\": false, \"output_format\": \"json\", \"parse_method\": \"DeepDOC\", \"preprocess\": [\"main_content\"], \"suffix\": [\"pdf\"], \"vlm\": {}}, \"spreadsheet\": {\"flatten_media_to_text\": false, \"output_format\": \"html\", \"parse_method\": \"DeepDOC\", \"preprocess\": [\"main_content\"], \"suffix\": [\"xls\", \"xlsx\", \"csv\"], \"vlm\": {}}, \"text&code\": {\"output_format\": \"json\", \"preprocess\": [\"main_content\"], \"suffix\": [\"txt\"]}}, \"Tokenizer:ColdCloudsDream\": {\"fields\": \"text\", \"filename_embd_weight\": 0.1, \"search_method\": [\"embedding\", \"full_text\"]}, \"QAChunker:TidyCloudsThink\": {}}", "resume": "{\"Extractor:ThreeDrinksAct\": {\"field_name\": \"metadata\", \"frequencyPenaltyEnabled\": true, \"frequency_penalty\": 0.7, \"llm_id\": \"THUDM/GLM-4.1V-9B-Thinking@SILICONFLOW\", \"maxTokensEnabled\": false, \"max_tokens\": 256, \"presencePenaltyEnabled\": true, \"presence_penalty\": 0.4, \"prompts\": [{\"content\": \"Content: {TitleChunker:FlatMiceFix@chunks}\", \"role\": \"user\"}], \"sys_prompt\": \"Act as a precise resume metadata extractor. Extract stable, chunk-supported metadata from the provided resume content.\\n\\nRules:\\n1. Use only information explicitly stated in the content. Do not infer, guess, normalize, or add missing facts.\\n2. The input may be only one chunk of a resume. Extract only what this content directly supports.\\n3. Use only these field names:\\ncandidate_name, gender, phone, email, city, location, nationality, linkedin, github, website, highest_degree, degree_levels, school_names, majors, graduation_years, work_experience_years, current_job_title, job_titles, company_names, job_experience, industries, target_job_titles, target_locations, employment_types, skills, certificates, awards, summary_tags\\n4. Ignore detailed responsibilities, project descriptions, achievement narratives, self-evaluation, and other low-value local details.\\n5. Keep values in the same language as the source text whenever possible.\\n6. Remove duplicates and keep only concise, high-value metadata.\\n7. Return only fields that are explicitly supported by the content. Do not return empty or unsupported fields.\\n\\nField guidance:\\n- highest_degree: highest explicit degree level mentioned\\n- degree_levels: all explicit degree levels mentioned\\n- school_names: explicit school, college, or university names\\n- majors: explicit fields of study\\n- graduation_years: explicit graduation years only\\n- work_experience_years: only if explicitly stated\\n- current_job_title: only if explicitly current or most recent\\n- job_titles: explicit role titles\\n- company_names: explicit employer names\\n- job_experience: concise structured work entries explicitly supported by the content, preferably including title, company, and time information when available\\n- industries: explicit industry names only\\n- target_job_titles: explicit desired roles only\\n- target_locations: explicit desired work locations only\\n- skills: concise, core, search-useful skills explicitly mentioned\\n- certificates: explicit certificate names only\\n- awards: explicit award names only\\n- summary_tags: short, high-value tags strictly supported by the content\\n\\nReturn only the extracted metadata. Do not output explanatory text.\", \"temperature\": 0.1, \"temperatureEnabled\": true, \"tenant_llm_id\": 29, \"topPEnabled\": true, \"top_p\": 0.3, \"auto_keywords\": 0, \"auto_questions\": 0, \"auto_tags\": 0, \"enable_metadata\": 0, \"metadata\": [], \"tag_file_id\": \"\"}, \"File\": {}, \"Parser:HipSignsRhyme\": {\"docx\": {\"flatten_media_to_text\": true, \"output_format\": \"json\", \"preprocess\": [\"main_content\"], \"suffix\": [\"docx\"], \"vlm\": {}}, \"pdf\": {\"flatten_media_to_text\": true, \"output_format\": \"json\", \"parse_method\": \"DeepDOC\", \"preprocess\": [\"main_content\"], \"suffix\": [\"pdf\"], \"vlm\": {}}, \"text&code\": {\"output_format\": \"json\", \"preprocess\": [\"main_content\"], \"suffix\": [\"txt\"]}}, \"TitleChunker:FlatMiceFix\": {\"hierarchy\": 1, \"include_heading_content\": false, \"levels\": [[\"^\\\\s*(?i:(?:\\\\d+[\\\\.\\\\)]\\\\s*)?(?:EDUCATION|ACADEMIC\\\\s*BACKGROUND|ACADEMIC\\\\s*HISTORY|EDUCATIONAL\\\\s*BACKGROUND|RELEVANT\\\\s*COURSEWORK|COURSEWORK|EXPERIENCE|WORK\\\\s*EXPERIENCE|PROFESSIONAL\\\\s*EXPERIENCE|RELEVANT\\\\s*EXPERIENCE|EMPLOYMENT\\\\s*HISTORY|CAREER\\\\s*HISTORY|INTERNSHIP\\\\s*EXPERIENCE|PROJECTS|PROJECT\\\\s*EXPERIENCE|ACADEMIC\\\\s*PROJECTS|PROFESSIONAL\\\\s*PROJECTS|SKILLS|TECHNICAL\\\\s*SKILLS|CORE\\\\s*COMPETENCIES|COMPETENCIES|QUALIFICATIONS|SUMMARY\\\\s*OF\\\\s*QUALIFICATIONS|CERTIFICATIONS|LICENSES|CERTIFICATES|AWARDS|HONORS|HONOURS|ACHIEVEMENTS|PUBLICATIONS|RESEARCH|RESEARCH\\\\s*EXPERIENCE|LEADERSHIP|LEADERSHIP\\\\s*EXPERIENCE|ACTIVITIES|EXTRACURRICULAR\\\\s*ACTIVITIES|ACTIVITIES\\\\s*(?:&|AND)\\\\s*SKILLS|INVOLVEMENT|CAMPUS\\\\s*INVOLVEMENT|VOLUNTEER\\\\s*EXPERIENCE|VOLUNTEERING|COMMUNITY\\\\s*SERVICE|LANGUAGES|INTERESTS|HOBBIES|PROFILE|PROFESSIONAL\\\\s*PROFILE|SUMMARY|PROFESSIONAL\\\\s*SUMMARY|CAREER\\\\s*SUMMARY|OBJECTIVE|CAREER\\\\s*OBJECTIVE|PERSONAL\\\\s*INFORMATION|CONTACT\\\\s*INFORMATION|ADDITIONAL\\\\s*INFORMATION|TRAINING))\\\\s*[::]?\\\\s*$\"], [\"^\\\\s*(?:\\\\d+[\\\\.、\\\\)]\\\\s*)?(?:教育背景|教育经历|学历背景|学术背景|技术背景|工作经历|工作经验|实习经历|项目经历|项目经验|科研经历|研究经历|校园经历|实践经历|专业经历|职业经历|技能|专业技能|技能特长|核心技能|技术栈|个人技能|工作技能|职业技能|技能与评价|技能与自我评价|工作技能与自我评价|职业技能与自我评价|证书|资格证书|职业资格|资质证书|获奖情况|获奖经历|荣誉|荣誉奖项|奖项|科研成果|论文发表|发表论文|领导经历|学生工作|校园活动|社团经历|活动经历|志愿经历|志愿服务|社会实践|语言能力|语言|自我评价|个人评价|自我总结|个人总结|个人优势|个人简介|个人信息|基本信息|联系方式|求职意向|应聘意向|职业目标|求职目标|兴趣爱好|兴趣特长|培训经历|其他信息|附加信息)\\\\s*[::]?\\\\s*$\"]], \"method\": \"hierarchy\"}, \"Tokenizer:KindHandsWin\": {\"fields\": \"text\", \"filename_embd_weight\": 0.1, \"search_method\": [\"embedding\", \"full_text\"]}}", "table": "{\"File\": {}, \"Parser:HipSignsRhyme\": {\"spreadsheet\": {\"flatten_media_to_text\": false, \"output_format\": \"html\", \"parse_method\": \"DeepDOC\", \"preprocess\": [\"main_content\"], \"suffix\": [\"xls\", \"xlsx\", \"csv\"], \"vlm\": {}, \"column_mode\": \"auto\", \"column_roles\": {}, \"column_names\": []}, \"text&code\": {\"output_format\": \"json\", \"preprocess\": [\"main_content\"], \"suffix\": [\"txt\"]}}, \"TableChunker:FastFoxesJump\": {}, \"Tokenizer:DeepLakesShine\": {\"fields\": \"text\", \"filename_embd_weight\": 0.1, \"search_method\": [\"embedding\", \"full_text\"]}}", - "knowledge_compiler": "{\"File\": {}, \"Compiler:KnownSwiftLions\": {\"language\": \"English\", \"llm_id\": \"\", \"variant\": \"structure\"}, \"Parser:HipSignsRhyme\": {\"setups\": {\"doc\": {\"output_format\": \"json\", \"preprocess\": [\"main_content\"], \"suffix\": [\"doc\"]}, \"docx\": {\"flatten_media_to_text\": false, \"output_format\": \"json\", \"preprocess\": [\"main_content\"], \"suffix\": [\"docx\"], \"vlm\": {}}, \"html\": {\"output_format\": \"json\", \"preprocess\": [\"main_content\"], \"suffix\": [\"htm\", \"html\"]}, \"markdown\": {\"flatten_media_to_text\": false, \"output_format\": \"json\", \"preprocess\": [\"main_content\"], \"suffix\": [\"md\", \"markdown\", \"mdx\"], \"vlm\": {}}, \"pdf\": {\"flatten_media_to_text\": false, \"output_format\": \"json\", \"parse_method\": \"DeepDOC\", \"preprocess\": [\"main_content\"], \"suffix\": [\"pdf\"], \"vlm\": {}}, \"spreadsheet\": {\"flatten_media_to_text\": false, \"output_format\": \"html\", \"parse_method\": \"DeepDOC\", \"preprocess\": [\"main_content\"], \"suffix\": [\"xls\", \"xlsx\", \"csv\"], \"vlm\": {}}, \"text&code\": {\"output_format\": \"json\", \"preprocess\": [\"main_content\"], \"suffix\": [\"txt\", \"py\", \"js\", \"java\", \"c\", \"cpp\", \"h\", \"php\", \"go\", \"ts\", \"sh\", \"cs\", \"kt\", \"sql\"]}}}, \"TokenChunker:SixApplesFall\": {\"children_delimiters\": [], \"chunk_token_size\": 512, \"delimiter_mode\": \"delimiter\", \"delimiters\": [\"\\n\", \"!\", \"?\", \"。\", \";\", \"!\", \"?\"], \"image_context_size\": 0, \"overlapped_percent\": 0, \"table_context_size\": 0}}", + "knowledge_compiler": "{\"File\": {}, \"Compiler:KnownSwiftLions\": {\"compilation_template_group_id\": \"\", \"llm_id\": \"\", \"plan\": false}, \"Parser:HipSignsRhyme\": {\"setups\": {\"doc\": {\"output_format\": \"json\", \"preprocess\": [\"main_content\"], \"suffix\": [\"doc\"]}, \"docx\": {\"flatten_media_to_text\": false, \"output_format\": \"json\", \"preprocess\": [\"main_content\"], \"suffix\": [\"docx\"], \"vlm\": {}}, \"html\": {\"output_format\": \"json\", \"preprocess\": [\"main_content\"], \"suffix\": [\"htm\", \"html\"]}, \"markdown\": {\"flatten_media_to_text\": false, \"output_format\": \"json\", \"preprocess\": [\"main_content\"], \"suffix\": [\"md\", \"markdown\", \"mdx\"], \"vlm\": {}}, \"pdf\": {\"flatten_media_to_text\": false, \"output_format\": \"json\", \"parse_method\": \"DeepDOC\", \"preprocess\": [\"main_content\"], \"suffix\": [\"pdf\"], \"vlm\": {}}, \"spreadsheet\": {\"flatten_media_to_text\": false, \"output_format\": \"html\", \"parse_method\": \"DeepDOC\", \"preprocess\": [\"main_content\"], \"suffix\": [\"xls\", \"xlsx\", \"csv\"], \"vlm\": {}}, \"text&code\": {\"output_format\": \"json\", \"preprocess\": [\"main_content\"], \"suffix\": [\"txt\", \"py\", \"js\", \"java\", \"c\", \"cpp\", \"h\", \"php\", \"go\", \"ts\", \"sh\", \"cs\", \"kt\", \"sql\"]}}}, \"TokenChunker:SixApplesFall\": {\"children_delimiters\": [], \"chunk_token_size\": 512, \"delimiter_mode\": \"delimiter\", \"delimiters\": [\"\\n\", \"!\", \"?\", \"。\", \";\", \"!\", \"?\"], \"image_context_size\": 0, \"overlapped_percent\": 0, \"table_context_size\": 0}}", } // Per-template test methods. Each resolves default component params from a diff --git a/web/src/locales/en.ts b/web/src/locales/en.ts index 29eafe2f96..8907f31ba8 100644 --- a/web/src/locales/en.ts +++ b/web/src/locales/en.ts @@ -1912,7 +1912,9 @@ Example: Virtual Hosted Style`, instruction: 'Instruction', globalRules: 'Global rules', globalRulesPlaceholder: 'Input global compilation rules', - plan: 'Plan (grouping wiki pages by topic via LLM)', + plan: 'Plan', + planTip: + 'Off: one wiki page per entity/concept. On: let the LLM decide to combine some entities/concepts into a single wiki page.', raptorTreeSettings: 'RAPTOR tree settings', summarizationPrompt: 'Summarization prompt', maxToken: 'Max token', diff --git a/web/src/locales/zh.ts b/web/src/locales/zh.ts index 22e7eb4839..6ecf82678e 100644 --- a/web/src/locales/zh.ts +++ b/web/src/locales/zh.ts @@ -1599,7 +1599,8 @@ NER:使用 spaCy NER 和基于规则的关键词提取来抽取实体和关系 instruction: 'Instruction', globalRules: '全局规则', globalRulesPlaceholder: '请输入全局编译规则', - plan: 'Plan (LLM 分组合并 wiki 页面)', + plan: 'Plan', + planTip: '关闭:每个实体或概念对应一个 wiki 页面。开启:让 LLM 决定将某些实体/概念组合到单个 wiki 页面。', raptorTreeSettings: 'RAPTOR 树设置', summarizationPrompt: '摘要提示词', maxToken: '最大 token 数', diff --git a/web/src/pages/agent/constant/pipeline.tsx b/web/src/pages/agent/constant/pipeline.tsx index 3b26171529..78032c1c2d 100644 --- a/web/src/pages/agent/constant/pipeline.tsx +++ b/web/src/pages/agent/constant/pipeline.tsx @@ -363,6 +363,7 @@ export const initialExtractorValues = { export const initialCompilationValues = { compilation_template_group_id: '', llm_id: '', + plan: false, outputs: { chunks: { type: 'Array', value: [] }, }, diff --git a/web/src/pages/agent/form/compilation-form/index.tsx b/web/src/pages/agent/form/compilation-form/index.tsx index 1cdbf0fc02..5b1dd00558 100644 --- a/web/src/pages/agent/form/compilation-form/index.tsx +++ b/web/src/pages/agent/form/compilation-form/index.tsx @@ -1,6 +1,8 @@ import { CompilationTemplateFormField } from '@/components/compilation-template-form-field'; import { LargeModelFormField } from '@/components/large-model-form-field'; +import { SwitchFormField } from '@/components/switch-fom-field'; import { Form } from '@/components/ui/form'; +import { useTranslate } from '@/hooks/common-hooks'; import { zodResolver } from '@hookform/resolvers/zod'; import { memo } from 'react'; import { useForm } from 'react-hook-form'; @@ -18,6 +20,7 @@ import { Output } from '../components/output'; export const FormSchema = z.object({ compilation_template_group_id: z.string().optional(), llm_id: z.string().optional(), + plan: z.boolean(), }); export type CompilationFormSchemaType = z.infer; @@ -31,6 +34,7 @@ const CompilationForm = ({ }: INextOperatorForm) => { const defaultValues = useFormValues(initialCompilationValues, node); const ownerTenantId = useOwnerTenantId(); + const { t } = useTranslate('setting'); const form = useForm({ defaultValues, @@ -49,6 +53,11 @@ const CompilationForm = ({ name="llm_id" ownerTenantId={ownerTenantId} > + {!hideOutputs && (