mirror of
https://github.com/infiniflow/ragflow.git
synced 2026-08-27 18:56:31 +08:00
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:
@@ -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".
|
||||
|
||||
@@ -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())
|
||||
|
||||
@@ -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
|
||||
}
|
||||
}
|
||||
|
||||
|
||||
@@ -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.
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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)
|
||||
}
|
||||
@@ -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)
|
||||
}
|
||||
@@ -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)
|
||||
}
|
||||
}
|
||||
|
||||
@@ -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 {
|
||||
|
||||
@@ -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)
|
||||
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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(
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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)
|
||||
|
||||
@@ -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)
|
||||
|
||||
@@ -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(),
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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
|
||||
}
|
||||
|
||||
@@ -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()
|
||||
|
||||
@@ -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)
|
||||
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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/长社',
|
||||
});
|
||||
});
|
||||
});
|
||||
@@ -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 {
|
||||
|
||||
Reference in New Issue
Block a user