From f37b6fc439800be90891ae8dcbd0b90368f6e92b Mon Sep 17 00:00:00 2001 From: buua436 Date: Mon, 24 Aug 2026 14:28:05 +0800 Subject: [PATCH] fix: account for checkpointed ingestion progress (#18667) Preserve completed component progress when an ingestion task resumes from a checkpoint. Checkpointed operators are skipped and do not emit lifecycle events again, which previously caused the document progress to remain below 100%. --- internal/agent/canvas/checkpoint_store.go | 13 + internal/handler/dataset_artifact.go | 7 +- .../knowledge_compiler/component_test.go | 8 +- .../component/knowledge_compiler/wiki/page.go | 11 + .../knowledge_compiler/wiki/prompt.go | 20 -- .../component/knowledge_compiler/wiki/wiki.go | 81 ++++--- .../knowledge_compiler/wiki/wiki_reduce.go | 189 --------------- .../wiki/wiki_reduce_test.go | 225 ------------------ .../knowledge_compiler/wiki/wiki_test.go | 108 ++++++++- .../service/burst_backpressure_test.go | 2 +- .../service/execute_task_ack_test.go | 28 +-- .../ingestion/service/execute_task_test.go | 10 +- internal/ingestion/service/heartbeat_test.go | 6 +- .../ingestion/service/ingestion_service.go | 27 ++- .../service/ingestor_lifecycle_test.go | 12 +- .../ingestion/service/process_message_test.go | 22 +- .../service/real_consumer_pipeline_test.go | 2 +- .../service/redelivery_counter_test.go | 4 +- internal/ingestion/service/run_task_test.go | 16 +- internal/ingestion/service/test_helpers.go | 10 + internal/router/router.go | 4 +- internal/service/dataset_artifact_service.go | 13 +- .../compilation/utils/parse-wiki-link.test.ts | 18 ++ .../compilation/utils/parse-wiki-link.ts | 4 +- 24 files changed, 301 insertions(+), 539 deletions(-) delete mode 100644 internal/ingestion/component/knowledge_compiler/wiki/wiki_reduce.go delete mode 100644 internal/ingestion/component/knowledge_compiler/wiki/wiki_reduce_test.go create mode 100644 web/src/pages/dataset/compilation/utils/parse-wiki-link.test.ts diff --git a/internal/agent/canvas/checkpoint_store.go b/internal/agent/canvas/checkpoint_store.go index 21dbfc8216..6a330e42f6 100644 --- a/internal/agent/canvas/checkpoint_store.go +++ b/internal/agent/canvas/checkpoint_store.go @@ -65,6 +65,19 @@ func NewRedisCheckPointStoreWithClient(client *redis.Client, ttl time.Duration) return &RedisCheckPointStore{client: client, ttl: ttl} } +// RedisCheckpointExists reports whether a pipeline checkpoint is present for +// id. It is used by ingestion progress handling to distinguish a fresh run +// from a resume: resumed nodes may not emit lifecycle events again, so their +// previous completed progress rows must be retained. +func RedisCheckpointExists(ctx context.Context, id string) (bool, error) { + rc := redis2.Get() + if rc == nil || rc.GetClient() == nil { + return false, errors.New("checkpoint store: redis client not initialized") + } + found, err := rc.GetClient().Exists(ctx, checkpointKeyPrefix+id).Result() + return found > 0, err +} + // Get implements eino's CheckPointStore.Get. Returns (nil, false, nil) when // the key does not exist (redis.Nil) so callers can distinguish "missing" // from "present-but-error". diff --git a/internal/handler/dataset_artifact.go b/internal/handler/dataset_artifact.go index 2d091b192f..4b1a282a3b 100644 --- a/internal/handler/dataset_artifact.go +++ b/internal/handler/dataset_artifact.go @@ -17,6 +17,7 @@ import ( "errors" "net/http" "strconv" + "strings" "ragflow/internal/common" "ragflow/internal/dao" @@ -126,6 +127,7 @@ func (h *DatasetArtifactHandler) ListArtifacts(c *gin.Context) { } // UpdateArtifact handles PUT /artifacts// — edit a wiki page. +// The slug may contain nested path segments, such as location/长社. func (h *DatasetArtifactHandler) UpdateArtifact(c *gin.Context) { user, tenantID, _ := h.datasetOwner(c, c.Param("dataset_id")) if tenantID == "" { @@ -133,7 +135,7 @@ func (h *DatasetArtifactHandler) UpdateArtifact(c *gin.Context) { } datasetID := c.Param("dataset_id") pageType := c.Param("page_type") - slug := c.Param("slug") + slug := strings.TrimPrefix(c.Param("slug"), "/") var req struct { ContentMd string `json:"content_md"` Title string `json:"title"` @@ -196,6 +198,7 @@ func (h *DatasetArtifactHandler) UpdateArtifact(c *gin.Context) { } // GetArtifact handles GET /artifacts// — single wiki page. +// The slug may contain nested path segments, such as location/长社. func (h *DatasetArtifactHandler) GetArtifact(c *gin.Context) { _, tenantID, _ := h.datasetOwner(c, c.Param("dataset_id")) if tenantID == "" { @@ -203,7 +206,7 @@ func (h *DatasetArtifactHandler) GetArtifact(c *gin.Context) { } datasetID := c.Param("dataset_id") pageType := c.Param("page_type") - slug := c.Param("slug") + slug := strings.TrimPrefix(c.Param("slug"), "/") detail, err := h.svc.GetWikiPage(c.Request.Context(), tenantID, datasetID, pageType, slug) if err != nil { common.ErrorWithCode(c, common.CodeDataError, err.Error()) diff --git a/internal/ingestion/component/knowledge_compiler/component_test.go b/internal/ingestion/component/knowledge_compiler/component_test.go index b5369a9e15..2636b3d94c 100644 --- a/internal/ingestion/component/knowledge_compiler/component_test.go +++ b/internal/ingestion/component/knowledge_compiler/component_test.go @@ -673,7 +673,7 @@ func TestKnowledgeCompiler_Wiki_UpdateMergesExistingPage(t *testing.T) { TenantID: tenantID, WikiPages: wikiStoreTestStub{page: &common.WikiPageCandidate{ ID: "existing-1", - Slug: "entity/alpha", + Slug: "entity/person/alpha", Title: "Alpha", PageType: "entity", Topic: "Alpha", @@ -714,7 +714,7 @@ func TestKnowledgeCompiler_Wiki_UpdateMergesExistingPage(t *testing.T) { if !ok { continue } - if cm["compile_kwd"] == "wiki_page" && cm["kc_kind"] == "page" && cm["slug_kwd"] == "entity/alpha" { + if cm["compile_kwd"] == "wiki_page" && cm["kc_kind"] == "page" && cm["slug_kwd"] == "entity/person/alpha" { page = cm break } @@ -840,7 +840,7 @@ func (wikiUpdateChat) Chat(_ context.Context, req common.ChatRequest) (*common.C }`}, nil case req.JSONMode && strings.Contains(req.UserPrompt, "Return a JSON compilation plan"): return &common.ChatResponse{Content: `{ - "pages":[{"action":"CREATE","slug":"entity/alpha","title":"Alpha","page_type":"entity","topic":"Alpha","entity_names":["Alpha"],"related_kb_pages":[],"priority":1,"lead":"Alpha overview","sections":[{"heading":"Overview","points":["Alpha overview"]}]}], + "pages":[{"action":"CREATE","slug":"entity/person/alpha","title":"Alpha","page_type":"entity","topic":"Alpha","entity_names":["Alpha"],"related_kb_pages":[],"priority":1,"lead":"Alpha overview","sections":[{"heading":"Overview","points":["Alpha overview"]}]}], "estimated_page_count":1, "compilation_notes":"ok" }`}, nil @@ -849,7 +849,7 @@ func (wikiUpdateChat) Chat(_ context.Context, req common.ChatRequest) (*common.C case strings.Contains(req.UserPrompt, "Existing page content"): return &common.ChatResponse{Content: "# Alpha\n\nAlpha launched a new process in 2026.\n\n## See also\n\n[[concept/alpha-protocol|Alpha Protocol]]"}, nil default: - return &common.ChatResponse{Content: `{"action":"UPDATE","slug":"entity/alpha","reason":"same entity"}`}, nil + return &common.ChatResponse{Content: `{"action":"UPDATE","slug":"entity/person/alpha","reason":"same entity"}`}, nil } } diff --git a/internal/ingestion/component/knowledge_compiler/wiki/page.go b/internal/ingestion/component/knowledge_compiler/wiki/page.go index 03b9e56fdf..a13b6d7c14 100644 --- a/internal/ingestion/component/knowledge_compiler/wiki/page.go +++ b/internal/ingestion/component/knowledge_compiler/wiki/page.go @@ -117,6 +117,17 @@ func buildPageProducts(tenantID, docID, page string, sourceChunkIDs []string) [] }}) } +// entityPageSlug returns the stable page identity for an extracted entity. +// Type and name use separate path segments so the identity is unambiguous. +func entityPageSlug(name, entityType string) string { + nameSlug := slugify(strings.Join(strings.Fields(name), " ")) + entityType = strings.Join(strings.Fields(entityType), " ") + if entityType == "" { + return "entity/" + nameSlug + } + return "entity/" + slugify(entityType) + "/" + nameSlug +} + // buildWikiPageProducts turns multiple wiki page drafts into page + section // products. Each page result becomes one page product and one section product // per heading in its Markdown body. diff --git a/internal/ingestion/component/knowledge_compiler/wiki/prompt.go b/internal/ingestion/component/knowledge_compiler/wiki/prompt.go index 95ef9d80e4..206df964ec 100644 --- a/internal/ingestion/component/knowledge_compiler/wiki/prompt.go +++ b/internal/ingestion/component/knowledge_compiler/wiki/prompt.go @@ -11,26 +11,6 @@ const wikiMapSystem = `You are a knowledge extraction engine. Extract structured const wikiPlanSystem = `You are a knowledge compilation planner. Given structured knowledge, produce a wiki page plan. Return ONLY valid JSON.` -const wikiReduceEntityDisambiguateSystem = `You are a knowledge canonicalization engine. Decide whether two named entities refer to the same real-world concept. Return ONLY valid JSON.` - -const wikiReduceEntityDisambiguateUserTemplate = `## Entity A -{entity_a} - -## Entity B -{entity_b} - -Return JSON: -{ - "merge": true, - "reason": "string" -} - -Rules: -- merge=true only when A and B are the same real-world entity (e.g. aliases, abbreviations, spelling variants of the same thing). -- merge=false when they are distinct concepts that merely co-occur. -- Prefer false when ambiguous. -- Return ONLY the JSON object.` - const wikiMapUserTemplate = `## Document context Document id: {doc_id} Batch contains {chunk_count} packed chunk(s). Each chunk is introduced by a diff --git a/internal/ingestion/component/knowledge_compiler/wiki/wiki.go b/internal/ingestion/component/knowledge_compiler/wiki/wiki.go index 717cbac419..c9bb1ca648 100644 --- a/internal/ingestion/component/knowledge_compiler/wiki/wiki.go +++ b/internal/ingestion/component/knowledge_compiler/wiki/wiki.go @@ -391,19 +391,6 @@ func (p *wikiPipeline) run() error { zap.Strings("concepts", conceptsDebug(reduced.Concepts)), zap.Int("claims", len(reduced.Claims)), zap.Strings("relations", relationsDebug(reduced.Relations))) - // Layer embedding + LLM disambiguation onto the exact-merged entities - // (REDUCE enhancement; concepts keep exact dedup). Degrades to a no-op when - // the embedder/chat seams are unavailable. - preDedup := reduced.Entities - p.reduced.Entities = p.dedupeEntities(preDedup) - appcommon.Debug("wiki: REDUCE embedding+LLM dedup done", - zap.String("dataset_id", p.datasetID), - zap.String("doc_id", p.runKey()), - zap.Int("pre_dedup_entities", len(preDedup)), - zap.Int("post_dedup_entities", len(p.reduced.Entities)), - zap.Strings("pre_dedup_names", entityNamesDebug(preDedup)), - zap.Strings("post_dedup_names", entityNamesDebug(p.reduced.Entities)), - zap.Strings("post_dedup_entities", entitiesDebug(p.reduced.Entities))) appcommon.Info("wiki: REDUCE done", zap.String("dataset_id", p.datasetID), zap.String("doc_id", p.runKey()), @@ -1528,11 +1515,13 @@ func reduceExtracts(extracts []wikiExtract) wikiExtract { } entities := map[entityKey]*wikiEntity{} concepts := map[conceptKey]*wikiConcept{} - seenRelations := map[string]bool{} + relationIndexes := map[string]int{} seenTopics := map[string]bool{} for _, ex := range extracts { for _, e := range ex.Entities { + e.Name = strings.Join(strings.Fields(e.Name), " ") + e.Type = strings.Join(strings.Fields(e.Type), " ") key := entityKey{Name: normKey(e.Name), Type: normKey(e.Type)} if cur, ok := entities[key]; ok { cur.Aliases = mergeStrings(cur.Aliases, e.Aliases) @@ -1543,8 +1532,6 @@ func reduceExtracts(extracts []wikiExtract) wikiExtract { continue } item := e - item.Name = strings.TrimSpace(item.Name) - item.Type = strings.TrimSpace(item.Type) item.Aliases = uniqueStrings(item.Aliases) item.SourceChunkIDs = uniqueStrings(item.SourceChunkIDs) entities[key] = &item @@ -1570,11 +1557,12 @@ func reduceExtracts(extracts []wikiExtract) wikiExtract { } for _, r := range ex.Relations { key := normKey(r.From) + "\x00" + normKey(r.Type) + "\x00" + normKey(r.To) - if seenRelations[key] { + if cur, ok := relationIndexes[key]; ok { + out.Relations[cur].SourceChunkIDs = mergeStrings(out.Relations[cur].SourceChunkIDs, r.SourceChunkIDs) continue } r.SourceChunkIDs = uniqueStrings(r.SourceChunkIDs) - seenRelations[key] = true + relationIndexes[key] = len(out.Relations) out.Relations = append(out.Relations, r) } for _, t := range ex.Topics { @@ -1597,10 +1585,10 @@ func reduceExtracts(extracts []wikiExtract) wikiExtract { keys = append(keys, k) } sort.Slice(keys, func(i, j int) bool { - if keys[i].Name == keys[j].Name { - return keys[i].Type < keys[j].Type + if keys[i].Type == keys[j].Type { + return keys[i].Name < keys[j].Name } - return keys[i].Name < keys[j].Name + return keys[i].Type < keys[j].Type }) for _, k := range keys { out.Entities = append(out.Entities, *entities[k]) @@ -1708,13 +1696,9 @@ func normalizeWikiPlan(plan wikiPlan, docID string, reduced wikiExtract) wikiPla func normalizeWikiPlanPages(pages []wikiPlanPage, reduced wikiExtract) []wikiPlanPage { // existing maps a plan slug to its index in out (exact-slug dedup); byTitle - // additionally collapses pages that share the same page_type + normalized - // title, so a single batch never yields two "吕布" pages whose slugs only - // differ in transliteration (lu-bu vs lv-bu). Scheme A (§5 of the duplicates - // research): page identity is page_type/slug, so the key is page-type-scoped - // (a same-title concept and topic may legitimately coexist) and the first - // page seen (LLM emission order) is kept, folding the duplicate's RelatedKB - // into the survivor so its outlinks are not silently dropped. + // additionally collapses non-canonical pages that share the same page_type + + // normalized title. Single-entity pages use their canonical slug instead so + // same-named entities of different types remain distinct. existing := map[string]int{} byTitle := map[string]int{} out := make([]wikiPlanPage, 0, len(pages)) @@ -1728,6 +1712,9 @@ func normalizeWikiPlanPages(pages []wikiPlanPage, reduced wikiExtract) []wikiPla continue } key := wikiTitleKey(page.PageType, page.Title) + if page.PageType == "entity" && len(page.EntityNames) == 1 { + key = page.PageType + "\x00" + page.Slug + } if idx, ok := byTitle[key]; ok { out[idx].RelatedKB = uniqueStrings(append(out[idx].RelatedKB, page.RelatedKB...)) continue @@ -1745,7 +1732,11 @@ func normalizeWikiPlanPages(pages []wikiPlanPage, reduced wikiExtract) []wikiPla } page = normalizeWikiPlanPage(page) idx := len(out) - byTitle[wikiTitleKey(page.PageType, page.Title)] = idx + key := wikiTitleKey(page.PageType, page.Title) + if page.PageType == "entity" && len(page.EntityNames) == 1 { + key = page.PageType + "\x00" + page.Slug + } + byTitle[key] = idx existing[page.Slug] = idx out = append(out, page) } @@ -1868,6 +1859,8 @@ func (p *wikiPipeline) buildModeAPlan() wikiPlan { fullSlugFor := func(name, pageType string) string { return pageType + "/" + normalizeWikiSlugHyphens(slugify(name)) } + entitySlugFor := func(e wikiEntity) string { return entityPageSlug(e.Name, e.Type) } + entitySlugsByName := map[string][]string{} slugToIndex := map[string]int{} var pages []wikiPlanPage addPage := func(slug, title, pageType string, entityNames []string) { @@ -1891,7 +1884,8 @@ func (p *wikiPipeline) buildModeAPlan() wikiPlan { if name == "" { continue } - slug := fullSlugFor(name, "entity") + slug := entitySlugFor(e) + entitySlugsByName[normKey(name)] = append(entitySlugsByName[normKey(name)], slug) addPage(slug, name, "entity", []string{name}) } for _, c := range reduced.Concepts { @@ -1909,8 +1903,13 @@ func (p *wikiPipeline) buildModeAPlan() wikiPlan { if from == "" || to == "" { continue } - fromSlug := fullSlugFor(from, "entity") - toSlug := fullSlugFor(to, "entity") + fromSlugs := uniqueStrings(entitySlugsByName[normKey(from)]) + toSlugs := uniqueStrings(entitySlugsByName[normKey(to)]) + if len(fromSlugs) != 1 || len(toSlugs) != 1 { + continue + } + fromSlug := fromSlugs[0] + toSlug := toSlugs[0] if fromSlug == toSlug { continue } @@ -1943,7 +1942,7 @@ func buildWikiFallbackPages(reduced wikiExtract) []wikiPlanPage { for _, e := range reduced.Entities { appendPage(wikiPlanPage{ Action: "CREATE", - Slug: "entity/" + slugify(e.Name), + Slug: entityPageSlug(e.Name, e.Type), Title: e.Name, PageType: "entity", Topic: e.Name, @@ -3079,6 +3078,22 @@ func l2Norm32(v []float32) float64 { return math.Sqrt(s) } +func cosine32(a, b []float32) float64 { + if len(a) != len(b) { + return 0 + } + na := l2Norm32(a) + nb := l2Norm32(b) + if na == 0 || nb == 0 { + return 0 + } + var dot float64 + for i := range a { + dot += float64(a[i]) * float64(b[i]) + } + return dot / (na * nb) +} + func maxInt(a, b int) int { if a > b { return a diff --git a/internal/ingestion/component/knowledge_compiler/wiki/wiki_reduce.go b/internal/ingestion/component/knowledge_compiler/wiki/wiki_reduce.go deleted file mode 100644 index 873e71c410..0000000000 --- a/internal/ingestion/component/knowledge_compiler/wiki/wiki_reduce.go +++ /dev/null @@ -1,189 +0,0 @@ -package wiki - -import ( - "strings" - - "ragflow/internal/ingestion/component/knowledge_compiler/common" -) - -// This file implements the REDUCE-stage canonical-entity enhancement that -// narrows the gap with Python's wiki.py canonicalization: -// -// - entities with distinct names but high embedding similarity are treated as -// ambiguous and sent to an LLM merge decision (collapsing near-duplicates); -// - concepts keep exact-term dedup, matching Python's current semantic (any -// embedding/LLM dedup for concepts must be a separate new capability with -// its own quality bar, not an alignment claim). -// -// The exact-key merge in reduceExtracts stays the deterministic baseline; this -// step layers embedding + LLM disambiguation on top. When the embedder or chat -// seam is unavailable, entities pass through unchanged (degrade gracefully). - -// wikiEntityMergeThreshold is the embedding-cosine similarity at or above which -// two distinct-name entities are considered ambiguous and sent to the LLM merge -// decision. It is deliberately high so only genuinely similar candidates reach -// the LLM. -const wikiEntityMergeThreshold = 0.85 - -// wikiEntityMergeMaxCalls caps how many LLM disambiguation calls a single -// REDUCE run may make, bounding the cost on entity-dense documents. -const wikiEntityMergeMaxCalls = 16 - -// wikiEntityMergeMaxCandidates caps how many candidate partners one entity is -// checked against to keep the pairwise scan bounded. -const wikiEntityMergeMaxCandidates = 8 - -// dedupeEntities returns a copy of in with ambiguous near-duplicate entities -// collapsed via LLM disambiguation. It is a no-op when fewer than two entities -// are present or when deps.Embed / deps.Chat are unavailable. -func (p *wikiPipeline) dedupeEntities(in []wikiEntity) []wikiEntity { - if len(in) < 2 || p.deps.Embed == nil || p.deps.Chat == nil { - return in - } - names := make([]string, len(in)) - for i, e := range in { - names[i] = e.Name - } - vecs, err := p.deps.Embed.Encode(p.ctx, names) - if err != nil || len(vecs) != len(in) { - return in - } - - // Canonical entity per input index: which index owns the final entity. - canon := make([]int, len(in)) - for i := range canon { - canon[i] = i - } - llmCalls := 0 - - // Greedy best-partner scan in input order (already deterministic after - // reduceExtracts sorts by name). - for i := 0; i < len(in) && llmCalls < wikiEntityMergeMaxCalls; i++ { - if canon[i] != i { - // Already merged into another canonical entity. - continue - } - bestIdx, bestSim := -1, -1.0 - checked := 0 - for j := 0; j < len(in) && checked < wikiEntityMergeMaxCandidates; j++ { - if i == j { - continue - } - if canon[j] != j { - // Consumed by an earlier merge; never a standalone partner. - continue - } - if normKey(in[i].Name) == normKey(in[j].Name) { - // Exact-name duplicates are already merged by reduceExtracts; - // never treat them as a pair here. - continue - } - // Same-type-only candidate filtering (Python canonicalizes entities - // within the same type). Two entities with provably different types - // are never ambiguous regardless of embedding similarity. An empty - // type is treated as compatible (cannot prove a difference). - if in[i].Type != "" && in[j].Type != "" && !strings.EqualFold(in[i].Type, in[j].Type) { - continue - } - checked++ - sim := cosine32(vecs[i], vecs[j]) - if sim >= wikiEntityMergeThreshold && sim > bestSim { - bestSim = sim - bestIdx = j - } - } - if bestIdx < 0 { - continue - } - // i and bestIdx are guaranteed standalone by the loop guards above. - // Count the call BEFORE issuing it so a persistent failure cannot drive - // unbounded external requests: the llmCalls budget is consumed even when - // the request fails. The outer loop's `llmCalls < max` guard then stops - // further iterations once the budget is exhausted. - llmCalls++ - merge, err := p.llmMergeEntityDecision(in[i], in[bestIdx]) - if err != nil { - // A failed disambiguation call should not abort the whole REDUCE; - // keep the entities separate and move on. - continue - } - if !merge { - continue - } - // Merge j into i: i is canonical, j is consumed. - canon[bestIdx] = i - in[i].Aliases = mergeStrings(in[i].Aliases, in[bestIdx].Aliases) - if in[i].Name != in[bestIdx].Name { - in[i].Aliases = mergeStrings(in[i].Aliases, []string{in[bestIdx].Name}) - } - in[i].SourceChunkIDs = mergeStrings(in[i].SourceChunkIDs, in[bestIdx].SourceChunkIDs) - if in[i].Type == "" { - in[i].Type = in[bestIdx].Type - } - } - - out := make([]wikiEntity, 0, len(in)) - for i := range in { - if canon[i] != i { - continue - } - out = append(out, in[i]) - } - return out -} - -// llmMergeEntityDecision asks the chat seam whether two distinct-name entities -// refer to the same real-world concept. Returns true to merge. -func (p *wikiPipeline) llmMergeEntityDecision(a, b wikiEntity) (bool, error) { - // Pass retryMax=0: dedupeEntities budgets llmCalls itself and counts one - // decision per call; letting GenJSON retry a failing chat seam would multiply - // each decision by (1 + jsonRetryMax) and blow through the budget (a - // persistent failure would issue 64 calls instead of <=16). - raw, err := common.GenJSON(p.ctx, p.deps.Chat, common.ChatRequest{ - LLMID: p.llmID, - SystemPrompt: wikiReduceEntityDisambiguateSystem, - UserPrompt: renderWikiTemplate(wikiReduceEntityDisambiguateUserTemplate, map[string]string{ - "entity_a": mustPrettyJSON(a), - "entity_b": mustPrettyJSON(b), - }), - }, 0) - if err != nil { - return false, err - } - v, ok := raw["merge"] - if !ok { - return false, nil - } - return toBoolValue(v), nil -} - -// toBoolValue interprets a loosely-typed boolean field returned by the LLM (the -// model may emit JSON true or a string like "true"/"yes"). -func toBoolValue(v any) bool { - switch x := v.(type) { - case bool: - return x - case string: - switch strings.ToLower(strings.TrimSpace(x)) { - case "true", "yes", "1", "same", "merge": - return true - } - case float64: - return x != 0 - } - return false -} - -// cosine32 computes the cosine similarity between two float32 vectors. -func cosine32(a, b []float32) float64 { - na := l2Norm32(a) - nb := l2Norm32(b) - if na == 0 || nb == 0 { - return 0 - } - var dot float64 - for i := 0; i < len(a) && i < len(b); i++ { - dot += float64(a[i]) * float64(b[i]) - } - return dot / (na * nb) -} diff --git a/internal/ingestion/component/knowledge_compiler/wiki/wiki_reduce_test.go b/internal/ingestion/component/knowledge_compiler/wiki/wiki_reduce_test.go deleted file mode 100644 index 6f621535fe..0000000000 --- a/internal/ingestion/component/knowledge_compiler/wiki/wiki_reduce_test.go +++ /dev/null @@ -1,225 +0,0 @@ -package wiki - -import ( - "context" - "errors" - "testing" - - "ragflow/internal/ingestion/component/knowledge_compiler/common" -) - -// TestDedupeEntities_NoSeamIsNoop verifies dedupeEntities degrades to a no-op -// when the embedder or chat seam is nil (M1-style unit safety). -func TestDedupeEntities_NoSeamIsNoop(t *testing.T) { - p := &wikiPipeline{ctx: context.Background()} - in := []wikiEntity{ - {Name: "Alpha", SourceChunkIDs: []string{"c1"}}, - {Name: "Alpha Corp", SourceChunkIDs: []string{"c2"}}, - } - got := p.dedupeEntities(in) - if len(got) != 2 { - t.Fatalf("got %d entities, want 2 (no-op without seams)", len(got)) - } -} - -// TestDedupeEntities_LLMMergesAmbiguousPair verifies two distinct-name entities -// with high embedding similarity are collapsed into one canonical entity via the -// LLM merge decision, with aliases and provenance merged. -func TestDedupeEntities_LLMMergesAmbiguousPair(t *testing.T) { - p := &wikiPipeline{ - ctx: context.Background(), - llmID: "llm1", - deps: common.Deps{ - Chat: reconcileChatStub{resp: `{"merge":true,"reason":"same company"}`}, - Embed: mergeEmbedStub{}, - }, - } - in := []wikiEntity{ - {Name: "Alpha Inc", Type: "org", SourceChunkIDs: []string{"c1"}}, - {Name: "Alpha Incorporated", Type: "org", SourceChunkIDs: []string{"c2"}}, - } - got := p.dedupeEntities(in) - if len(got) != 1 { - t.Fatalf("got %d entities, want 1 (LLM merge)", len(got)) - } - if got[0].Name != "Alpha Inc" { - t.Fatalf("canonical name = %q, want Alpha Inc", got[0].Name) - } - if len(got[0].SourceChunkIDs) != 2 { - t.Fatalf("provenance = %#v, want 2 chunk ids", got[0].SourceChunkIDs) - } - if len(got[0].Aliases) == 0 { - t.Fatalf("aliases not merged: %#v", got[0].Aliases) - } -} - -// TestDedupeEntities_LLMRejectsDistinct verifies the LLM rejecting a merge keeps -// both entities distinct. -func TestDedupeEntities_LLMRejectsDistinct(t *testing.T) { - p := &wikiPipeline{ - ctx: context.Background(), - llmID: "llm1", - deps: common.Deps{ - Chat: reconcileChatStub{resp: `{"merge":false,"reason":"distinct products"}`}, - Embed: mergeEmbedStub{}, - }, - } - in := []wikiEntity{ - {Name: "Alpha", SourceChunkIDs: []string{"c1"}}, - {Name: "Beta", SourceChunkIDs: []string{"c2"}}, - } - got := p.dedupeEntities(in) - if len(got) != 2 { - t.Fatalf("got %d entities, want 2 (LLM rejected merge)", len(got)) - } -} - -// TestDedupeEntities_ExactNameIsNotAmbiguous verifies entities with identical -// normalized names (already collapsed by reduceExtracts before this stage) are -// not treated as ambiguous: they pass through untouched and no LLM call is made -// for the exact-name pair. -func TestDedupeEntities_ExactNameIsNotAmbiguous(t *testing.T) { - p := &wikiPipeline{ - ctx: context.Background(), - llmID: "llm1", - deps: common.Deps{ - Chat: reconcileChatStub{resp: `{"merge":true}`}, - Embed: mergeEmbedStub{}, - }, - } - in := []wikiEntity{ - {Name: "Alpha", SourceChunkIDs: []string{"c1"}}, - {Name: "Alpha", SourceChunkIDs: []string{"c2"}}, - } - got := p.dedupeEntities(in) - // Exact-name duplicates are out of scope for the embedding step; both are - // kept unchanged (they would already be one entity after reduceExtracts). - if len(got) != 2 { - t.Fatalf("got %d entities, want 2 (exact-name pairs are not ambiguous)", len(got)) - } -} - -// TestDedupeEntities_ConceptStaysExact validates the REDUCE boundary: concept -// dedup must remain exact (no embedding/LLM). This test guards that the entity -// enhancement never touches concepts. -func TestReduceExtracts_ConceptsStayExact(t *testing.T) { - reduced := reduceExtracts([]wikiExtract{ - {Concepts: []wikiConcept{{Term: "RAG", Definition: "d1", SourceChunkIDs: []string{"c1"}}}}, - {Concepts: []wikiConcept{{Term: "Retrieval Augmented Generation", Definition: "d2", SourceChunkIDs: []string{"c2"}}}}, - }) - if len(reduced.Concepts) != 2 { - t.Fatalf("concepts = %d, want 2 (exact-term dedup must not collapse distinct terms)", len(reduced.Concepts)) - } -} - -// TestDedupeEntities_FailingChatIsBudgeted locks F4: a chat seam that -// persistently fails must not drive unbounded external calls. The llmCalls -// budget is consumed before the request, so the loop stops after at most -// wikiEntityMergeMaxCalls attempts. -func TestDedupeEntities_FailingChatIsBudgeted(t *testing.T) { - calls := 0 - p := &wikiPipeline{ - ctx: context.Background(), - llmID: "llm1", - deps: common.Deps{ - Chat: chatFunc(func(_ context.Context, _ common.ChatRequest) (*common.ChatResponse, error) { - calls++ - return nil, errors.New("llm down") - }), - Embed: mergeEmbedStub{}, - }, - } - // 20 entities with distinct names but identical embeddings => every pair is - // ambiguous and would trigger an LLM call. - in := make([]wikiEntity, 0, 20) - for i := 0; i < 20; i++ { - in = append(in, wikiEntity{Name: "Entity " + itoa(i), Type: "person"}) - } - got := p.dedupeEntities(in) - if calls > wikiEntityMergeMaxCalls { - t.Fatalf("chat calls = %d, want <= %d despite persistent failures", calls, wikiEntityMergeMaxCalls) - } - // No merges happen because every call fails, so all 20 entities survive. - if len(got) != 20 { - t.Fatalf("entities = %d, want 20 (no merges on failure)", len(got)) - } -} - -// TestDedupeEntities_CrossTypeHighSimDoesNotCallLLM locks F5: entities with -// provably different types must never be treated as ambiguous, even with -// identical embeddings, so no LLM call is made for them. -func TestDedupeEntities_CrossTypeHighSimDoesNotCallLLM(t *testing.T) { - calls := 0 - p := &wikiPipeline{ - ctx: context.Background(), - llmID: "llm1", - deps: common.Deps{ - Chat: chatFunc(func(_ context.Context, _ common.ChatRequest) (*common.ChatResponse, error) { - calls++ - return &common.ChatResponse{Content: `{"merge":true}`}, nil - }), - Embed: mergeEmbedStub{}, - }, - } - // Both embed to [1,1,1] (identical vectors, cosine 1.0) but types differ. - in := []wikiEntity{ - {Name: "Alpha", Type: "person"}, - {Name: "Beta Corp", Type: "org"}, - } - got := p.dedupeEntities(in) - if calls != 0 { - t.Fatalf("chat calls = %d, want 0 (cross-type pairs must not reach the LLM)", calls) - } - if len(got) != 2 { - t.Fatalf("entities = %d, want 2 (cross-type entities must stay distinct)", len(got)) - } -} - -// mergeEmbedStub returns embeddings where identical names share a vector and -// distinct names are far apart (cosine ~0), so it can drive both the ambiguous -// and distinct test paths deterministically. -type mergeEmbedStub struct{} - -func (mergeEmbedStub) Encode(_ context.Context, texts []string) ([][]float32, error) { - out := make([][]float32, len(texts)) - for i, text := range texts { - // "Alpha Inc" and "Alpha Incorporated" both contain "alpha" -> same - // vector; "Beta" differs. - if containsFold(text, "beta") { - out[i] = []float32{1, 0, 0} - continue - } - out[i] = []float32{1, 1, 1} - } - return out, nil -} - -func (mergeEmbedStub) Dimensions() int { return 3 } - -func containsFold(s, sub string) bool { - return len(s) >= len(sub) && (len(sub) == 0 || indexFold(s, sub) >= 0) -} - -func indexFold(s, sub string) int { - if sub == "" { - return 0 - } - ls := toLowerASCII(s) - lsub := toLowerASCII(sub) - for i := 0; i+len(lsub) <= len(ls); i++ { - if ls[i:i+len(lsub)] == lsub { - return i - } - } - return -1 -} - -func toLowerASCII(s string) string { - b := []byte(s) - for i := range b { - if b[i] >= 'A' && b[i] <= 'Z' { - b[i] += 'a' - 'A' - } - } - return string(b) -} diff --git a/internal/ingestion/component/knowledge_compiler/wiki/wiki_test.go b/internal/ingestion/component/knowledge_compiler/wiki/wiki_test.go index 3402354d60..11b12dfcec 100644 --- a/internal/ingestion/component/knowledge_compiler/wiki/wiki_test.go +++ b/internal/ingestion/component/knowledge_compiler/wiki/wiki_test.go @@ -2,6 +2,7 @@ package wiki import ( "context" + "slices" "strings" "sync" "testing" @@ -24,8 +25,8 @@ func TestReduceExtracts_MergesProvenance(t *testing.T) { if len(reduced.Entities) != 1 { t.Fatalf("entities=%d, want 1", len(reduced.Entities)) } - if ids := reduced.Entities[0].SourceChunkIDs; len(ids) != 2 { - t.Fatalf("entity provenance = %#v, want 2 chunk ids", ids) + if ids := reduced.Entities[0].SourceChunkIDs; !slices.Equal(ids, []string{"c1", "c2"}) { + t.Fatalf("entity provenance = %#v, want [c1 c2]", ids) } if len(reduced.Claims) != 2 { t.Fatalf("claims=%d, want 2", len(reduced.Claims)) @@ -209,6 +210,20 @@ func TestNormalizeWikiPlanPages_DoesNotMergeAcrossTypes(t *testing.T) { } } +func TestNormalizeWikiPlanPages_DoesNotMergeTypedEntitiesWithSameTitle(t *testing.T) { + pages := []wikiPlanPage{ + {Slug: "entity/fruit/苹果", Title: "苹果", PageType: "entity", EntityNames: []string{"苹果"}}, + {Slug: "entity/company/苹果", Title: "苹果", PageType: "entity", EntityNames: []string{"苹果"}}, + } + got := normalizeWikiPlanPages(pages, wikiExtract{Entities: []wikiEntity{ + {Name: "苹果", Type: "fruit"}, + {Name: "苹果", Type: "company"}, + }}) + if len(got) != 2 { + t.Fatalf("normalizeWikiPlanPages = %#v, want both typed entity pages", got) + } +} + func TestMergePlanCandidates_DeduplicatesWithoutLLMMerge(t *testing.T) { p := &wikiPipeline{ docID: "doc-1", @@ -418,3 +433,92 @@ func TestMergeWikiPageContent_PreservesShortExistingPage(t *testing.T) { t.Fatalf("merged page dropped incoming content: %q", merged) } } + +func TestReduceExtracts_MergesDuplicateEntitySlug(t *testing.T) { + reduced := reduceExtracts([]wikiExtract{ + {Entities: []wikiEntity{{Name: "曹操", Type: "person", Aliases: []string{"孟德"}, SourceChunkIDs: []string{"c1"}}}}, + {Entities: []wikiEntity{{Name: "曹操", Type: "person", SourceChunkIDs: []string{"c2"}}}}, + }) + if len(reduced.Entities) != 1 { + t.Fatalf("entities = %d, want 1", len(reduced.Entities)) + } + if got := reduced.Entities[0].SourceChunkIDs; !slices.Equal(got, []string{"c1", "c2"}) { + t.Fatalf("source chunk ids = %#v, want [c1 c2]", got) + } + if got := reduced.Entities[0].Aliases; len(got) != 1 || got[0] != "孟德" { + t.Fatalf("aliases = %#v, want [孟德]", got) + } +} + +func TestReduceExtracts_DifferentEntityTypesKeepDifferentSlugs(t *testing.T) { + reduced := reduceExtracts([]wikiExtract{ + {Entities: []wikiEntity{{Name: "苹果", Type: "fruit"}}}, + {Entities: []wikiEntity{{Name: "苹果", Type: "company"}}}, + }) + if len(reduced.Entities) != 2 { + t.Fatalf("entities = %d, want 2", len(reduced.Entities)) + } + if got := entityPageSlug(reduced.Entities[0].Name, reduced.Entities[0].Type); got == entityPageSlug(reduced.Entities[1].Name, reduced.Entities[1].Type) { + t.Fatalf("different entity types have the same slug %q", got) + } +} + +func TestReduceExtracts_EntityIdentityDoesNotCollideAtHyphenBoundary(t *testing.T) { + reduced := reduceExtracts([]wikiExtract{{Entities: []wikiEntity{ + {Name: "bar-baz", Type: "foo"}, + {Name: "baz", Type: "foo-bar"}, + }}}) + if len(reduced.Entities) != 2 { + t.Fatalf("entities = %#v, want two distinct identities", reduced.Entities) + } + first := entityPageSlug("bar-baz", "foo") + second := entityPageSlug("baz", "foo-bar") + if first == second { + t.Fatalf("entity slugs collide: %q", first) + } +} + +func TestReduceExtracts_NormalizesEntityWhitespace(t *testing.T) { + reduced := reduceExtracts([]wikiExtract{{Entities: []wikiEntity{ + {Name: "John Smith", Type: "person"}, + {Name: "John Smith", Type: "person"}, + }}}) + if len(reduced.Entities) != 1 { + t.Fatalf("entities = %#v, want whitespace-equivalent names merged", reduced.Entities) + } +} + +func TestReduceExtracts_DoesNotMergeSimilarNames(t *testing.T) { + reduced := reduceExtracts([]wikiExtract{ + {Entities: []wikiEntity{{Name: "Alpha", Type: "org"}}}, + {Entities: []wikiEntity{{Name: "Alpha Incorporated", Type: "org"}}}, + }) + if len(reduced.Entities) != 2 { + t.Fatalf("entities = %d, want 2; REDUCE must not perform semantic merging", len(reduced.Entities)) + } +} + +func TestReduceExtracts_MergesDuplicateRelationProvenance(t *testing.T) { + reduced := reduceExtracts([]wikiExtract{ + {Relations: []wikiRelation{{From: "A", To: "B", Type: "knows", SourceChunkIDs: []string{"c1"}}}}, + {Relations: []wikiRelation{{From: "A", To: "B", Type: "knows", SourceChunkIDs: []string{"c2"}}}}, + }) + if len(reduced.Relations) != 1 || !slices.Equal(reduced.Relations[0].SourceChunkIDs, []string{"c1", "c2"}) { + t.Fatalf("relations = %#v, want one relation with both source chunks", reduced.Relations) + } +} + +func TestEntityPageSlugIncludesType(t *testing.T) { + if got, want := entityPageSlug("曹操", "person"), "entity/person/曹操"; got != want { + t.Fatalf("entityPageSlug = %q, want %q", got, want) + } + if got, want := entityPageSlug("曹操", ""), "entity/曹操"; got != want { + t.Fatalf("entityPageSlug without type = %q, want %q", got, want) + } +} + +func TestCosine32RejectsDifferentDimensions(t *testing.T) { + if got := cosine32([]float32{1}, []float32{1, 1}); got != 0 { + t.Fatalf("cosine32 unequal dimensions = %v, want 0", got) + } +} diff --git a/internal/ingestion/service/burst_backpressure_test.go b/internal/ingestion/service/burst_backpressure_test.go index ad77fd7ecf..2da6a308ab 100644 --- a/internal/ingestion/service/burst_backpressure_test.go +++ b/internal/ingestion/service/burst_backpressure_test.go @@ -76,7 +76,7 @@ func TestProcessMessage_BurstNoTaskLossUnderBackpressure(t *testing.T) { taskIDs := seedBurstTasks(t, db, burstTaskCount) - ingestor := NewIngestor("test", 1, []string{"pdf"}) + ingestor := newUnitIngestor("test", 1, []string{"pdf"}) // Slow worker: model a saturated pipeline so the channel is full during // the burst, forcing the backpressure path. ingestor.runDocumentTask = func(ctx context.Context, _ *entity.IngestionTask) error { diff --git a/internal/ingestion/service/execute_task_ack_test.go b/internal/ingestion/service/execute_task_ack_test.go index 960de9aaec..9f160e2046 100644 --- a/internal/ingestion/service/execute_task_ack_test.go +++ b/internal/ingestion/service/execute_task_ack_test.go @@ -54,7 +54,7 @@ func TestExecuteTask_AcksMessageOnCompletion(t *testing.T) { ) ctx := t.Context() - ingestor := NewIngestor("test", 1, []string{"pdf"}) + ingestor := newUnitIngestor("test", 1, []string{"pdf"}) ingestor.runDocumentTask = func(ctx context.Context, _ *entity.IngestionTask) error { return nil } @@ -88,7 +88,7 @@ func TestExecuteTask_AcksMessageOnFailure(t *testing.T) { ) ctx := t.Context() - ingestor := NewIngestor("test", 1, []string{"pdf"}) + ingestor := newUnitIngestor("test", 1, []string{"pdf"}) ingestor.runDocumentTask = func(ctx context.Context, _ *entity.IngestionTask) error { return errors.New("boom") } @@ -122,7 +122,7 @@ func TestExecuteTask_AcksMessageOnContextCancel(t *testing.T) { testutil.WithTenantID("tenant-1"), ) - ingestor := NewIngestor("test", 1, []string{"pdf"}) + ingestor := newUnitIngestor("test", 1, []string{"pdf"}) var runCalled bool ingestor.runDocumentTask = func(ctx context.Context, _ *entity.IngestionTask) error { runCalled = true @@ -158,7 +158,7 @@ func TestExecuteTask_HeartbeatsInProgressDuringLongTask(t *testing.T) { ) ctx := t.Context() - ingestor := NewIngestor("test", 1, []string{"pdf"}) + ingestor := newUnitIngestor("test", 1, []string{"pdf"}) ingestor.heartbeatInterval = 5 * time.Millisecond started := make(chan struct{}) @@ -202,7 +202,7 @@ func TestExecuteTask_HeartbeatsInProgressDuringLongTask(t *testing.T) { // succeed; a second claim while the first worker is still processing must // fail. This is the local guard that catches MQ redeliveries. func TestClaimTask_FirstTrueThenFalse(t *testing.T) { - ingestor := NewIngestor("test", 1, []string{"pdf"}) + ingestor := newUnitIngestor("test", 1, []string{"pdf"}) if !ingestor.claimTask("task-1") { t.Fatal("first claim should succeed") @@ -222,7 +222,7 @@ func TestClaimTask_FirstTrueThenFalse(t *testing.T) { // TestClaimTask_AfterReleaseCanReclaim: after a worker finishes and releases // the task, a fresh claim (e.g. on restart) must succeed again. func TestClaimTask_AfterReleaseCanReclaim(t *testing.T) { - ingestor := NewIngestor("test", 1, []string{"pdf"}) + ingestor := newUnitIngestor("test", 1, []string{"pdf"}) ingestor.claimTask("task-1") ingestor.releaseTask("task-1") @@ -244,7 +244,7 @@ func TestExecuteTask_ReleasesTaskFromCurrentTasks(t *testing.T) { testutil.WithTenantID("tenant-1"), ) - ingestor := NewIngestor("test", 1, []string{"pdf"}) + ingestor := newUnitIngestor("test", 1, []string{"pdf"}) ingestor.runDocumentTask = func(ctx context.Context, _ *entity.IngestionTask) error { return nil } @@ -265,7 +265,7 @@ func TestExecuteTask_ReleasesTaskFromCurrentTasks(t *testing.T) { // TestSettleMessage_AckOnTerminal: body returns true -> Ack, no Nack. func TestSettleMessage_AckOnTerminal(t *testing.T) { - ingestor := NewIngestor("test", 1, []string{"pdf"}) + ingestor := newUnitIngestor("test", 1, []string{"pdf"}) handle := &fakeTaskHandle{} ctx := t.Context() taskCtx := newAckTaskCtx(context.Background(), "task-1", "doc-1", handle) @@ -279,7 +279,7 @@ func TestSettleMessage_AckOnTerminal(t *testing.T) { // TestSettleMessage_NackOnNonTerminal: body returns false -> Nack, no Ack. func TestSettleMessage_NackOnNonTerminal(t *testing.T) { - ingestor := NewIngestor("test", 1, []string{"pdf"}) + ingestor := newUnitIngestor("test", 1, []string{"pdf"}) handle := &fakeTaskHandle{} ctx := t.Context() taskCtx := newAckTaskCtx(context.Background(), "task-1", "doc-1", handle) @@ -305,7 +305,7 @@ func TestSettleMessage_RecoversPanicAndAcksWhenTaskTerminal(t *testing.T) { testutil.WithTenantID("tenant-1"), ) - ingestor := NewIngestor("test", 1, []string{"pdf"}) + ingestor := newUnitIngestor("test", 1, []string{"pdf"}) handle := &fakeTaskHandle{} ctx := t.Context() taskCtx := newAckTaskCtx(context.Background(), taskID, docID, handle) @@ -340,7 +340,7 @@ func TestSettleMessage_RecoversPanicAndAcksWhenTaskTerminal(t *testing.T) { // TestAckOrNack_AckOnTerminal: terminal=true -> Ack called, Nack not called. func TestAckOrNack_AckOnTerminal(t *testing.T) { - ingestor := NewIngestor("test", 1, []string{"pdf"}) + ingestor := newUnitIngestor("test", 1, []string{"pdf"}) handle := &fakeTaskHandle{} taskCtx := newAckTaskCtx(context.Background(), "task-1", "doc-1", handle) @@ -353,7 +353,7 @@ func TestAckOrNack_AckOnTerminal(t *testing.T) { // TestAckOrNack_NackOnNonTerminal: terminal=false -> Nack called, Ack not called. func TestAckOrNack_NackOnNonTerminal(t *testing.T) { - ingestor := NewIngestor("test", 1, []string{"pdf"}) + ingestor := newUnitIngestor("test", 1, []string{"pdf"}) handle := &fakeTaskHandle{} taskCtx := newAckTaskCtx(context.Background(), "task-1", "doc-1", handle) @@ -366,7 +366,7 @@ func TestAckOrNack_NackOnNonTerminal(t *testing.T) { // TestAckOrNack_NoOpWhenNoHandle: nil handle -> no ack/nack, no panic. func TestAckOrNack_NoOpWhenNoHandle(t *testing.T) { - ingestor := NewIngestor("test", 1, []string{"pdf"}) + ingestor := newUnitIngestor("test", 1, []string{"pdf"}) taskCtx := taskpkg.NewTaskContextForScheduling( context.Background(), &entity.IngestionTask{ID: "task-1", DocumentID: "doc-1", DatasetID: "kb-1", Status: common.RUNNING}, @@ -392,7 +392,7 @@ func TestSettleMessage_DBTruthOverridesBodyReturn(t *testing.T) { testutil.WithTenantID("tenant-1"), ) - ingestor := NewIngestor("test", 1, []string{"pdf"}) + ingestor := newUnitIngestor("test", 1, []string{"pdf"}) handle := &fakeTaskHandle{} taskCtx := newAckTaskCtx(context.Background(), taskID, docID, handle) diff --git a/internal/ingestion/service/execute_task_test.go b/internal/ingestion/service/execute_task_test.go index d7c9582a98..7b22b2bf01 100644 --- a/internal/ingestion/service/execute_task_test.go +++ b/internal/ingestion/service/execute_task_test.go @@ -37,7 +37,7 @@ func TestExecuteTask_CheckpointParseFailureDoesNotKillProcess(t *testing.T) { t.Fatalf("create bad task log: %v", err) } - ingestor := NewIngestor("test", 1, []string{"pdf"}) + ingestor := newUnitIngestor("test", 1, []string{"pdf"}) // Replace runDocumentTask to ensure it doesn't get called var runDocumentTaskCalled bool ingestor.runDocumentTask = func(ctx context.Context, ingestionTask *entity.IngestionTask) error { @@ -89,7 +89,7 @@ func TestDefaultRunDocumentTask_BothPipelineAndParserMissing(t *testing.T) { t.Fatalf("clear parser_id: %v", err) } - ingestor := NewIngestor("test", 1, []string{"pdf"}) + ingestor := newUnitIngestor("test", 1, []string{"pdf"}) err := ingestor.defaultRunDocumentTask(context.Background(), &entity.IngestionTask{ ID: taskID, DocumentID: docID, @@ -122,7 +122,7 @@ func TestDefaultRunDocumentTask_ParserIDWithoutPipelineID(t *testing.T) { testutil.WithTaskID("task-1"), ) - ingestor := NewIngestor("test", 1, []string{"pdf"}) + ingestor := newUnitIngestor("test", 1, []string{"pdf"}) err := ingestor.defaultRunDocumentTask(context.Background(), &entity.IngestionTask{ ID: taskID, DocumentID: docID, @@ -152,7 +152,7 @@ func TestExecuteTask_RunsDocumentTask(t *testing.T) { testutil.WithTenantID("tenant-1"), ) - ingestor := NewIngestor("test", 1, []string{"pdf"}) + ingestor := newUnitIngestor("test", 1, []string{"pdf"}) var runDocumentTaskCalled bool var gotTaskID string var gotProgress []float64 @@ -212,7 +212,7 @@ func TestExecuteTask_CancelBeforePipeline(t *testing.T) { testutil.WithTenantID("tenant-1"), ) - ingestor := NewIngestor("test", 1, []string{"pdf"}) + ingestor := newUnitIngestor("test", 1, []string{"pdf"}) ingestor.cancelCheck = func(ctx context.Context, taskID string) bool { return true } var runDocumentTaskCalled bool diff --git a/internal/ingestion/service/heartbeat_test.go b/internal/ingestion/service/heartbeat_test.go index 243306ec15..fe7a82b98a 100644 --- a/internal/ingestion/service/heartbeat_test.go +++ b/internal/ingestion/service/heartbeat_test.go @@ -25,7 +25,7 @@ func (h *controllableHandle) InProgress() error { return h.inProgre // TestStartHeartbeat_TicksInProgressUntilStop: with a short interval the // heartbeat goroutine calls InProgress repeatedly; stop() halts it. func TestStartHeartbeat_TicksInProgressUntilStop(t *testing.T) { - ingestor := NewIngestor("test", 1, []string{"pdf"}) + ingestor := newUnitIngestor("test", 1, []string{"pdf"}) ingestor.heartbeatInterval = 2 * time.Millisecond handle := &fakeTaskHandle{} @@ -45,7 +45,7 @@ func TestStartHeartbeat_TicksInProgressUntilStop(t *testing.T) { // concurrent InProgress on the same message. Regression guard for the // close-without-wait heartbeat shutdown (problem 1). func TestStartHeartbeat_StopWaitsForInFlightInProgress(t *testing.T) { - ingestor := NewIngestor("test", 1, []string{"pdf"}) + ingestor := newUnitIngestor("test", 1, []string{"pdf"}) ingestor.heartbeatInterval = time.Millisecond started := make(chan struct{}) @@ -89,7 +89,7 @@ func TestStartHeartbeat_StopWaitsForInFlightInProgress(t *testing.T) { // TestStartHeartbeat_NoOpWhenNoHandle: with no MQ handle (standalone/test path) // startHeartbeat returns a no-op stop and starts no goroutine. func TestStartHeartbeat_NoOpWhenNoHandle(t *testing.T) { - ingestor := NewIngestor("test", 1, []string{"pdf"}) + ingestor := newUnitIngestor("test", 1, []string{"pdf"}) ingestor.heartbeatInterval = time.Millisecond taskCtx := taskpkg.NewTaskContextForScheduling( diff --git a/internal/ingestion/service/ingestion_service.go b/internal/ingestion/service/ingestion_service.go index d0a5b24dec..71983a9f49 100644 --- a/internal/ingestion/service/ingestion_service.go +++ b/internal/ingestion/service/ingestion_service.go @@ -27,6 +27,7 @@ import ( "sync/atomic" "time" + "ragflow/internal/agent/canvas" "ragflow/internal/common" "ragflow/internal/dao" "ragflow/internal/engine" @@ -106,6 +107,10 @@ type Ingestor struct { // cancel-flag lookup that mirrors Python's has_canceled(). Tests may // override this to simulate cancel without Redis. cancelCheck func(ctx context.Context, taskID string) bool + + // checkpointExists distinguishes a fresh run from a checkpoint resume. + // Tests inject this dependency so unit tests do not require Redis. + checkpointExists func(ctx context.Context, taskID string) (bool, error) } func NewIngestor(name string, maxConcurrency int32, supportedTypes []string) *Ingestor { @@ -131,6 +136,7 @@ func NewIngestor(name string, maxConcurrency int32, supportedTypes []string) *In } ingestor.runDocumentTask = ingestor.defaultRunDocumentTask ingestor.cancelCheck = ingestor.defaultCancelCheck + ingestor.checkpointExists = canvas.RedisCheckpointExists ingestor.kcConcurrency = maxConcurrency // parallel dataset-level compile workers default to the task width return ingestor } @@ -615,14 +621,31 @@ func (e *Ingestor) runTask(ctx context.Context, task *entity.IngestionTask) bool } return ok } - if err := e.ingestionTaskSvc.ClearComponentProgress(ctx, task.ID); err != nil { - common.Error(fmt.Sprintf("Failed to clear previous component progress for task %s", task.ID), err) + checkpointExists := e.checkpointExists + if checkpointExists == nil { + checkpointExists = canvas.RedisCheckpointExists + } + resumeCheckpoint, checkpointErr := checkpointExists(ctx, task.ID) + if checkpointErr != nil { + common.Error(fmt.Sprintf("Failed to check checkpoint for task %s", task.ID), checkpointErr) ok := e.markFailed(ctx, task.ID) if ok { e.recordTerminalPipelineLog(ctx, task, string(entity.TaskStatusFail)) } return ok } + if !resumeCheckpoint { + if err := e.ingestionTaskSvc.ClearComponentProgress(ctx, task.ID); err != nil { + common.Error(fmt.Sprintf("Failed to clear previous component progress for task %s", task.ID), err) + ok := e.markFailed(ctx, task.ID) + if ok { + e.recordTerminalPipelineLog(ctx, task, string(entity.TaskStatusFail)) + } + return ok + } + } else { + common.Info(fmt.Sprintf("Preserving component progress for checkpoint resume of task %s", task.ID)) + } // This is a new run (IncrementRunCount succeeded). Any Redis cancel flag // that exists now is stale — a leftover from a previous run whose diff --git a/internal/ingestion/service/ingestor_lifecycle_test.go b/internal/ingestion/service/ingestor_lifecycle_test.go index cb95f882a4..01b65c1f3b 100644 --- a/internal/ingestion/service/ingestor_lifecycle_test.go +++ b/internal/ingestion/service/ingestor_lifecycle_test.go @@ -34,7 +34,7 @@ import ( // pool and activeWorkers would exceed concurrency after the second call. func TestStartWorkerPool_StartOnceIdempotent(t *testing.T) { const concurrency int32 = 3 - ingestor := NewIngestor("test-idempotent", concurrency, nil) + ingestor := newUnitIngestor("test-idempotent", concurrency, nil) ingestor.startWorkerPool() // Wait for all workers to enter their loop (they block on the select @@ -66,7 +66,7 @@ func TestStartWorkerPool_StartOnceIdempotent(t *testing.T) { // for all worker goroutines to exit without hanging. func TestStop_GracefulShutdown(t *testing.T) { const concurrency int32 = 2 - ingestor := NewIngestor("test-shutdown", concurrency, nil) + ingestor := newUnitIngestor("test-shutdown", concurrency, nil) // Start workers; they will block on the task channel since nothing is pushed. ingestor.startWorkerPool() @@ -91,7 +91,7 @@ func TestStop_GracefulShutdown(t *testing.T) { // Stop. Without this, the admin graceful-shutdown path is dead (cmd blocks // forever on the receive). func TestStop_ClosesShutdownCh(t *testing.T) { - ingestor := NewIngestor("test-shutdown-ch", 1, nil) + ingestor := newUnitIngestor("test-shutdown-ch", 1, nil) ingestor.Stop(context.Background()) select { case <-ingestor.ShutdownCh: @@ -112,7 +112,7 @@ func TestStop_TimesOutWhenWorkerStuck(t *testing.T) { _, _, docID, taskID := testutil.SeedTestData(t, db, testutil.WithPipelineID("flow-1")) const concurrency int32 = 1 - ingestor := NewIngestor("test-stuck", concurrency, []string{"pdf"}) + ingestor := newUnitIngestor("test-stuck", concurrency, []string{"pdf"}) ingestor.startWorkerPool() // runDocumentTask blocks on release and ignores ctx, simulating a @@ -168,7 +168,7 @@ func TestStop_TimesOutWhenWorkerStuck(t *testing.T) { // long DB query). Without BP3, the initial cancelCheck call runs // synchronously and pollCancel cannot observe done until it returns. func TestPollCancel_ExitsWhenDoneClosed(t *testing.T) { - ingestor := NewIngestor("test", 1, []string{"pdf"}) + ingestor := newUnitIngestor("test", 1, []string{"pdf"}) // Block cancelCheck until released — simulate a stuck DB call. blocking := make(chan struct{}) @@ -220,7 +220,7 @@ func TestStart_FullPathReturnsAndStartsWorkers(t *testing.T) { t.Cleanup(func() { engine.SetMessageQueueEngine(previousEngine) }) const concurrency int32 = 2 - ing := NewIngestor("test-start-fullpath", concurrency, nil) + ing := newUnitIngestor("test-start-fullpath", concurrency, nil) t.Cleanup(func() { ing.Stop(context.Background()) }) done := make(chan error, 1) diff --git a/internal/ingestion/service/process_message_test.go b/internal/ingestion/service/process_message_test.go index 56cda0456c..671db5d30c 100644 --- a/internal/ingestion/service/process_message_test.go +++ b/internal/ingestion/service/process_message_test.go @@ -28,7 +28,7 @@ func TestProcessMessage_MemoryTaskDispatches(t *testing.T) { cleanup := testutil.ReplaceDBForTest(t, db) defer cleanup() - ingestor := NewIngestor("test", 1, []string{"pdf"}) + ingestor := newUnitIngestor("test", 1, []string{"pdf"}) // Memory extractor must be enabled for the memory branch to enqueue. ingestor.SetMemoryMessageService(service.NewMemoryMessageService(nil)) @@ -77,7 +77,7 @@ func TestProcessMessage_MemoryTaskDisabledAcks(t *testing.T) { cleanup := testutil.ReplaceDBForTest(t, db) defer cleanup() - ingestor := NewIngestor("test", 1, []string{"pdf"}) // memorySvc nil by default + ingestor := newUnitIngestor("test", 1, []string{"pdf"}) // memorySvc nil by default handle := newFakeHandle("mem-task-2", common.TaskTypeMemory) ingestor.processMessage(handle) @@ -108,7 +108,7 @@ func TestExecuteMemoryTaskAlreadyFailedAcks(t *testing.T) { t.Fatalf("insert already-failed task: %v", err) } - ingestor := NewIngestor("test", 1, []string{"pdf"}) + ingestor := newUnitIngestor("test", 1, []string{"pdf"}) // Real memory service (non-nil memories) so HandleSaveToMemoryTask gets past // the nil-guard and reaches the progress==-1 "already failed" branch. ingestor.SetMemoryMessageService(service.NewMemoryMessageService(service.NewMemoryService())) @@ -145,7 +145,7 @@ func TestExecuteMemoryTaskTransientFailureNacks(t *testing.T) { cleanup := testutil.ReplaceDBForTest(t, db) defer cleanup() - ingestor := NewIngestor("test", 1, []string{"pdf"}) + ingestor := newUnitIngestor("test", 1, []string{"pdf"}) ingestor.SetMemoryMessageService(service.NewMemoryMessageService(service.NewMemoryService())) handle := &fakeTaskHandle{msg: common.TaskMessage{TaskID: "mem-task-x", TaskType: common.TaskTypeMemory}} @@ -177,7 +177,7 @@ func TestProcessMessage_NonIngestionTaskAcks(t *testing.T) { cleanup := testutil.ReplaceDBForTest(t, db) defer cleanup() - ingestor := NewIngestor("test", 1, []string{"pdf"}) + ingestor := newUnitIngestor("test", 1, []string{"pdf"}) handle := newFakeHandle("task-1", "not-ingestion") ingestor.processMessage(handle) @@ -196,7 +196,7 @@ func TestProcessMessage_TaskNotFoundAcks(t *testing.T) { cleanup := testutil.ReplaceDBForTest(t, db) defer cleanup() - ingestor := NewIngestor("test", 1, []string{"pdf"}) + ingestor := newUnitIngestor("test", 1, []string{"pdf"}) // No task seeded in DB — StartRunning returns ErrTaskNotFound. handle := newFakeHandle("no-such-task", common.TaskTypeIngestionTask) @@ -234,7 +234,7 @@ func TestProcessMessage_AlreadyCompletedAcks(t *testing.T) { t.Fatalf("set COMPLETED: %v", err) } - ingestor := NewIngestor("test", 1, []string{"pdf"}) + ingestor := newUnitIngestor("test", 1, []string{"pdf"}) handle := newFakeHandle(taskID, common.TaskTypeIngestionTask) ingestor.processMessage(handle) @@ -267,7 +267,7 @@ func TestProcessMessage_ClaimFailsAcks(t *testing.T) { defer cleanup() _, _, _, taskID := testutil.SeedTestData(t, db, testutil.WithPipelineID("flow-1")) - ingestor := NewIngestor("test", 1, []string{"pdf"}) + ingestor := newUnitIngestor("test", 1, []string{"pdf"}) // Pre-claim the task so processMessage sees a claim conflict. ingestor.claimTask(taskID) @@ -291,7 +291,7 @@ func TestProcessMessage_ClaimSucceedsEnqueues(t *testing.T) { defer cleanup() _, _, _, taskID := testutil.SeedTestData(t, db, testutil.WithPipelineID("flow-1")) - ingestor := NewIngestor("test", 1, []string{"pdf"}) + ingestor := newUnitIngestor("test", 1, []string{"pdf"}) handle := newFakeHandle(taskID, common.TaskTypeIngestionTask) ingestor.processMessage(handle) @@ -322,7 +322,7 @@ func TestProcessMessage_ChannelFullBlocksUntilSlot(t *testing.T) { _, _, _, taskID := testutil.SeedTestData(t, db, testutil.WithPipelineID("flow-1")) // maxConcurrency=2 → channel cap=4. Fill it completely. - ingestor := NewIngestor("test", 2, []string{"pdf"}) + ingestor := newUnitIngestor("test", 2, []string{"pdf"}) for i := 0; i < cap(ingestor.taskChan); i++ { ingestor.taskChan <- taskpkg.NewTaskContextForScheduling(nil, &entity.IngestionTask{ID: "filler"}) } @@ -387,7 +387,7 @@ func TestProcessMessage_StartRunningErrorNacks(t *testing.T) { t.Fatalf("drop table: %v", err) } - ingestor := NewIngestor("test", 1, []string{"pdf"}) + ingestor := newUnitIngestor("test", 1, []string{"pdf"}) handle := newFakeHandle(taskID, common.TaskTypeIngestionTask) ingestor.processMessage(handle) diff --git a/internal/ingestion/service/real_consumer_pipeline_test.go b/internal/ingestion/service/real_consumer_pipeline_test.go index 8848e5f249..beed984c24 100644 --- a/internal/ingestion/service/real_consumer_pipeline_test.go +++ b/internal/ingestion/service/real_consumer_pipeline_test.go @@ -91,7 +91,7 @@ func TestRealConsumer_PipelineMessageRoutesToExecuteTask(t *testing.T) { t.Fatalf("task status after UpdateStatusIfCurrent = %s, want %s", task.Status, common.RUNNING) } - ingestor := NewIngestor("queue-test", 1, []string{"pdf"}) + ingestor := newUnitIngestor("queue-test", 1, []string{"pdf"}) var routedToPipeline bool taskCtx := taskpkg.NewTaskContextForScheduling( context.Background(), diff --git a/internal/ingestion/service/redelivery_counter_test.go b/internal/ingestion/service/redelivery_counter_test.go index 87335263a9..9f3f1b9c8f 100644 --- a/internal/ingestion/service/redelivery_counter_test.go +++ b/internal/ingestion/service/redelivery_counter_test.go @@ -92,7 +92,7 @@ func TestRunTask_RedeliveryOfCompletedTaskCountsOnce(t *testing.T) { defer cleanup() _, kbID, docID, taskID := testutil.SeedTestData(t, db, testutil.WithPipelineID("flow-1")) - ingestor := NewIngestor("test", 1, []string{"pdf"}) + ingestor := newUnitIngestor("test", 1, []string{"pdf"}) ingestor.runDocumentTask = applyResult(ingestor, docID, kbID) // First delivery: parse succeeds, counters applied once, task -> COMPLETED. @@ -115,7 +115,7 @@ func TestRunTask_RedeliveryAfterIncompleteRunCountsOnce(t *testing.T) { defer cleanup() _, kbID, docID, taskID := testutil.SeedTestData(t, db, testutil.WithPipelineID("flow-1")) - ingestor := NewIngestor("test", 1, []string{"pdf"}) + ingestor := newUnitIngestor("test", 1, []string{"pdf"}) ingestor.runDocumentTask = applyResult(ingestor, docID, kbID) // Prior run: counters applied, but the task never completed (crash before diff --git a/internal/ingestion/service/run_task_test.go b/internal/ingestion/service/run_task_test.go index bfc362f7e2..9ff3dbdd7e 100644 --- a/internal/ingestion/service/run_task_test.go +++ b/internal/ingestion/service/run_task_test.go @@ -21,7 +21,7 @@ func TestRunTask_ContextCancelledBeforeCheckpoint(t *testing.T) { defer cleanup() _, _, _, taskID := testutil.SeedTestData(t, db, testutil.WithPipelineID("flow-1")) - ingestor := NewIngestor("test", 1, []string{"pdf"}) + ingestor := newUnitIngestor("test", 1, []string{"pdf"}) var runDocCalled bool ingestor.runDocumentTask = func(ctx context.Context, _ *entity.IngestionTask) error { runDocCalled = true @@ -79,7 +79,7 @@ func TestRunTask_CorruptedRunCountSkipped(t *testing.T) { t.Fatalf("insert bad log: %v", err) } - ingestor := NewIngestor("test", 1, []string{"pdf"}) + ingestor := newUnitIngestor("test", 1, []string{"pdf"}) var runDocCalled bool ingestor.runDocumentTask = func(ctx context.Context, _ *entity.IngestionTask) error { runDocCalled = true @@ -115,7 +115,7 @@ func TestRunTask_RunDocumentTaskFailureMarksFailed(t *testing.T) { defer cleanup() _, _, docID, taskID := testutil.SeedTestData(t, db, testutil.WithPipelineID("flow-1")) - ingestor := NewIngestor("test", 1, []string{"pdf"}) + ingestor := newUnitIngestor("test", 1, []string{"pdf"}) ingestor.runDocumentTask = func(ctx context.Context, _ *entity.IngestionTask) error { return errors.New("boom") } @@ -153,7 +153,7 @@ func TestRunTask_PipelineCancelledMarksStopped(t *testing.T) { t.Fatalf("set task STOPPING: %v", err) } - ingestor := NewIngestor("test", 1, []string{"pdf"}) + ingestor := newUnitIngestor("test", 1, []string{"pdf"}) ingestor.runDocumentTask = func(ctx context.Context, _ *entity.IngestionTask) error { return context.Canceled } @@ -186,7 +186,7 @@ func TestRunTask_ComponentTimeoutMarksFailed(t *testing.T) { defer cleanup() _, _, docID, taskID := testutil.SeedTestData(t, db, testutil.WithPipelineID("flow-1")) - ingestor := NewIngestor("test", 1, []string{"pdf"}) + ingestor := newUnitIngestor("test", 1, []string{"pdf"}) ingestor.runDocumentTask = func(ctx context.Context, _ *entity.IngestionTask) error { return context.DeadlineExceeded } @@ -227,7 +227,7 @@ func TestRunTask_AlreadyCompletedAcksNotRedelivers(t *testing.T) { t.Fatalf("set task COMPLETED: %v", err) } - ingestor := NewIngestor("test", 1, []string{"pdf"}) + ingestor := newUnitIngestor("test", 1, []string{"pdf"}) ingestor.runDocumentTask = func(ctx context.Context, _ *entity.IngestionTask) error { return nil } @@ -263,7 +263,7 @@ func TestRunTask_PipelineSucceedsConcurrentStopSettlesStopped(t *testing.T) { defer cleanup() _, _, docID, taskID := testutil.SeedTestData(t, db, testutil.WithPipelineID("flow-1")) - ingestor := NewIngestor("test", 1, []string{"pdf"}) + ingestor := newUnitIngestor("test", 1, []string{"pdf"}) ingestor.runDocumentTask = func(ctx context.Context, task *entity.IngestionTask) error { // Simulate the user pressing Stop mid-pipeline: RUNNING->STOPPING. if _, err := ingestor.ingestionTaskSvc.RequestStop(ctx, task.ID); err != nil { @@ -298,7 +298,7 @@ func TestRunTask_SuccessfulCompletion(t *testing.T) { defer cleanup() _, _, docID, taskID := testutil.SeedTestData(t, db, testutil.WithPipelineID("flow-1")) - ingestor := NewIngestor("test", 1, []string{"pdf"}) + ingestor := newUnitIngestor("test", 1, []string{"pdf"}) ingestor.runDocumentTask = func(ctx context.Context, _ *entity.IngestionTask) error { return nil } diff --git a/internal/ingestion/service/test_helpers.go b/internal/ingestion/service/test_helpers.go index aa6b9440d6..05e887d1ce 100644 --- a/internal/ingestion/service/test_helpers.go +++ b/internal/ingestion/service/test_helpers.go @@ -16,6 +16,8 @@ package service +import "context" + // ────────────────────────────────────────────────────────── // Ingestor Test Helpers // ────────────────────────────────────────────────────────── @@ -30,6 +32,14 @@ func NewTestIngestor() *Ingestor { // IngestorOption configures a test Ingestor. type IngestorOption func(*Ingestor) +func newUnitIngestor(name string, maxConcurrency int32, supportedTypes []string) *Ingestor { + ingestor := NewIngestor(name, maxConcurrency, supportedTypes) + ingestor.checkpointExists = func(context.Context, string) (bool, error) { + return false, nil + } + return ingestor +} + // SetupTestIngestor creates a new test Ingestor with the given options. func SetupTestIngestor(t testingT, opts ...IngestorOption) *Ingestor { t.Helper() diff --git a/internal/router/router.go b/internal/router/router.go index 9aa64e2813..55752573f3 100644 --- a/internal/router/router.go +++ b/internal/router/router.go @@ -362,8 +362,8 @@ func (r *Router) Setup(engine *gin.Engine) { datasets.GET("/:dataset_id/artifacts/topics", r.datasetArtifactHandler.ListArtifactTopics) datasets.GET("/:dataset_id/artifacts/alteration", r.datasetArtifactHandler.GetArtifactAlteration) datasets.GET("/:dataset_id/artifacts/graph", r.datasetArtifactHandler.GetArtifactGraph) - datasets.GET("/:dataset_id/artifacts/:page_type/:slug", r.datasetArtifactHandler.GetArtifact) - datasets.PUT("/:dataset_id/artifacts/:page_type/:slug", r.datasetArtifactHandler.UpdateArtifact) + datasets.GET("/:dataset_id/artifacts/:page_type/*slug", r.datasetArtifactHandler.GetArtifact) + datasets.PUT("/:dataset_id/artifacts/:page_type/*slug", r.datasetArtifactHandler.UpdateArtifact) datasets.GET("/:dataset_id/artifacts/structure", r.datasetArtifactHandler.ListStructures) datasets.DELETE("/:dataset_id/artifacts/structure", r.datasetArtifactHandler.DeleteStructures) diff --git a/internal/service/dataset_artifact_service.go b/internal/service/dataset_artifact_service.go index d374bf18e9..ea531b4491 100644 --- a/internal/service/dataset_artifact_service.go +++ b/internal/service/dataset_artifact_service.go @@ -433,11 +433,11 @@ func (s *DatasetArtifactService) GetWikiGraph(ctx context.Context, tenantID, dat // so the frontend can build artifact// links that // round-trip. entity_type_kwd stores "wiki_" + page_type (e.g. // "wiki_topic"); strip the prefix. slug_kwd stores the full - // "/" form; expose the trailing bare slug. + // "/" form; preserve nested slug segments. fullSlug := firstStringValue(c["slug_kwd"]) bareSlug := fullSlug pageType := strings.TrimPrefix(firstStringValue(c["entity_type_kwd"]), "wiki_") - if idx := strings.LastIndex(bareSlug, "/"); idx >= 0 { + if idx := strings.IndexByte(bareSlug, '/'); idx >= 0 { pageType = bareSlug[:idx] bareSlug = bareSlug[idx+1:] } @@ -949,13 +949,12 @@ func firstStringValue(v interface{}) string { return "" } -// bareWikiSlug strips the "/" prefix from a full wiki slug -// ("topic/yellow-turban-rebellion" -> "yellow-turban-rebellion"), matching the -// bare-slug form the graph UI keys nodes/relations on. Slugs with no prefix are -// returned unchanged. +// bareWikiSlug strips only the first path segment from a full wiki slug +// ("entity/location/长社" -> "location/长社"). Nested type/name segments are +// preserved so distinct typed entities remain distinguishable in graph links. func bareWikiSlug(slug string) string { s := strings.TrimSpace(slug) - if idx := strings.LastIndex(s, "/"); idx >= 0 && idx < len(s)-1 { + if idx := strings.IndexByte(s, '/'); idx >= 0 && idx < len(s)-1 { return s[idx+1:] } return s diff --git a/web/src/pages/dataset/compilation/utils/parse-wiki-link.test.ts b/web/src/pages/dataset/compilation/utils/parse-wiki-link.test.ts new file mode 100644 index 0000000000..f20ae98347 --- /dev/null +++ b/web/src/pages/dataset/compilation/utils/parse-wiki-link.test.ts @@ -0,0 +1,18 @@ +import { parseWikiLinkHref } from './parse-wiki-link'; + +describe('parseWikiLinkHref', () => { + it('preserves nested typed slugs in artifact links', () => { + expect( + parseWikiLinkHref( + 'artifact/fb9bfae2a00b43e59ecaea5e86d90c91/entity/person/张角', + ), + ).toEqual({ pageType: 'entity', slug: 'person/张角' }); + }); + + it('preserves nested typed slugs in simple links', () => { + expect(parseWikiLinkHref('entity/location/长社')).toEqual({ + pageType: 'entity', + slug: 'location/长社', + }); + }); +}); diff --git a/web/src/pages/dataset/compilation/utils/parse-wiki-link.ts b/web/src/pages/dataset/compilation/utils/parse-wiki-link.ts index 7accc647b8..828cf94b1c 100644 --- a/web/src/pages/dataset/compilation/utils/parse-wiki-link.ts +++ b/web/src/pages/dataset/compilation/utils/parse-wiki-link.ts @@ -19,7 +19,7 @@ export function parseWikiLinkHref( // Prefer the artifact/{datasetId}/{pageType}/{slug} form. const artifactMatch = normalized.match( - /(?:^|\/)artifact\/[^/]+\/(entity|concept|topic)\/([^/\s"']+)/, + /(?:^|\/)artifact\/[^/]+\/(entity|concept|topic)\/([^\s"']+)/, ); if (artifactMatch) { return { @@ -30,7 +30,7 @@ export function parseWikiLinkHref( // Fallback to a plain {pageType}/{slug} form. const simpleMatch = normalized.match( - /(?:^|\/)(entity|concept|topic)\/([^/\s"']+)/, + /(?:^|\/)(entity|concept|topic)\/([^\s"']+)/, ); if (simpleMatch) { return {