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%.
This commit is contained in:
buua436
2026-08-24 14:28:05 +08:00
committed by GitHub
parent 01d83f3679
commit f37b6fc439
24 changed files with 301 additions and 539 deletions

View File

@@ -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".

View File

@@ -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/<page_type>/<slug> — 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/<page_type>/<slug> — 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())

View File

@@ -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
}
}

View File

@@ -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.

View File

@@ -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

View File

@@ -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

View File

@@ -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)
}

View File

@@ -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)
}

View File

@@ -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)
}
}

View File

@@ -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 {

View File

@@ -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)

View File

@@ -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

View File

@@ -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(

View File

@@ -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

View File

@@ -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)

View File

@@ -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)

View File

@@ -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(),

View File

@@ -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

View File

@@ -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
}

View File

@@ -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()

View File

@@ -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)

View File

@@ -433,11 +433,11 @@ func (s *DatasetArtifactService) GetWikiGraph(ctx context.Context, tenantID, dat
// so the frontend can build artifact/<page_type>/<slug> links that
// round-trip. entity_type_kwd stores "wiki_" + page_type (e.g.
// "wiki_topic"); strip the prefix. slug_kwd stores the full
// "<page_type>/<slug>" form; expose the trailing bare slug.
// "<page_type>/<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 "<page_type>/" 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

View File

@@ -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/长社',
});
});
});

View File

@@ -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 {