diff --git a/AGENTS.md b/AGENTS.md index 733bd0ec42..37f9f28a51 100644 --- a/AGENTS.md +++ b/AGENTS.md @@ -77,6 +77,7 @@ Rules: ## Working Rules - Before editing, inspect the nearest code path that actually owns the behavior. +- When handling review comments, independently verify each substantive claim against the current code or tests before accepting, rejecting, or acting on it. - Keep changes small and local unless the task is explicitly a broader refactor. - Prefer one implementation path instead of preserving old and new versions side by side. - Preserve behavior with focused tests when the behavior is still valid; do not keep tests that protect obsolete behavior. diff --git a/agent/templates/compiler.json b/agent/templates/compiler.json index f38526883c..437524b766 100644 --- a/agent/templates/compiler.json +++ b/agent/templates/compiler.json @@ -5,7 +5,8 @@ "obj": { "component_name": "Compiler", "params": { - "compilation_template_group_id": "c3aa748c8b2111f191f3047c16ec874f", + "compilation_template_group_id": "", + "llm_id": "", "outputs": { "chunks": { "type": "Array", @@ -408,7 +409,8 @@ { "data": { "form": { - "compilation_template_group_id": "c3aa748c8b2111f191f3047c16ec874f", + "compilation_template_group_id": "", + "llm_id": "", "outputs": { "chunks": { "type": "Array", diff --git a/cmd/ragflow_server.go b/cmd/ragflow_server.go index ea77ed5fdb..61a55280cf 100644 --- a/cmd/ragflow_server.go +++ b/cmd/ragflow_server.go @@ -30,6 +30,7 @@ import ( agenttool "ragflow/internal/agent/tool" "ragflow/internal/channels" "ragflow/internal/handler" + "ragflow/internal/ingestion/knowledge_compile" ingestion "ragflow/internal/ingestion/service" "ragflow/internal/mcp" "ragflow/internal/router" @@ -41,6 +42,7 @@ import ( "ragflow/internal/service/file" "ragflow/internal/service/nav" "ragflow/internal/service/nlp" + "ragflow/internal/service/wikisearch" "ragflow/internal/storage" "ragflow/internal/syncer" "ragflow/internal/tokenizer" @@ -554,7 +556,15 @@ func runIngestor(ctx context.Context, cancel context.CancelFunc, args *serverArg // writes available_int=0 compiled chunks; they just won't be merged until // the consumer is available. globalConfig := server.GetConfig() - ingestor := ingestion.NewIngestor(*args.name, 2, []string{"pdf", "docx", "txt"}) + ingestorCfg := globalConfig.GetIngestorConfig() + const maxIngestorConcurrency = int32(1<<30 - 1) + if ingestorCfg.MaxConcurrentWorkers > int(maxIngestorConcurrency) { + return fmt.Errorf("ingestor max_concurrent_workers %d exceeds maximum %d", ingestorCfg.MaxConcurrentWorkers, maxIngestorConcurrency) + } + // Apply the configured compiler pool size (no-op when 0; the pool keeps its + // vCPU default, overridable via KC_COMPILE_CONCURRENCY). + knowledge_compile.SetCompilerConcurrency(ingestorCfg.CompilerPoolSize) + ingestor := ingestion.NewIngestor(*args.name, int32(ingestorCfg.MaxConcurrentWorkers), []string{"pdf", "docx", "txt"}) ingestor.SetKnowledgeCompileModelConfig( globalConfig.GetDefaultChatModel().Name, globalConfig.GetDefaultEmbeddingModel().Name, @@ -850,6 +860,14 @@ func startServer(ctx context.Context) { // on demand so Search/UpsertDoc can embed queries/summaries automatically. nav.SetNavService(nlp.NewNavService(service.NewNavEmbedder(modelProviderService, ""))) + // Install the compiled-wiki search service. It is backed directly by the + // document engine: QueryPages filters the tenant-scoped index to + // compile_kwd="wiki_page" (+ supported kinds) so ordinary source chunks are + // never relabeled as wiki pages, and BackfillChunks fetches original chunks + // by id. When the engine is unavailable the service degrades to empty so the + // agent falls back to hybrid search (no failing call). + wikisearch.SetService(wikisearch.NewEngineService(engine.Get())) + // Initialize router r := router.NewRouter(authHandler, userHandler, diff --git a/conf/service_conf.yaml b/conf/service_conf.yaml index 25073a0fa3..59f7de50e2 100644 --- a/conf/service_conf.yaml +++ b/conf/service_conf.yaml @@ -69,6 +69,7 @@ otel: ingestor: mq_type: 'nats' max_concurrent_workers: 1 + compiler_pool_size: 0 file_syncer: max_concurrent_syncs: 1 sync_interval: 3 diff --git a/docker/service_conf.yaml.template b/docker/service_conf.yaml.template index 5b77399425..531f29a725 100644 --- a/docker/service_conf.yaml.template +++ b/docker/service_conf.yaml.template @@ -87,6 +87,7 @@ otel: ingestor: mq_type: 'nats' max_concurrent_workers: 1 + compiler_pool_size: 0 file_syncer: max_concurrent_syncs: 1 sync_interval: 3 diff --git a/internal/agent/harness/agentic_rag.go b/internal/agent/harness/agentic_rag.go index 24b419199d..d843ff96ac 100644 --- a/internal/agent/harness/agentic_rag.go +++ b/internal/agent/harness/agentic_rag.go @@ -46,7 +46,18 @@ type AgenticState struct { // - medium+ (decompose_and_search / agentic_research / deep_research): // pre_search grounds the planner, then decompose-and-search runs until a // sufficiency verdict stops it. +// +// RunAgenticRAG drives the agentic-search graph. It computes the route itself +// and delegates to RunAgenticRAGWithRoute so production runners that need a +// route-aware search strategy (e.g. prefer wiki_query on a wiki suggestion) can +// reuse the same flow with a pre-computed route. func RunAgenticRAG(ctx context.Context, db *gorm.DB, question, keywords, modeLabel string, search SearchFn) AnswerResult { + return RunAgenticRAGWithRoute(ctx, db, question, keywords, modeLabel, RouteNode(ctx, db, question, modeLabel), search) +} + +// RunAgenticRAGWithRoute is the route-aware core of RunAgenticRAG. It performs +// pre_search → planner → orchestrator → formalize_answer with the given route. +func RunAgenticRAGWithRoute(ctx context.Context, db *gorm.DB, question, keywords, modeLabel string, route RouteDecision, search SearchFn) AnswerResult { state := &AgenticState{ Question: strings.TrimSpace(question), Keywords: keywords, @@ -56,8 +67,8 @@ func RunAgenticRAG(ctx context.Context, db *gorm.DB, question, keywords, modeLab return AnswerResult{FinalAnswer: emptyResultMessage, Empty: true} } - // ── route ── - state.Route = RouteNode(ctx, db, state.Question, modeLabel) + // ── route (pre-computed by the caller) ── + state.Route = route // ── pre_search (decomposition modes only) ── if state.Route.RequiresDecomposition { diff --git a/internal/agent/harness/production.go b/internal/agent/harness/production.go index 2a6f3a9139..0b1f739585 100644 --- a/internal/agent/harness/production.go +++ b/internal/agent/harness/production.go @@ -27,24 +27,34 @@ import ( "gorm.io/gorm" "ragflow/internal/agent/tool" + "ragflow/internal/common" "ragflow/internal/service/nav" + "ragflow/internal/service/wikisearch" ) // ProductionRunner wires the real agentic-search tools (hybrid_search, -// dataset_navigation_by_tree) into the RunAgenticRAG flow, so the tools are -// actually invoked rather than merely registered. This is the production -// counterpart to the unit-testable SearchFn seam. +// dataset_navigation_by_tree, wiki_query) into the RunAgenticRAG flow, so the +// tools are actually invoked rather than merely registered. This is the +// production counterpart to the unit-testable SearchFn seam. type ProductionRunner struct { db *gorm.DB tenantID string datasetIDs []string searchTool einotool.InvokableTool navSvc nav.NavService // defaults to nav.GetNavService() when nil + wikiSvc wikisearch.Service + // webTool is an optional, already-configured web search tool. When nil the + // runner never exposes web fallback (P8: no web provider configured => the + // agent does not attempt web search and no failing tool call is made). + webTool einotool.InvokableTool } // NewProductionRunner builds a ProductionRunner backed by the real tools. The // dataset-nav router (harness.NavigateDatasetByTree) resolves its NavService -// lazily via nav.GetNavService(). +// lazily via nav.GetNavService(). When a web provider is configured (a Tavily +// API key is present), the runner also wires the web fallback tool so +// high/ultra modes can fill an empty KB result from the web; otherwise no web +// tool is attached and no web call is ever attempted (P8/R2). func NewProductionRunner(db *gorm.DB, tenantID string, datasetIDs []string) (*ProductionRunner, error) { searchBase, err := tool.BuildByName("hybrid_search", nil) if err != nil { @@ -54,7 +64,11 @@ func NewProductionRunner(db *gorm.DB, tenantID string, datasetIDs []string) (*Pr if !ok { return nil, fmt.Errorf("hybrid_search is not invokable") } - return &ProductionRunner{db: db, tenantID: tenantID, datasetIDs: datasetIDs, searchTool: search}, nil + r := &ProductionRunner{db: db, tenantID: tenantID, datasetIDs: datasetIDs, searchTool: search} + if common.GetEnv(common.EnvTavilyApiKey) != "" { + r.webTool = tool.NewTavilyTool() + } + return r, nil } // newProductionRunnerWithTools builds a ProductionRunner with an injected @@ -64,19 +78,61 @@ func newProductionRunnerWithTools(db *gorm.DB, tenantID string, datasetIDs []str return &ProductionRunner{db: db, tenantID: tenantID, datasetIDs: datasetIDs, searchTool: searchTool, navSvc: navSvc} } -// Run executes the agentic-search graph with the real tools. It returns the -// final answer. +// Run executes the agentic-search graph with the real tools. It computes the +// route once and uses it to pick a search strategy: when the route suggests a +// wiki compilation and the bound KBs actually carry wiki artifacts, the runner +// tries wiki_query first and falls back to general hybrid search on an empty +// result; otherwise it uses hybrid search. Web fallback is only reachable when a +// web provider is configured (P8). Returns the final answer. func (r *ProductionRunner) Run(ctx context.Context, question, keywords, modeLabel string) AnswerResult { if r.searchTool == nil { log.Printf("agentic_rag: production runner not fully wired (search tool missing)") return AnswerResult{FinalAnswer: emptyResultMessage, Empty: true} } - // The router tool returns a doc list; feed it as the search DocScope. + route := RouteNode(ctx, r.db, question, modeLabel) + + // Base hybrid search, optionally scoped by the nav router for decomposition + // modes. + searchFn := r.hybridSearchFn(ctx, question, keywords, modeLabel) + + // P8/R4: web fallback is phase-gated — only wired for modes whose + // AvailableTools actually include web_search (high/ultra), AND only when a + // web provider is configured. Low/medium never trigger external web requests + // from an empty KB result. Unconfigured => no web tool call is ever attempted. + if modeAllowsWeb(modeLabel) { + searchFn = r.webFallbackFn(searchFn) + } + + // P5: prefer wiki when the route suggests it AND the bound KBs carry the + // artifact; fall back to hybrid on empty/absent wiki results. + if route.SuggestsCompilation == "wiki" && r.wikiAvailable(ctx) { + searchFn = r.wikiPreferredSearchFn(searchFn) + } + return RunAgenticRAGWithRoute(ctx, r.db, question, keywords, modeLabel, route, searchFn) +} + +// modeAllowsWeb reports whether the mode's AvailableTools include web_search, so +// web fallback is only reachable in the modes that are supposed to have it +// (high/ultra). Unknown modes are treated as not allowing web. +func modeAllowsWeb(modeLabel string) bool { + mode, ok := GetMode(modeLabel) + if !ok { + return false + } + for _, name := range mode.AvailableTools { + if name == "web_search" { + return true + } + } + return false +} + +// hybridSearchFn builds the base hybrid search closure (optionally doc-scoped +// for decomposition modes). +func (r *ProductionRunner) hybridSearchFn(ctx context.Context, question, keywords, modeLabel string) SearchFn { searchFn := func(ctx context.Context, query, kws string) ([]map[string]interface{}, []map[string]interface{}) { return r.search(ctx, query, kws, nil) } - - // For decomposition modes, route the doc scope first via the nav tool. if mode, _ := GetMode(modeLabel); mode.RequiresDecomposition { docs := r.routeDocs(ctx, question, keywords) if len(docs) > 0 { @@ -85,7 +141,55 @@ func (r *ProductionRunner) Run(ctx context.Context, question, keywords, modeLabe } } } - return RunAgenticRAG(ctx, r.db, question, keywords, modeLabel, searchFn) + return searchFn +} + +// wikiPreferredSearchFn wraps the hybrid searchFn so that each search first asks +// the compiled wiki for the query and only falls back to hybrid when the wiki +// returns nothing (or the wiki backend is unavailable). This is the P5 route +// consumption: a wiki suggestion selects the wiki path without discarding the +// hybrid fallback. +func (r *ProductionRunner) wikiPreferredSearchFn(hybrid SearchFn) SearchFn { + return func(ctx context.Context, query, kws string) ([]map[string]interface{}, []map[string]interface{}) { + chunks, aggs := r.wikiSearch(ctx, query, kws) + if len(chunks) > 0 { + return chunks, aggs + } + return hybrid(ctx, query, kws) + } +} + +// wikiAvailable reports whether the bound datasets carry searchable wiki +// artifacts, so the runner only selects the wiki path when it can actually serve. +func (r *ProductionRunner) wikiAvailable(ctx context.Context) bool { + ws := r.wikiSvc + if ws == nil { + ws = wikisearch.GetService() + } + if ws == nil { + return false + } + return ws.AvailableFor(ctx, r.tenantID, r.datasetIDs) +} + +// wikiSearch invokes the wiki_query tool against the compiled wiki. It returns +// empty chunks (never a hard error) when the service is unavailable or yields +// nothing, so the caller falls back to hybrid search. +func (r *ProductionRunner) wikiSearch(ctx context.Context, query, keywords string) ([]map[string]interface{}, []map[string]interface{}) { + ws := r.wikiSvc + if ws == nil { + ws = wikisearch.GetService() + } + if ws == nil || !ws.AvailableFor(ctx, r.tenantID, r.datasetIDs) { + return nil, nil + } + res, err := ws.QueryPages(ctx, r.tenantID, r.datasetIDs, query, keywords, 12) + if err != nil || len(res.Chunks) == 0 { + return nil, nil + } + // P7: backfill the original source chunks referenced by the compiled page + // hits, deduped and bounded, so the answer can cite raw evidence. + return r.expandCompiledEvidence(ctx, res.Chunks, res.DocAggs) } // search invokes the hybrid_search tool and normalizes its chunk output. @@ -108,6 +212,204 @@ func (r *ProductionRunner) search(ctx context.Context, query, keywords string, d return res.Chunks, nil } +// expandCompiledEvidence backfills the ORIGINAL source chunks a compiled-page +// hit was built from (P7/R3). It collects the page hits' source_chunk_ids +// (bounded per page and in total), then asks the concrete wiki service to fetch +// them BY ID — scoped to the tenant + datasets — so the answer can cite raw +// evidence. When the page hits carry no source ids, the service is unavailable, +// or none of the ids resolve, the page results are kept as-is (safe degradation; +// nothing is fabricated). +func (r *ProductionRunner) expandCompiledEvidence(ctx context.Context, chunks, aggs []map[string]interface{}) ([]map[string]interface{}, []map[string]interface{}) { + if len(chunks) == 0 { + return chunks, aggs + } + ws := r.wikiSvc + if ws == nil { + ws = wikisearch.GetService() + } + if ws == nil { + return chunks, aggs + } + const maxEvidencePerPage = 4 + const maxEvidenceTotal = 12 + + // Collect bounded source-chunk ids from the page hits (deduped, in page + // order), grouped by dataset so the backfill stays within each KB's scope. + var sourceIDs []string + seen := map[string]bool{} + datasets := map[string]bool{} + for _, c := range chunks { + if len(sourceIDs) >= maxEvidenceTotal { + break + } + ids := stringSlice(c["source_chunk_ids"]) + count := 0 + for _, id := range ids { + if count >= maxEvidencePerPage { + break + } + if id == "" || seen[id] { + continue + } + seen[id] = true + count++ + sourceIDs = append(sourceIDs, id) + if ds := stringValue(c["dataset_id"]); ds != "" { + datasets[ds] = true + } + if len(sourceIDs) >= maxEvidenceTotal { + break + } + } + } + if len(sourceIDs) == 0 { + return chunks, aggs + } + // Scope the backfill to the page hits' datasets (fall back to all bound + // datasets when the page hits carry none). Build a fresh slice: never mutate + // r.datasetIDs. + scope := make([]string, 0, len(r.datasetIDs)) + if len(datasets) == 0 { + scope = append(scope, r.datasetIDs...) + } else { + for ds := range datasets { + scope = append(scope, ds) + } + } + evidence, err := ws.BackfillChunks(ctx, r.tenantID, scope, sourceIDs) + if err != nil || len(evidence) == 0 { + return chunks, aggs + } + + // Stable merge: page results first (in retrieval order), then the backfilled + // evidence, deduped by chunk key. + merged := append([]map[string]interface{}(nil), chunks...) + keys := map[string]bool{} + for _, c := range chunks { + if k := chunkKey(c); k != "" { + keys[k] = true + } + } + for _, e := range evidence { + k := chunkKey(e) + if k != "" && !keys[k] { + keys[k] = true + merged = append(merged, e) + } + } + // Doc aggs: union the page doc aggs with the evidence docs. + dseen := map[string]bool{} + for _, d := range aggs { + if id, _ := d["doc_id"].(string); id != "" { + dseen[id] = true + } + } + for _, e := range evidence { + id := stringValue(e["doc_id"]) + if id == "" { + continue + } + if !dseen[id] { + dseen[id] = true + aggs = append(aggs, map[string]interface{}{"doc_id": id, "doc_name": stringValue(e["docnm_kwd"])}) + } + } + return merged, aggs +} + +// webFallbackFn wraps a SearchFn so that, when the KB search returns nothing, a +// configured web provider is invoked to fill the gap (P8). It is only used when +// webTool is non-nil; otherwise it returns the hybrid path unchanged and no web +// tool call is ever attempted (no failing call when unconfigured). +func (r *ProductionRunner) webFallbackFn(hybrid SearchFn) SearchFn { + if r.webTool == nil { + return hybrid + } + return func(ctx context.Context, query, kws string) ([]map[string]interface{}, []map[string]interface{}) { + chunks, aggs := hybrid(ctx, query, kws) + if len(chunks) > 0 { + return chunks, aggs + } + raw, err := r.webTool.InvokableRun(ctx, mustJSON(map[string]interface{}{"query": query, "keywords": kws})) + if err != nil { + return nil, nil + } + var res struct { + Chunks []map[string]interface{} `json:"chunks"` + Results []map[string]interface{} `json:"results"` + } + if err := json.Unmarshal([]byte(raw), &res); err != nil { + return nil, nil + } + // Normalize web evidence into the same agentic evidence shape as KB + // chunks. Accept both the agent "chunks" envelope and the Tavily + // "results" envelope (tavily.go returns {"results":[...]}); each result + // contributes content + a doc_id reference so the answer can retain the + // source URL. + src := res.Chunks + if len(src) == 0 { + src = res.Results + } + out := make([]map[string]interface{}, 0, len(src)) + for _, c := range src { + url := firstNonEmpty(stringValue(c["url"]), stringValue(c["link"]), stringValue(c["source"])) + if url == "" { + continue + } + content := firstNonEmpty(stringValue(c["content"]), stringValue(c["raw_content"]), stringValue(c["text"])) + if content == "" { + continue + } + docID := stringValue(c["doc_id"]) + if docID == "" { + docID = url + "|" + stringValue(c["source"]) + } + out = append(out, map[string]interface{}{ + "chunk_id": docID, "content_with_weight": content, + "doc_id": docID, "docnm_kwd": firstNonEmpty(stringValue(c["title"]), stringValue(c["source"])), + "dataset_id": stringValue(c["dataset_id"]), "url": url, "source": "web", + }) + } + if len(out) == 0 { + return nil, nil + } + return out, nil + } +} + +func stringValue(v interface{}) string { + if s, ok := v.(string); ok { + return s + } + return "" +} + +func firstNonEmpty(ss ...string) string { + for _, s := range ss { + if strings.TrimSpace(s) != "" { + return s + } + } + return "" +} + +func stringSlice(v interface{}) []string { + if raw, ok := v.([]string); ok { + return raw + } + arr, ok := v.([]interface{}) + if !ok { + return nil + } + out := make([]string, 0, len(arr)) + for _, item := range arr { + if s, ok := item.(string); ok { + out = append(out, s) + } + } + return out +} + // routeDocs derives the doc scope via the canonical dataset-nav router // (harness.NavigateDatasetByTree — the full LLM two-round selection). It routes // across ALL bound datasets and merges the doc ids, so every KB contributes its diff --git a/internal/agent/harness/production_wiki_test.go b/internal/agent/harness/production_wiki_test.go new file mode 100644 index 0000000000..31c9344814 --- /dev/null +++ b/internal/agent/harness/production_wiki_test.go @@ -0,0 +1,276 @@ +package harness + +import ( + "context" + "strings" + "sync" + "testing" + + "gorm.io/gorm" + + "ragflow/internal/agent/component" + "ragflow/internal/service/wikisearch" +) + +// fakeWikiSvcHarness is a deterministic wikisearch.Service double for the +// production runner. +type fakeWikiSvcHarness struct { + available bool + // pages are pre-shaped page chunks (map form) returned by QueryPages. + pages []map[string]interface{} + // backfill maps chunk id -> evidence content for BackfillChunks. + backfill map[string]string + calls []string + queryCount int + mu sync.Mutex +} + +func (f *fakeWikiSvcHarness) AvailableFor(_ context.Context, _ string, _ []string) bool { + return f.available +} + +func (f *fakeWikiSvcHarness) QueryPages(_ context.Context, _ string, _ []string, query, _ string, _ int) (wikisearch.SearchResult, error) { + f.mu.Lock() + f.queryCount++ + f.calls = append(f.calls, query) + f.mu.Unlock() + res := wikisearch.SearchResult{Chunks: append([]map[string]interface{}(nil), f.pages...), DocAggs: []map[string]interface{}{}} + for _, c := range res.Chunks { + if docID, _ := c["doc_id"].(string); docID != "" { + res.DocAggs = append(res.DocAggs, map[string]interface{}{"doc_id": docID, "doc_name": c["docnm_kwd"]}) + } + } + return res, nil +} + +func (f *fakeWikiSvcHarness) seenQueries() []string { + f.mu.Lock() + defer f.mu.Unlock() + return append([]string(nil), f.calls...) +} + +// backfill is the fake's per-id evidence map (chunk id -> content). +func (f *fakeWikiSvcHarness) BackfillChunks(_ context.Context, _ string, _ []string, chunkIDs []string) ([]map[string]interface{}, error) { + out := make([]map[string]interface{}, 0, len(chunkIDs)) + for _, id := range chunkIDs { + content, ok := f.backfill[id] + if !ok { + continue + } + out = append(out, map[string]interface{}{ + "chunk_id": id, "content_with_weight": content, "doc_id": "d1", "docnm_kwd": "Doc", "dataset_id": "kb1", + }) + } + return out, nil +} + +// wikiRouteChat returns a route JSON suggesting wiki for the route stage and a +// plain final answer otherwise. +type wikiRouteChat struct{} + +func (wikiRouteChat) Invoke(_ context.Context, _ *gorm.DB, req component.ChatInvokeRequest) (*component.ChatInvokeResponse, error) { + // The route stage carries the route prompt (with "suggests_compilation") as + // a system message; detect it across any message (system or user). + isRoute := false + for _, m := range req.Messages { + if strings.Contains(m.Content, "suggests_compilation") { + isRoute = true + break + } + } + if isRoute { + return &component.ChatInvokeResponse{Content: `{"question_type":"analytical","requires_decomposition":false,"suggests_compilation":"wiki"}`}, nil + } + return &component.ChatInvokeResponse{Content: "final wiki answer"}, nil +} + +func installWikiRouteChat(t *testing.T) { + t.Helper() + component.SetDefaultChatInvoker(wikiRouteChat{}) + t.Cleanup(func() { component.SetDefaultChatInvoker(nil) }) +} + +// TestProductionRunner_WikiPreferred_WhenSuggested drives a low-mode run with a +// wiki suggestion and an available wiki service, and asserts the wiki service is +// queried (P5: route suggestion selects the wiki path). +func TestProductionRunner_WikiPreferred_WhenSuggested(t *testing.T) { + installWikiRouteChat(t) + hybrid := &fakeInvokableTool{name: "hybrid_search", fn: func(_ context.Context, _ string) string { + return `{"chunks":[]}` + }} + wikiSvc := &fakeWikiSvcHarness{available: true, pages: []map[string]interface{}{ + {"chunk_id": "wiki/entity/alpha", "content_with_weight": "# Alpha", "doc_id": "kb1", "docnm_kwd": "Alpha", "wiki_slug_kwd": "entity/alpha", "dataset_id": "kb1"}, + }} + runner := newProductionRunnerWithTools(nil, "t1", []string{"kb1"}, hybrid, nil) + runner.wikiSvc = wikiSvc + res := runner.Run(context.Background(), "What is Alpha?", "", "low") + if res.FinalAnswer != "final wiki answer" { + t.Errorf("final answer = %q, want wiki chat output", res.FinalAnswer) + } + if len(wikiSvc.seenQueries()) == 0 { + t.Errorf("wiki service was not queried despite a wiki suggestion") + } +} + +// TestProductionRunner_WikiEmpty_FallsBackToHybrid asserts that when the wiki +// service yields nothing, the runner falls back to hybrid search (P5 must not +// discard the general retrieval fallback). +func TestProductionRunner_WikiEmpty_FallsBackToHybrid(t *testing.T) { + installWikiRouteChat(t) + hybrid := &fakeInvokableTool{name: "hybrid_search", fn: func(_ context.Context, _ string) string { + return `{"chunks":[{"chunk_id":"c1","content_with_weight":"hybrid evidence"}]}` + }} + wikiSvc := &fakeWikiSvcHarness{available: true} // no pages + runner := newProductionRunnerWithTools(nil, "t1", []string{"kb1"}, hybrid, nil) + runner.wikiSvc = wikiSvc + res := runner.Run(context.Background(), "What is Alpha?", "", "low") + if len(wikiSvc.seenQueries()) == 0 { + t.Errorf("wiki service should have been attempted") + } + if !strings.Contains(hybrid.args(), `"query":"What is Alpha?"`) { + t.Errorf("hybrid fallback was not invoked after empty wiki; args=%s", hybrid.args()) + } + if res.FinalAnswer == "" { + t.Errorf("expected a final answer from the hybrid fallback") + } +} + +// TestProductionRunner_NoWikiWithoutSuggestion asserts the wiki service is NOT +// queried when the route does not suggest wiki. +func TestProductionRunner_NoWikiWithoutSuggestion(t *testing.T) { + // route chat returns no suggestion for a generic route call + installRouteChat(t) // routeChat returns plain text -> route falls back (no suggestion) + hybrid := &fakeInvokableTool{name: "hybrid_search", fn: func(_ context.Context, _ string) string { + return `{"chunks":[{"chunk_id":"c1","content_with_weight":"evidence"}]}` + }} + wikiSvc := &fakeWikiSvcHarness{available: true, pages: []map[string]interface{}{ + {"chunk_id": "wiki/s", "content_with_weight": "c", "doc_id": "kb1", "docnm_kwd": "T", "wiki_slug_kwd": "s", "dataset_id": "kb1"}, + }} + runner := newProductionRunnerWithTools(nil, "t1", []string{"kb1"}, hybrid, nil) + runner.wikiSvc = wikiSvc + runner.Run(context.Background(), "What is Alpha?", "", "low") + if len(wikiSvc.seenQueries()) != 0 { + t.Errorf("wiki service must not be queried without a wiki suggestion; calls=%v", wikiSvc.seenQueries()) + } +} + +// TestProductionRunner_WebFallback_UnconfiguredIsNoop asserts that with no web +// tool configured the runner never attempts a web call (P8 gate). +func TestProductionRunner_WebFallback_UnconfiguredIsNoop(t *testing.T) { + webTool := &fakeInvokableTool{name: "web_search", fn: func(_ context.Context, _ string) string { + return `{"chunks":[{"chunk_id":"w1","content_with_weight":"web evidence"}]}` + }} + hybrid := &fakeInvokableTool{name: "hybrid_search", fn: func(_ context.Context, _ string) string { + return `{"chunks":[]}` + }} + // webTool is NOT set on the runner -> the runner must not invoke it. + runner := newProductionRunnerWithTools(nil, "t1", []string{"kb1"}, hybrid, nil) + // With webTool nil, webFallbackFn returns the hybrid path unchanged, so the + // web tool is never called. + fn := runner.webFallbackFn(func(ctx context.Context, q, k string) ([]map[string]interface{}, []map[string]interface{}) { + return nil, nil + }) + chunks, _ := fn(context.Background(), "q", "k") + if len(chunks) != 0 { + t.Errorf("expected no chunks from the no-web fallback path") + } + if webTool.args() != "" { + t.Errorf("web tool must not be called when unconfigured; args=%s", webTool.args()) + } +} + +// TestProductionRunner_WebFallback_TavilyResultsNormalized asserts R2: the web +// fallback consumes the Tavily `{"results":[...]}` envelope (tavily.go contract) +// and normalizes each result into an agent evidence chunk with a doc_id +// reference, not the agent `chunks` shape. +func TestProductionRunner_WebFallback_TavilyResultsNormalized(t *testing.T) { + webTool := &fakeInvokableTool{name: "web_search", fn: func(_ context.Context, _ string) string { + return `{"results":[{"title":"Alpha docs","url":"https://example.com/x","content":"web evidence body","source":"example.com"}]}` + }} + runner := newProductionRunnerWithTools(nil, "t1", []string{"kb1"}, nil, nil) + runner.webTool = webTool + fn := runner.webFallbackFn(func(ctx context.Context, q, k string) ([]map[string]interface{}, []map[string]interface{}) { + return nil, nil + }) + chunks, _ := fn(context.Background(), "q", "k") + if len(chunks) != 1 { + t.Fatalf("chunks = %d, want 1 normalized Tavily result", len(chunks)) + } + if chunks[0]["content_with_weight"] != "web evidence body" { + t.Errorf("content = %v, want web evidence body", chunks[0]["content_with_weight"]) + } + if chunks[0]["doc_id"] == "" || !strings.Contains(chunks[0]["doc_id"].(string), "https://example.com/x") { + t.Errorf("doc_id must reference the source url: %v", chunks[0]["doc_id"]) + } + if chunks[0]["source"] != "web" { + t.Errorf("source = %v, want web", chunks[0]["source"]) + } + if webTool.args() == "" { + t.Errorf("web tool was not invoked when hybrid was empty") + } +} + +// TestModeAllowsWeb asserts R4 gating: only modes whose AvailableTools include +// web_search (high/ultra) allow web fallback; low/medium/unknown do not. +func TestModeAllowsWeb(t *testing.T) { + if !modeAllowsWeb("high") || !modeAllowsWeb("ultra") { + t.Errorf("high/ultra must allow web search") + } + for _, m := range []string{"low", "medium", "fast", "unknown"} { + if modeAllowsWeb(m) { + t.Errorf("mode %q must NOT allow web search", m) + } + } +} + +// TestProductionRunner_CompiledEvidenceExpansion asserts P7/R3: page hits +// carrying source_chunk_ids get the ORIGINAL chunks fetched by id (via the wiki +// service BackfillChunks) appended after the page results, deduped. +func TestProductionRunner_CompiledEvidenceExpansion(t *testing.T) { + wikiSvc := &fakeWikiSvcHarness{available: true, backfill: map[string]string{"c1": "raw chunk 1", "c2": "raw chunk 2"}} + runner := newProductionRunnerWithTools(nil, "t1", []string{"kb1"}, nil, nil) + runner.wikiSvc = wikiSvc + chunks := []map[string]interface{}{ + { + "chunk_id": "wiki/s", "content_with_weight": "# Page", "doc_id": "d1", "docnm_kwd": "Page", "dataset_id": "kb1", + "source_chunk_ids": []interface{}{"c1", "c2"}, + }, + } + merged, aggs := runner.expandCompiledEvidence(context.Background(), chunks, []map[string]interface{}{}) + if len(merged) != 3 { + t.Fatalf("merged = %d, want 3 (page + 2 backfilled evidence chunks)", len(merged)) + } + if merged[0]["chunk_id"] != "wiki/s" || merged[1]["chunk_id"] != "c1" || merged[2]["chunk_id"] != "c2" { + t.Errorf("merge order wrong: %v", merged) + } + if merged[1]["content_with_weight"] != "raw chunk 1" { + t.Errorf("evidence chunk content = %v, want raw chunk 1 (by-id backfill)", merged[1]["content_with_weight"]) + } + // Doc aggs must include the evidence doc. + found := false + for _, d := range aggs { + if d["doc_id"] == "d1" { + found = true + } + } + if !found { + t.Errorf("evidence doc not unioned into doc aggs: %v", aggs) + } +} + +// TestProductionRunner_CompiledEvidence_NoServiceIsNoop asserts P7/R3 degrades +// safely: when the wiki service is unavailable, page hits (even with +// source_chunk_ids) keep the page results unchanged and fabricate nothing. +func TestProductionRunner_CompiledEvidence_NoServiceIsNoop(t *testing.T) { + runner := newProductionRunnerWithTools(nil, "t1", []string{"kb1"}, nil, nil) // wikiSvc nil + chunks := []map[string]interface{}{ + {"chunk_id": "wiki/s", "content_with_weight": "# Page", "dataset_id": "kb1", "source_chunk_ids": []interface{}{"c1", "c2"}}, + } + merged, _ := runner.expandCompiledEvidence(context.Background(), chunks, []map[string]interface{}{}) + if len(merged) != 1 { + t.Fatalf("merged = %d, want 1 (no service => no fabricated evidence)", len(merged)) + } + if merged[0]["chunk_id"] != "wiki/s" { + t.Errorf("page result changed: %v", merged[0]) + } +} diff --git a/internal/agent/harness/route.go b/internal/agent/harness/route.go index 5459cfd0d5..9e5ca228c1 100644 --- a/internal/agent/harness/route.go +++ b/internal/agent/harness/route.go @@ -50,9 +50,10 @@ Output format (JSON): ` type routeResult struct { - QuestionType string `json:"question_type"` - RequiresDecomp *bool `json:"requires_decomposition"` - Reasoning string `json:"reasoning"` + QuestionType string `json:"question_type"` + RequiresDecomp *bool `json:"requires_decomposition"` + SuggestsCompilation string `json:"suggests_compilation"` + Reasoning string `json:"reasoning"` } // RouteNode mirrors Python route_node. It classifies the question into a @@ -104,10 +105,28 @@ func decide(question, modeLabel string, res routeResult) RouteDecision { QuestionType: qType, RequiresDecomposition: mode.RequiresDecomposition && needDecomp, ExecutionStrategy: mode.Strategy, + SuggestsCompilation: normalizeCompilationSuggestion(res.SuggestsCompilation), Reasoning: res.Reasoning, } } +// normalizeCompilationSuggestion maps the LLM's free-text suggestion to a +// canonical compiled-artifact key: "", "toc", "graph", or "wiki". Anything else +// (including "null") collapses to "" so the production runner never routes on an +// unknown artifact name. +func normalizeCompilationSuggestion(s string) string { + switch strings.ToLower(strings.TrimSpace(s)) { + case "toc": + return "toc" + case "graph", "knowledge_graph", "kg": + return "graph" + case "wiki", "compiled": + return "wiki" + default: + return "" + } +} + func fallbackRoute(question, modeLabel, reason string) RouteDecision { return RouteDecision{ Question: question, ThinkingMode: modeLabel, QuestionType: "factual", diff --git a/internal/agent/harness/route_test.go b/internal/agent/harness/route_test.go index 03d9170079..da9eec04a7 100644 --- a/internal/agent/harness/route_test.go +++ b/internal/agent/harness/route_test.go @@ -74,6 +74,48 @@ func TestRouteNode_EmptyQuestionFallsBack(t *testing.T) { } } +// TestRouteNode_SuggestsCompilation asserts the route preserves a normalized +// compiled-artifact suggestion (P5) so the production runner can prefer the wiki +// tool. +func TestRouteNode_SuggestsCompilation(t *testing.T) { + installChat(t, `{"question_type":"analytical","requires_decomposition":true,"suggests_compilation":"wiki"}`) + r := RouteNode(context.Background(), nil, "What does the domain say about X?", "medium") + if r.SuggestsCompilation != "wiki" { + t.Errorf("suggests_compilation = %q, want wiki", r.SuggestsCompilation) + } +} + +// TestNormalizeCompilationSuggestion asserts free-text suggestions map to the +// canonical keys and unknown/null collapse to "". +func TestNormalizeCompilationSuggestion(t *testing.T) { + cases := map[string]string{ + "wiki": "wiki", + "WIKI": "wiki", + "compiled": "wiki", + "graph": "graph", + "kg": "graph", + "toc": "toc", + "null": "", + "": "", + "spaghetti": "", + } + for in, want := range cases { + if got := normalizeCompilationSuggestion(in); got != want { + t.Errorf("normalizeCompilationSuggestion(%q) = %q, want %q", in, got, want) + } + } +} + +// TestRouteNode_WikiSuggestionSurvivesFence asserts a fenced JSON route still +// carries the wiki suggestion through decide(). +func TestRouteNode_WikiSuggestionSurvivesFence(t *testing.T) { + installChat(t, "```json\n{\"question_type\":\"procedural\",\"requires_decomposition\":false,\"suggests_compilation\":\"graph\"}\n```") + r := RouteNode(context.Background(), nil, "How is this structured?", "medium") + if r.SuggestsCompilation != "graph" { + t.Errorf("suggests_compilation = %q, want graph", r.SuggestsCompilation) + } +} + // TestPlannerNode_DirectMode asserts a non-decomposed route yields one coarse // claim without calling the LLM. func TestPlannerNode_DirectMode(t *testing.T) { diff --git a/internal/agent/harness/types.go b/internal/agent/harness/types.go index 3b9c56d357..2f3393d1c8 100644 --- a/internal/agent/harness/types.go +++ b/internal/agent/harness/types.go @@ -26,7 +26,13 @@ type RouteDecision struct { QuestionType string // factual | comparative | analytical | procedural | exploratory | verification | summarization RequiresDecomposition bool ExecutionStrategy string // direct_search | decompose_and_search | agentic_research | deep_research - Reasoning string + // SuggestsCompilation is the compiled-artifact type the route suggests using + // for retrieval: "" (none) | "toc" | "graph" | "wiki". Mirrors Python's + // route `suggests_compilation`. It is preserved so the production runner can + // prefer a compiled wiki/graph tool when the bound KBs actually carry the + // artifact, and fall back to general hybrid search otherwise. + SuggestsCompilation string + Reasoning string } // ClaimTarget mirrors Python ClaimTarget. diff --git a/internal/agent/tool/registry.go b/internal/agent/tool/registry.go index 89fafeefd7..f5143b9ec2 100644 --- a/internal/agent/tool/registry.go +++ b/internal/agent/tool/registry.go @@ -40,6 +40,7 @@ var registry = map[string]Factory{ "hybrid_search": noConfig("hybrid_search", func() einotool.BaseTool { return NewAgenticSearchTool(toolHybridSearch) }), "vector_search": noConfig("vector_search", func() einotool.BaseTool { return NewAgenticSearchTool(toolVectorSearch) }), "bm25_search": noConfig("bm25_search", func() einotool.BaseTool { return NewAgenticSearchTool(toolBM25Search) }), + "wiki_query": noConfig("wiki_query", func() einotool.BaseTool { return NewWikiQueryTool() }), "deepl": noConfig("deepl", func() einotool.BaseTool { return NewDeepLTool() }), "duckduckgo": buildDuckDuckGoTool, "email": buildEmailTool, diff --git a/internal/agent/tool/wiki_query.go b/internal/agent/tool/wiki_query.go new file mode 100644 index 0000000000..bd02c9a165 --- /dev/null +++ b/internal/agent/tool/wiki_query.go @@ -0,0 +1,115 @@ +// +// Copyright 2026 The InfiniFlow Authors. All Rights Reserved. +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. +// + +package tool + +import ( + "context" + "encoding/json" + "fmt" + "strings" + + einotool "github.com/cloudwego/eino/components/tool" + "github.com/cloudwego/eino/schema" + + "ragflow/internal/service/wikisearch" +) + +// WikiQueryTool is the wiki_query agent tool (Python harness/tools/exploration.py +// wiki_query). It hybrid-searches the compiled wiki/artifact pages of the bound +// datasets and returns each page's rendered markdown as a chunk, narrowed by +// keywords. Input keeps query + keywords so the LLM's tool schema matches the +// other search tools. +// +// The tool is scoped to the calling tenant and bound datasets, so results never +// leak across tenants/KBs. When the wiki-search service is not configured, or the +// bound datasets have no wiki artifacts, it returns an empty result so the agent +// falls back to general hybrid search. +type WikiQueryTool struct { + service wikisearch.Service // nil => resolve lazily via wikisearch.GetService() + topN int +} + +// NewWikiQueryTool returns the wiki_query tool. The service is resolved lazily +// from the wikisearch singleton unless overridden for tests. +func NewWikiQueryTool() *WikiQueryTool { + return &WikiQueryTool{topN: 12} +} + +// newWikiQueryToolWithService builds a tool with an injected service, for tests. +func newWikiQueryToolWithService(service wikisearch.Service) *WikiQueryTool { + return &WikiQueryTool{service: service, topN: 12} +} + +type wikiQueryArgs struct { + Query string `json:"query"` + Keywords string `json:"keywords,omitempty"` +} + +func (w *WikiQueryTool) Info(_ context.Context) (*schema.ToolInfo, error) { + return &schema.ToolInfo{ + Name: "wiki_query", + Desc: "Search the compiled wiki of the bound knowledge base(s). Returns rendered wiki page content as passages.", + ParamsOneOf: schema.NewParamsOneOfByParams(map[string]*schema.ParameterInfo{ + "query": {Type: schema.String, Required: true, Desc: "The search query."}, + "keywords": {Type: schema.String, Desc: "Comma-separated keywords to narrow results."}, + }), + }, nil +} + +// InvokableRun executes a wiki lookup. It never returns a hard error for an +// empty/unconfigured backend so the agent can fall back to hybrid search. +func (w *WikiQueryTool) InvokableRun(ctx context.Context, argumentsInJSON string, _ ...einotool.Option) (string, error) { + var args wikiQueryArgs + if err := json.Unmarshal([]byte(argumentsInJSON), &args); err != nil { + return "", fmt.Errorf("wiki_query: parse arguments: %w", err) + } + svc := w.service + if svc == nil { + svc = wikisearch.GetService() + } + tenantID := canvasTenantID(ctx) + datasetIDs := canvasDatasetIDs(ctx, nil) + if svc == nil || tenantID == "" || len(datasetIDs) == 0 { + return emptyWikiResult(), nil + } + if !svc.AvailableFor(ctx, tenantID, datasetIDs) { + return emptyWikiResult(), nil + } + topN := w.topN + if topN <= 0 { + topN = 12 + } + res, err := svc.QueryPages(ctx, tenantID, datasetIDs, strings.TrimSpace(args.Query), args.Keywords, topN) + if err != nil { + return emptyWikiResult(), nil + } + if res.Chunks == nil { + res.Chunks = []map[string]interface{}{} + } + if res.DocAggs == nil { + res.DocAggs = []map[string]interface{}{} + } + out, err := json.Marshal(map[string]interface{}{"answer": "", "chunks": res.Chunks, "doc_aggs": res.DocAggs}) + if err != nil { + return emptyWikiResult(), nil + } + return string(out), nil +} + +func emptyWikiResult() string { + return `{"answer":"","chunks":[],"doc_aggs":[]}` +} diff --git a/internal/agent/tool/wiki_query_test.go b/internal/agent/tool/wiki_query_test.go new file mode 100644 index 0000000000..11e6911187 --- /dev/null +++ b/internal/agent/tool/wiki_query_test.go @@ -0,0 +1,125 @@ +package tool + +import ( + "context" + "encoding/json" + "testing" + + "ragflow/internal/agent/runtime" + "ragflow/internal/service/wikisearch" +) + +// fakeWikiService is a deterministic wikisearch.Service double. +type fakeWikiService struct { + available map[string]bool // datasetID -> has artifact + pages func(query string) []map[string]interface{} + backfill map[string]string // chunk id -> content + callCount int +} + +func (f *fakeWikiService) AvailableFor(_ context.Context, _ string, datasetIDs []string) bool { + for _, ds := range datasetIDs { + if f.available[ds] { + return true + } + } + return false +} + +func (f *fakeWikiService) QueryPages(_ context.Context, _ string, _ []string, query, _ string, topN int) (wikisearch.SearchResult, error) { + f.callCount++ + if f.pages == nil || len(f.pages(query)) == 0 { + return wikisearch.SearchResult{Chunks: []map[string]interface{}{}, DocAggs: []map[string]interface{}{}}, nil + } + res := wikisearch.SearchResult{Chunks: append([]map[string]interface{}(nil), f.pages(query)...), DocAggs: []map[string]interface{}{}} + seen := map[string]bool{} + for _, c := range res.Chunks { + if docID, _ := c["doc_id"].(string); docID != "" && !seen[docID] { + seen[docID] = true + res.DocAggs = append(res.DocAggs, map[string]interface{}{"doc_id": docID, "doc_name": c["docnm_kwd"]}) + } + } + return res, nil +} + +func (f *fakeWikiService) BackfillChunks(_ context.Context, _ string, _ []string, chunkIDs []string) ([]map[string]interface{}, error) { + out := make([]map[string]interface{}, 0, len(chunkIDs)) + for _, id := range chunkIDs { + content, ok := f.backfill[id] + if !ok { + continue + } + out = append(out, map[string]interface{}{"chunk_id": id, "content_with_weight": content, "doc_id": "d1", "docnm_kwd": "Doc"}) + } + return out, nil +} + +// wikiToolRun runs the wiki_query tool with a single-dataset canvas context +// (the tool derives tenant + dataset scope from canvas state). +func wikiToolRun(t *testing.T, svc wikisearch.Service, tenant string, kb string, query string) map[string]interface{} { + t.Helper() + state := runtime.NewCanvasState("run-1", "task-1") + state.Sys["tenant_id"] = tenant + state.Sys["dataset_id"] = kb + ctx := runtime.WithState(context.Background(), state) + tool := newWikiQueryToolWithService(svc) + args, _ := json.Marshal(map[string]interface{}{"query": query}) + raw, err := tool.InvokableRun(ctx, string(args)) + if err != nil { + t.Fatalf("wiki_query InvokableRun err = %v", err) + } + var out map[string]interface{} + if err := json.Unmarshal([]byte(raw), &out); err != nil { + t.Fatalf("bad wiki_query output: %v", err) + } + return out +} + +func TestWikiQueryTool_ReturnsPages(t *testing.T) { + svc := &fakeWikiService{ + available: map[string]bool{"kb1": true}, + pages: func(q string) []map[string]interface{} { + return []map[string]interface{}{{"chunk_id": "wiki/entity/alpha", "content_with_weight": "# Alpha", "doc_id": "kb1", "docnm_kwd": "Alpha", "wiki_slug_kwd": "entity/alpha", "dataset_id": "kb1"}} + }, + } + out := wikiToolRun(t, svc, "t1", "kb1", "alpha") + chunks, _ := out["chunks"].([]interface{}) + if len(chunks) != 1 { + t.Fatalf("chunks = %d, want 1", len(chunks)) + } + first := chunks[0].(map[string]interface{}) + if first["content_with_weight"] != "# Alpha" { + t.Errorf("page content = %v, want # Alpha", first["content_with_weight"]) + } + if first["wiki_slug_kwd"] != "entity/alpha" { + t.Errorf("slug = %v, want entity/alpha", first["wiki_slug_kwd"]) + } +} + +func TestWikiQueryTool_EmptyWhenNoArtifact(t *testing.T) { + svc := &fakeWikiService{available: map[string]bool{"kb2": true}} + out := wikiToolRun(t, svc, "t1", "kb1", "alpha") // kb1 has no artifact + if chunks, _ := out["chunks"].([]interface{}); len(chunks) != 0 { + t.Fatalf("chunks = %d, want 0 (kb1 has no wiki artifact)", len(chunks)) + } +} + +func TestWikiQueryTool_EmptyWhenNoService(t *testing.T) { + out := wikiToolRun(t, nil, "t1", "kb1", "alpha") + if chunks, _ := out["chunks"].([]interface{}); len(chunks) != 0 { + t.Fatalf("chunks = %d, want 0 (no service configured)", len(chunks)) + } +} + +func TestWikiQueryTool_ScopeRespected(t *testing.T) { + svc := &fakeWikiService{ + available: map[string]bool{"kb1": true}, + pages: func(q string) []map[string]interface{} { + return []map[string]interface{}{{"chunk_id": "wiki/s", "content_with_weight": "c", "doc_id": "kb1", "docnm_kwd": "T", "wiki_slug_kwd": "s", "dataset_id": "kb1"}} + }, + } + out := wikiToolRun(t, svc, "t1", "kb1", "alpha") + if _, ok := out["chunks"]; !ok { + t.Fatalf("missing chunks key") + } +} diff --git a/internal/common/parser_config.go b/internal/common/parser_config.go index adc1793b96..df3c53588f 100644 --- a/internal/common/parser_config.go +++ b/internal/common/parser_config.go @@ -1,129 +1,5 @@ package common -import "strings" - -// InjectExtractorLLMID finds all Extractor component entries (keys prefixed -// with "extractor:" or "extractor_") in parserConfig and sets their llm_id -// to the given value. Returns whether any entry was updated. -func InjectExtractorLLMID(parserConfig map[string]interface{}, llmID string) bool { - if parserConfig == nil || llmID == "" { - return false - } - updated := false - for cid, raw := range parserConfig { - compMap, ok := raw.(map[string]interface{}) - if !ok { - continue - } - cidLower := strings.ToLower(cid) - if strings.HasPrefix(cidLower, "extractor:") || strings.HasPrefix(cidLower, "extractor_") { - if current, ok := compMap["llm_id"].(string); !ok || current == "" { - compMap["llm_id"] = llmID - updated = true - } - } - } - return updated -} - -// InjectExtractorEnableMetadata enables auto-metadata (enable_metadata) extraction -// on every Extractor node when the dataset has enable_metadata on and a -// non-empty field set (metadata and/or built_in_metadata). The dataset-level -// enable_metadata flag is authoritative (mirrors Python task_executor.py:519, -// which reads parser_config directly and never consults a per-node flag): a -// shipped DSL that defaults enable_metadata to 0 is still turned on. Only a -// node the user already turned ON (enable_metadata truthy) is left untouched, -// so an explicit per-node enablement keeps its own config. The field schema is -// taken from parserConfig["metadata"] and parserConfig["built_in_metadata"] -// (combined). Returns whether any entry was updated. -func InjectExtractorEnableMetadata(parserConfig map[string]interface{}) bool { - if parserConfig == nil { - return false - } - if !isTruthy(parserConfig["enable_metadata"]) { - return false - } - fields := metadataFieldDefs(parserConfig) - if len(fields) == 0 { - return false - } - updated := false - for cid, raw := range parserConfig { - compMap, ok := raw.(map[string]interface{}) - if !ok { - continue - } - cidLower := strings.ToLower(cid) - if !strings.HasPrefix(cidLower, "extractor:") && !strings.HasPrefix(cidLower, "extractor_") { - continue - } - // The dataset-level enable_metadata flag is authoritative (mirrors - // Python task_executor.py:519, which reads parser_config directly and - // never consults a per-node flag). Only a node the user already turned - // ON (truthy) is left alone; a shipped DSL that defaults the field to - // 0 must still be enabled by the dataset flag, otherwise auto-metadata - // could never turn on for any of the built-in pipelines. - if isTruthy(compMap["enable_metadata"]) { - continue - } - compMap["enable_metadata"] = 1 - compMap["metadata"] = fields - updated = true - } - return updated -} - -// metadataFieldDefs combines parserConfig["metadata"] and -// parserConfig["built_in_metadata"] into the field list injected as -// metadata (each entry keeps key/type/description/enum, mirroring the -// stored shape from dataset/helpers.go normalizeMetadataConfigFields). -// -// It returns []any (i.e. []interface{}) rather than []map[string]interface{} -// because the injected value is handed to NewExtractorComponent, which reads -// params["metadata"].([]any). A []map[string]interface{} value would fail that -// type assertion (Go slice types are not covariant) and the field schema would -// be silently dropped, so auto-metadata never reached ExtractorParam.Metadata. -func metadataFieldDefs(parserConfig map[string]interface{}) []any { - var out []any - for _, key := range []string{"metadata", "built_in_metadata"} { - raw, ok := parserConfig[key].([]interface{}) - if !ok { - continue - } - for _, item := range raw { - m, ok := item.(map[string]interface{}) - if !ok { - continue - } - k, _ := m["key"].(string) - if strings.TrimSpace(k) == "" { - continue - } - out = append(out, m) - } - } - return out -} - -// isTruthy reports whether a parserConfig flag (e.g. enable_metadata) is on. -// It tolerates bool, numeric >0 and the strings "true"/"1" so storage -// representation differences don't silently disable the feature. -func isTruthy(v interface{}) bool { - switch t := v.(type) { - case bool: - return t - case string: - return t == "true" || t == "1" || t == "True" || t == "TRUE" - case float64: - return t > 0 - case int: - return t > 0 - case int64: - return t > 0 - } - return false -} - // deepCopyMap duplicates a JSON-like map so later merges do not mutate shared defaults. func deepCopyMap(source map[string]interface{}) map[string]interface{} { if source == nil { diff --git a/internal/common/parser_config_test.go b/internal/common/parser_config_test.go deleted file mode 100644 index e6d1195eba..0000000000 --- a/internal/common/parser_config_test.go +++ /dev/null @@ -1,135 +0,0 @@ -package common - -import "testing" - -func TestInjectExtractorLLMID_SkipWhenUUID(t *testing.T) { - uuid := "9e819c2442b14f9dab46062916e29195" - pc := map[string]interface{}{ - "Extractor:A": map[string]interface{}{ - "llm_id": uuid, - }, - } - InjectExtractorLLMID(pc, "Qwen/Qwen3-8B@siliconflow") - id := pc["Extractor:A"].(map[string]interface{})["llm_id"].(string) - if id != uuid { - t.Fatalf("expected UUID preserved, got %q", id) - } -} - -func TestInjectExtractorLLMID_SkipWhenComposite(t *testing.T) { - composite := "Qwen/Qwen3-8B@siliconflow" - pc := map[string]interface{}{ - "Extractor:B": map[string]interface{}{ - "llm_id": composite, - }, - } - InjectExtractorLLMID(pc, "DeepSeek@siliconflow") - id := pc["Extractor:B"].(map[string]interface{})["llm_id"].(string) - if id != composite { - t.Fatalf("expected composite preserved, got %q", id) - } -} - -func TestInjectExtractorLLMID_InjectWhenEmpty(t *testing.T) { - defaultLLM := "Qwen/Qwen3-8B@siliconflow" - pc := map[string]interface{}{ - "Extractor:C": map[string]interface{}{}, - } - InjectExtractorLLMID(pc, defaultLLM) - id := pc["Extractor:C"].(map[string]interface{})["llm_id"].(string) - if id != defaultLLM { - t.Fatalf("expected %q injected, got %q", defaultLLM, id) - } -} - -func TestInjectExtractorLLMID_NoExtractor(t *testing.T) { - pc := map[string]interface{}{ - "Parser:X": map[string]interface{}{"llm_id": ""}, - } - InjectExtractorLLMID(pc, "default@provider") - if _, ok := pc["Parser:X"]; !ok { - t.Fatal("expected Parser:X still present") - } -} - -func TestInjectExtractorEnableMetadata_Disabled(t *testing.T) { - pc := map[string]interface{}{ - "Extractor:A": map[string]interface{}{}, - } - if InjectExtractorEnableMetadata(pc) { - t.Fatal("expected no update when enable_metadata is off") - } - if _, ok := pc["Extractor:A"].(map[string]interface{})["enable_metadata"]; ok { - t.Fatal("enable_metadata should not be set") - } -} - -func TestInjectExtractorEnableMetadata_InjectsAndMerges(t *testing.T) { - pc := map[string]interface{}{ - "enable_metadata": true, - "metadata": []interface{}{ - map[string]interface{}{"key": "author", "type": "string", "description": "doc author", "enum": []interface{}{"Alice", "Bob"}}, - }, - "built_in_metadata": []interface{}{ - map[string]interface{}{"key": "year", "type": "number"}, - }, - "Extractor:A": map[string]interface{}{}, - "Parser:X": map[string]interface{}{}, - } - if !InjectExtractorEnableMetadata(pc) { - t.Fatal("expected update") - } - ext := pc["Extractor:A"].(map[string]interface{}) - if got, _ := ext["enable_metadata"].(int); got != 1 { - t.Fatalf("expected enable_metadata=1, got %v", ext["enable_metadata"]) - } - fields, ok := ext["metadata"].([]any) - if !ok || len(fields) != 2 { - t.Fatalf("expected 2 merged metadata, got %#v", ext["metadata"]) - } - if _, ok := pc["Parser:X"].(map[string]interface{})["enable_metadata"]; ok { - t.Fatal("Parser node must not be touched") - } -} - -func TestInjectExtractorEnableMetadata_RespectsExplicitEnabled(t *testing.T) { - // A node the user already turned ON keeps its own config (no clobber). - pc := map[string]interface{}{ - "enable_metadata": true, - "metadata": []interface{}{map[string]interface{}{"key": "author"}}, - "Extractor:A": map[string]interface{}{"enable_metadata": 1}, - } - if InjectExtractorEnableMetadata(pc) { - t.Fatal("expected no update when user explicitly enabled enable_metadata") - } -} - -func TestInjectExtractorEnableMetadata_OverridesDefaultZero(t *testing.T) { - // Shipped DSLs default enable_metadata to 0; the dataset flag must still - // turn them on (otherwise auto-metadata could never activate). - pc := map[string]interface{}{ - "enable_metadata": true, - "metadata": []interface{}{map[string]interface{}{"key": "author"}}, - "Extractor:A": map[string]interface{}{"enable_metadata": 0}, - } - if !InjectExtractorEnableMetadata(pc) { - t.Fatal("expected update: dataset flag must override default enable_metadata=0") - } - ext := pc["Extractor:A"].(map[string]interface{}) - if got, _ := ext["enable_metadata"].(int); got != 1 { - t.Fatalf("expected enable_metadata=1 after override, got %v", ext["enable_metadata"]) - } - if _, ok := ext["metadata"].([]any); !ok { - t.Fatalf("expected metadata field schema injected, got %#v", ext["metadata"]) - } -} - -func TestInjectExtractorEnableMetadata_NoFields(t *testing.T) { - pc := map[string]interface{}{ - "enable_metadata": true, - "Extractor:A": map[string]interface{}{}, - } - if InjectExtractorEnableMetadata(pc) { - t.Fatal("expected no update when no fields configured") - } -} diff --git a/internal/entity/knowledge_compile_doc.go b/internal/entity/knowledge_compile_doc.go index 84f144b0c6..1c38685919 100644 --- a/internal/entity/knowledge_compile_doc.go +++ b/internal/entity/knowledge_compile_doc.go @@ -17,6 +17,16 @@ package entity import "time" +// Dataset-level compile lifecycle states. Shared source of truth for the +// scheduler (which writes State on the knowledge_compile_docs row) and the +// dataset compilation-status API (which reads it back). +const ( + DatasetStateIdle = "idle" // no scheduling row / nothing to do + DatasetStatePending = "pending" // backlog non-empty, awaiting claim + DatasetStateRunning = "running" // a worker holds the lease and is merging + DatasetStateCompleted = "completed" // backlog drained to empty +) + // KnowledgeCompileDataset is the MySQL scheduling row for the dataset-level // post-processing consumer (knowledge_compile_design.md §11.4, Option E). It is // the scheduling system of record: backlog_doc_ids holds the not-yet-processed @@ -42,8 +52,20 @@ type KnowledgeCompileDataset struct { ClaimToken string `gorm:"column:claim_token;size:64;not null;default:''" json:"claim_token"` ClaimExpiresAt *time.Time `gorm:"column:claim_expires_at;default:null" json:"claim_expires_at"` Priority int `gorm:"column:priority;not null;default:0" json:"priority"` - CreatedAt time.Time `gorm:"column:created_at;autoCreateTime" json:"created_at"` - UpdatedAt time.Time `gorm:"column:updated_at;autoUpdateTime" json:"updated_at"` + // State is the dataset-level compile lifecycle state surfaced to the API: + // idle | pending | running | completed. It is written by the scheduler and + // consumer; the API never derives it from the backlog alone. Default is a + // scalar string literal on a varchar column, which MySQL allows (Error 1101 + // only affects TEXT/BLOB, not varchar). + State string `gorm:"column:state;size:16;not null;default:'idle'" json:"state"` + // ErrorMsg is the most recent failure/retry diagnostic. It is TEXT with no + // DDL default (MySQL 8.0.13+ rejects a literal default on TEXT, Error 1101); + // the application always writes it explicitly when set. + ErrorMsg string `gorm:"column:error_msg;type:text;not null" json:"error_msg"` + // LastCompletedAt records the last time the backlog drained to empty. + LastCompletedAt *time.Time `gorm:"column:last_completed_at;default:null" json:"last_completed_at"` + CreatedAt time.Time `gorm:"column:created_at;autoCreateTime" json:"created_at"` + UpdatedAt time.Time `gorm:"column:updated_at;autoUpdateTime" json:"updated_at"` } // TableName pins the scheduling table name. diff --git a/internal/entity/models/bedrock.go b/internal/entity/models/bedrock.go index 5d461cfb3e..604e7c60eb 100644 --- a/internal/entity/models/bedrock.go +++ b/internal/entity/models/bedrock.go @@ -499,10 +499,6 @@ func mapChatConfigToInference(cfg *ChatConfig) *bedrockInferenceConfig { } inf := &bedrockInferenceConfig{} hasField := false - if cfg.MaxTokens != nil { - inf.MaxTokens = cfg.MaxTokens - hasField = true - } if cfg.Temperature != nil { inf.Temperature = cfg.Temperature hasField = true diff --git a/internal/entity/models/bedrock_test.go b/internal/entity/models/bedrock_test.go index c3a51cc860..968636d200 100644 --- a/internal/entity/models/bedrock_test.go +++ b/internal/entity/models/bedrock_test.go @@ -207,8 +207,8 @@ func TestMapChatConfigToInferenceForwardsAllFields(t *testing.T) { if inf == nil { t.Fatal("expected non-nil inferenceConfig") } - if inf.MaxTokens == nil || *inf.MaxTokens != 4096 { - t.Errorf("maxTokens=%v", inf.MaxTokens) + if inf.MaxTokens != nil { + t.Errorf("maxTokens should be omitted, got %v", inf.MaxTokens) } if inf.Temperature == nil || *inf.Temperature != 0.5 { t.Errorf("temperature=%v", inf.Temperature) diff --git a/internal/entity/models/model.go b/internal/entity/models/model.go index 0d7b4e186f..17069d9095 100644 --- a/internal/entity/models/model.go +++ b/internal/entity/models/model.go @@ -23,6 +23,7 @@ import ( "os" "path/filepath" "sort" + "strconv" "strings" ) @@ -167,6 +168,7 @@ type Model struct { Class *string `json:"class"` MaxDimension *int `json:"max_dimension"` // used by embedding models Dimensions []int `json:"dimensions"` + BatchSize *int `json:"batch_size"` // max texts per Embed request; used by embedding models Alias []string `json:"alias"` Rank *int `json:"rank"` ModelTypeMap map[string]bool @@ -418,6 +420,9 @@ func (pm *ProviderManager) ListAllModels() ([]map[string]interface{}, error) { if len(model.Dimensions) > 0 { modelData["dimensions"] = model.Dimensions } + if model.BatchSize != nil { + modelData["batch_size"] = *model.BatchSize + } if model.Thinking != nil { modelData["thinking"] = "supported" } @@ -444,6 +449,34 @@ func (pm *ProviderManager) GetModelByNameOrAlias(modelName string) *Model { return nil } +// DefaultEmbeddingBatchSize is the fallback per-request embedding input count +// when a model reports no batch_size capability (mirrors Python's +// settings.EMBEDDING_BATCH_SIZE). +const DefaultEmbeddingBatchSize = 16 + +// GetEmbeddingBatchSize returns the max texts per Embed request for the named +// model. It consults the model provider capability (batch_size, added to +// all_models.json by #17877/#17878) and falls back to +// DefaultEmbeddingBatchSize when unset or unknown. The +// TOKENIZER_EMBEDDING_BATCH_SIZE env var, when valid, overrides everything for +// local tuning. +func GetEmbeddingBatchSize(modelName string) int { + if v := os.Getenv("TOKENIZER_EMBEDDING_BATCH_SIZE"); v != "" { + if n, err := strconv.Atoi(v); err == nil && n > 0 { + return n + } + } + if modelName != "" { + if pm := GetProviderManager(); pm != nil { + if m := pm.GetModelByNameOrAlias(modelName); m != nil && m.BatchSize != nil && *m.BatchSize > 0 { + return *m.BatchSize + } + } + } + return DefaultEmbeddingBatchSize +} + +// 2. Show specific provider information (including base_url) func (pm *ProviderManager) GetProviderByName(providerName string) (map[string]interface{}, error) { provider := pm.FindProvider(providerName) @@ -489,6 +522,9 @@ func (pm *ProviderManager) ListModels(providerName string) ([]map[string]interfa "max_dimension": model.MaxDimension, "dimensions": model.Dimensions, } + if model.BatchSize != nil { + modelData["batch_size"] = *model.BatchSize + } if model.Thinking != nil { modelData["thinking"] = "supported" } diff --git a/internal/entity/models/replicate.go b/internal/entity/models/replicate.go index 3c6a658721..93b653d14d 100644 --- a/internal/entity/models/replicate.go +++ b/internal/entity/models/replicate.go @@ -164,9 +164,6 @@ func replicateInputFromMessages(messages []Message, chatModelConfig *ChatConfig) input["system_prompt"] = systemPrompt } if chatModelConfig != nil { - if chatModelConfig.MaxTokens != nil { - input["max_new_tokens"] = *chatModelConfig.MaxTokens - } if chatModelConfig.Temperature != nil { input["temperature"] = *chatModelConfig.Temperature } diff --git a/internal/entity/models/replicate_test.go b/internal/entity/models/replicate_test.go index d4e80bbdbc..8ea68352b1 100644 --- a/internal/entity/models/replicate_test.go +++ b/internal/entity/models/replicate_test.go @@ -105,8 +105,8 @@ func TestReplicateOfficialChatHappyPath(t *testing.T) { if input["system_prompt"] != "be helpful" { t.Errorf("system_prompt=%v", input["system_prompt"]) } - if input["max_new_tokens"] != float64(128) { - t.Errorf("max_new_tokens=%v", input["max_new_tokens"]) + if _, ok := input["max_new_tokens"]; ok { + t.Errorf("max_new_tokens should be omitted, got %v", input["max_new_tokens"]) } // Stop is deliberately filtered out because Replicate model // inputs are model-specific and upstream support is undefined. diff --git a/internal/entity/models/siliconflow.go b/internal/entity/models/siliconflow.go index 5dac0642a6..4b9b1cae91 100644 --- a/internal/entity/models/siliconflow.go +++ b/internal/entity/models/siliconflow.go @@ -139,9 +139,6 @@ type siliconflowEmbeddingResponse struct { } `json:"usage"` } -// siliconflowMaxBatchSize is the per-request input limit documented at -const siliconflowMaxBatchSize = 32 - // Embed embeds a list of texts into embeddings func (s *SiliconflowModel) Embed(ctx context.Context, modelName *string, texts []string, apiConfig *APIConfig, embeddingConfig *EmbeddingConfig, modelUsage *common.ModelUsage) ([]EmbeddingData, error) { if err := s.baseModel.APIConfigCheck(apiConfig); err != nil { @@ -151,8 +148,16 @@ func (s *SiliconflowModel) Embed(ctx context.Context, modelName *string, texts [ if len(texts) == 0 { return []EmbeddingData{}, nil } - if len(texts) > siliconflowMaxBatchSize { - return nil, fmt.Errorf("siliconflow supports a maximum of %d inputs per request", siliconflowMaxBatchSize) + // Per-request input cap: resolved from the provider capability (batch_size + // added to all_models.json by #17877/#17878) and falling back to a safe + // default. This defends against callers that bypass batch splitting upstream. + var modelNameStr string + if modelName != nil { + modelNameStr = *modelName + } + maxBatch := GetEmbeddingBatchSize(modelNameStr) + if len(texts) > maxBatch { + return nil, fmt.Errorf("siliconflow supports a maximum of %d inputs per request", maxBatch) } if modelName == nil || *modelName == "" { diff --git a/internal/entity/models/types.go b/internal/entity/models/types.go index bfcc887df9..1305c8c72f 100644 --- a/internal/entity/models/types.go +++ b/internal/entity/models/types.go @@ -219,10 +219,11 @@ type ParseFileConfig struct { // EmbeddingModel wraps a ModelDriver with embedding-specific configuration type EmbeddingModel struct { - ModelDriver ModelDriver - ModelName *string - APIConfig *APIConfig - MaxTokens int // Max input tokens for the embedding model, used for text truncation + ModelDriver ModelDriver + ModelName *string + APIConfig *APIConfig + MaxTokens int // Max input tokens for the embedding model, used for text truncation + MaxBatchSize *int // Max texts per Embed request; nil means "resolve from provider capability at use site" } // NewEmbeddingModel creates a new EmbeddingModel @@ -235,6 +236,22 @@ func NewEmbeddingModel(driver ModelDriver, modelName *string, apiConfig *APIConf } } +// ResolveBatchSize returns the max texts per Embed request for this embedding +// model. It prefers an explicit MaxBatchSize set at construction time and falls +// back to the provider capability (all_models.json batch_size, added by +// #17877/#17878) via GetEmbeddingBatchSize, which itself defaults to +// DefaultEmbeddingBatchSize. +func (m *EmbeddingModel) ResolveBatchSize() int { + if m != nil && m.MaxBatchSize != nil && *m.MaxBatchSize > 0 { + return *m.MaxBatchSize + } + var name string + if m != nil && m.ModelName != nil { + name = *m.ModelName + } + return GetEmbeddingBatchSize(name) +} + // RerankModel wraps a ModelDriver with rerank-specific configuration type RerankModel struct { ModelDriver ModelDriver diff --git a/internal/entity/models/zhipu-ai.go b/internal/entity/models/zhipu-ai.go index 98920be0fe..75bebb9ad7 100644 --- a/internal/entity/models/zhipu-ai.go +++ b/internal/entity/models/zhipu-ai.go @@ -55,7 +55,7 @@ func (z *ZhipuAIModel) Name() string { return "zhipu" } -// ChatWithMessages sends multiple messages with roles and returns response +// ChatWithMessages sends multiple messages with roles and returns response. func (z *ZhipuAIModel) ChatWithMessages(ctx context.Context, modelName string, messages []Message, apiConfig *APIConfig, chatModelConfig *ChatConfig, modelUsage *common.ModelUsage) (*ChatResponse, error) { if err := z.baseModel.APIConfigCheck(apiConfig); err != nil { return nil, err diff --git a/internal/handler/compilation_status_test.go b/internal/handler/compilation_status_test.go new file mode 100644 index 0000000000..b76b490685 --- /dev/null +++ b/internal/handler/compilation_status_test.go @@ -0,0 +1,195 @@ +// +// Copyright 2026 The InfiniFlow Authors. All Rights Reserved. +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +package handler + +import ( + "encoding/json" + "net/http" + "net/http/httptest" + "net/url" + "testing" + + "github.com/gin-gonic/gin" + "github.com/glebarez/sqlite" + "gorm.io/gorm" + + "ragflow/internal/common" + "ragflow/internal/dao" + "ragflow/internal/entity" + dataset "ragflow/internal/service/dataset" +) + +// setupCompilationStatusHandlerDB migrates the minimal schema for the +// GET /datasets/:id/compilation/status handler and pushes it onto dao.DB. +func setupCompilationStatusHandlerDB(t *testing.T) *gorm.DB { + t.Helper() + db, err := gorm.Open(sqlite.Open("file:"+url.QueryEscape(t.Name())+"?mode=memory&cache=shared"), &gorm.Config{ + TranslateError: true, + }) + if err != nil { + t.Fatalf("failed to open sqlite: %v", err) + } + if err := db.AutoMigrate( + &entity.Knowledgebase{}, + &entity.KnowledgeCompileDataset{}, + ); err != nil { + t.Fatalf("failed to migrate test schema: %v", err) + } + origDB := dao.DB + dao.DB = db + t.Cleanup(func() { dao.DB = origDB }) + return db +} + +func insertCompilationStatusHandlerKB(t *testing.T, kbID, ownerID string) { + t.Helper() + status := string(entity.StatusValid) + kb := &entity.Knowledgebase{ + ID: kbID, + TenantID: ownerID, + Name: "compile-status-handler-kb", + EmbdID: "BAAI/bge-large-zh-v1.5@Builtin", + CreatedBy: ownerID, + Permission: string(entity.TenantPermissionMe), + Status: &status, + } + if err := dao.DB.Create(kb).Error; err != nil { + t.Fatalf("insert kb: %v", err) + } +} + +func newCompilationStatusHandlerRouter() *gin.Engine { + gin.SetMode(gin.TestMode) + h := NewDatasetsHandler(dataset.NewDatasetService(), nil) + r := gin.New() + r.GET("/api/v1/datasets/:dataset_id/compilation/status", func(c *gin.Context) { + c.Set("user", &entity.User{ID: "user-1"}) + h.GetCompilationStatus(c) + }) + return r +} + +type compilationStatusResponse struct { + Code int `json:"code"` + Message string `json:"message"` + Data map[string]interface{} `json:"data"` +} + +func getCompilationStatus(t *testing.T, r *gin.Engine, datasetID string) (int, compilationStatusResponse) { + t.Helper() + resp := httptest.NewRecorder() + req := httptest.NewRequest(http.MethodGet, + "/api/v1/datasets/"+datasetID+"/compilation/status", nil) + r.ServeHTTP(resp, req) + var body compilationStatusResponse + if err := json.Unmarshal(resp.Body.Bytes(), &body); err != nil { + t.Fatalf("unmarshal response: %v body=%s", err, resp.Body.String()) + } + return resp.Code, body +} + +// TestCompilationStatusHandler_NoRowIdle verifies a dataset with no scheduling +// row returns idle with zero counts. +func TestCompilationStatusHandler_NoRowIdle(t *testing.T) { + db := setupCompilationStatusHandlerDB(t) + insertCompilationStatusHandlerKB(t, "kb-status-idle", "user-1") + _ = db + + status, body := getCompilationStatus(t, newCompilationStatusHandlerRouter(), "kb-status-idle") + if status != http.StatusOK { + t.Fatalf("status=%d want 200", status) + } + if body.Code != int(common.CodeSuccess) { + t.Fatalf("code=%d message=%q", body.Code, body.Message) + } + if body.Data["state"] != entity.DatasetStateIdle { + t.Fatalf("state=%v want idle", body.Data["state"]) + } + if n, _ := body.Data["inflight"].(float64); n != 0 { + t.Fatalf("inflight=%v want 0", body.Data["inflight"]) + } + if n, _ := body.Data["backlog"].(float64); n != 0 { + t.Fatalf("backlog=%v want 0", body.Data["backlog"]) + } +} + +// TestCompilationStatusHandler_FullOutput locks the JSON contract for a row +// with state, inflight/backlog counts, and error diagnostic. +func TestCompilationStatusHandler_FullOutput(t *testing.T) { + db := setupCompilationStatusHandlerDB(t) + insertCompilationStatusHandlerKB(t, "kb-status-full", "user-1") + + row := entity.KnowledgeCompileDataset{ + DatasetID: "kb-status-full", + TenantID: "user-1", + BacklogDocIDs: `[{"doc_id":"d2","event_type":"completed","seq":2}]`, + InflightDocIDs: `[{"doc_id":"d1","event_type":"completed","seq":1}]`, + State: entity.DatasetStatePending, + ErrorMsg: "merge failed: boom", + } + if err := db.Create(&row).Error; err != nil { + t.Fatalf("insert scheduling row: %v", err) + } + + status, body := getCompilationStatus(t, newCompilationStatusHandlerRouter(), "kb-status-full") + if status != http.StatusOK { + t.Fatalf("status=%d want 200", status) + } + if body.Code != int(common.CodeSuccess) { + t.Fatalf("code=%d message=%q", body.Code, body.Message) + } + if body.Data["state"] != entity.DatasetStatePending { + t.Fatalf("state=%v want pending", body.Data["state"]) + } + if n, _ := body.Data["inflight"].(float64); n != 1 { + t.Fatalf("inflight=%v want 1", body.Data["inflight"]) + } + if n, _ := body.Data["backlog"].(float64); n != 1 { + t.Fatalf("backlog=%v want 1", body.Data["backlog"]) + } + if body.Data["error"] != "merge failed: boom" { + t.Fatalf("error=%v want %q", body.Data["error"], "merge failed: boom") + } +} + +// TestCompilationStatusHandler_Unauthorized verifies a user who does not own the +// dataset is rejected with a data error (HTTP 200 + non-zero code, matching the +// handler's ErrorWithCode contract). +func TestCompilationStatusHandler_Unauthorized(t *testing.T) { + db := setupCompilationStatusHandlerDB(t) + // KB is owned by user-1; the router sets user to user-1, so this test must + // exercise the case where the KB belongs to a different owner. We insert the + // KB under a different owner tenant than the request user by re-pointing the + // KB owner to "other-owner". + insertCompilationStatusHandlerKB(t, "kb-status-forbidden", "other-owner") + if err := db.Create(&entity.KnowledgeCompileDataset{ + DatasetID: "kb-status-forbidden", + TenantID: "other-owner", + BacklogDocIDs: "[]", + InflightDocIDs: "[]", + State: entity.DatasetStateRunning, + }).Error; err != nil { + t.Fatalf("insert scheduling row: %v", err) + } + + _, body := getCompilationStatus(t, newCompilationStatusHandlerRouter(), "kb-status-forbidden") + if body.Code != int(common.CodeDataError) { + t.Fatalf("code=%d want %d", body.Code, common.CodeDataError) + } + if body.Message != "no authorization" { + t.Fatalf("message=%q want %q", body.Message, "no authorization") + } +} diff --git a/internal/handler/components_testpkg/components_handler_test.go b/internal/handler/components_testpkg/components_handler_test.go index 59c9f2df47..1fefbb7650 100644 --- a/internal/handler/components_testpkg/components_handler_test.go +++ b/internal/handler/components_testpkg/components_handler_test.go @@ -134,8 +134,8 @@ func TestComponentsHandler_NoFilter(t *testing.T) { // TestComponentsHandler_FilterIngestion verifies the // ?category=ingestion filter returns the ingestion components -// (Extractor, File, Parser, Tokenizer + 9 chunker variants). Names -// must be sorted ascending (plan §4 task 1 stable output). +// (Compiler, Extractor, File, Parser, Tokenizer + 9 chunker variants). +// Names must be sorted ascending (plan §4 task 1 stable output). func TestComponentsHandler_FilterIngestion(t *testing.T) { eng := newComponentsTestRig(t) w := doRequest(t, eng, "/api/v1/components?category=ingestion") @@ -146,7 +146,7 @@ func TestComponentsHandler_FilterIngestion(t *testing.T) { _, _, data := decodeEnvelope(t, w.Body.Bytes()) wantNames := []string{ - "extractor", "file", "grouptitlechunker", "hierarchytitlechunker", + "compiler", "extractor", "file", "grouptitlechunker", "hierarchytitlechunker", "onechunker", "parser", "presentationchunker", "qachunker", "tablechunker", "tagchunker", "titlechunker", "tokenchunker", "tokenizer", } @@ -172,7 +172,7 @@ func TestComponentsHandler_FilterMultiple(t *testing.T) { _, _, data := decodeEnvelope(t, w.Body.Bytes()) wantNames := []string{ - "extractor", "file", "grouptitlechunker", "hierarchytitlechunker", + "compiler", "extractor", "file", "grouptitlechunker", "hierarchytitlechunker", "onechunker", "parser", "presentationchunker", "qachunker", "tablechunker", "tagchunker", "titlechunker", "tokenchunker", "tokenizer", } @@ -273,7 +273,7 @@ func TestComponentsHandler_CaseInsensitive(t *testing.T) { } _, _, data := decodeEnvelope(t, w.Body.Bytes()) wantNames := []string{ - "extractor", "file", "grouptitlechunker", "hierarchytitlechunker", + "compiler", "extractor", "file", "grouptitlechunker", "hierarchytitlechunker", "onechunker", "parser", "presentationchunker", "qachunker", "tablechunker", "tagchunker", "titlechunker", "tokenchunker", "tokenizer", } diff --git a/internal/handler/dataset.go b/internal/handler/dataset.go index 37feaec5e4..1fbaf8509e 100644 --- a/internal/handler/dataset.go +++ b/internal/handler/dataset.go @@ -1008,114 +1008,28 @@ func (h *DatasetsHandler) AggregateTags(c *gin.Context) { common.SuccessWithData(c, result, "success") } -// RunIndex Run an indexing task (graph/raptor/mindmap) for a dataset. -func (h *DatasetsHandler) RunIndex(c *gin.Context) { +// GetCompilationStatus returns the dataset-level knowledge-compile lifecycle +// state (scheduler contract for API_PROXY_SCHEME=go/hybrid). It replaces the +// Python-era TraceIndex task-progress endpoint for the Go backend. +func (h *DatasetsHandler) GetCompilationStatus(c *gin.Context) { user, errorCode, errorMessage := GetUser(c) if errorCode != common.CodeSuccess { common.ErrorWithCode(c, errorCode, errorMessage) return } - datasetID := strings.TrimSpace(c.Param("dataset_id")) if datasetID == "" { common.ResponseWithCodeData(c, common.CodeDataError, nil, "dataset_id is required") return } - userID := strings.TrimSpace(user.ID) - if userID == "" { - common.ResponseWithCodeData(c, common.CodeDataError, nil, "user_id is required") - return - } - ctx := c.Request.Context() - indexType := strings.ToLower(strings.TrimSpace(c.Query("type"))) - data, code, err := h.datasetsService.RunIndex(ctx, userID, datasetID, indexType) + st, code, err := h.datasetsService.GetDatasetCompilationStatus(ctx, userID, datasetID) if err != nil { common.ErrorWithCode(c, code, err.Error()) return } - - common.SuccessWithData(c, data, "success") -} - -// TraceIndex Trace an indexing task (graph/raptor/mindmap) for a dataset. -func (h *DatasetsHandler) TraceIndex(c *gin.Context) { - user, errorCode, errorMessage := GetUser(c) - if errorCode != common.CodeSuccess { - common.ErrorWithCode(c, errorCode, errorMessage) - return - } - - datasetID := strings.TrimSpace(c.Param("dataset_id")) - if datasetID == "" { - common.ResponseWithCodeData(c, common.CodeDataError, nil, "dataset_id is required") - return - } - - userID := strings.TrimSpace(user.ID) - if userID == "" { - common.ResponseWithCodeData(c, common.CodeDataError, nil, "user_id is required") - return - } - - ctx := c.Request.Context() - - indexType := strings.ToLower(strings.TrimSpace(c.Query("type"))) - result, code, err := h.datasetsService.TraceIndex(ctx, datasetID, userID, indexType) - if err != nil { - common.ErrorWithCode(c, code, err.Error()) - return - } - if result == nil { - common.SuccessWithData(c, map[string]interface{}{}, "success") - return - } - - common.SuccessWithData(c, result, "success") -} - -// DeleteIndex Delete an indexing task (graph/raptor/mindmap) for a dataset. -func (h *DatasetsHandler) DeleteIndex(c *gin.Context) { - user, errorCode, errorMessage := GetUser(c) - if errorCode != common.CodeSuccess { - common.ErrorWithCode(c, errorCode, errorMessage) - return - } - - datasetID := strings.TrimSpace(c.Param("dataset_id")) - if datasetID == "" { - common.ResponseWithCodeData(c, common.CodeDataError, nil, "dataset_id is required") - return - } - - userID := strings.TrimSpace(user.ID) - if userID == "" { - common.ResponseWithCodeData(c, common.CodeDataError, nil, "user_id is required") - return - } - - indexType := strings.ToLower(strings.TrimSpace(c.Param("index_type"))) - if indexType == "" { - indexType = strings.ToLower(strings.TrimSpace(c.Query("type"))) - } - - wipeArg := strings.ToLower(strings.TrimSpace(c.DefaultQuery("wipe", "true"))) - wipe := true - switch wipeArg { - case "false", "0", "no", "off": - wipe = false - } - - ctx := c.Request.Context() - - code, err := h.datasetsService.DeleteIndex(ctx, userID, datasetID, indexType, wipe) - if err != nil { - common.ErrorWithCode(c, code, err.Error()) - return - } - - common.SuccessWithData(c, map[string]interface{}{}, "success") + common.SuccessWithData(c, st, "success") } // ListMetadataFlattened handles GET /api/v1/datasets/metadata/flattened. diff --git a/internal/handler/dataset_artifact.go b/internal/handler/dataset_artifact.go index 1c898f6469..dcddfddda7 100644 --- a/internal/handler/dataset_artifact.go +++ b/internal/handler/dataset_artifact.go @@ -119,7 +119,9 @@ func (h *DatasetArtifactHandler) ListArtifacts(c *gin.Context) { common.ErrorWithCode(c, common.CodeDataError, err.Error()) return } - common.SuccessWithData(c, gin.H{"total": total, "pages": items}, "success") + // Python's list_wiki_pages returns {total, items}; align the Go port so the + // shared frontend (which reads data.items) stays compatible. + common.SuccessWithData(c, gin.H{"total": total, "items": items}, "success") } // UpdateArtifact handles PUT /artifacts// — edit a wiki page. @@ -239,7 +241,9 @@ func (h *DatasetArtifactHandler) ListArtifactTopics(c *gin.Context) { common.ErrorWithCode(c, common.CodeDataError, err.Error()) return } - common.SuccessWithData(c, gin.H{"total": total, "topics": items}, "success") + // Python's list_wiki_topics returns {total, items}; align the Go port so the + // shared frontend (which reads data.items) stays compatible. + common.SuccessWithData(c, gin.H{"total": total, "items": items}, "success") } // GetArtifactAlteration handles GET /artifacts/alteration — wiki alteration summary. diff --git a/internal/ingestion/component/knowledge_compiler/common/deps.go b/internal/ingestion/component/knowledge_compiler/common/deps.go index 0f37d2c4be..2c4d00fefa 100644 --- a/internal/ingestion/component/knowledge_compiler/common/deps.go +++ b/internal/ingestion/component/knowledge_compiler/common/deps.go @@ -106,10 +106,11 @@ type Deps struct { Redis RedisClient // optional (datasetnav) TenantID string DatasetID string - // LLMMaxLength is the chat model's context window in tokens. RAPTOR uses it - // to truncate each cluster's texts so the summary prompt fits the window - // (mirrors Python self._llm_model.max_length). - LLMMaxLength int + // ModelContextLen is the chat model's context window in tokens + // (content_length). The prompt-budget helpers (wikiMapMaxTokens, + // deriveWikiPlanBudget, buildClusterContent) use it to size the input/output + // quotas (mirrors Python self._llm_model.max_length). + ModelContextLen int } // DepsResolver resolves the per-run Deps from a tenant/llm/embedding triple. diff --git a/internal/ingestion/component/knowledge_compiler/common/jsonchat.go b/internal/ingestion/component/knowledge_compiler/common/jsonchat.go index e6f868b91e..d5937d0a41 100644 --- a/internal/ingestion/component/knowledge_compiler/common/jsonchat.go +++ b/internal/ingestion/component/knowledge_compiler/common/jsonchat.go @@ -4,6 +4,7 @@ import ( "context" "encoding/json" "fmt" + "log" "regexp" "strings" ) @@ -34,6 +35,13 @@ func GenJSON(ctx context.Context, chat ChatInvoker, req ChatRequest) (map[string return m, nil } } + // Diagnostic: surface how far parsing got so a formatting failure is not + // opaque. Each candidate is reported with its own unmarshal error so we can + // tell an unfenced/truncated payload from a genuine syntax error. + for i, candidate := range jsonCandidates(resp.Content) { + _, err := tryUnmarshalJSONErr(candidate) + log.Printf("knowledge_compiler: GenJSON candidate[%d] len=%d parse_err=%v body=%q", i, len(candidate), err, truncate(candidate, 300)) + } return nil, fmt.Errorf("knowledge_compiler: LLM response is not parseable JSON: %q", truncate(resp.Content, 200)) } @@ -54,15 +62,20 @@ func jsonCandidates(s string) []string { } func tryUnmarshalJSON(s string) (map[string]any, bool) { + m, err := tryUnmarshalJSONErr(s) + return m, err == nil +} + +func tryUnmarshalJSONErr(s string) (map[string]any, error) { s = strings.TrimSpace(s) if s == "" { - return nil, false + return nil, fmt.Errorf("empty candidate") } var m map[string]any if err := json.Unmarshal([]byte(s), &m); err != nil { - return nil, false + return nil, err } - return m, true + return m, nil } func truncate(s string, n int) string { diff --git a/internal/ingestion/component/knowledge_compiler/component.go b/internal/ingestion/component/knowledge_compiler/component.go index b88f762b93..e1e1f25ac9 100644 --- a/internal/ingestion/component/knowledge_compiler/component.go +++ b/internal/ingestion/component/knowledge_compiler/component.go @@ -9,9 +9,11 @@ import ( "encoding/json" "fmt" "log" + "sort" "strings" "ragflow/internal/agent/runtime" + "ragflow/internal/ingestion/component/globals" "ragflow/internal/ingestion/component/knowledge_compiler/common" "ragflow/internal/ingestion/component/knowledge_compiler/mindmap" "ragflow/internal/ingestion/component/knowledge_compiler/structure" @@ -37,7 +39,12 @@ var chunkerOutputs = map[string]string{ "_ERROR": "Set only on validation failure.", } -const componentNameKnowledgeCompiler = "KnowledgeCompiler" +// componentNameCompiler is the canonical, unified component name for the +// knowledge-compilation flow. It matches the Python side +// (rag/flow/compiler/compiler.py registers component_name = "Compiler"), so a +// canvas saved by the Python frontend and Go's built-in ingestion templates +// both reference the node as "Compiler" and resolve to the same component. +const componentNameCompiler = "Compiler" // KnowledgeCompilerComponent is the runtime.Component surface. Param is set at // construction from the DSL; per-call overrides flow through the inputs map. @@ -60,10 +67,6 @@ func NewKnowledgeCompilerComponent(name string, params map[string]any) (runtime. func (c *KnowledgeCompilerComponent) Inputs() map[string]string { return map[string]string{ "chunks": "List of map[string]any from upstream chunker/parser; each must carry id + text/content_with_weight.", - "llm_id": "Optional per-call LLM id override.", - "embedding_model": "Optional per-call embedding model override.", - "tenant_id": "Optional tenant scope (defaults to resolver context).", - "dataset_id": "Optional dataset scope (wiki historical dedup).", "historical_candidates": "Optional []common.Candidate override for historical dedup (test/offline).", } } @@ -85,14 +88,13 @@ func (c *KnowledgeCompilerComponent) Outputs() map[string]string { func (c *KnowledgeCompilerComponent) Invoke(ctx context.Context, db *gorm.DB, inputs map[string]any) (map[string]any, error) { _ = db param := c.Param - if v, ok := inputs["llm_id"].(string); ok && v != "" { - param.LLMID = v - } - if v, ok := inputs["embedding_model"].(string); ok && v != "" { - param.EmbeddingModel = v - } - tenantID, _ := inputs["tenant_id"].(string) - datasetID, _ := inputs["dataset_id"].(string) + // Resolve the run-level tenant scope from the shared CanvasState.Globals + // bag first (seeded by the pipeline at run start), falling back to the + // component's own input map. Mirrors parser.go: it keeps the tenant id from + // being lost when the upstream output map narrows it, which would otherwise + // leave the template-group lookup with an empty tenant and fail loudly. + tenantID := globals.GlobalOrInput(ctx, inputs, "tenant_id", "") + datasetID := globals.GlobalOrInput(ctx, inputs, "dataset_id", "") // Resolve the compilation template spec(s). Priority: // compilation_template_id > compilation_template_group_id. The variant is @@ -415,10 +417,13 @@ func kindOrVariant(p common.Product) string { // variantCompileKWD maps each Go variant to the compile_kwd discriminator value // Python writes into ES (rag/advanced_rag/knowlege_compile). It is the primary // key that distinguishes compiled knowledge units from ordinary chunks and -// routes retrieval-side filters (e.g. "compile_kwd": ["artifact_page"]). +// routes retrieval-side filters. The wiki value MUST be "wiki_page" (Python's +// canonical WIKI_PAGE_COMPILE_KWD in wiki.py:1661 / wiki_incremental.py:44 / +// dataset_wiki_generator.py:108) so Go-produced wiki pages are visible to the +// artifact API (dataset_artifact_service.go reads compile_kwd="wiki_page"). var variantCompileKWD = map[common.Variant]string{ common.VariantStructure: "structure", - common.VariantWiki: "artifact_page", + common.VariantWiki: "wiki_page", common.VariantTree: "tree", common.VariantMindmap: "mindmap", } @@ -470,6 +475,14 @@ func productsToChunkDocs(products []common.Product) ([]schema.ChunkDoc, error) { if v := metaString(p.Meta, "compile_kwd"); v != "" { compileKWD = v } + // Wiki sub-parts: sections get their own compile_kwd so that a page + // search on compile_kwd="wiki_page" returns pages only (page.go emits + // both kind:"page" and kind:"section" rows under VariantWiki). This is + // the schema-backed page/section discriminator: "wiki_page" == page, + // "wiki_section" == a page sub-section. + if p.Variant == common.VariantWiki && metaString(p.Meta, "kind") == "section" && compileKWD == "wiki_page" { + compileKWD = "wiki_section" + } if compileKWD == "" { compileKWD = string(p.Variant) } @@ -574,12 +587,25 @@ func applyVariantColumns(doc *schema.ChunkDoc, p common.Product) error { case common.VariantWiki: // One artifact_page row per wiki page; section rows reuse the same // page-level columns so retrieval-side filters work uniformly. - if v := metaString(p.Meta, "slug"); v != "" { - if err := doc.SetExtraValue("slug_kwd", v); err != nil { - return err - } - if err := doc.SetExtraValue("artifact_slug_kwd", v); err != nil { - return err + // Match the Python writer contract (api/db/db_models.py slug_kwd): + // slug_kwd stores the full "/" form, so retrieval + // filters (GetWikiPage) can reconstruct it directly. page_type is also + // stored separately for topic grouping. + if pageType := metaString(p.Meta, "page_type"); pageType != "" { + if slug := metaString(p.Meta, "slug"); slug != "" { + // Normalize to the full "/" form (Python writer + // contract). Idempotent: a slug that already carries the prefix + // (some producers emit pageType/slug directly) is left as-is. + fullSlug := slug + if !strings.Contains(slug, "/") { + fullSlug = pageType + "/" + slug + } + if err := doc.SetExtraValue("slug_kwd", fullSlug); err != nil { + return err + } + if err := doc.SetExtraValue("artifact_slug_kwd", fullSlug); err != nil { + return err + } } } if v := metaString(p.Meta, "title"); v != "" { @@ -741,7 +767,21 @@ func metaStringSlice(m map[string]any, key string) []string { // headless / manual chaining reads them from the component output map, so they // must be forwarded when present. func mergeChunks(inputs map[string]any, compiled []schema.ChunkDoc) map[string]any { - raw, _ := inputs["chunks"].([]any) + // Accept both the []any and []map[string]any chunk carriers (the chunker + // emits the latter; buildInputs already handles both). Without this, the + // original source chunks would be dropped when the carrier is []map[string]any. + var raw []any + switch v := inputs["chunks"].(type) { + case []any: + raw = v + case []map[string]any: + raw = make([]any, 0, len(v)) + for _, m := range v { + raw = append(raw, m) + } + default: + log.Printf("knowledge_compiler: mergeChunks: unexpected chunks type %T", inputs["chunks"]) + } merged := make([]any, 0, len(raw)+len(compiled)) for _, r := range raw { merged = append(merged, r) @@ -772,6 +812,16 @@ func mergeChunks(inputs map[string]any, compiled []schema.ChunkDoc) map[string]a // serialization shape, and the one place where inputs are validated, defaulted, // and enriched (e.g. extracting each chunk's pre-computed embedding) before any // LLM/embedding work begins. +// mapKeys returns the sorted keys of m, for diagnostics logging. +func mapKeys(m map[string]any) []string { + keys := make([]string, 0, len(m)) + for k := range m { + keys = append(keys, k) + } + sort.Strings(keys) + return keys +} + func buildInputs(inputs map[string]any, param common.Param) (common.Inputs, error) { in := common.Inputs{ LLMID: param.LLMID, @@ -781,32 +831,45 @@ func buildInputs(inputs map[string]any, param common.Param) (common.Inputs, erro if d, ok := inputs["doc_id"].(string); ok && d != "" { in.DocID = d } - if raw, ok := inputs["chunks"].([]any); ok { - for _, r := range raw { - m, ok := r.(map[string]any) + // The upstream pipeline hands chunks over as a []any of map[string]any in + // some paths and as a []map[string]any in others (the chunker emits the + // latter). Accept both so the knowledge compiler never silently drops the + // whole upstream output on a type mismatch. + var raw []map[string]any + switch v := inputs["chunks"].(type) { + case []any: + raw = make([]map[string]any, 0, len(v)) + for _, item := range v { + m, ok := item.(map[string]any) if !ok { continue } - ch := common.Chunk{Meta: m} - if id, ok := m["id"].(string); ok { - ch.ID = id - } - if t, ok := m["text"].(string); ok { - ch.Text = t - } - if cw, ok := m["content_with_weight"].(string); ok { - ch.Content = cw - } - // Reuse the embedding the upstream pipeline already computed on the - // chunk (stored under q__vec); variants fall back to embedding - // on demand when it is absent. A chunk must carry exactly one vector. - vec, err := common.VectorFromChunkMap(m, 0) - if err != nil { - return in, err - } - ch.Vector = vec - in.Chunks = append(in.Chunks, ch) + raw = append(raw, m) } + case []map[string]any: + raw = v + default: + log.Printf("knowledge_compiler: buildInputs: unexpected chunks type %T", inputs["chunks"]) + } + log.Printf("knowledge_compiler: buildInputs: accepted %d chunk(s) from inputs[chunks]", len(raw)) + for _, m := range raw { + ch := common.Chunk{Meta: m} + if id, ok := m["id"].(string); ok { + ch.ID = id + } + if t, ok := m["text"].(string); ok { + ch.Text = t + } + if cw, ok := m["content_with_weight"].(string); ok { + ch.Content = cw + } + // Reuse the embedding the upstream pipeline already computed on the + // chunk (stored under q__vec); variants fall back to embedding + // on demand when it is absent. A chunk must carry exactly one vector. + if vec, err := common.VectorFromChunkMap(m, 0); err == nil { + ch.Vector = vec + } + in.Chunks = append(in.Chunks, ch) } if hc, ok := inputs["historical_candidates"].([]common.Candidate); ok { in.HistoricalCandidates = hc @@ -824,16 +887,17 @@ func buildInputs(inputs map[string]any, param common.Param) (common.Inputs, erro } func init() { - runtime.MustRegister(componentNameKnowledgeCompiler, runtime.CategoryIngestion, - NewKnowledgeCompilerComponent, runtime.Metadata{ - Version: "0.1.0", - Inputs: map[string]string{ - "chunks": "Upstream chunker/parser output chunks (id + text/content_with_weight).", - "llm_id": "Optional LLM id override.", - "embedding_model": "Optional embedding model override.", - "tenant_id": "Optional tenant scope.", - "dataset_id": "Optional dataset scope (wiki historical dedup).", - }, - Outputs: chunkerOutputs, - }) + // Register under the single unified name "Compiler" (matching the Python + // side) so both Python-saved canvases and Go's built-in ingestion templates + // resolve to the same component without name translation. + meta := runtime.Metadata{ + Version: "0.1.0", + Inputs: map[string]string{ + "chunks": "Upstream chunker/parser output chunks (id + text/content_with_weight).", + "historical_candidates": "Optional historical dedup candidates for offline/test runs.", + }, + Outputs: chunkerOutputs, + } + runtime.MustRegister(componentNameCompiler, runtime.CategoryIngestion, + NewKnowledgeCompilerComponent, meta) } diff --git a/internal/ingestion/component/knowledge_compiler/component_test.go b/internal/ingestion/component/knowledge_compiler/component_test.go index c46a8847af..a9aaa5418a 100644 --- a/internal/ingestion/component/knowledge_compiler/component_test.go +++ b/internal/ingestion/component/knowledge_compiler/component_test.go @@ -12,6 +12,8 @@ import ( "sync" "testing" + "ragflow/internal/agent/runtime" + "ragflow/internal/ingestion/component/globals" "ragflow/internal/ingestion/component/knowledge_compiler/common" "ragflow/internal/service/nav" @@ -173,7 +175,7 @@ func TestKnowledgeCompiler_Structure_EndToEnd(t *testing.T) { installMockDeps(t) installVariantTemplateResolver(t, "structure") - c, err := NewKnowledgeCompilerComponent("KnowledgeCompiler", map[string]any{ + c, err := NewKnowledgeCompilerComponent("Compiler", map[string]any{ "compilation_template_id": "tpl-structure", "llm_id": "llm1", "embedding_model": "emb1", }) if err != nil { @@ -232,7 +234,7 @@ func TestKnowledgeCompiler_Structure_EndToEnd(t *testing.T) { } func TestKnowledgeCompiler_UnknownVariant(t *testing.T) { - c, err := NewKnowledgeCompilerComponent("KnowledgeCompiler", map[string]any{"compilation_template_id": "nope"}) + c, err := NewKnowledgeCompilerComponent("Compiler", map[string]any{"compilation_template_id": "nope"}) if err != nil { t.Fatalf("NewKnowledgeCompilerComponent: %v", err) } @@ -378,7 +380,7 @@ func TestKnowledgeCompiler_Alias_Mindmap(t *testing.T) { installMockDeps(t) // "mind_map" is the deprecated alias for "mindmap"; both resolve to the // implemented mindmap variant and must run (not ErrUnknownVariant / stub). - c, err := NewKnowledgeCompilerComponent("KnowledgeCompiler", map[string]any{ + c, err := NewKnowledgeCompilerComponent("Compiler", map[string]any{ "compilation_template_id": "mind_map", "llm_id": "llm1", "embedding_model": "emb1", }) if err != nil { @@ -408,7 +410,7 @@ func runVariant(t *testing.T, variant string, extra map[string]any) []map[string for k, v := range extra { params[k] = v } - c, err := NewKnowledgeCompilerComponent("KnowledgeCompiler", params) + c, err := NewKnowledgeCompilerComponent("Compiler", params) if err != nil { t.Fatalf("NewKnowledgeCompilerComponent(%s): %v", variant, err) } @@ -531,7 +533,7 @@ func TestKnowledgeCompiler_EmitsChunks(t *testing.T) { installMockDeps(t) installVariantTemplateResolver(t, "structure") - c, err := NewKnowledgeCompilerComponent("KnowledgeCompiler", map[string]any{ + c, err := NewKnowledgeCompilerComponent("Compiler", map[string]any{ "compilation_template_id": "tpl-structure", "llm_id": "llm1", "embedding_model": "emb1", }) if err != nil { @@ -588,7 +590,7 @@ func TestKnowledgeCompiler_TemplateIDsAndProvenance(t *testing.T) { installMockDeps(t) installVariantTemplateResolver(t, "structure") - c, err := NewKnowledgeCompilerComponent("KnowledgeCompiler", map[string]any{ + c, err := NewKnowledgeCompilerComponent("Compiler", map[string]any{ "compilation_template_id": "tpl-structure", "llm_id": "llm1", "embedding_model": "emb1", @@ -672,7 +674,7 @@ func (m constEmbedder) Encode(_ context.Context, texts []string) ([][]float32, e func TestKnowledgeCompiler_Tree_DegenerateNoInfiniteLoop(t *testing.T) { installProseDeps(t) installVariantTemplateResolver(t, "tree") - c, err := NewKnowledgeCompilerComponent("KnowledgeCompiler", map[string]any{ + c, err := NewKnowledgeCompilerComponent("Compiler", map[string]any{ "compilation_template_id": "tpl-tree", "llm_id": "llm1", "embedding_model": "emb1", "extra": map[string]any{"tree_order": 4}, }) @@ -727,7 +729,7 @@ func TestKnowledgeCompiler_Wiki_HistoricalDedupDropsDuplicates(t *testing.T) { t.Cleanup(func() { common.SetDepsResolver(nil) }) installVariantTemplateResolver(t, "wiki") - c, err := NewKnowledgeCompilerComponent("KnowledgeCompiler", map[string]any{ + c, err := NewKnowledgeCompilerComponent("Compiler", map[string]any{ "compilation_template_id": "tpl-wiki", "llm_id": "llm1", "embedding_model": "emb1", }) if err != nil { @@ -786,7 +788,7 @@ func TestKnowledgeCompiler_Wiki_UpdateMergesExistingPage(t *testing.T) { t.Cleanup(func() { common.SetDepsResolver(nil) }) installVariantTemplateResolver(t, "wiki") - c, err := NewKnowledgeCompilerComponent("KnowledgeCompiler", map[string]any{ + c, err := NewKnowledgeCompilerComponent("Compiler", map[string]any{ "compilation_template_id": "tpl-wiki", "llm_id": "llm1", "embedding_model": "emb1", }) if err != nil { @@ -813,7 +815,7 @@ func TestKnowledgeCompiler_Wiki_UpdateMergesExistingPage(t *testing.T) { if !ok { continue } - if cm["compile_kwd"] == "artifact_page" && cm["kc_kind"] == "page" && cm["slug_kwd"] == "entity/alpha" { + if cm["compile_kwd"] == "wiki_page" && cm["kc_kind"] == "page" && cm["slug_kwd"] == "entity/alpha" { page = cm break } @@ -867,7 +869,7 @@ func TestKnowledgeCompiler_Wiki_HistoricalDedupScopedByDataset(t *testing.T) { t.Cleanup(func() { common.SetDepsResolver(nil) }) installVariantTemplateResolver(t, "wiki") - c, err := NewKnowledgeCompilerComponent("KnowledgeCompiler", map[string]any{ + c, err := NewKnowledgeCompilerComponent("Compiler", map[string]any{ "compilation_template_id": "tpl-wiki", "llm_id": "llm1", "embedding_model": "emb1", "enable_historical_dedup": true, }) @@ -982,7 +984,7 @@ func TestKnowledgeCompiler_Structure_FencedJSONNotDropped(t *testing.T) { t.Cleanup(func() { common.SetDepsResolver(nil) }) installVariantTemplateResolver(t, "structure") - c, err := NewKnowledgeCompilerComponent("KnowledgeCompiler", map[string]any{ + c, err := NewKnowledgeCompilerComponent("Compiler", map[string]any{ "compilation_template_id": "tpl-structure", "llm_id": "llm1", "embedding_model": "emb1", }) if err != nil { @@ -1022,7 +1024,7 @@ func TestKnowledgeCompiler_Structure_MalformedJSONFailsLoud(t *testing.T) { t.Cleanup(func() { common.SetDepsResolver(nil) }) installVariantTemplateResolver(t, "structure") - c, err := NewKnowledgeCompilerComponent("KnowledgeCompiler", map[string]any{ + c, err := NewKnowledgeCompilerComponent("Compiler", map[string]any{ "compilation_template_id": "tpl-structure", "llm_id": "llm1", "embedding_model": "emb1", }) if err != nil { @@ -1048,7 +1050,7 @@ func TestKnowledgeCompiler_PassThroughEnvelope(t *testing.T) { installMockDeps(t) installVariantTemplateResolver(t, "structure") - c, err := NewKnowledgeCompilerComponent("KnowledgeCompiler", map[string]any{ + c, err := NewKnowledgeCompilerComponent("Compiler", map[string]any{ "compilation_template_id": "tpl-structure", "llm_id": "llm1", "embedding_model": "emb1", }) if err != nil { @@ -1116,7 +1118,7 @@ func TestKnowledgeCompiler_GroupIDsResolvedToTemplateIDs(t *testing.T) { // compilation_template_group_id (not the obsolete plural list) selects the // group; compilation_template_id is absent so the group path is taken. - c, err := NewKnowledgeCompilerComponent("KnowledgeCompiler", map[string]any{ + c, err := NewKnowledgeCompilerComponent("Compiler", map[string]any{ "compilation_template_group_id": "grp1", "llm_id": "llm1", "embedding_model": "emb1", @@ -1177,7 +1179,7 @@ func TestKnowledgeCompiler_GroupIDsWithoutResolverFailsLoud(t *testing.T) { installMockDeps(t) common.SetGroupResolver(nil) // ensure no resolver is installed - c, err := NewKnowledgeCompilerComponent("KnowledgeCompiler", map[string]any{ + c, err := NewKnowledgeCompilerComponent("Compiler", map[string]any{ "compilation_template_group_id": "grp1", "llm_id": "llm1", "embedding_model": "emb1", @@ -1232,6 +1234,148 @@ var testGroupResolver common.GroupResolver = func(ctx context.Context, db *gorm. return groupIDs, nil } +// TestKnowledgeCompiler_RegistryResolvesUnifiedName locks the unified-name +// contract: the knowledge-compilation node is registered under "Compiler" +// (matching the Python side rag/flow/compiler/compiler.py component_name), so +// both a Python-saved canvas and Go's built-in ingestion templates resolve to +// the same KnowledgeCompilerComponent through runtime.DefaultRegistry. +func TestKnowledgeCompiler_RegistryResolvesUnifiedName(t *testing.T) { + factory, category, _, ok := runtime.DefaultRegistry.Lookup("Compiler") + if !ok { + t.Fatal("runtime registry has no component \"Compiler\"; the Python canvas and Go templates both use this name") + } + if category != runtime.CategoryIngestion { + t.Fatalf("component \"Compiler\" category = %q, want %q", category, runtime.CategoryIngestion) + } + c, err := factory("Compiler", map[string]any{"compilation_template_id": "tree", "llm_id": "llm1", "embedding_model": "emb1"}) + if err != nil { + t.Fatalf("factory(\"Compiler\"): %v", err) + } + if _, ok := c.(*KnowledgeCompilerComponent); !ok { + t.Fatalf("factory(\"Compiler\") produced %T, want *KnowledgeCompilerComponent", c) + } +} + +// TestKnowledgeCompiler_TenantFromGlobals locks the tenant-resolution contract: +// in the production canvas run the run-level tenant_id lives in the shared +// CanvasState.Globals bag (seeded by the pipeline at run start), not necessarily +// in the KnowledgeCompiler's own input map. The component must resolve it through +// globals.GlobalOrInput; otherwise the template/group lookup gets an empty tenant +// and fails with "compilation_template_group ... not found for tenant". +func TestKnowledgeCompiler_TenantFromGlobals(t *testing.T) { + // Attach a CanvasState to the context and seed the run-level tenant id into + // the global bag, as the pipeline does at run start. + ctx := runtime.WithState(context.Background(), runtime.NewCanvasState("run-id", "sess-id")) + globals.SeedIngestionGlobals(ctx, map[string]any{"tenant_id": "tenant-from-globals"}) + + // Install a template resolver that records the tenant id it is called with. + var gotTenant string + prev := testTemplateResolver + common.SetTemplateResolver(func(ctx context.Context, db *gorm.DB, tenantID, templateID string) (common.TemplateInfo, error) { + gotTenant = tenantID + return common.TemplateInfo{ID: templateID, Kind: "structure", Config: map[string]any{}}, nil + }) + t.Cleanup(func() { common.SetTemplateResolver(prev) }) + + c, err := NewKnowledgeCompilerComponent("Compiler", map[string]any{"compilation_template_id": "tpl-x"}) + if err != nil { + t.Fatalf("construct: %v", err) + } + + // Invoke without tenant_id in the input map; the tenant must come from the + // global bag. The template resolution (and thus gotTenant) happens early in + // Invoke, before the variant's LLM/embedding deps are exercised — which is + // all this test needs to assert. + _, _ = c.Invoke(ctx, nil, map[string]any{ + "llm_id": "llm1", + "chunks": []any{map[string]any{"id": "c1", "content_with_weight": "alpha beta", "text": "alpha beta"}}, + "embedding_model": "emb1", + }) + if gotTenant != "tenant-from-globals" { + t.Fatalf("template resolver saw tenant %q, want %q (tenant_id must be read from CanvasState.Globals)", gotTenant, "tenant-from-globals") + } +} + +// TestKnowledgeCompiler_BuildInputsAcceptsMapSliceChunks locks the chunk-carrier +// contract: the upstream chunker hands chunks over as []map[string]any (see the +// chunk map shape {"id","text","ck_type","doc_type_kwd","tk_nums"} observed from +// the running pipeline), not as []any of map. buildInputs must accept both +// shapes; a strict []any assertion alone would silently drop every chunk and +// leave the knowledge compiler with empty input (compiling nothing yet still +// reporting success). +func TestKnowledgeCompiler_BuildInputsAcceptsMapSliceChunks(t *testing.T) { + in, err := buildInputs(map[string]any{ + "chunks": []map[string]any{ + {"id": "c1", "text": "《三国演义》", "ck_type": "text"}, + {"id": "c2", "text": "滚滚长江东逝水", "ck_type": "text"}, + }, + }, common.Param{}) + if err != nil { + t.Fatalf("buildInputs: %v", err) + } + if len(in.Chunks) != 2 { + t.Fatalf("buildInputs produced %d chunks, want 2 (upstream sends []map[string]any)", len(in.Chunks)) + } + if in.Chunks[0].ID != "c1" || in.Chunks[0].Text != "《三国演义》" { + t.Fatalf("chunk[0] = %+v, want id=c1 text=《三国演义》", in.Chunks[0]) + } + + // The legacy []any-of-map shape must still work. + in2, err := buildInputs(map[string]any{ + "chunks": []any{map[string]any{"id": "x", "text": "t"}}, + }, common.Param{}) + if err != nil { + t.Fatalf("buildInputs ([]any): %v", err) + } + if len(in2.Chunks) != 1 || in2.Chunks[0].ID != "x" { + t.Fatalf("buildInputs ([]any) produced %+v, want 1 chunk id=x", in2.Chunks) + } +} + +// TestProductsToChunkDocs_PageVsSectionCompileKWD locks the page/section +// discriminator: a wiki page product is stamped compile_kwd="wiki_page" and a +// wiki section product compile_kwd="wiki_section", so a page search on +// compile_kwd="wiki_page" (engine_service / kcWikiPageStore) returns pages only. +func TestProductsToChunkDocs_PageVsSectionCompileKWD(t *testing.T) { + page := common.Product{ + ID: "page-id", DocID: "d1", TenantID: "t1", Variant: common.VariantWiki, + Content: "# Alpha\n\nBody", ParentID: "", + Meta: map[string]any{"kind": "page", "slug": "entity/alpha", "title": "Alpha", "page_type": "entity", "source_chunk_ids": []string{"c1"}}, + } + section := common.Product{ + ID: "section-id", DocID: "d1", TenantID: "t1", Variant: common.VariantWiki, + Content: "Section body", ParentID: "page-id", + Meta: map[string]any{"kind": "section", "slug": "overview", "page_slug": "entity/alpha", "section_level": 1, "source_chunk_ids": []string{"c1"}}, + } + docs, err := productsToChunkDocs([]common.Product{page, section}) + if err != nil { + t.Fatalf("productsToChunkDocs: %v", err) + } + var pageKWD, sectionKWD string + var sectionParent string + for _, d := range docs { + // Product.Meta is preserved under the kc_* round-trip keys; the page/ + // section kind lives at "kc_kind". + kind, _ := d.GetExtraString("kc_kind") + if kind == "page" { + pageKWD, _ = d.GetExtraString("compile_kwd") + } + if kind == "section" { + sectionKWD, _ = d.GetExtraString("compile_kwd") + sectionParent, _ = d.GetExtraString("parent_kwd") + } + } + if pageKWD != "wiki_page" { + t.Errorf("page compile_kwd = %q, want wiki_page", pageKWD) + } + if sectionKWD != "wiki_section" { + t.Errorf("section compile_kwd = %q, want wiki_section (schema-backed page/section discriminator)", sectionKWD) + } + if sectionParent != "page-id" { + t.Errorf("section parent_kwd = %q, want page-id", sectionParent) + } +} + // TestMain installs the stub resolvers for the variant unit tests. func TestMain(m *testing.M) { common.SetTemplateResolver(testTemplateResolver) diff --git a/internal/ingestion/component/knowledge_compiler/golden_test.go b/internal/ingestion/component/knowledge_compiler/golden_test.go index 09ed27dcc5..7ffa3b16d2 100644 --- a/internal/ingestion/component/knowledge_compiler/golden_test.go +++ b/internal/ingestion/component/knowledge_compiler/golden_test.go @@ -79,7 +79,7 @@ func runVariantChunksWithInputs(t *testing.T, variant string, extra, inputsExtra for k, v := range extra { params[k] = v } - c, err := NewKnowledgeCompilerComponent("KnowledgeCompiler", params) + c, err := NewKnowledgeCompilerComponent("Compiler", params) if err != nil { t.Fatalf("NewKnowledgeCompilerComponent(%s): %v", variant, err) } diff --git a/internal/ingestion/component/knowledge_compiler/pool_wiring.go b/internal/ingestion/component/knowledge_compiler/pool_wiring.go index 6eb4ec017a..3d4a92b7f8 100644 --- a/internal/ingestion/component/knowledge_compiler/pool_wiring.go +++ b/internal/ingestion/component/knowledge_compiler/pool_wiring.go @@ -20,6 +20,7 @@ import ( "ragflow/internal/ingestion/component/knowledge_compiler/mindmap" "ragflow/internal/ingestion/component/knowledge_compiler/structure" + "ragflow/internal/ingestion/component/knowledge_compiler/wiki" "ragflow/internal/ingestion/knowledge_compile" ) @@ -35,4 +36,5 @@ func init() { } structure.SetBatchSubmitter(submit) mindmap.SetBatchSubmitter(submit) + wiki.SetBatchSubmitter(submit) } diff --git a/internal/ingestion/component/knowledge_compiler/tree/raptor.go b/internal/ingestion/component/knowledge_compiler/tree/raptor.go index 697b1a845e..5813e2e460 100644 --- a/internal/ingestion/component/knowledge_compiler/tree/raptor.go +++ b/internal/ingestion/component/knowledge_compiler/tree/raptor.go @@ -195,7 +195,7 @@ func buildTree(ctx context.Context, deps common.Deps, llmID, tenantID, docID str // text to a per-chunk token budget so the cluster fits the LLM context // (Python: len_per_chunk = (max_length - max_token) / len(texts); // truncate(t, len_per_chunk), raptor.py:389-390). - content := buildClusterContent(texts, task.pointIdxs, deps.LLMMaxLength, maxToken) + content := buildClusterContent(texts, task.pointIdxs, deps.ModelContextLen, maxToken) system := raptorSystemHelper + strings.Replace(taskPrompt, "{cluster_content}", content, 1) summary, err := summarizeTexts(ctx, deps, llmID, system, raptorTitleInstruction, maxToken) if err != nil { @@ -292,7 +292,7 @@ func buildTree(ctx context.Context, deps common.Deps, llmID, tenantID, docID str log.Printf("tree: no top-level summaries produced, skipping root node") return nil } - rootContent := buildClusterContent(topLevelTexts, allIndices(len(topLevelTexts)), deps.LLMMaxLength, maxToken) + rootContent := buildClusterContent(topLevelTexts, allIndices(len(topLevelTexts)), deps.ModelContextLen, maxToken) rootSummary, err := summarizeTexts(ctx, deps, llmID, raptorSystemHelper+strings.Replace(taskPrompt, "{cluster_content}", rootContent, 1), raptorTitleInstruction, maxToken) @@ -389,14 +389,14 @@ func titleOf(summary string) string { // (max_length - max_token) / len(texts); truncate(t, len_per_chunk)). The token // budget uses the cl100k_base encoder, mirroring Python's truncate (token-level, // not character-level). -func buildClusterContent(texts []string, idxs []int, llmMaxLength, maxToken int) string { +func buildClusterContent(texts []string, idxs []int, modelContextLen, maxToken int) string { if len(idxs) == 0 { return "" } - if llmMaxLength <= 0 { - llmMaxLength = common.DefaultLLMContextLength + if modelContextLen <= 0 { + modelContextLen = common.DefaultLLMContextLength } - per := (llmMaxLength - maxToken) / len(idxs) + per := (modelContextLen - maxToken) / len(idxs) if per < 1 { per = 1 } diff --git a/internal/ingestion/component/knowledge_compiler/wiki/page_test.go b/internal/ingestion/component/knowledge_compiler/wiki/page_test.go index 5b8c563c7d..e4dcab3216 100644 --- a/internal/ingestion/component/knowledge_compiler/wiki/page_test.go +++ b/internal/ingestion/component/knowledge_compiler/wiki/page_test.go @@ -202,6 +202,7 @@ func TestTransformWikiLinks(t *testing.T) { "See [[concept/beta|Beta]] and [[entity/alpha]]. Also [Beta](artifact/kb/concept/beta).", "kb", map[string]string{"entity/alpha": "Alpha", "concept/beta": "Beta"}, + map[string]string{"entity/alpha": "entity", "concept/beta": "concept"}, ) if !strings.Contains(rendered, "(artifact/kb/concept/beta)") || !strings.Contains(rendered, "(artifact/kb/entity/alpha)") { t.Fatalf("rendered links not rewritten: %q", rendered) @@ -211,6 +212,26 @@ func TestTransformWikiLinks(t *testing.T) { } } +func TestTransformWikiLinksBareSlugGetsPageType(t *testing.T) { + // A bare wikilink (no page_type prefix) resolves its page_type from the + // plan map so the rendered link is clickable by the frontend parser. + rendered, outlinks := transformWikiLinks( + "See [[董卓]] and [[刘备|Liu Bei]].", + "kb", + map[string]string{"董卓": "董卓", "刘备": "刘备"}, + map[string]string{"董卓": "entity", "刘备": "entity"}, + ) + if !strings.Contains(rendered, "(artifact/kb/entity/董卓)") { + t.Fatalf("bare slug link missing page_type: %q", rendered) + } + if !strings.Contains(rendered, "(artifact/kb/entity/刘备)") { + t.Fatalf("bare slug link missing page_type: %q", rendered) + } + if len(outlinks) != 2 { + t.Fatalf("outlinks = %#v, want 2", outlinks) + } +} + func TestSlugify(t *testing.T) { cases := []struct{ in, want string }{ {"Hello World", "hello-world"}, diff --git a/internal/ingestion/component/knowledge_compiler/wiki/prompt.go b/internal/ingestion/component/knowledge_compiler/wiki/prompt.go index 6da9738c75..43f9cac2c9 100644 --- a/internal/ingestion/component/knowledge_compiler/wiki/prompt.go +++ b/internal/ingestion/component/knowledge_compiler/wiki/prompt.go @@ -15,6 +15,26 @@ const wikiPlanSystem = `You are a knowledge compilation planner. Given structure const wikiRefineSystem = `You are a technical writer. Write a complete wiki page from the plan, evidence checklist, and source text. Preserve factual density, keep the source language, and return only markdown.` +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 @@ -134,23 +154,16 @@ Return a JSON compilation plan with one or more page entries: } Rules: +- Return at most {max_pages} page entries for this batch. - Prefer one page per high-signal entity or concept when the batch supports it. +- Merge minor or weakly-supported facts into broader topic pages instead of emitting tiny standalone pages. - Use page_type=entity for entity pages, page_type=concept for concept pages, and page_type=topic for cross-cutting themes. - entity_names must name the entities and concepts that justify the page. - related_kb_pages should list other slugs from the same plan that the page should cross-link to. +- Keep lead concise (one sentence) and keep sections compact (no more than 4 sections, no more than 3 short points per section). - Keep the page in the source language. - Return ONLY the JSON object.` -const wikiPlanMergeUserTemplate = `## Knowledge base context -Document id: {doc_id} - -## Partial plans -{candidates} - -Merge the partial plans into one final compilation plan with the same JSON shape as above. -Preserve distinct pages when they cover different entities or concepts; drop only near-duplicate slugs. -Return ONLY the JSON object.` - const wikiPlanReconcileSystem = `You are a wiki page reconciliation engine. Compare a planned wiki page with existing wiki pages and decide whether the planned page should UPDATE one of them or CREATE a new page. Return only valid JSON.` const wikiPlanReconcileUserTemplate = `## Planned page diff --git a/internal/ingestion/component/knowledge_compiler/wiki/wiki.go b/internal/ingestion/component/knowledge_compiler/wiki/wiki.go index 2cb95570af..1837d390e9 100644 --- a/internal/ingestion/component/knowledge_compiler/wiki/wiki.go +++ b/internal/ingestion/component/knowledge_compiler/wiki/wiki.go @@ -21,6 +21,55 @@ import ( "ragflow/internal/ingestion/component/knowledge_compiler/structure" ) +// batchSubmitter fans out the MAP-stage extraction jobs on the process-wide +// knowledge-compilation pool. It is injected by the knowledge_compiler wiring +// so every variant shares one vCPU-sized concurrency bound; when nil the +// batches run sequentially (the historic default). +var batchSubmitter func(ctx context.Context, jobs []func() error) error + +// SetBatchSubmitter installs the shared-pool fan-out used by Run's MAP stage. +// Pass nil to revert to serial execution. +func SetBatchSubmitter(submit func(ctx context.Context, jobs []func() error) error) { + batchSubmitter = submit +} + +// runBatches mirrors the other compiler variants: concurrent under the wired +// global compiler pool, or serial when no submitter is set. The first error is +// returned after all jobs settle; the global pool is never StopWait'd. +func runBatches(ctx context.Context, jobs []func() error) error { + if len(jobs) == 0 { + return nil + } + if batchSubmitter != nil { + return batchSubmitter(ctx, jobs) + } + for _, j := range jobs { + if err := j(); err != nil { + return err + } + } + return nil +} + +// wikiMapTokenBudget is the input-token budget per MAP extraction batch. It is +// intentionally well below the chat model's context window so the LLM has +// generous room to emit the entity/concept/claim/relation/topic JSON without +// hitting the output-token limit and truncating the payload. +const wikiMapTokenBudget = 2048 + +// wikiMapMaxTokens derives the extraction output budget from the model's +// context length and the per-batch input budget: once the batch has consumed +// wikiMapTokenBudget input tokens, the rest of the window is handed to the +// output — but never below the input budget itself, so a small-input batch can +// still get a proportionally large extraction payload. modelContextLen is the +// model's total context window in tokens (0 means unknown). +func wikiMapMaxTokens(modelContextLen int) int { + if modelContextLen <= 0 { + modelContextLen = common.DefaultLLMContextLength + } + return max(modelContextLen-wikiMapTokenBudget, wikiMapTokenBudget) +} + type wikiPipeline struct { ctx context.Context deps common.Deps @@ -35,6 +84,12 @@ type wikiPipeline struct { reduced wikiExtract plan wikiPlan pages []wikiPageResult + // planBudget is the resolved global page budget for the current planning + // run (target approx + hard cap). + planBudget wikiPlanBudget + // planCapacityExcluded counts the planned pages dropped to fit the global + // hard cap; testable and reported for observability. + planCapacityExcluded int } type wikiExtract struct { @@ -100,8 +155,33 @@ type wikiPlanPage struct { Priority int `json:"priority"` Lead string `json:"lead"` Sections []wikiPlanSection `json:"sections"` + // MentionCount is an internal (non-serialized) signal used for the + // deterministic page selection when the merged plan exceeds the global hard + // cap. It is computed from the reduced extract, not read from JSON. + MentionCount int `json:"-"` } +// Reconciliation thresholds. These are a deliberate Go-specific refinement of +// the Python wiki.py contract, NOT a byte-for-byte alignment: +// +// Python `_wiki_reconcile_with_kb` (wiki.py:1900-1957) queries KNN with +// `extra_options={"similarity": update_threshold}` (update_threshold=0.95), so +// candidates below that threshold are normally filtered out at retrieval and +// the item becomes CREATE directly. Its MAYBE band ([maybe=0.60, update=0.95)) +// is only reachable when a backend still returns low-score candidates despite +// the similarity filter. +// +// Go's `FindSimilarPages` returns top-K candidates without a similarity floor, +// so it genuinely sees low-score candidates Python never does. To exploit that +// richer signal we keep a real two-band decision: +// +// - Score >= update_threshold (0.92) -> direct UPDATE +// - Score < maybe_threshold (0.78) -> CREATE (no match) +// - Score in [maybe, update) -> title/topic/entity overlap +// heuristic first (direct UPDATE), else MAYBE resolved by the LLM. +// +// The title/topic/entity overlap straight-to-UPDATE shortcut is a Go-only +// enhancement; it is NOT a Python-aligned behavior. Keep it documented as such. const ( wikiPlanUpdateThreshold = 0.92 wikiPlanMaybeThreshold = 0.78 @@ -199,6 +279,10 @@ func (p *wikiPipeline) run() error { return err } p.reduced = reduceExtracts(p.mapExtracts) + // 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. + p.reduced.Entities = p.dedupeEntities(p.reduced.Entities) plan, err := p.runPlan() if err != nil { return err @@ -213,27 +297,66 @@ func (p *wikiPipeline) run() error { } func (p *wikiPipeline) runMap() error { - batches := common.PackBatches(p.inputs.Chunks, 4096, p.deps.Tokenizer) - for _, batch := range batches { - if err := p.ctx.Err(); err != nil { - return err - } - extract, err := p.mapBatch(batch) - if err != nil { - return err - } - p.mapExtracts = append(p.mapExtracts, extract) + // Keep each batch small enough that the LLM's entity/relation JSON output + // for the batch stays well under the model's output-token limit. 2048 input + // tokens per batch (was a hard-coded 4096) leaves generous headroom for the + // extraction payload; oversize batches caused truncated, unparseable JSON + // (unexpected end of JSON input) on the real pipeline. + batches := common.PackBatches(p.inputs.Chunks, wikiMapTokenBudget, p.deps.Tokenizer) + extracts, err := runMapBatches(p.ctx, batches, p.mapBatch) + if err != nil { + return err } + p.mapExtracts = append(p.mapExtracts, extracts...) return nil } +func runMapBatches( + ctx context.Context, + batches [][]common.Chunk, + mapBatch func([]common.Chunk) (wikiExtract, error), +) ([]wikiExtract, error) { + if len(batches) == 0 { + return nil, nil + } + extracts := make([]wikiExtract, len(batches)) + jobs := make([]func() error, 0, len(batches)) + for i, batch := range batches { + i, batch := i, batch + jobs = append(jobs, func() error { + if err := ctx.Err(); err != nil { + return err + } + extract, err := mapBatch(batch) + if err != nil { + return err + } + // Distinct slice index per batch keeps results stable without locks. + extracts[i] = extract + return nil + }) + } + if err := runBatches(ctx, jobs); err != nil { + return nil, err + } + return extracts, nil +} + func (p *wikiPipeline) mapBatch(batch []common.Chunk) (wikiExtract, error) { parserConfig, _ := p.inputs.VariantSpecific["parser_config"].(map[string]any) user, _ := buildWikiMapPrompt(p.docID, batch, parserConfig, p.param.Language) + // Give the extraction step a generous output budget so the entity/relation + // JSON is not silently truncated by the model's default output cap (that + // produced "unexpected end of JSON input" from GenJSON). The output budget is + // tied to the per-batch input budget: once the batch consumes + // wikiMapTokenBudget tokens of the model's context, the remainder is left + // for the extraction payload (and never less than the input budget itself). + mt := wikiMapMaxTokens(p.deps.ModelContextLen) raw, err := common.GenJSON(p.ctx, p.deps.Chat, common.ChatRequest{ LLMID: p.llmID, SystemPrompt: wikiMapSystem, UserPrompt: user, + MaxTokens: &mt, }) if err != nil { return wikiExtract{}, err @@ -246,31 +369,85 @@ func (p *wikiPipeline) runPlan() (wikiPlan, error) { if len(batches) == 0 { batches = []wikiExtract{p.reduced} } - plans := make([]wikiPlan, 0, len(batches)) + totalItems := 0 + for _, b := range batches { + totalItems += wikiExtractItemCount(b) + } + p.planBudget = deriveWikiPlanBudget(p.deps.ModelContextLen, totalItems) + // Quota allocation must use the achievable cap (min(Target, Max)): when the + // model's output capacity is smaller than the item-count-derived target, the + // planner must be asked for at most Max pages so the sum of per-batch + // max_pages never exceeds the capacity that can actually be emitted. Using + // Target here would re-introduce the truncated-JSON risk the budget exists + // to eliminate. + quotas := allocatePlanQuotas(batches, p.planBudget.Cap()) + + // approvedReduced is the set of items that actually got a non-zero quota. + // It is what the fallback/normalization may reference so zero-quota items + // can never leak back into the plan via buildWikiFallbackPages. + approved := wikiExtract{} + plans := make([]wikiPlan, len(batches)) + jobs := make([]func() error, 0, len(batches)) for i, batch := range batches { - plan, err := p.runPlanBatch(batch, i+1, len(batches)) - if err != nil { - return wikiPlan{}, err + i, batch := i, batch + quota := quotas[i] + if quota <= 0 { + // Zero-quota batch: no planner call and no fallback page. It is + // intentionally left as the zero wikiPlan{} so the merge sees no + // pages from it. + continue } - plans = append(plans, plan) + approved.Entities = append(approved.Entities, batch.Entities...) + approved.Concepts = append(approved.Concepts, batch.Concepts...) + approved.Claims = append(approved.Claims, batch.Claims...) + approved.Relations = append(approved.Relations, batch.Relations...) + approved.Topics = append(approved.Topics, batch.Topics...) + jobs = append(jobs, func() error { + if err := p.ctx.Err(); err != nil { + return err + } + plan, err := p.runPlanBatch(batch, i+1, len(batches), quota) + if err != nil { + return err + } + // Distinct slice index keeps results stable without locks. + plans[i] = plan + return nil + }) } - if len(plans) == 1 { - plan := normalizeWikiPlan(plans[0], p.docID, p.reduced) - return p.reconcilePlan(plan) + if err := runBatches(p.ctx, jobs); err != nil { + return wikiPlan{}, err } - plan, err := p.mergePlanCandidates(plans) + + merged := p.mergePlanCandidates(plans, approved) + var excluded int + merged.Pages, excluded = truncatePlanPagesByCap(merged.Pages, p.planBudget.Max, approved) + p.planCapacityExcluded += excluded + merged.Pages = normalizeWikiPlanPageLinks(merged.Pages) + reconciled, err := p.reconcilePlan(merged) if err != nil { return wikiPlan{}, err } - return p.reconcilePlan(plan) + // reconcilePlan rewrites page.Slug to the matched existing page's slug, which + // can re-introduce self-referential or plan-absent related links; normalize + // once more so stale links don't reach the stored pages. + reconciled.Pages = normalizeWikiPlanPageLinks(reconciled.Pages) + return reconciled, nil } +// runRefine fans the REFINE stage out to page-level jobs on the shared compiler +// pool, mirroring the P1 error model: all jobs are submitted, awaited, and the +// first error is returned; each job checks ctx before starting. Results are +// written to pre-allocated per-index slots so the output is in normalized-plan +// order regardless of completion order (no concurrent append to a shared slice). +// When no submitter is wired, jobs run serially (historic default). func (p *wikiPipeline) runRefine() ([]wikiPageResult, error) { pages := normalizeWikiPlanPages(p.plan.Pages, p.reduced) if len(pages) == 0 { return nil, nil } pageTitles := map[string]string{} + slugToPageType := map[string]string{} allPlanSlugs := make([]string, 0, len(pages)) for _, page := range pages { if page.Slug == "" { @@ -278,108 +455,155 @@ func (p *wikiPipeline) runRefine() ([]wikiPageResult, error) { } allPlanSlugs = append(allPlanSlugs, page.Slug) pageTitles[page.Slug] = page.Title + // Every planned page carries its type (entity/concept/...); the link + // renderer needs it so artifact/// links are + // clickable (frontend parseWikiLinkHref only matches entity|concept). + if pt := strings.TrimSpace(page.PageType); pt != "" { + slugToPageType[page.Slug] = pt + } } entityLookup := buildWikiEntityLookup(p.reduced.Entities) conceptLookup := buildWikiConceptLookup(p.reduced.Concepts) - results := make([]wikiPageResult, 0, len(pages)) - for _, planItem := range pages { - if p.ctx.Err() != nil { - return nil, p.ctx.Err() - } - evidence := assembleWikiPageEvidence(planItem, p.reduced.Claims, entityLookup, conceptLookup) - sourceChunkIDs := collectWikiEvidenceChunkIDs(evidence) - sourceContext := buildSourceContext(p.inputs.Chunks, sourceChunkIDs) - if strings.TrimSpace(sourceContext) == "" { - sourceContext = buildSourceContext(p.inputs.Chunks, p.reduced.sourceChunkIDs()) - } - available := make([]string, 0, len(allPlanSlugs)) - for _, slug := range allPlanSlugs { - if slug != planItem.Slug { - available = append(available, "- [["+slug+"]]") + + results := make([]wikiPageResult, len(pages)) + jobs := make([]func() error, 0, len(pages)) + for i, planItem := range pages { + i, planItem := i, planItem + jobs = append(jobs, func() error { + if err := p.ctx.Err(); err != nil { + return err } - } - if len(available) == 0 { - available = []string{"(none — this is the only page)"} - } - var existing *common.WikiPageCandidate - var err error - if strings.EqualFold(planItem.Action, "UPDATE") && p.deps.WikiPages != nil { - existing, err = p.deps.WikiPages.GetPageBySlug(p.ctx, p.tenantID, p.datasetID, planItem.Slug) + res, err := p.runRefinePage(planItem, allPlanSlugs, pageTitles, slugToPageType, entityLookup, conceptLookup) if err != nil { - return nil, err + return err } - } - existingSection := "" - existingRaw := "" - if existing != nil { - existingRaw = firstNonEmpty(existing.ContentMDRaw, existing.ContentMD) - if strings.TrimSpace(existingRaw) != "" { - existingSection = "## Existing page content (UPDATE — integrate new evidence into this)\n\n" + existingRaw + "\n" - } - } - user := renderWikiTemplate(wikiRefineWriterUserTemplate, map[string]string{ - "action": firstNonEmpty(planItem.Action, "CREATE"), - "slug": planItem.Slug, - "title": firstNonEmpty(planItem.Title, planItem.Slug), - "page_type": firstNonEmpty(planItem.PageType, "concept"), - "all_plan_slugs": strings.Join(available, "\n"), - "existing_section": existingSection, - "source_context": sourceContext, - "evidence_count": fmt.Sprintf("%d", len(evidence)), - "evidence_blocks": formatWikiEvidenceBlocks(evidence), - }) - resp, err := p.deps.Chat.Chat(p.ctx, common.ChatRequest{ - LLMID: p.llmID, - SystemPrompt: buildWikiRefineWriterSystem(""), - UserPrompt: user, - }) - if err != nil { - return nil, err - } - if resp == nil { - return nil, fmt.Errorf("knowledge_compiler: wiki refine returned no response") - } - contentRaw := strings.TrimSpace(firstNonEmpty(resp.Content)) - if contentRaw == "" { - contentRaw = "# " + firstNonEmpty(planItem.Title, planItem.Slug) + "\n\n(Page generation produced no content.)" - } - if strings.TrimSpace(existingRaw) != "" { - contentRaw, err = p.mergeWikiPageContent(existingRaw, contentRaw, planItem.Slug) - if err != nil { - return nil, err - } - } - contentRendered, outlinks := transformWikiLinks(contentRaw, firstNonEmpty(p.datasetID, p.docID), pageTitles) - sourceDocIDs := collectWikiSourceDocIDs(p.inputs.Chunks, sourceChunkIDs, p.docID) - summary := firstParagraph(contentRendered) - if summary == "" { - summary = firstNonEmpty(planItem.Title, planItem.Slug) - } - topic := firstNonEmpty(planItem.Topic, planItem.Title, planItem.Slug) - results = append(results, wikiPageResult{ - Slug: planItem.Slug, - Title: firstNonEmpty(planItem.Title, planItem.Slug), - PageType: firstNonEmpty(planItem.PageType, "concept"), - Topic: topic, - Action: firstNonEmpty(planItem.Action, "CREATE"), - EntityNames: uniqueStrings(planItem.EntityNames), - RelatedKBPages: uniqueStrings(planItem.RelatedKB), - ContentRaw: contentRaw, - Content: contentRendered, - Summary: summary, - Outlinks: outlinks, - SourceChunkIDs: sourceChunkIDs, - SourceDocIDs: sourceDocIDs, + results[i] = res + return nil }) } + if err := runBatches(p.ctx, jobs); err != nil { + return nil, err + } return results, nil } -func (p *wikiPipeline) runPlanBatch(batch wikiExtract, batchIndex, batchTotal int) (wikiPlan, error) { +// runRefinePage generates one page result from a normalized plan page. UPDATE +// merge, evidence assembly, and source-context building are unchanged; this is +// the per-page unit that runRefine fans out. +func (p *wikiPipeline) runRefinePage( + planItem wikiPlanPage, + allPlanSlugs []string, + pageTitles map[string]string, + slugToPageType map[string]string, + entityLookup map[string]wikiExtractItem, + conceptLookup map[string]wikiExtractItem, +) (wikiPageResult, error) { + evidence := assembleWikiPageEvidence(planItem, p.reduced.Claims, entityLookup, conceptLookup) + sourceChunkIDs := collectWikiEvidenceChunkIDs(evidence) + sourceContext := buildSourceContext(p.inputs.Chunks, sourceChunkIDs) + if strings.TrimSpace(sourceContext) == "" { + sourceContext = buildSourceContext(p.inputs.Chunks, p.reduced.sourceChunkIDs()) + } + available := make([]string, 0, len(allPlanSlugs)) + for _, slug := range allPlanSlugs { + if slug != planItem.Slug { + available = append(available, "- [["+slug+"]]") + } + } + if len(available) == 0 { + available = []string{"(none — this is the only page)"} + } + var existing *common.WikiPageCandidate + var err error + if strings.EqualFold(planItem.Action, "UPDATE") && p.deps.WikiPages != nil { + existing, err = p.deps.WikiPages.GetPageBySlug(p.ctx, p.tenantID, p.datasetID, planItem.Slug) + if err != nil { + return wikiPageResult{}, err + } + } + existingSection := "" + existingRaw := "" + if existing != nil { + existingRaw = firstNonEmpty(existing.ContentMDRaw, existing.ContentMD) + if strings.TrimSpace(existingRaw) != "" { + existingSection = "## Existing page content (UPDATE — integrate new evidence into this)\n\n" + existingRaw + "\n" + } + } + user := renderWikiTemplate(wikiRefineWriterUserTemplate, map[string]string{ + "action": firstNonEmpty(planItem.Action, "CREATE"), + "slug": planItem.Slug, + "title": firstNonEmpty(planItem.Title, planItem.Slug), + "page_type": firstNonEmpty(planItem.PageType, "concept"), + "all_plan_slugs": strings.Join(available, "\n"), + "existing_section": existingSection, + "source_context": sourceContext, + "evidence_count": fmt.Sprintf("%d", len(evidence)), + "evidence_blocks": formatWikiEvidenceBlocks(evidence), + }) + resp, err := p.deps.Chat.Chat(p.ctx, common.ChatRequest{ + LLMID: p.llmID, + SystemPrompt: buildWikiRefineWriterSystem(""), + UserPrompt: user, + }) + if err != nil { + return wikiPageResult{}, err + } + if resp == nil { + return wikiPageResult{}, fmt.Errorf("knowledge_compiler: wiki refine returned no response") + } + contentRaw := strings.TrimSpace(firstNonEmpty(resp.Content)) + if contentRaw == "" { + contentRaw = "# " + firstNonEmpty(planItem.Title, planItem.Slug) + "\n\n(Page generation produced no content.)" + } + if strings.TrimSpace(existingRaw) != "" { + contentRaw, err = p.mergeWikiPageContent(existingRaw, contentRaw, planItem.Slug) + if err != nil { + return wikiPageResult{}, err + } + } + contentRendered, outlinks := transformWikiLinks(contentRaw, firstNonEmpty(p.datasetID, p.docID), pageTitles, slugToPageType) + sourceDocIDs := collectWikiSourceDocIDs(p.inputs.Chunks, sourceChunkIDs, p.docID) + summary := firstParagraph(contentRendered) + if summary == "" { + summary = firstNonEmpty(planItem.Title, planItem.Slug) + } + topic := firstNonEmpty(planItem.Topic, planItem.Title, planItem.Slug) + return wikiPageResult{ + Slug: planItem.Slug, + Title: firstNonEmpty(planItem.Title, planItem.Slug), + PageType: firstNonEmpty(planItem.PageType, "concept"), + Topic: topic, + Action: firstNonEmpty(planItem.Action, "CREATE"), + EntityNames: uniqueStrings(planItem.EntityNames), + RelatedKBPages: uniqueStrings(planItem.RelatedKB), + ContentRaw: contentRaw, + Content: contentRendered, + Summary: summary, + Outlinks: outlinks, + SourceChunkIDs: sourceChunkIDs, + SourceDocIDs: sourceDocIDs, + }, nil +} + +// maxPagesForBatch is the per-batch planner ceiling. It is the allocated quota +// (a fraction of the global target) directly: the quota is already bounded by +// target <= max <= output-token capacity, so no separate static cap is needed. +// A zero/negative quota yields 1 so runPlanBatch is never told "0 pages". The +// old static wikiPlanMaxPagesPerBatch cap is gone: capping a large quota at 8 +// would make the P0 target unreachable for a single high-quota batch. +func maxPagesForBatch(quota int) int { + if quota < 1 { + return 1 + } + return quota +} + +func (p *wikiPipeline) runPlanBatch(batch wikiExtract, batchIndex, batchTotal, quota int) (wikiPlan, error) { user := renderWikiTemplate(wikiPlanBatchUserTemplate, map[string]string{ "doc_id": p.docID, "batch_index": fmt.Sprintf("%d", batchIndex), "batch_total": fmt.Sprintf("%d", batchTotal), + "max_pages": fmt.Sprintf("%d", maxPagesForBatch(quota)), "entities": mustJSON(batch.Entities), "concepts": mustJSON(batch.Concepts), "claims": mustJSON(batch.Claims), @@ -397,20 +621,52 @@ func (p *wikiPipeline) runPlanBatch(batch wikiExtract, batchIndex, batchTotal in return parseWikiPlan(raw, p.docID, batch), nil } -func (p *wikiPipeline) mergePlanCandidates(plans []wikiPlan) (wikiPlan, error) { - user := renderWikiTemplate(wikiPlanMergeUserTemplate, map[string]string{ - "doc_id": p.docID, - "candidates": mustPrettyJSON(plans), - }) - raw, err := common.GenJSON(p.ctx, p.deps.Chat, common.ChatRequest{ - LLMID: p.llmID, - SystemPrompt: wikiPlanSystem, - UserPrompt: user, - }) - if err != nil { - return wikiPlan{}, err +// mergePlanCandidates merges per-batch plans into one plan. reduced is the item +// set the merged plan is allowed to reference (fallback/normalization); in the +// quota-filtered PLAN path this is the approved (non-zero-quota) set so items +// from skipped batches can never leak back in via fallback pages. +func (p *wikiPipeline) mergePlanCandidates(plans []wikiPlan, reduced wikiExtract) wikiPlan { + merged := wikiPlan{} + mergedEntities := map[string]bool{} + mergedRelated := map[string]bool{} + for _, plan := range plans { + if merged.Title == "" { + merged.Title = strings.TrimSpace(plan.Title) + } + if merged.Slug == "" { + merged.Slug = strings.TrimSpace(plan.Slug) + } + if merged.Lead == "" { + merged.Lead = strings.TrimSpace(plan.Lead) + } + if merged.PageType == "" { + merged.PageType = strings.TrimSpace(plan.PageType) + } + if merged.Topic == "" { + merged.Topic = strings.TrimSpace(plan.Topic) + } + if len(merged.Sections) == 0 && len(plan.Sections) > 0 { + merged.Sections = append([]wikiPlanSection(nil), plan.Sections...) + } + merged.Pages = append(merged.Pages, plan.Pages...) + for _, name := range plan.Entities { + name = strings.TrimSpace(name) + if name != "" && !mergedEntities[name] { + mergedEntities[name] = true + merged.Entities = append(merged.Entities, name) + } + } + for _, slug := range plan.Related { + slug = strings.TrimSpace(slug) + if slug != "" && !mergedRelated[slug] { + mergedRelated[slug] = true + merged.Related = append(merged.Related, slug) + } + } } - return normalizeWikiPlan(parseWikiPlan(raw, p.docID, p.reduced), p.docID, p.reduced), nil + merged = normalizeWikiPlan(merged, p.docID, reduced) + merged.Pages = normalizeWikiPlanPageLinks(merged.Pages) + return merged } func (p *wikiPipeline) reconcilePlan(plan wikiPlan) (wikiPlan, error) { @@ -461,6 +717,9 @@ func (p *wikiPipeline) reconcilePlan(plan wikiPlan) (wikiPlan, error) { return plan, nil } +// reconcilePlanPage decides UPDATE / CREATE for one planned page against the +// existing wiki-page store. See the threshold block above for how the band and +// the overlap heuristic relate to Python's wiki.py contract. func (p *wikiPipeline) reconcilePlanPage(page wikiPlanPage, queryVec []float32) (*common.WikiPageCandidate, error) { if p.deps.WikiPages == nil { return nil, nil @@ -943,6 +1202,34 @@ func normalizeWikiPlanPages(pages []wikiPlanPage, reduced wikiExtract) []wikiPla return out } +func normalizeWikiPlanPageLinks(pages []wikiPlanPage) []wikiPlanPage { + if len(pages) == 0 { + return nil + } + valid := make(map[string]bool, len(pages)) + for _, page := range pages { + if slug := strings.TrimSpace(page.Slug); slug != "" { + valid[slug] = true + } + } + out := make([]wikiPlanPage, 0, len(pages)) + for _, page := range pages { + related := make([]string, 0, len(page.RelatedKB)) + seen := map[string]bool{} + for _, slug := range page.RelatedKB { + slug = strings.TrimSpace(slug) + if slug == "" || slug == page.Slug || !valid[slug] || seen[slug] { + continue + } + seen[slug] = true + related = append(related, slug) + } + page.RelatedKB = related + out = append(out, page) + } + return out +} + func normalizeWikiPlanPage(page wikiPlanPage) wikiPlanPage { page.Action = strings.ToUpper(strings.TrimSpace(page.Action)) if page.Action == "" { @@ -1589,7 +1876,24 @@ var ( wikiArtifactMarkdownLink = regexp.MustCompile(`\[([^\]]+)\]\(([^)]+)\)`) ) -func transformWikiLinks(content, kbID string, pageTitles map[string]string) (string, []string) { +// pageTypeOf resolves the page_type for a slug so rendered internal links use +// the artifact/// form the frontend wiki-link parser +// requires (it only matches entity|concept). Unknown links fall back to "page". +func pageTypeOf(slug string, slugToPageType map[string]string) string { + if pt := strings.TrimSpace(slugToPageType[slug]); pt != "" { + return pt + } + if strings.Contains(slug, "/") { + // Slugs may already carry a / prefix; reuse it. + first := strings.SplitN(slug, "/", 2)[0] + if first == "entity" || first == "concept" || first == "topic" { + return first + } + } + return "page" +} + +func transformWikiLinks(content, kbID string, pageTitles, slugToPageType map[string]string) (string, []string) { kbID = strings.TrimSpace(kbID) seen := map[string]bool{} var outlinks []string @@ -1635,13 +1939,29 @@ func transformWikiLinks(content, kbID string, pageTitles map[string]string) (str } return "" } + // link renders artifact/// so the frontend wiki-link + // parser (artifact///, page_type in {entity,concept}) + // can resolve and navigate to the target page. Slugs may already carry a + // / prefix, in which case that prefix is reused as the + // page_type and the name is emitted as the bare slug. + link := func(label, slug string) string { + bareSlug := slug + pageType := pageTypeOf(slug, slugToPageType) + if idx := strings.Index(slug, "/"); idx >= 0 && slug[:idx] == pageType { + bareSlug = slug[idx+1:] + } + if pageType == "" { + pageType = "page" + } + return "[" + label + "](artifact/" + kbID + "/" + pageType + "/" + bareSlug + ")" + } rewriteMD := func(label, href string) string { slug := artifactSlug(href) if slug == "" { return "[" + label + "](" + href + ")" } track(slug) - return "[" + displayText(label, slug) + "](artifact/" + kbID + "/" + slug + ")" + return link(displayText(label, slug), slug) } out := wikiArtifactMarkdownLink.ReplaceAllStringFunc(content, func(match string) string { sub := wikiArtifactMarkdownLink.FindStringSubmatch(match) @@ -1657,7 +1977,7 @@ func transformWikiLinks(content, kbID string, pageTitles map[string]string) (str } slug := strings.TrimSpace(sub[1]) track(slug) - return "[" + sub[2] + "](artifact/" + kbID + "/" + slug + ")" + return link(sub[2], slug) }) out = wikiWikilinkSimpleRe.ReplaceAllStringFunc(out, func(match string) string { sub := wikiWikilinkSimpleRe.FindStringSubmatch(match) @@ -1666,7 +1986,7 @@ func transformWikiLinks(content, kbID string, pageTitles map[string]string) (str } slug := strings.TrimSpace(sub[1]) track(slug) - return "[" + displayText(slug, slug) + "](artifact/" + kbID + "/" + slug + ")" + return link(displayText(slug, slug), slug) }) return out, outlinks } @@ -1874,15 +2194,6 @@ func (p *wikiPipeline) maybeSourceIDs() []string { return out } -func (p *wikiPipeline) runMapBatch(batch []common.Chunk) error { - extract, err := p.mapBatch(batch) - if err != nil { - return err - } - p.mapExtracts = append(p.mapExtracts, extract) - return nil -} - // dedupHistorical drops products that are near-duplicates of existing historical // artifacts, implementing cross-run dedup for the wiki variant. This remains a // read-only historical lookup and does not store any wiki intermediate state. diff --git a/internal/ingestion/component/knowledge_compiler/wiki/wiki_budget.go b/internal/ingestion/component/knowledge_compiler/wiki/wiki_budget.go new file mode 100644 index 0000000000..c1807d3b4b --- /dev/null +++ b/internal/ingestion/component/knowledge_compiler/wiki/wiki_budget.go @@ -0,0 +1,259 @@ +package wiki + +import "sort" + +// This file implements the PLAN-stage page-budget controls that align the Go +// wiki variant with Python's wiki.py: +// +// - a global target_page_count derived from item count (clamp(8, total//3, 60)); +// - a dynamic max_page_count derived from the model's context window that acts +// as an unbreakable hard cap (output-token capacity vs page-token estimate); +// - per-batch page quotas distributed by largest-remainder so the quota sum +// equals the global target and no batch silently multiplies the page count; +// - a deterministic, mention-grounded truncation that selects the top pages +// under the global cap and reports how many were excluded. +// +// The provider never receives an explicit output max_tokens on the wiki path +// (see P0 in tasks/2026-08-04-wiki-go-python-gap-alignment-plan.md): output +// scale is controlled purely through max_pages in the prompt + the in-code +// truncation below. + +// Python alignment constants (wiki.py:1670-1674, 1783-1787). +const ( + wikiPlanMaxOutputTokens = 4096 + wikiPlanOutputSafetyTokens = 256 + wikiPlanPageTokenEstimate = 48 + wikiPlanTargetPageCountMin = 8 + wikiPlanTargetPageCountMax = 60 +) + +// wikiTargetPageCount mirrors Python _wiki_target_page_count: +// clamp(8, total//3, 60). +func wikiTargetPageCount(totalItems int) int { + if totalItems <= 0 { + return wikiPlanTargetPageCountMin + } + if n := totalItems / 3; n < wikiPlanTargetPageCountMin { + return wikiPlanTargetPageCountMin + } else if n > wikiPlanTargetPageCountMax { + return wikiPlanTargetPageCountMax + } else { + return n + } +} + +// wikiPlanBudget is the resolved page budget for one planning run. +type wikiPlanBudget struct { + // Target is the approximate page count the planner should aim for. It is + // only approximate: batches are allocated quotas that sum to it, but the + // merged result may differ slightly. Target is NOT a capacity guarantee. + Target int + // Max is the unbreakable global hard cap derived from output-token + // capacity. The merged, slug-deduped page list is truncated to at most Max + // pages regardless of what the batches produced. Max may be below Target + // when the model's output capacity is smaller than the item-count-derived + // target; that is deliberate (never ask a small-window model for more pages + // than its output can hold). + Max int +} + +// Cap is the page budget the planner is actually allowed to emit. It is the +// achievable bound min(Target, Max): a capacity-limited model must never be +// asked for more pages than its output can hold, so the cap (not the +// item-count-derived Target) drives per-batch quota allocation. +func (b wikiPlanBudget) Cap() int { + if b.Max < b.Target { + return b.Max + } + return b.Target +} + +// deriveWikiPlanBudget computes the global page budget from the model's context +// window and the reduced item count, mirroring Python's +// output_tokens / output_page_capacity / max_page_count derivation +// (wiki.py:2066-2078). modelContextLen is the chat model's context window in +// tokens (0 means unknown). +func deriveWikiPlanBudget(modelContextLen, totalItems int) wikiPlanBudget { + target := wikiTargetPageCount(totalItems) + + if modelContextLen <= 0 { + modelContextLen = 8192 + } + // output_tokens = min(4096, max(1024, int(model_context * 0.4))). + outputTokens := modelContextLen * 2 / 5 // 0.4 + if outputTokens < 1024 { + outputTokens = 1024 + } + if outputTokens > wikiPlanMaxOutputTokens { + outputTokens = wikiPlanMaxOutputTokens + } + + capacity := (outputTokens - wikiPlanOutputSafetyTokens) / wikiPlanPageTokenEstimate + if capacity < 1 { + capacity = 1 + } + maxCount := capacity + if n := target + 8; n < maxCount { + maxCount = n + } + if n := target * 2; n < maxCount { + maxCount = n + } + // Max is the unbreakable global hard cap. It is NOT raised back up to + // Target when output-token capacity is small: a small-window model must + // never be asked to emit more pages than its output capacity permits, or we + // reintroduce truncated-JSON risk. When capacity < Target, Max simply lands + // below Target and the achievable page count is capacity-bound. + return wikiPlanBudget{Target: target, Max: maxCount} +} + +// wikiExtractItemCount counts the planning items in one reduced extract. It is +// the unit used for proportional quota allocation. +func wikiExtractItemCount(e wikiExtract) int { + return len(e.Entities) + len(e.Concepts) + len(e.Claims) + len(e.Relations) + len(e.Topics) +} + +// allocatePlanQuotas distributes totalTarget pages across batches proportionally +// to each batch's item count using the largest-remainder method, padding by +// remainder in original batch order. When the number of batches exceeds the +// target, small batches naturally receive a zero quota (their floor rounds to +// zero and no remainder remains for them). +// +// Invariant: when len(batches) <= totalTarget, the returned quotas sum exactly +// to totalTarget; when len(batches) > totalTarget they sum to totalTarget but +// some entries are zero. In all cases no quota exceeds totalTarget, so the +// global target is never duplicated per batch. +func allocatePlanQuotas(batches []wikiExtract, totalTarget int) []int { + if len(batches) == 0 { + return nil + } + if len(batches) == 1 { + return []int{totalTarget} + } + items := make([]int, len(batches)) + total := 0 + for i, b := range batches { + items[i] = wikiExtractItemCount(b) + total += items[i] + } + if total <= 0 { + total = len(batches) + } + quotas := make([]int, len(batches)) + remaining := totalTarget + for i := range batches { + q := items[i] * totalTarget / total + quotas[i] = q + remaining -= q + } + // Largest-remainder: hand out the leftover pages to batches with the + // largest fractional remainder, breaking ties by original index (stable + // sort preserves first-seen order). + type remItem struct { + remainder int + idx int + } + rems := make([]remItem, len(batches)) + for i := range batches { + rems[i] = remItem{remainder: items[i] * totalTarget % total, idx: i} + } + sort.SliceStable(rems, func(a, b int) bool { + if rems[a].remainder == rems[b].remainder { + return rems[a].idx < rems[b].idx + } + return rems[a].remainder > rems[b].remainder + }) + for i := 0; i < len(rems) && remaining > 0; i++ { + quotas[rems[i].idx]++ + remaining-- + } + return quotas +} + +// pageMentionCount estimates how strongly a planned page is grounded in the +// reduced extract by counting the distinct source chunks that mention its +// entities, concepts, or subject claims. It is used as the deterministic +// priority when pages must be dropped to fit the global hard cap. +func pageMentionCount(page wikiPlanPage, reduced wikiExtract) int { + names := map[string]bool{} + for _, n := range page.EntityNames { + if k := normKey(n); k != "" { + names[k] = true + } + } + if len(names) == 0 { + if t := normKey(page.Title); t != "" { + names[t] = true + } + if topic := normKey(page.Topic); topic != "" { + names[topic] = true + } + } + chunks := map[string]bool{} + for _, e := range reduced.Entities { + if !names[normKey(e.Name)] { + continue + } + for _, c := range e.SourceChunkIDs { + chunks[c] = true + } + } + for _, c := range reduced.Concepts { + if !names[normKey(c.Term)] { + continue + } + for _, cid := range c.SourceChunkIDs { + chunks[cid] = true + } + } + for _, c := range reduced.Claims { + if !names[normKey(c.Subject)] { + continue + } + for _, cid := range c.SourceChunkIDs { + chunks[cid] = true + } + } + return len(chunks) +} + +// truncatePlanPagesByCap keeps at most maxPageCount planned pages, selecting by +// deterministic priority (mention count descending, then priority ascending, +// then slug), and returns the number of pages excluded by the cap. The output +// preserves the input order (original priority/slug order after normalize) so +// downstream slug-dedup and link normalization stay stable. It never fabricates +// a fallback page to fill the budget. +func truncatePlanPagesByCap(pages []wikiPlanPage, maxPageCount int, reduced wikiExtract) ([]wikiPlanPage, int) { + if maxPageCount < 0 { + maxPageCount = 0 + } + if len(pages) <= maxPageCount { + return pages, 0 + } + type scored struct { + idx int + pg wikiPlanPage + mc int + } + scoredPages := make([]scored, len(pages)) + for i, pg := range pages { + scoredPages[i] = scored{idx: i, pg: pg, mc: pageMentionCount(pg, reduced)} + } + sort.SliceStable(scoredPages, func(a, b int) bool { + if scoredPages[a].mc != scoredPages[b].mc { + return scoredPages[a].mc > scoredPages[b].mc + } + if scoredPages[a].pg.Priority != scoredPages[b].pg.Priority { + return scoredPages[a].pg.Priority < scoredPages[b].pg.Priority + } + return scoredPages[a].pg.Slug < scoredPages[b].pg.Slug + }) + selected := scoredPages[:maxPageCount] + // Restore original order by index so output order is deterministic. + sort.SliceStable(selected, func(a, b int) bool { return selected[a].idx < selected[b].idx }) + out := make([]wikiPlanPage, 0, maxPageCount) + for _, s := range selected { + out = append(out, s.pg) + } + return out, len(pages) - maxPageCount +} diff --git a/internal/ingestion/component/knowledge_compiler/wiki/wiki_budget_test.go b/internal/ingestion/component/knowledge_compiler/wiki/wiki_budget_test.go new file mode 100644 index 0000000000..2a26cb2c36 --- /dev/null +++ b/internal/ingestion/component/knowledge_compiler/wiki/wiki_budget_test.go @@ -0,0 +1,559 @@ +package wiki + +import ( + "context" + "errors" + "strings" + "sync" + "testing" + "time" + + "ragflow/internal/ingestion/component/knowledge_compiler/common" +) + +func TestWikiTargetPageCount_Clamp(t *testing.T) { + cases := []struct { + total int + want int + }{ + {0, 8}, // default floor + {1, 8}, // below floor + {24, 8}, // 24//3 = 8 + {60, 20}, // 60//3 = 20 + {180, 60}, // 180//3 = 60 (cap) + {500, 60}, // above cap + } + for _, c := range cases { + if got := wikiTargetPageCount(c.total); got != c.want { + t.Fatalf("wikiTargetPageCount(%d) = %d, want %d", c.total, got, c.want) + } + } +} + +// TestDeriveWikiPlanBudget_MaxReflectsOutputCapacity locks the corrected P0 +// contract: Max is the unbreakable output-capacity bound and is NOT raised back +// up to Target. A small-window model must never be asked for more pages than its +// output capacity permits. +func TestDeriveWikiPlanBudget_MaxReflectsOutputCapacity(t *testing.T) { + // Tiny window (modelLen=1024): output_tokens = max(1024, 1024*0.4=409) = + // 1024; capacity = (1024-256)//48 = 16. For a large item count + // (target=60), Max must stay at 16 (capacity-bound), NOT be raised to 60. + b := deriveWikiPlanBudget(1024, 1000) + if b.Target != 60 { + t.Fatalf("Target = %d, want 60", b.Target) + } + if b.Max != 16 { + t.Fatalf("Max = %d, want 16 (capacity-bound, must not re-raise to Target 60)", b.Max) + } + // A tiny item count with the same window: target = 8, max = min(16, 16, 16) + // = 16. + b = deriveWikiPlanBudget(1024, 1) + if b.Max != 16 { + t.Fatalf("Max = %d, want 16", b.Max) + } + // A roomy window: Max = min(capacity, target+8, target*2). For total=1000 + // (target 60) and window 8192: output=3276, capacity=62 -> max=min(62,68,120)=62. + b = deriveWikiPlanBudget(8192, 1000) + if b.Max != 62 { + t.Fatalf("Max = %d, want 62", b.Max) + } +} + +func TestDeriveWikiPlanBudget_OutputCapacityBounds(t *testing.T) { + // With a 8192 model: output_tokens = min(4096, max(1024, 8192*0.4=3276)) + // = 3276; capacity = (3276-256)//48 = 62. For total=1000 target=60, + // max = min(62, max(68, 120)) = 62. Max must equal 62 and be >= target 60. + b := deriveWikiPlanBudget(8192, 1000) + if b.Target != 60 { + t.Fatalf("Target = %d, want 60", b.Target) + } + want := 62 + if b.Max != want { + t.Fatalf("Max = %d, want %d (output-token capacity)", b.Max, want) + } +} + +func TestAllocatePlanQuotas_SumsToTarget(t *testing.T) { + batches := []wikiExtract{ + {Entities: make([]wikiEntity, 5)}, + {Concepts: make([]wikiConcept, 5)}, + {Claims: make([]wikiClaim, 5)}, + } + quotas := allocatePlanQuotas(batches, 10) + sum := 0 + for _, q := range quotas { + sum += q + } + if sum != 10 { + t.Fatalf("quota sum = %d, want 10 (got %v)", sum, quotas) + } + if len(quotas) != 3 { + t.Fatalf("len(quotas) = %d, want 3", len(quotas)) + } +} + +func TestAllocatePlanQuotas_LargestRemainderOrdered(t *testing.T) { + // 7 items in batch0, 3 in batch1, target=10: + // floors: 7 and 3; remainders 0 and 0 -> [7,3]. + batches := []wikiExtract{ + {Entities: make([]wikiEntity, 7)}, + {Concepts: make([]wikiConcept, 3)}, + } + quotas := allocatePlanQuotas(batches, 10) + if quotas[0] != 7 || quotas[1] != 3 { + t.Fatalf("quotas = %v, want [7 3]", quotas) + } + + // 7,2,1 target=10: floors 7,2,1 rem=0 -> [7,2,1]. + batches = []wikiExtract{ + {Entities: make([]wikiEntity, 7)}, + {Concepts: make([]wikiConcept, 2)}, + {Claims: make([]wikiClaim, 1)}, + } + quotas = allocatePlanQuotas(batches, 10) + if quotas[0] != 7 || quotas[1] != 2 || quotas[2] != 1 { + t.Fatalf("quotas = %v, want [7 2 1]", quotas) + } +} + +func TestAllocatePlanQuotas_ZeroForOverflowingBatches(t *testing.T) { + // More batches than target: some batches must get a zero quota and none may + // exceed the target. + target := 4 + batches := make([]wikiExtract, 8) + for i := range batches { + batches[i] = wikiExtract{Entities: []wikiEntity{{Name: "e"}}} + } + quotas := allocatePlanQuotas(batches, target) + sum := 0 + zero := 0 + for _, q := range quotas { + sum += q + if q == 0 { + zero++ + } + } + if sum != target { + t.Fatalf("quota sum = %d, want %d", sum, target) + } + if zero == 0 { + t.Fatalf("expected at least one zero quota with %d batches > target %d", len(batches), target) + } + for _, q := range quotas { + if q > target { + t.Fatalf("quota %d exceeds target %d", q, target) + } + } +} + +func TestTruncatePlanPagesByCap_SelectsByMentionCount(t *testing.T) { + reduced := wikiExtract{ + Entities: []wikiEntity{ + {Name: "High", SourceChunkIDs: []string{"a", "b", "c", "d"}}, + {Name: "Low", SourceChunkIDs: []string{"a"}}, + }, + } + pages := []wikiPlanPage{ + {Slug: "entity/low", Title: "Low", EntityNames: []string{"Low"}, Priority: 1}, + {Slug: "entity/high", Title: "High", EntityNames: []string{"High"}, Priority: 2}, + } + kept, excluded := truncatePlanPagesByCap(pages, 1, reduced) + if excluded != 1 { + t.Fatalf("excluded = %d, want 1", excluded) + } + if len(kept) != 1 || kept[0].Slug != "entity/high" { + t.Fatalf("kept = %#v, want entity/high", kept) + } +} + +func TestTruncatePlanPagesByCap_NoCapNoDrop(t *testing.T) { + pages := []wikiPlanPage{ + {Slug: "a", Priority: 1}, + {Slug: "b", Priority: 2}, + } + kept, excluded := truncatePlanPagesByCap(pages, 5, wikiExtract{}) + if excluded != 0 || len(kept) != 2 { + t.Fatalf("got kept=%d excluded=%d, want 2/0", len(kept), excluded) + } +} + +func TestTruncatePlanPagesByCap_PreservesInputOrder(t *testing.T) { + reduced := wikiExtract{ + Entities: []wikiEntity{ + {Name: "X", SourceChunkIDs: []string{"a"}}, + {Name: "Y", SourceChunkIDs: []string{"a", "b"}}, + }, + } + // Cap is large enough to keep everything; input order must be preserved. + pages := []wikiPlanPage{ + {Slug: "z", Title: "Z", EntityNames: []string{"X"}, Priority: 2}, + {Slug: "a", Title: "A", EntityNames: []string{"Y"}, Priority: 1}, + } + kept, _ := truncatePlanPagesByCap(pages, 5, reduced) + if len(kept) != 2 || kept[0].Slug != "z" || kept[1].Slug != "a" { + t.Fatalf("kept = %#v, want input order [z a]", kept) + } +} + +// TestRunPlan_PromptMaxPagesNeverExceedsCap locks the capacity-limited quota +// fix: when the model's output capacity is smaller than the item-derived target +// (e.g. ModelContextLen=1024, target 60, Max 16), the sum of the per-batch +// "at most N page entries" values placed in the planner prompts must never +// exceed Max. This prevents the truncated-JSON risk from re-appearing. +func TestRunPlan_PromptMaxPagesNeverExceedsCap(t *testing.T) { + previous := batchSubmitter + defer SetBatchSubmitter(previous) + + SetBatchSubmitter(func(ctx context.Context, jobs []func() error) error { + for _, j := range jobs { + if err := ctx.Err(); err != nil { + return err + } + if err := j(); err != nil { + return err + } + } + return ctx.Err() + }) + + var mu sync.Mutex + var maxPagesSeen []int + big := strings.Repeat("x", 5000) + // 12 large entities each pack as their own (or small) batch, giving multiple + // batches. total items >= 36 => target clamps to 60; ModelContextLen=1024 => + // output capacity 16 => Max = min(16, 68, 120) = 16 => Cap = 16. + entities := make([]wikiEntity, 0, 12) + for i := 0; i < 12; i++ { + entities = append(entities, wikiEntity{Name: "Ent " + itoa(i) + big}) + } + p := &wikiPipeline{ + ctx: context.Background(), + deps: common.Deps{ + ModelContextLen: 1024, + Chat: chatFunc(func(_ context.Context, req common.ChatRequest) (*common.ChatResponse, error) { + if n := extractMaxPages(req.UserPrompt); n >= 0 { + mu.Lock() + maxPagesSeen = append(maxPagesSeen, n) + mu.Unlock() + } + return &common.ChatResponse{Content: `{"pages":[]}`}, nil + }), + }, + reduced: wikiExtract{Entities: entities}, + docID: "doc-1", + } + if _, err := p.runPlan(); err != nil { + t.Fatalf("runPlan err = %v", err) + } + if len(maxPagesSeen) == 0 { + t.Fatalf("no planning prompt captured max_pages") + } + sum := 0 + for _, n := range maxPagesSeen { + sum += n + } + if sum > p.planBudget.Max { + t.Fatalf("sum of per-batch max_pages = %d, want <= Max %d (target 60)", sum, p.planBudget.Max) + } +} + +// extractMaxPages parses the "at most N page entries" instruction from a plan +// prompt, returning -1 when absent. +func extractMaxPages(prompt string) int { + const marker = "at most " + idx := strings.Index(prompt, marker) + if idx < 0 { + return -1 + } + rest := prompt[idx+len(marker):] + j := 0 + for j < len(rest) && rest[j] >= '0' && rest[j] <= '9' { + j++ + } + if j == 0 { + return -1 + } + n := 0 + for _, c := range rest[:j] { + n = n*10 + int(c-'0') + } + return n +} + +// TestMergePlanCandidates_FallbackOnlyUsesApprovedItems locks F3: the fallback +// page set is built from the approved (non-zero-quota) item set only, so items +// from skipped zero-quota batches can never leak back into the plan. +func TestMergePlanCandidates_FallbackOnlyUsesApprovedItems(t *testing.T) { + p := &wikiPipeline{docID: "doc-1"} + approved := wikiExtract{ + Entities: []wikiEntity{{Name: "Approved", SourceChunkIDs: []string{"c1"}}}, + } + // All approved batches returned no pages; the merged plan must fall back to + // approved items only. + merged := p.mergePlanCandidates(nil, approved) + if len(merged.Pages) == 0 { + t.Fatalf("expected at least one fallback page") + } + hasApproved := false + for _, pg := range merged.Pages { + for _, n := range pg.EntityNames { + if strings.Contains(n, "Skipped") { + t.Fatalf("fallback leaked zero-quota item %q", n) + } + if strings.Contains(n, "Approved") { + hasApproved = true + } + } + } + if !hasApproved { + t.Fatalf("fallback missing approved item") + } +} + +// TestRunPlan_TruncatesToGlobalHardCap drives runPlan through a planner that +// returns more pages than the derived global max_page_count, and asserts the +// merged page list is truncated to the hard cap with the excluded count +// recorded. This is the P0 acceptance criterion that the final page count never +// exceeds max_page_count after slug dedup + global cap. +func TestRunPlan_TruncatesToGlobalHardCap(t *testing.T) { + previous := batchSubmitter + defer SetBatchSubmitter(previous) + + SetBatchSubmitter(func(ctx context.Context, jobs []func() error) error { + for _, j := range jobs { + if err := j(); err != nil { + return err + } + } + return nil + }) + + // Planner returns 30 pages. With one entity and ModelContextLen unset, + // target = clamp(8, 1//3, 60) = 8, and max = min(capacity=62, 16, 16) = 16. + pages := make([]map[string]any, 0, 30) + for i := 0; i < 30; i++ { + pages = append(pages, map[string]any{ + "action": "CREATE", + "slug": "entity/item-" + itoa(i), + "title": "Item " + itoa(i), + "page_type": "entity", + "topic": "Item", + "entity_names": []any{"Entity"}, + "priority": i + 1, + }) + } + payload := map[string]any{"pages": pages} + p := &wikiPipeline{ + ctx: context.Background(), + deps: common.Deps{ + Chat: reconcileChatStub{resp: mustJSON(payload)}, + }, + reduced: wikiExtract{ + Entities: []wikiEntity{{Name: "Entity", SourceChunkIDs: []string{"c1"}}}, + }, + docID: "doc-1", + } + plan, err := p.runPlan() + if err != nil { + t.Fatalf("runPlan err = %v", err) + } + if got := len(plan.Pages); got != 16 { + t.Fatalf("plan pages = %d, want 16 (global hard cap)", got) + } + if got := p.planCapacityExcluded; got != 14 { + t.Fatalf("planCapacityExcluded = %d, want 14", got) + } +} + +func itoa(i int) string { + if i == 0 { + return "0" + } + neg := i < 0 + if neg { + i = -i + } + var b []byte + for i > 0 { + b = append([]byte{byte('0' + i%10)}, b...) + i /= 10 + } + if neg { + b = append([]byte{'-'}, b...) + } + return string(b) +} + +// batchPlanChatStub returns one page per planning batch based on which entity +// name is present in the batch prompt. It lets a fake submitter drive each +// batch's planner call with a distinct, deterministic result. +type batchPlanChatStub struct{} + +func (batchPlanChatStub) Chat(_ context.Context, req common.ChatRequest) (*common.ChatResponse, error) { + var title string + switch { + case strings.Contains(req.UserPrompt, "Alpha"): + title = "Alpha" + case strings.Contains(req.UserPrompt, "Beta"): + title = "Beta" + default: + title = "Gamma" + } + return &common.ChatResponse{Content: `{"pages":[{"action":"CREATE","slug":"entity/` + slugify(title) + `","title":"` + title + `","page_type":"entity","topic":"` + title + `","entity_names":["` + title + `"],"priority":1}]}`}, nil +} + +// TestRunPlan_ParallelBatchesMergeInOrder drives runPlan through a submitter +// that completes batches out of order (batch1 finishes before batch0) and +// asserts the merged plan preserves the original batch order deterministically. +// This exercises the P1 invariant that jobs write only their own index and the +// merge reads slots in order. +func TestRunPlan_ParallelBatchesMergeInOrder(t *testing.T) { + previous := batchSubmitter + defer SetBatchSubmitter(previous) + + SetBatchSubmitter(func(ctx context.Context, jobs []func() error) error { + var wg sync.WaitGroup + for i, j := range jobs { + i, j := i, j + wg.Add(1) + go func() { + defer wg.Done() + if i == 0 { + time.Sleep(30 * time.Millisecond) // batch0 completes last + } + j() + }() + } + wg.Wait() + return ctx.Err() + }) + + // Three entities sized so Alpha+Beta pack into batch1 and Gamma falls into + // batch2 (token budget 3500). + big := strings.Repeat("x", 7000) + p := &wikiPipeline{ + ctx: context.Background(), + deps: common.Deps{ + Chat: batchPlanChatStub{}, + }, + reduced: wikiExtract{ + Entities: []wikiEntity{ + {Name: "Alpha" + big}, + {Name: "Beta"}, + {Name: "Gamma" + big}, + }, + }, + docID: "doc-1", + } + plan, err := p.runPlan() + if err != nil { + t.Fatalf("runPlan err = %v", err) + } + // Batch1 (Alpha) must appear before batch2 (Gamma) in the merged plan. + if len(plan.Pages) < 2 { + t.Fatalf("plan pages = %d, want >= 2", len(plan.Pages)) + } + if plan.Pages[0].Title != "Alpha" { + t.Fatalf("merged pages[0].Title = %q, want Alpha (batch order preserved)", plan.Pages[0].Title) + } + if plan.Pages[1].Title != "Gamma" { + t.Fatalf("merged pages[1].Title = %q, want Gamma", plan.Pages[1].Title) + } +} + +// TestRunPlan_ParallelBatchesFirstError verifies the P1 error model: the first +// batch error is returned after all submitted jobs settle. +func TestRunPlan_ParallelBatchesFirstError(t *testing.T) { + previous := batchSubmitter + defer SetBatchSubmitter(previous) + + SetBatchSubmitter(func(ctx context.Context, jobs []func() error) error { + var wg sync.WaitGroup + errs := make(chan error, len(jobs)) + for _, j := range jobs { + j := j + wg.Add(1) + go func() { + defer wg.Done() + errs <- j() + }() + } + wg.Wait() + close(errs) + for err := range errs { + if err != nil { + return err + } + } + return ctx.Err() + }) + + big := strings.Repeat("x", 7000) + boom := errors.New("planning failed") + p := &wikiPipeline{ + ctx: context.Background(), + deps: common.Deps{ + Chat: failPlanChatStub{err: boom}, + }, + reduced: wikiExtract{ + Entities: []wikiEntity{ + {Name: "Alpha" + big}, + {Name: "Beta"}, + {Name: "Gamma" + big}, + }, + }, + docID: "doc-1", + } + if _, err := p.runPlan(); err != boom { + t.Fatalf("runPlan err = %v, want boom", err) + } +} + +// failPlanChatStub fails every planning call with a fixed error. +type failPlanChatStub struct { + err error +} + +func (f failPlanChatStub) Chat(_ context.Context, _ common.ChatRequest) (*common.ChatResponse, error) { + return nil, f.err +} + +// TestRunPlan_CancelledCtxAborts verifies that a cancelled context aborts the +// planning fan-out and surfaces the context error. +func TestRunPlan_CancelledCtxAborts(t *testing.T) { + previous := batchSubmitter + defer SetBatchSubmitter(previous) + + SetBatchSubmitter(func(ctx context.Context, jobs []func() error) error { + for _, j := range jobs { + if err := ctx.Err(); err != nil { + return err + } + if err := j(); err != nil { + return err + } + } + return ctx.Err() + }) + + big := strings.Repeat("x", 7000) + ctx, cancel := context.WithCancel(context.Background()) + cancel() + p := &wikiPipeline{ + ctx: ctx, + deps: common.Deps{ + Chat: batchPlanChatStub{}, + }, + reduced: wikiExtract{ + Entities: []wikiEntity{ + {Name: "Alpha" + big}, + {Name: "Beta"}, + {Name: "Gamma" + big}, + }, + }, + docID: "doc-1", + } + if _, err := p.runPlan(); err == nil { + t.Fatalf("runPlan err = nil, want context cancelled") + } +} diff --git a/internal/ingestion/component/knowledge_compiler/wiki/wiki_reduce.go b/internal/ingestion/component/knowledge_compiler/wiki/wiki_reduce.go new file mode 100644 index 0000000000..d506d630dd --- /dev/null +++ b/internal/ingestion/component/knowledge_compiler/wiki/wiki_reduce.go @@ -0,0 +1,185 @@ +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) { + 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), + }), + }) + if err != nil { + return false, err + } + v, ok := raw["merge"] + if !ok { + return false, nil + } + return toBoolValue(v), nil +} + +// toBoolValue interprets a loosely-typed boolean field returned by the LLM (the +// model may emit JSON true or a string like "true"/"yes"). +func toBoolValue(v any) bool { + switch x := v.(type) { + case bool: + return x + case string: + switch strings.ToLower(strings.TrimSpace(x)) { + case "true", "yes", "1", "same", "merge": + return true + } + case float64: + return x != 0 + } + return false +} + +// cosine32 computes the cosine similarity between two float32 vectors. +func cosine32(a, b []float32) float64 { + na := l2Norm32(a) + nb := l2Norm32(b) + if na == 0 || nb == 0 { + return 0 + } + var dot float64 + for i := 0; i < len(a) && i < len(b); i++ { + dot += float64(a[i]) * float64(b[i]) + } + return dot / (na * nb) +} diff --git a/internal/ingestion/component/knowledge_compiler/wiki/wiki_reduce_test.go b/internal/ingestion/component/knowledge_compiler/wiki/wiki_reduce_test.go new file mode 100644 index 0000000000..6f621535fe --- /dev/null +++ b/internal/ingestion/component/knowledge_compiler/wiki/wiki_reduce_test.go @@ -0,0 +1,225 @@ +package wiki + +import ( + "context" + "errors" + "testing" + + "ragflow/internal/ingestion/component/knowledge_compiler/common" +) + +// TestDedupeEntities_NoSeamIsNoop verifies dedupeEntities degrades to a no-op +// when the embedder or chat seam is nil (M1-style unit safety). +func TestDedupeEntities_NoSeamIsNoop(t *testing.T) { + p := &wikiPipeline{ctx: context.Background()} + in := []wikiEntity{ + {Name: "Alpha", SourceChunkIDs: []string{"c1"}}, + {Name: "Alpha Corp", SourceChunkIDs: []string{"c2"}}, + } + got := p.dedupeEntities(in) + if len(got) != 2 { + t.Fatalf("got %d entities, want 2 (no-op without seams)", len(got)) + } +} + +// TestDedupeEntities_LLMMergesAmbiguousPair verifies two distinct-name entities +// with high embedding similarity are collapsed into one canonical entity via the +// LLM merge decision, with aliases and provenance merged. +func TestDedupeEntities_LLMMergesAmbiguousPair(t *testing.T) { + p := &wikiPipeline{ + ctx: context.Background(), + llmID: "llm1", + deps: common.Deps{ + Chat: reconcileChatStub{resp: `{"merge":true,"reason":"same company"}`}, + Embed: mergeEmbedStub{}, + }, + } + in := []wikiEntity{ + {Name: "Alpha Inc", Type: "org", SourceChunkIDs: []string{"c1"}}, + {Name: "Alpha Incorporated", Type: "org", SourceChunkIDs: []string{"c2"}}, + } + got := p.dedupeEntities(in) + if len(got) != 1 { + t.Fatalf("got %d entities, want 1 (LLM merge)", len(got)) + } + if got[0].Name != "Alpha Inc" { + t.Fatalf("canonical name = %q, want Alpha Inc", got[0].Name) + } + if len(got[0].SourceChunkIDs) != 2 { + t.Fatalf("provenance = %#v, want 2 chunk ids", got[0].SourceChunkIDs) + } + if len(got[0].Aliases) == 0 { + t.Fatalf("aliases not merged: %#v", got[0].Aliases) + } +} + +// TestDedupeEntities_LLMRejectsDistinct verifies the LLM rejecting a merge keeps +// both entities distinct. +func TestDedupeEntities_LLMRejectsDistinct(t *testing.T) { + p := &wikiPipeline{ + ctx: context.Background(), + llmID: "llm1", + deps: common.Deps{ + Chat: reconcileChatStub{resp: `{"merge":false,"reason":"distinct products"}`}, + Embed: mergeEmbedStub{}, + }, + } + in := []wikiEntity{ + {Name: "Alpha", SourceChunkIDs: []string{"c1"}}, + {Name: "Beta", SourceChunkIDs: []string{"c2"}}, + } + got := p.dedupeEntities(in) + if len(got) != 2 { + t.Fatalf("got %d entities, want 2 (LLM rejected merge)", len(got)) + } +} + +// TestDedupeEntities_ExactNameIsNotAmbiguous verifies entities with identical +// normalized names (already collapsed by reduceExtracts before this stage) are +// not treated as ambiguous: they pass through untouched and no LLM call is made +// for the exact-name pair. +func TestDedupeEntities_ExactNameIsNotAmbiguous(t *testing.T) { + p := &wikiPipeline{ + ctx: context.Background(), + llmID: "llm1", + deps: common.Deps{ + Chat: reconcileChatStub{resp: `{"merge":true}`}, + Embed: mergeEmbedStub{}, + }, + } + in := []wikiEntity{ + {Name: "Alpha", SourceChunkIDs: []string{"c1"}}, + {Name: "Alpha", SourceChunkIDs: []string{"c2"}}, + } + got := p.dedupeEntities(in) + // Exact-name duplicates are out of scope for the embedding step; both are + // kept unchanged (they would already be one entity after reduceExtracts). + if len(got) != 2 { + t.Fatalf("got %d entities, want 2 (exact-name pairs are not ambiguous)", len(got)) + } +} + +// TestDedupeEntities_ConceptStaysExact validates the REDUCE boundary: concept +// dedup must remain exact (no embedding/LLM). This test guards that the entity +// enhancement never touches concepts. +func TestReduceExtracts_ConceptsStayExact(t *testing.T) { + reduced := reduceExtracts([]wikiExtract{ + {Concepts: []wikiConcept{{Term: "RAG", Definition: "d1", SourceChunkIDs: []string{"c1"}}}}, + {Concepts: []wikiConcept{{Term: "Retrieval Augmented Generation", Definition: "d2", SourceChunkIDs: []string{"c2"}}}}, + }) + if len(reduced.Concepts) != 2 { + t.Fatalf("concepts = %d, want 2 (exact-term dedup must not collapse distinct terms)", len(reduced.Concepts)) + } +} + +// TestDedupeEntities_FailingChatIsBudgeted locks F4: a chat seam that +// persistently fails must not drive unbounded external calls. The llmCalls +// budget is consumed before the request, so the loop stops after at most +// wikiEntityMergeMaxCalls attempts. +func TestDedupeEntities_FailingChatIsBudgeted(t *testing.T) { + calls := 0 + p := &wikiPipeline{ + ctx: context.Background(), + llmID: "llm1", + deps: common.Deps{ + Chat: chatFunc(func(_ context.Context, _ common.ChatRequest) (*common.ChatResponse, error) { + calls++ + return nil, errors.New("llm down") + }), + Embed: mergeEmbedStub{}, + }, + } + // 20 entities with distinct names but identical embeddings => every pair is + // ambiguous and would trigger an LLM call. + in := make([]wikiEntity, 0, 20) + for i := 0; i < 20; i++ { + in = append(in, wikiEntity{Name: "Entity " + itoa(i), Type: "person"}) + } + got := p.dedupeEntities(in) + if calls > wikiEntityMergeMaxCalls { + t.Fatalf("chat calls = %d, want <= %d despite persistent failures", calls, wikiEntityMergeMaxCalls) + } + // No merges happen because every call fails, so all 20 entities survive. + if len(got) != 20 { + t.Fatalf("entities = %d, want 20 (no merges on failure)", len(got)) + } +} + +// TestDedupeEntities_CrossTypeHighSimDoesNotCallLLM locks F5: entities with +// provably different types must never be treated as ambiguous, even with +// identical embeddings, so no LLM call is made for them. +func TestDedupeEntities_CrossTypeHighSimDoesNotCallLLM(t *testing.T) { + calls := 0 + p := &wikiPipeline{ + ctx: context.Background(), + llmID: "llm1", + deps: common.Deps{ + Chat: chatFunc(func(_ context.Context, _ common.ChatRequest) (*common.ChatResponse, error) { + calls++ + return &common.ChatResponse{Content: `{"merge":true}`}, nil + }), + Embed: mergeEmbedStub{}, + }, + } + // Both embed to [1,1,1] (identical vectors, cosine 1.0) but types differ. + in := []wikiEntity{ + {Name: "Alpha", Type: "person"}, + {Name: "Beta Corp", Type: "org"}, + } + got := p.dedupeEntities(in) + if calls != 0 { + t.Fatalf("chat calls = %d, want 0 (cross-type pairs must not reach the LLM)", calls) + } + if len(got) != 2 { + t.Fatalf("entities = %d, want 2 (cross-type entities must stay distinct)", len(got)) + } +} + +// mergeEmbedStub returns embeddings where identical names share a vector and +// distinct names are far apart (cosine ~0), so it can drive both the ambiguous +// and distinct test paths deterministically. +type mergeEmbedStub struct{} + +func (mergeEmbedStub) Encode(_ context.Context, texts []string) ([][]float32, error) { + out := make([][]float32, len(texts)) + for i, text := range texts { + // "Alpha Inc" and "Alpha Incorporated" both contain "alpha" -> same + // vector; "Beta" differs. + if containsFold(text, "beta") { + out[i] = []float32{1, 0, 0} + continue + } + out[i] = []float32{1, 1, 1} + } + return out, nil +} + +func (mergeEmbedStub) Dimensions() int { return 3 } + +func containsFold(s, sub string) bool { + return len(s) >= len(sub) && (len(sub) == 0 || indexFold(s, sub) >= 0) +} + +func indexFold(s, sub string) int { + if sub == "" { + return 0 + } + ls := toLowerASCII(s) + lsub := toLowerASCII(sub) + for i := 0; i+len(lsub) <= len(ls); i++ { + if ls[i:i+len(lsub)] == lsub { + return i + } + } + return -1 +} + +func toLowerASCII(s string) string { + b := []byte(s) + for i := range b { + if b[i] >= 'A' && b[i] <= 'Z' { + b[i] += 'a' - 'A' + } + } + return string(b) +} diff --git a/internal/ingestion/component/knowledge_compiler/wiki/wiki_refine_test.go b/internal/ingestion/component/knowledge_compiler/wiki/wiki_refine_test.go new file mode 100644 index 0000000000..a8bf19bedf --- /dev/null +++ b/internal/ingestion/component/knowledge_compiler/wiki/wiki_refine_test.go @@ -0,0 +1,156 @@ +package wiki + +import ( + "context" + "errors" + "strings" + "sync" + "testing" + "time" + + "ragflow/internal/ingestion/component/knowledge_compiler/common" +) + +// refineChatStub returns per-page markdown keyed by the page title in the +// writer prompt so each page's result is distinct and deterministic. +type refineChatStub struct{} + +func (refineChatStub) Chat(_ context.Context, req common.ChatRequest) (*common.ChatResponse, error) { + title := "Page" + for _, cand := range []string{"Alpha", "Beta", "Gamma"} { + if strings.Contains(req.UserPrompt, cand) { + title = cand + break + } + } + return &common.ChatResponse{Content: "# " + title + "\n\nContent for " + title + ".\n"}, nil +} + +func refinePipeline() *wikiPipeline { + return &wikiPipeline{ + ctx: context.Background(), + tenantID: "t1", + datasetID: "kb1", + llmID: "llm1", + docID: "doc-1", + deps: common.Deps{ + Chat: refineChatStub{}, + }, + reduced: wikiExtract{ + Entities: []wikiEntity{{Name: "Alpha", SourceChunkIDs: []string{"c1"}}}, + Claims: []wikiClaim{{Statement: "Alpha exists", Subject: "Alpha", SourceChunkIDs: []string{"c1"}}}, + }, + inputs: common.Inputs{ + Chunks: []common.Chunk{{ID: "c1", Text: "Alpha content", Meta: map[string]any{"doc_id": "doc-1"}}}, + }, + } +} + +func TestRunRefine_ParallelPagesKeepPlanOrder(t *testing.T) { + previous := batchSubmitter + defer SetBatchSubmitter(previous) + + SetBatchSubmitter(func(ctx context.Context, jobs []func() error) error { + var wg sync.WaitGroup + for i, j := range jobs { + i, j := i, j + wg.Add(1) + go func() { + defer wg.Done() + if i == 0 { + time.Sleep(30 * time.Millisecond) // page 0 completes last + } + j() + }() + } + wg.Wait() + return ctx.Err() + }) + + p := refinePipeline() + p.plan = wikiPlan{ + Pages: []wikiPlanPage{ + {Action: "CREATE", Slug: "entity/alpha", Title: "Alpha", PageType: "entity", Topic: "Alpha", EntityNames: []string{"Alpha"}, Priority: 1}, + {Action: "CREATE", Slug: "entity/beta", Title: "Beta", PageType: "entity", Topic: "Beta", EntityNames: []string{"Beta"}, Priority: 2}, + }, + } + got, err := p.runRefine() + if err != nil { + t.Fatalf("runRefine err = %v", err) + } + if len(got) != 2 { + t.Fatalf("got %d pages, want 2", len(got)) + } + if got[0].Title != "Alpha" || got[1].Title != "Beta" { + t.Fatalf("page order = [%s, %s], want [Alpha, Beta] (plan order preserved)", got[0].Title, got[1].Title) + } + if !strings.Contains(got[0].Content, "Content for Alpha") { + t.Fatalf("page0 content missing: %q", got[0].Content) + } +} + +func TestRunRefine_FirstErrorAborts(t *testing.T) { + previous := batchSubmitter + defer SetBatchSubmitter(previous) + + boom := errors.New("refine failed") + SetBatchSubmitter(func(ctx context.Context, jobs []func() error) error { + var wg sync.WaitGroup + errs := make(chan error, len(jobs)) + for _, j := range jobs { + j := j + wg.Add(1) + go func() { + defer wg.Done() + errs <- j() + }() + } + wg.Wait() + close(errs) + for err := range errs { + if err != nil { + return err + } + } + return ctx.Err() + }) + + p := refinePipeline() + p.plan = wikiPlan{Pages: []wikiPlanPage{ + {Action: "CREATE", Slug: "entity/alpha", Title: "Alpha", Priority: 1}, + }} + p.deps.Chat = chatFunc(func(_ context.Context, _ common.ChatRequest) (*common.ChatResponse, error) { + return nil, boom + }) + if _, err := p.runRefine(); err != boom { + t.Fatalf("runRefine err = %v, want boom", err) + } +} + +func TestRunRefine_CancelledCtxAborts(t *testing.T) { + previous := batchSubmitter + defer SetBatchSubmitter(previous) + + SetBatchSubmitter(func(ctx context.Context, jobs []func() error) error { + for _, j := range jobs { + if err := ctx.Err(); err != nil { + return err + } + if err := j(); err != nil { + return err + } + } + return ctx.Err() + }) + + ctx, cancel := context.WithCancel(context.Background()) + cancel() + p := refinePipeline() + p.ctx = ctx + p.plan = wikiPlan{Pages: []wikiPlanPage{ + {Action: "CREATE", Slug: "entity/alpha", Title: "Alpha", Priority: 1}, + }} + if _, err := p.runRefine(); err == nil { + t.Fatalf("runRefine err = nil, want context cancelled") + } +} diff --git a/internal/ingestion/component/knowledge_compiler/wiki/wiki_test.go b/internal/ingestion/component/knowledge_compiler/wiki/wiki_test.go index 5f233ba648..2cf03c3b07 100644 --- a/internal/ingestion/component/knowledge_compiler/wiki/wiki_test.go +++ b/internal/ingestion/component/knowledge_compiler/wiki/wiki_test.go @@ -3,7 +3,9 @@ package wiki import ( "context" "strings" + "sync" "testing" + "time" "ragflow/internal/ingestion/component/knowledge_compiler/common" ) @@ -44,6 +46,77 @@ func TestPackWikiPlanBatches_SplitsLargeInput(t *testing.T) { } } +// TestWikiMapMaxTokens_OutputBudgetTracksInputBudget locks the input/output +// budget coupling: the extraction MaxTokens must leave at least the whole +// wikiMapTokenBudget input budget of headroom and, with a roomy model, give the +// output the rest of the context window after the batch's input is reserved. +func TestWikiMapMaxTokens_OutputBudgetTracksInputBudget(t *testing.T) { + // Unknown model context -> default window (DefaultLLMContextLength). Output + // gets the whole window minus the input budget. + got := wikiMapMaxTokens(0) + if want := common.DefaultLLMContextLength - wikiMapTokenBudget; got != want { + t.Fatalf("wikiMapMaxTokens(0) = %d, want %d", got, want) + } + // A model window that barely fits one batch must still grant at least the + // input budget of output space (never starve the output). + if got := wikiMapMaxTokens(2048); got != wikiMapTokenBudget { + t.Fatalf("wikiMapMaxTokens(2048) = %d, want %d (floor at input budget)", got, wikiMapTokenBudget) + } + // A roomy model: output = window - input budget. + if got := wikiMapMaxTokens(16384); got != 16384-wikiMapTokenBudget { + t.Fatalf("wikiMapMaxTokens(16384) = %d, want %d", got, 16384-wikiMapTokenBudget) + } +} + +func TestRunMapBatches_PreservesBatchOrderWithSubmitter(t *testing.T) { + previous := batchSubmitter + defer SetBatchSubmitter(previous) + + SetBatchSubmitter(func(ctx context.Context, jobs []func() error) error { + var wg sync.WaitGroup + errs := make(chan error, len(jobs)) + for _, job := range jobs { + job := job + wg.Add(1) + go func() { + defer wg.Done() + errs <- job() + }() + } + wg.Wait() + close(errs) + for err := range errs { + if err != nil { + return err + } + } + return ctx.Err() + }) + + batches := [][]common.Chunk{ + {{ID: "slow", Text: "slow"}}, + {{ID: "fast-1", Text: "fast-1"}}, + {{ID: "fast-2", Text: "fast-2"}}, + } + got, err := runMapBatches(context.Background(), batches, func(batch []common.Chunk) (wikiExtract, error) { + if batch[0].ID == "slow" { + time.Sleep(25 * time.Millisecond) + } + return wikiExtract{Topics: []string{batch[0].ID}}, nil + }) + if err != nil { + t.Fatalf("runMapBatches err = %v", err) + } + if len(got) != len(batches) { + t.Fatalf("runMapBatches len = %d, want %d", len(got), len(batches)) + } + for i, want := range []string{"slow", "fast-1", "fast-2"} { + if len(got[i].Topics) != 1 || got[i].Topics[0] != want { + t.Fatalf("runMapBatches[%d] = %#v, want topic %q", i, got[i], want) + } + } +} + func TestBuildSourceContext_SelectsKnownChunks(t *testing.T) { ctx := buildSourceContext([]common.Chunk{ {ID: "c1", Text: "alpha text"}, @@ -71,6 +144,61 @@ func TestNormalizeWikiPlanPages_FallbacksToEntitiesAndConcepts(t *testing.T) { } } +func TestMergePlanCandidates_DeduplicatesWithoutLLMMerge(t *testing.T) { + p := &wikiPipeline{ + docID: "doc-1", + reduced: wikiExtract{ + Entities: []wikiEntity{{Name: "Alpha"}}, + }, + } + merged := p.mergePlanCandidates([]wikiPlan{ + { + Title: "Alpha", + Pages: []wikiPlanPage{ + { + Slug: "entity/alpha", + Title: "Alpha", + PageType: "entity", + Topic: "Alpha", + EntityNames: []string{"Alpha"}, + RelatedKB: []string{"entity/beta", "missing", "entity/alpha"}, + Priority: 2, + }, + }, + }, + { + Pages: []wikiPlanPage{ + { + Slug: "entity/beta", + Title: "Beta", + PageType: "entity", + Topic: "Beta", + EntityNames: []string{"Beta"}, + RelatedKB: []string{"entity/alpha"}, + Priority: 1, + }, + { + Slug: "entity/alpha", + Title: "Alpha duplicate", + PageType: "entity", + Topic: "Alpha", + EntityNames: []string{"Alpha"}, + Priority: 3, + }, + }, + }, + }, p.reduced) + if len(merged.Pages) != 2 { + t.Fatalf("merged pages = %d, want 2", len(merged.Pages)) + } + if merged.Pages[0].Slug != "entity/beta" || merged.Pages[1].Slug != "entity/alpha" { + t.Fatalf("merged page order = %#v", merged.Pages) + } + if got := merged.Pages[1].RelatedKB; len(got) != 1 || got[0] != "entity/beta" { + t.Fatalf("alpha related links = %#v, want [entity/beta]", got) + } +} + type reconcileChatStub struct { resp string } @@ -134,6 +262,48 @@ func TestReconcilePlanPage_MaybeUsesLLMDecision(t *testing.T) { } } +// TestReconcilePlanPage_OverlapHeuristicSkipsLLM locks the Go-only enhancement: +// a candidate whose score is inside [maybe, update) but whose topic matches the +// planned page's topic is promoted straight to UPDATE without an LLM round. +func TestReconcilePlanPage_OverlapHeuristicSkipsLLM(t *testing.T) { + called := false + p := &wikiPipeline{ + ctx: context.Background(), + tenantID: "t1", + datasetID: "kb1", + llmID: "llm1", + deps: common.Deps{ + Chat: chatFunc(func(_ context.Context, _ common.ChatRequest) (*common.ChatResponse, error) { + called = true + return &common.ChatResponse{Content: `{"action":"CREATE"}`}, nil + }), + Embed: reconcileEmbedStub{}, + WikiPages: wikiStoreStub{similar: []common.WikiPageCandidate{{Slug: "topic/alpha", Title: "Alpha topic", Topic: "Alpha", Score: 0.85}}}, + }, + } + got, err := p.reconcilePlanPage(wikiPlanPage{ + Slug: "topic/alpha-new", + Title: "Alpha Topic", + PageType: "topic", + Topic: "Alpha", + }, []float32{0.1, 0.2, 0.3}) + if err != nil { + t.Fatalf("reconcilePlanPage err = %v", err) + } + if got == nil || got.Slug != "topic/alpha" { + t.Fatalf("reconcilePlanPage = %#v, want topic/alpha (topic overlap promotes to UPDATE)", got) + } + if called { + t.Fatalf("overlap heuristic should not invoke the LLM") + } +} + +type chatFunc func(context.Context, common.ChatRequest) (*common.ChatResponse, error) + +func (f chatFunc) Chat(ctx context.Context, req common.ChatRequest) (*common.ChatResponse, error) { + return f(ctx, req) +} + func TestReconcilePlanPage_LowScoreSkipsLLM(t *testing.T) { p := &wikiPipeline{ ctx: context.Background(), diff --git a/internal/ingestion/component/tokenizer.go b/internal/ingestion/component/tokenizer.go index f8001736f6..ba91cdaf98 100644 --- a/internal/ingestion/component/tokenizer.go +++ b/internal/ingestion/component/tokenizer.go @@ -82,10 +82,8 @@ import ( "encoding/json" "fmt" "log" - "os" "regexp" "slices" - "strconv" "strings" "go.uber.org/zap" @@ -100,19 +98,6 @@ import ( const ComponentNameTokenizer = "Tokenizer" -// embeddingBatchSize returns the embedding batch size, matching Python's -// settings.EMBEDDING_BATCH_SIZE. Reads TOKENIZER_EMBEDDING_BATCH_SIZE env -// var; defaults to 16. Invalid / non-positive values fall back to the -// default (diff Tokenizer Omission-3). -func embeddingBatchSize() int { - if v := os.Getenv("TOKENIZER_EMBEDDING_BATCH_SIZE"); v != "" { - if n, err := strconv.Atoi(v); err == nil && n > 0 { - return n - } - } - return 16 -} - // titleExtRE strips a trailing file-extension (e.g. ".pdf") from the // upstream document name before tokenizing it. Mirrors the python // `re.sub(r"\.[a-zA-Z]+$", "", name)` in tokenizer.py:137. @@ -134,6 +119,7 @@ type EmbeddingResult struct { // Embedder is the testability seam for the embedding branch. type Embedder interface { MaxTokens() int + BatchSize() int Encode(ctx context.Context, texts []string) ([]EmbeddingResult, error) } @@ -447,8 +433,12 @@ func (c *TokenizerComponent) embedChunks(ctx context.Context, tenantID, kbID, em } contentResults := make([]EmbeddingResult, 0, len(texts)) - for start := 0; start < len(texts); start += embeddingBatchSize() { - end := start + embeddingBatchSize() + batchSize := embedder.BatchSize() + if batchSize <= 0 { + return nil, 0, fmt.Errorf("tokenizer: embedder reported non-positive batch size %d", batchSize) + } + for start := 0; start < len(texts); start += batchSize { + end := start + batchSize if end > len(texts) { end = len(texts) } diff --git a/internal/ingestion/component/tokenizer_test.go b/internal/ingestion/component/tokenizer_test.go index 1dbe579168..f6ba5acb42 100644 --- a/internal/ingestion/component/tokenizer_test.go +++ b/internal/ingestion/component/tokenizer_test.go @@ -414,6 +414,8 @@ type countMismatchedEmbedder struct{ want int } func (c *countMismatchedEmbedder) MaxTokens() int { return 2048 } +func (c *countMismatchedEmbedder) BatchSize() int { return 16 } + func (c *countMismatchedEmbedder) Encode(ctx context.Context, texts []string) ([]EmbeddingResult, error) { out := make([]EmbeddingResult, c.want) for i := range out { diff --git a/internal/ingestion/component/tokenizer_unit_test.go b/internal/ingestion/component/tokenizer_unit_test.go index f82fc434f2..5c60d96cb7 100644 --- a/internal/ingestion/component/tokenizer_unit_test.go +++ b/internal/ingestion/component/tokenizer_unit_test.go @@ -30,6 +30,7 @@ import ( "time" "ragflow/internal/agent/runtime" + "ragflow/internal/entity/models" "ragflow/internal/ingestion/component/schema" "ragflow/internal/tokenizer" ) @@ -56,6 +57,10 @@ func (s *stubEmbedder) MaxTokens() int { return s.maxTokens } +func (s *stubEmbedder) BatchSize() int { + return 16 +} + func (s *stubEmbedder) Encode(ctx context.Context, texts []string) ([]EmbeddingResult, error) { s.calls.Add(1) copied := append([]string(nil), texts...) @@ -436,20 +441,22 @@ func TestTruncateForEmbedding_UnconfiguredClampsToDefault(t *testing.T) { // TestEmbeddingBatchSizeEnvVar covers Tokenizer Omission-3: the batch size // must be configurable via TOKENIZER_EMBEDDING_BATCH_SIZE env var, matching -// Python's configurable settings.EMBEDDING_BATCH_SIZE. +// Python's configurable settings.EMBEDDING_BATCH_SIZE. The resolved batch size +// is now provided by models.GetEmbeddingBatchSize (env override -> provider +// capability -> default), surfaced through Embedder.BatchSize(). func TestEmbeddingBatchSizeEnvVar(t *testing.T) { - if got := embeddingBatchSize(); got != 16 { - t.Errorf("embeddingBatchSize() default = %d, want 16", got) + if got := models.GetEmbeddingBatchSize(""); got != models.DefaultEmbeddingBatchSize { + t.Errorf("GetEmbeddingBatchSize() default = %d, want %d", got, models.DefaultEmbeddingBatchSize) } os.Setenv("TOKENIZER_EMBEDDING_BATCH_SIZE", "32") t.Cleanup(func() { os.Unsetenv("TOKENIZER_EMBEDDING_BATCH_SIZE") }) - if got := embeddingBatchSize(); got != 32 { - t.Errorf("embeddingBatchSize() after env = %d, want 32", got) + if got := models.GetEmbeddingBatchSize(""); got != 32 { + t.Errorf("GetEmbeddingBatchSize() after env = %d, want 32", got) } // Invalid value falls back to default. os.Setenv("TOKENIZER_EMBEDDING_BATCH_SIZE", "bad") - if got := embeddingBatchSize(); got != 16 { - t.Errorf("embeddingBatchSize() invalid env = %d, want 16", got) + if got := models.GetEmbeddingBatchSize(""); got != models.DefaultEmbeddingBatchSize { + t.Errorf("GetEmbeddingBatchSize() invalid env = %d, want %d", got, models.DefaultEmbeddingBatchSize) } } diff --git a/internal/ingestion/knowledge_compile/consumer.go b/internal/ingestion/knowledge_compile/consumer.go index 1d8bb6c3df..2a663c7c87 100644 --- a/internal/ingestion/knowledge_compile/consumer.go +++ b/internal/ingestion/knowledge_compile/consumer.go @@ -21,8 +21,11 @@ import ( "sync" "time" + "ragflow/internal/common" "ragflow/internal/engine" kccommon "ragflow/internal/ingestion/component/knowledge_compiler/common" + + "go.uber.org/zap" ) // Consumer is the dataset-level post-processing worker (§11.5). Multiple @@ -183,10 +186,19 @@ func (c *Consumer) processClaim(ctx context.Context, cr ClaimResult) { // must leave the claimed batch in the backlog for reclamation/retry rather // than silently dropping it (C5: never ack what we failed to merge). if batchErr != nil { + common.Error("knowledge_compile: batch processing failed, leaving batch for retry", + batchErr, + zap.String("dataset_id", datasetID), + zap.Int("entries", len(cr.Entries))) + if err := c.scheduler.SetError(ctx, datasetID, cr.Token, batchErr.Error()); err != nil { + common.Warn("knowledge_compile: failed to record error_msg", + zap.String("dataset_id", datasetID), zap.Error(err)) + } return } if _, err := c.scheduler.Ack(ctx, datasetID, cr.Token, cr.Entries); err != nil { - _ = err + common.Warn("knowledge_compile: ack failed", + zap.String("dataset_id", datasetID), zap.Error(err)) } } @@ -195,6 +207,10 @@ func (c *Consumer) processClaim(ctx context.Context, cr ClaimResult) { // returns an error if any reader/dedup/writer step fails so the caller can // leave the batch for reclamation instead of acking dropped work. func (c *Consumer) processBatch(ctx context.Context, tenant, kb string, entries []BacklogEntry) error { + common.Info("knowledge_compile: processing claimed batch", + zap.String("dataset_id", kb), + zap.String("tenant_id", tenant), + zap.Int("entries", len(entries))) c.mu.Lock() if c.tombs == nil { c.tombs = map[string]map[string]uint64{} @@ -355,7 +371,7 @@ func (c *Consumer) processBatch(ctx context.Context, tenant, kb string, entries // fan it out across the shared global compilerPool (vCPU-sized). Output order // is irrelevant: merged rows are upserted by their idempotent dataset-level // id, and each candidate lands in exactly one group / the unmatched set. - jobs := make([]compilerJob, 0, len(candidates)) + jobs := make([]CompilerJob, 0, len(candidates)) for _, cand := range candidates { cand := cand jobs = append(jobs, func() error { @@ -437,5 +453,10 @@ func (c *Consumer) processBatch(ctx context.Context, tenant, kb string, entries delete(c.tombs[kb], docID) } c.mu.Unlock() + common.Info("knowledge_compile: batch merge complete", + zap.String("dataset_id", kb), + zap.Int("completed_docs", len(completed)), + zap.Int("deleted_docs", len(deleted)), + zap.Int("merged_rows_written", len(mergedFinal))) return nil } diff --git a/internal/ingestion/knowledge_compile/consumer_test.go b/internal/ingestion/knowledge_compile/consumer_test.go index 2c3c9d9035..29f0cec042 100644 --- a/internal/ingestion/knowledge_compile/consumer_test.go +++ b/internal/ingestion/knowledge_compile/consumer_test.go @@ -249,3 +249,154 @@ func TestSchedulerReclaimExpired(t *testing.T) { t.Fatalf("expected 1 entry after reclaim, got %d", len(cr2.Entries)) } } + +// rowCounts snapshots a fake scheduling row's inflight/backlog entry counts. +type rowCounts struct { + inflight int + backlog int +} + +func fakeRowCounts(sch *FakeScheduler, datasetID string) (rowCounts, string) { + sch.mu.Lock() + defer sch.mu.Unlock() + r := sch.rows[datasetID] + if r == nil { + return rowCounts{}, "" + } + return rowCounts{inflight: len(r.inflight), backlog: len(r.backlog)}, r.state +} + +// TestSchedulerStateMachineLocksStateAndCounts locks the full lifecycle state +// machine (plan v4.1 §9.2) on the FakeScheduler, which mirrors the MySQL +// scheduler's transitions: +// +// Publish -> pending; Claim -> running (+error cleared); +// SetError (failed batch left for retry) keeps running + records error; +// lease expiry -> reclaimOne -> pending (inflight moved back to backlog); +// Ack with backlog drained -> completed. +func TestSchedulerStateMachineLocksStateAndCounts(t *testing.T) { + sch := NewFakeScheduler() + + // Publish two docs: state=pending, backlog=2, inflight=0. + if err := sch.Publish(context.Background(), "t1", "kb1", "d1", string(EventTypeCompleted), 1); err != nil { + t.Fatalf("publish d1: %v", err) + } + if err := sch.Publish(context.Background(), "t1", "kb1", "d2", string(EventTypeCompleted), 2); err != nil { + t.Fatalf("publish d2: %v", err) + } + if c, s := fakeRowCounts(sch, "kb1"); s != DatasetStatePending || c.backlog != 2 || c.inflight != 0 { + t.Fatalf("after publish: want state=pending backlog=2 inflight=0, got state=%s %+v", s, c) + } + + // Claim: state=running, backlog moves to inflight (batch=2), error cleared. + cr, ok, err := sch.Claim(context.Background(), "kb1") + if err != nil || !ok { + t.Fatalf("claim: ok=%v err=%v", ok, err) + } + if len(cr.Entries) != 2 { + t.Fatalf("expected 2-entry claim, got %d", len(cr.Entries)) + } + if c, s := fakeRowCounts(sch, "kb1"); s != DatasetStateRunning || c.backlog != 0 || c.inflight != 2 { + t.Fatalf("after claim: want state=running backlog=0 inflight=2, got state=%s %+v", s, c) + } + + // Failed batch left for retry: SetError records a diagnostic, state stays running. + if err := sch.SetError(context.Background(), "kb1", cr.Token, "boom"); err != nil { + t.Fatalf("set error: %v", err) + } + sch.mu.Lock() + gotErr := sch.rows["kb1"].errorMsg + sch.mu.Unlock() + if gotErr != "boom" { + t.Fatalf("expected errorMsg=boom, got %q", gotErr) + } + if _, s := fakeRowCounts(sch, "kb1"); s != DatasetStateRunning { + t.Fatalf("failed batch must stay running (left for retry), got state=%s", s) + } + + // Lease expires: reclaimOne moves inflight back to backlog, clears lease -> pending. + past := time.Now().Add(-time.Hour) + sch.mu.Lock() + sch.rows["kb1"].expires = &past + sch.mu.Unlock() + // Lock the reclaim transition in isolation (before any re-claim): the fake's + // reclaim helper is the same code path TryClaim uses, mirroring reclaimOne. + sch.mu.Lock() + if id := sch.fakeReclaimExpired(time.Now()); id != "kb1" { + sch.mu.Unlock() + t.Fatalf("expected kb1 to be reclaimed, got %q", id) + } + sch.mu.Unlock() + if c, s := fakeRowCounts(sch, "kb1"); s != DatasetStatePending || c.backlog != 2 || c.inflight != 0 { + t.Fatalf("after reclaim: want state=pending backlog=2 inflight=0, got state=%s %+v", s, c) + } + + // Claim again then Ack the drained batch -> completed, counts zeroed. + cr2, ok3, err := sch.Claim(context.Background(), "kb1") + if err != nil || !ok3 { + t.Fatalf("re-claim: ok=%v err=%v", ok3, err) + } + if _, err := sch.Ack(context.Background(), "kb1", cr2.Token, cr2.Entries); err != nil { + t.Fatalf("ack: %v", err) + } + if c, s := fakeRowCounts(sch, "kb1"); s != DatasetStateCompleted || c.backlog != 0 || c.inflight != 0 { + t.Fatalf("after ack drain: want state=completed backlog=0 inflight=0, got state=%s %+v", s, c) + } +} + +// TestSchedulerSetErrorScopedToClaimToken locks the concurrency guard on +// SetError: a failed batch is diagnosed only while its own claim token is still +// live. If worker A's lease expires and is reclaimed, worker B re-claims and +// completes; a late SetError from A (stale token) must NOT overwrite the row's +// diagnostic, otherwise a completed state would be misread as failed. +func TestSchedulerSetErrorScopedToClaimToken(t *testing.T) { + sch := NewFakeScheduler() + if err := sch.Publish(context.Background(), "t1", "kb1", "d1", string(EventTypeCompleted), 1); err != nil { + t.Fatalf("publish: %v", err) + } + + // A claims and begins processing (running). + crA, ok, err := sch.Claim(context.Background(), "kb1") + if err != nil || !ok { + t.Fatalf("claim A: ok=%v err=%v", ok, err) + } + + // A's lease expires before it finishes; the sweeper reclaims the inflight + // batch back to backlog (pending) and B takes over. + past := time.Now().Add(-time.Hour) + sch.mu.Lock() + sch.rows["kb1"].expires = &past + sch.mu.Unlock() + if id := sch.fakeReclaimExpired(time.Now()); id != "kb1" { + t.Fatalf("expected kb1 reclaimed, got %q", id) + } + crB, okB, err := sch.Claim(context.Background(), "kb1") + if err != nil || !okB { + t.Fatalf("claim B: ok=%v err=%v", okB, err) + } + if crA.Token == crB.Token { + t.Fatalf("expected distinct claim tokens, got %q", crA.Token) + } + + // B succeeds and drains the backlog to completed. + if _, err := sch.Ack(context.Background(), "kb1", crB.Token, crB.Entries); err != nil { + t.Fatalf("ack B: %v", err) + } + if c, s := fakeRowCounts(sch, "kb1"); s != DatasetStateCompleted || c.backlog != 0 || c.inflight != 0 { + t.Fatalf("after B ack: want completed empty, got state=%s %+v", s, c) + } + + // A's late failure arrives with its now-stale token: it must be ignored. + if err := sch.SetError(context.Background(), "kb1", crA.Token, "stale failure from A"); err != nil { + t.Fatalf("stale set error: %v", err) + } + sch.mu.Lock() + gotErr := sch.rows["kb1"].errorMsg + sch.mu.Unlock() + if gotErr != "" { + t.Fatalf("stale SetError overwrote diagnostic: got errorMsg=%q want empty", gotErr) + } + if _, s := fakeRowCounts(sch, "kb1"); s != DatasetStateCompleted { + t.Fatalf("stale SetError changed state: got %q want completed", s) + } +} diff --git a/internal/ingestion/knowledge_compile/dedup.go b/internal/ingestion/knowledge_compile/dedup.go index 6c0cec3993..6cc775769c 100644 --- a/internal/ingestion/knowledge_compile/dedup.go +++ b/internal/ingestion/knowledge_compile/dedup.go @@ -79,7 +79,7 @@ func NewLLMDeduper(chat kccommon.ChatInvoker, embed kccommon.Embedder, llmID str // never blocks on a single job, so a stopped pool returns an error instead of // hanging DecideBatch. decider.SetSubmitter(func(ctx context.Context, fn func() error) error { - return SubmitCompilerJobs(ctx, []compilerJob{fn}) + return SubmitCompilerJobs(ctx, []CompilerJob{fn}) }) return &llmDeduper{group: structure.NewGroupedDeduper(decider), decider: decider, embed: embed} } diff --git a/internal/ingestion/knowledge_compile/pool.go b/internal/ingestion/knowledge_compile/pool.go index b6f5dd31a7..6b527dffd3 100644 --- a/internal/ingestion/knowledge_compile/pool.go +++ b/internal/ingestion/knowledge_compile/pool.go @@ -24,10 +24,11 @@ import ( "ragflow/internal/utility" ) -// compilerJob is one unit of knowledge-compilation work (an I/O- or -// LLM-bounded task) executed on the shared global pool. It is a type alias for -// func() error so callers can pass plain []func() error slices without a cast. -type compilerJob = func() error +// CompilerJob is one unit of knowledge-compilation work (an I/O- or +// LLM-bounded task) executed on the shared global pool. It is an exported type +// alias for func() error so callers (including lower-level packages such as the +// knowledge_compiler wiring) can pass plain []func() error slices without a cast. +type CompilerJob = func() error // compilerPool is the process-wide bounded worker pool that drives cross-doc // concurrency for every knowledge-compilation stage: the DocEngine KNN pass in @@ -43,10 +44,10 @@ type compilerJob = func() error // (KNN / write / delete) or LLM-bounded (merge decisions) rather than // CPU-bounded, so the degree of useful parallelism is capped by the number of // available cores rather than by a hand-tuned constant. -var compilerPool = utility.NewWorkerPool[compilerJob, struct{}]( +var compilerPool = utility.NewWorkerPool[CompilerJob, struct{}]( compilerConcurrency(), compilerConcurrency()*4, - func(_ context.Context, j compilerJob) (struct{}, error) { return struct{}{}, j() }, + func(_ context.Context, j CompilerJob) (struct{}, error) { return struct{}{}, j() }, ) // compilerConcurrency resolves the global pool size. It defaults to the host @@ -82,11 +83,11 @@ func SetCompilerConcurrency(n int) { // then Wait on each in a second pass on the calling goroutine. This keeps the // fan-out bounded by the shared pool's worker count while avoiding len(jobs) // short-lived goroutines. -func runCompilerJobs(ctx context.Context, jobs []compilerJob) error { +func runCompilerJobs(ctx context.Context, jobs []CompilerJob) error { if len(jobs) == 0 { return nil } - futures := make([]utility.WorkerPoolFuture[compilerJob, struct{}], 0, len(jobs)) + futures := make([]utility.WorkerPoolFuture[CompilerJob, struct{}], 0, len(jobs)) var firstErr error for _, j := range jobs { f, err := compilerPool.Submit(ctx, j) @@ -121,7 +122,7 @@ func runCompilerJobs(ctx context.Context, jobs []compilerJob) error { // SubmitCompilerJob runs a single job on the global pool and waits for it, // returning its error. Used to inject bounded parallelism into lower-level // packages (e.g. structure.LLMMergeDecider) without creating an import cycle. -func SubmitCompilerJob(ctx context.Context, fn compilerJob) error { +func SubmitCompilerJob(ctx context.Context, fn CompilerJob) error { f, err := compilerPool.Submit(ctx, fn) if err != nil { return err @@ -138,10 +139,10 @@ func SubmitCompilerJob(ctx context.Context, fn compilerJob) error { // the one process-wide compiler pool. Implementations must submit every job to // the shared pool, wait for all to finish, and return the first non-nil error // (without StopWait-ing the global pool). -type CompilerBatchSubmitter func(ctx context.Context, jobs []compilerJob) error +type CompilerBatchSubmitter func(ctx context.Context, jobs []CompilerJob) error // SubmitCompilerJobs fans out a batch of jobs on the global pool and returns the // first error. This is the CompilerBatchSubmitter handed to variant packages. -func SubmitCompilerJobs(ctx context.Context, jobs []compilerJob) error { +func SubmitCompilerJobs(ctx context.Context, jobs []CompilerJob) error { return runCompilerJobs(ctx, jobs) } diff --git a/internal/ingestion/knowledge_compile/reader.go b/internal/ingestion/knowledge_compile/reader.go index 3b0332a85e..a8deeee6ae 100644 --- a/internal/ingestion/knowledge_compile/reader.go +++ b/internal/ingestion/knowledge_compile/reader.go @@ -68,6 +68,16 @@ var compiledSelectFields = []string{ "slug_kwd", "type", } +// wikiSelectFields are the additional columns a wiki page carries (beyond +// compiledSelectFields) that must survive the doc→merge round-trip so the +// dataset-level merged rows keep the fields the artifact API and page renderers +// depend on (page_type_kwd/topic_kwd/title_kwd/...). +var wikiSelectFields = []string{ + "page_type_kwd", "topic_kwd", "title_kwd", + "entity_names_kwd", "summary_with_weight", + "related_kb_pages_kwd", "outlinks_kwd", "section_level_int", +} + // loadDocProductsLimit is the per-page size used when scrolling a single // document's compiled rows. A document can compile more than this many rows, so // LoadDocProducts pages until the engine returns fewer than a full page. @@ -92,7 +102,7 @@ func (r engineReader) LoadDocProducts(ctx context.Context, tenant, kb, docID str IndexNames: []string{fmt.Sprintf("ragflow_%s", tenant)}, KbIDs: []string{kb}, Filter: map[string]interface{}{"doc_id": docID}, - SelectFields: compiledSelectFields, + SelectFields: append(append([]string(nil), compiledSelectFields...), wikiSelectFields...), Limit: loadDocProductsLimit, Offset: offset, }) @@ -148,8 +158,36 @@ func productFromChunkMap(c map[string]interface{}, tenant string) (kccommon.Prod meta["kind"] = "relation" } if v, ok := c["slug_kwd"].(string); ok && v != "" { + // slug_kwd is the full "/" form (Python writer + // contract); reconstruct it verbatim so the round-trip stays full-form. meta["slug"] = v } + // Restore wiki page fields so the merged product (and hence the dataset-level + // merged row) retains the metadata the artifact API and page renderers read. + if v, ok := c["page_type_kwd"].(string); ok && v != "" { + meta["page_type"] = v + } + if v, ok := c["topic_kwd"].(string); ok && v != "" { + meta["topic"] = v + } + if v, ok := c["title_kwd"].(string); ok && v != "" { + meta["title"] = v + } + if v, ok := c["summary_with_weight"].(string); ok && v != "" { + meta["summary"] = v + } + if v := metaStringSlice(c, "entity_names_kwd"); len(v) > 0 { + meta["entity_names"] = v + } + if v := metaStringSlice(c, "related_kb_pages_kwd"); len(v) > 0 { + meta["related_kb_pages"] = v + } + if v := metaStringSlice(c, "outlinks_kwd"); len(v) > 0 { + meta["outlinks"] = v + } + if v, ok := metaInt(c, "section_level_int"); ok { + meta["section_level"] = v + } if v, ok := c["type"].(string); ok && v != "" { meta["type"] = v } @@ -196,9 +234,10 @@ func (r engineReader) SearchSimilar(ctx context.Context, tenant, kb string, vari IndexNames: []string{fmt.Sprintf("ragflow_%s", tenant)}, KbIDs: []string{kb}, Limit: topN, - SelectFields: []string{"id", "doc_id", "kb_id", "content_with_weight", "kc_payload", + SelectFields: append([]string{"id", "doc_id", "kb_id", "content_with_weight", "kc_payload", "name_kwd", "entity_type_kwd", "from_entity_kwd", "to_entity_kwd", "slug_kwd", "type", "source_chunk_ids", "source_doc_ids", "kc_merged", "compile_kwd"}, + wikiSelectFields...), Filter: map[string]interface{}{ "kc_merged": 1, "compile_kwd": string(variant), diff --git a/internal/ingestion/knowledge_compile/scheduler.go b/internal/ingestion/knowledge_compile/scheduler.go index 323311877c..03e9f15be4 100644 --- a/internal/ingestion/knowledge_compile/scheduler.go +++ b/internal/ingestion/knowledge_compile/scheduler.go @@ -23,9 +23,11 @@ import ( "sync" "time" + "ragflow/internal/common" "ragflow/internal/engine" "ragflow/internal/entity" + "go.uber.org/zap" "gorm.io/gorm" "gorm.io/gorm/clause" ) @@ -41,6 +43,15 @@ var ErrClaimTokenMismatch = errors.New("knowledge_compile: claim token mismatch" // scheduling truth. const notifySubject = "notify.kc.workers" +// Dataset compile lifecycle states. Defined in entity so the scheduler, the +// dataset service (status API) and tests share one source of truth. +const ( + DatasetStateIdle = entity.DatasetStateIdle + DatasetStatePending = entity.DatasetStatePending + DatasetStateRunning = entity.DatasetStateRunning + DatasetStateCompleted = entity.DatasetStateCompleted +) + // BacklogEntry is one scheduling unit appended to a KB's backlog (Option E // §11.4). It carries the doc id plus the original event kind/seq so the // consumer can re-apply the same out-of-order / tombstone guards as the @@ -95,6 +106,14 @@ type Claimer interface { // only when backlog is also empty. Ack(ctx context.Context, datasetID, token string, batch []BacklogEntry) (backlogRemaining int, err error) + // SetError records a diagnostic message for a failed batch without changing + // the lifecycle state (the batch is left for retry, so state stays running). + // It is best-effort for observability. token must be the claim token of the + // batch that failed: the update only applies while that exact claim is still + // live, so a stale worker whose lease was reclaimed cannot overwrite the + // status of the worker that took over. + SetError(ctx context.Context, datasetID, token, errMsg string) error + // SubscribeNotify returns a channel of dataset ids pushed by Publish, or nil // when the implementation has no push wake-up (callers fall back to polling). SubscribeNotify(ctx context.Context) (<-chan string, error) @@ -197,11 +216,23 @@ func (s *mysqlScheduler) Publish(ctx context.Context, tenantID, datasetID, docID backlog := parseEntries(row.BacklogDocIDs) backlog = append(backlog, entry) row.BacklogDocIDs = marshalEntries(backlog) + // Surface a pending state unless a worker is already running (a live lease + // means the consumer is mid-merge on this KB; a pending transition would + // wrongly hide that). completed -> pending when new work arrives. + if row.State != DatasetStateRunning { + row.State = DatasetStatePending + } return tx.Save(&row).Error }) if err != nil { return fmt.Errorf("knowledge_compile: publish backlog %s: %w", datasetID, err) } + common.Info("knowledge_compile: published backlog entry", + zap.String("dataset_id", datasetID), + zap.String("tenant_id", tenantID), + zap.String("doc_id", docID), + zap.String("event_type", eventType), + zap.Uint64("seq", seq)) return s.notify(ctx, datasetID) } @@ -238,9 +269,17 @@ func (s *mysqlScheduler) claimRow(ctx context.Context, tx *gorm.DB, datasetID st row.ClaimToken = generateHolder() exp := now.Add(s.leaseTTL) row.ClaimExpiresAt = &exp + row.State = DatasetStateRunning + row.ErrorMsg = "" if err := tx.Save(&row).Error; err != nil { return ClaimResult{}, false, err } + common.Info("knowledge_compile: claimed dataset batch", + zap.String("dataset_id", row.DatasetID), + zap.String("tenant_id", row.TenantID), + zap.Int("batch_size", len(batch)), + zap.Int("backlog_remaining", len(backlog)-n), + zap.String("token", row.ClaimToken)) return ClaimResult{DatasetID: row.DatasetID, TenantID: row.TenantID, Entries: batch, Token: row.ClaimToken}, true, nil } @@ -286,18 +325,61 @@ func (s *mysqlScheduler) Ack(ctx context.Context, datasetID, token string, batch row.ClaimToken = "" row.ClaimExpiresAt = nil } + remaining = len(parseEntries(row.BacklogDocIDs)) + // Backlog still has work -> pending; drained to empty -> completed. + if remaining > 0 { + row.State = DatasetStatePending + } else { + row.State = DatasetStateCompleted + now := time.Now() + row.LastCompletedAt = &now + } if err := tx.Save(&row).Error; err != nil { return err } - remaining = len(parseEntries(row.BacklogDocIDs)) return nil }) if err != nil { return 0, err } + common.Info("knowledge_compile: acked dataset batch", + zap.String("dataset_id", datasetID), + zap.Int("batch_size", len(batch)), + zap.Int("backlog_remaining", remaining)) return remaining, nil } +// SetError records a best-effort diagnostic on the dataset row when a batch +// merge fails (consumer failure path). It does not change the lifecycle state: +// a failed batch is left in backlog for retry, so the row stays running/pending +// and the error message is surfaced by the status API for diagnosis. +// +// token scopes the write to the exact claim that failed: the update only lands +// while that claim token is still live on the row. If the original worker's +// lease expired and another worker took over (a new claim token), a stale +// SetError from the old worker is a no-op and cannot overwrite the new worker's +// status/diagnostic. +func (s *mysqlScheduler) SetError(ctx context.Context, datasetID, token, errMsg string) error { + if s.db == nil { + return nil + } + msg := errMsg + if len(msg) > 4000 { + msg = msg[:4000] + } + res := s.db.WithContext(ctx).Model(&entity.KnowledgeCompileDataset{}). + Where("dataset_id = ? AND claim_token = ?", datasetID, token). + Update("error_msg", msg) + if res.Error != nil { + return res.Error + } + if res.RowsAffected == 0 { + common.Warn("knowledge_compile: set_error ignored (claim token mismatch / lease reclaimed)", + zap.String("dataset_id", datasetID)) + } + return nil +} + func (s *mysqlScheduler) TouchClaim(ctx context.Context, datasetID, token string, ttl time.Duration) (bool, error) { if s.db == nil { return false, nil @@ -308,6 +390,10 @@ func (s *mysqlScheduler) TouchClaim(ctx context.Context, datasetID, token string if res.Error != nil { return false, res.Error } + if res.RowsAffected == 0 { + common.Warn("knowledge_compile: claim touch failed (lease lost or token mismatch)", + zap.String("dataset_id", datasetID), zap.String("token", token)) + } return res.RowsAffected > 0, nil } @@ -405,9 +491,14 @@ func (s *mysqlScheduler) reclaimOne(ctx context.Context, tx *gorm.DB, now time.T cur.ClaimOwner = "" cur.ClaimToken = "" cur.ClaimExpiresAt = nil + cur.State = DatasetStatePending if err := tx.Save(&cur).Error; err != nil { return "", false, err } + common.Warn("knowledge_compile: reclaimed expired inflight lease", + zap.String("dataset_id", cur.DatasetID), + zap.String("tenant_id", cur.TenantID), + zap.Int("reclaimed_entries", len(inflight))) return cur.DatasetID, true, nil } return "", false, nil @@ -420,7 +511,13 @@ func (s *mysqlScheduler) notify(ctx context.Context, datasetID string) error { return nil } payload, _ := json.Marshal(map[string]string{"dataset_id": datasetID}) - return s.mq.PublishKnowledgeCompile(notifySubject, payload) + if err := s.mq.PublishKnowledgeCompile(notifySubject, payload); err != nil { + common.Warn("knowledge_compile: publish notify failed (workers will poll)", + zap.String("dataset_id", datasetID), zap.Error(err)) + return err + } + common.Info("knowledge_compile: published worker notify", zap.String("dataset_id", datasetID)) + return nil } func (s *mysqlScheduler) SubscribeNotify(ctx context.Context) (<-chan string, error) { @@ -461,6 +558,8 @@ type fakeRow struct { owner string token string expires *time.Time + state string + errorMsg string } // FakeScheduler is an in-memory Publisher + Claimer used by tests. It mirrors @@ -498,6 +597,10 @@ func (f *FakeScheduler) Publish(_ context.Context, tenantID, datasetID, docID, e r.tenant = tenantID } r.backlog = append(r.backlog, BacklogEntry{DocID: docID, EventType: eventType, Seq: seq}) + // Mirror the MySQL scheduler: surface pending unless a worker is running. + if r.state != DatasetStateRunning { + r.state = DatasetStatePending + } select { case f.notifyCh <- datasetID: default: @@ -528,6 +631,8 @@ func (f *FakeScheduler) Claim(_ context.Context, datasetID string) (ClaimResult, r.token = generateHolder() exp := now.Add(f.leaseTTL) r.expires = &exp + r.state = DatasetStateRunning + r.errorMsg = "" return ClaimResult{DatasetID: datasetID, TenantID: r.tenant, Entries: batch, Token: r.token}, true, nil } @@ -545,6 +650,12 @@ func (f *FakeScheduler) Ack(_ context.Context, datasetID, token string, batch [] if len(r.inflight) == 0 { r.owner, r.token, r.expires = "", "", nil } + // Mirror the MySQL scheduler: backlog still has work -> pending; drained -> completed. + if len(r.backlog) > 0 { + r.state = DatasetStatePending + } else { + r.state = DatasetStateCompleted + } return len(r.backlog), nil } @@ -560,6 +671,39 @@ func (f *FakeScheduler) TouchClaim(_ context.Context, datasetID, token string, _ return true, nil } +func (f *FakeScheduler) SetError(_ context.Context, datasetID, token, errMsg string) error { + f.mu.Lock() + defer f.mu.Unlock() + r, ok := f.rows[datasetID] + if !ok || r.token != token { + return nil + } + r.errorMsg = errMsg + return nil +} + +// fakeReclaimExpired finds one dataset with an expired lease and non-empty +// inflight, moves the inflight batch back to backlog, clears the lease, and +// marks the row pending — mirroring the MySQL scheduler's reclaimOne. It returns +// the reclaimed dataset id, or "" when nothing is expired. The lock must already +// be held by the caller. +func (f *FakeScheduler) fakeReclaimExpired(now time.Time) string { + for id, r := range f.rows { + if r.owner != "" && r.expires != nil && r.expires.After(now) { + continue + } + if len(r.inflight) > 0 { + r.backlog = append(r.backlog, r.inflight...) + r.inflight = nil + r.owner, r.token, r.expires = "", "", nil + // reclaimOne: inflight -> backlog, lease cleared -> pending. + r.state = DatasetStatePending + return id + } + } + return "" +} + // TryClaim mirrors the production flow: claim a ready dataset, otherwise reclaim // an expired lease and claim it. func (f *FakeScheduler) TryClaim(ctx context.Context) (ClaimResult, bool, error) { @@ -573,18 +717,7 @@ func (f *FakeScheduler) TryClaim(ctx context.Context) (ClaimResult, bool, error) } } if readyID == "" { - for id, r := range f.rows { - if r.owner != "" && r.expires != nil && r.expires.After(now) { - continue - } - if len(r.inflight) > 0 { - r.backlog = append(r.backlog, r.inflight...) - r.inflight = nil - r.owner, r.token, r.expires = "", "", nil - expiredID = id - break - } - } + expiredID = f.fakeReclaimExpired(now) } f.mu.Unlock() if readyID != "" { diff --git a/internal/ingestion/knowledge_compile/service.go b/internal/ingestion/knowledge_compile/service.go index 7faedcd0f4..8141605b9c 100644 --- a/internal/ingestion/knowledge_compile/service.go +++ b/internal/ingestion/knowledge_compile/service.go @@ -101,7 +101,7 @@ func defaultDeduperFactory(tenant string) (Deduper, error) { if err != nil { return nil, err } - return NewLLMDeduper(deps.Chat, deps.Embed, defaultLLMID, 0.99, deps.LLMMaxLength), nil + return NewLLMDeduper(deps.Chat, deps.Embed, defaultLLMID, 0.99, deps.ModelContextLen), nil } func generateHolder() string { diff --git a/internal/ingestion/knowledge_compile/writer.go b/internal/ingestion/knowledge_compile/writer.go index 126e13d71c..45e6fded28 100644 --- a/internal/ingestion/knowledge_compile/writer.go +++ b/internal/ingestion/knowledge_compile/writer.go @@ -20,6 +20,7 @@ import ( "crypto/sha256" "encoding/hex" "fmt" + "strings" "ragflow/internal/engine" "ragflow/internal/engine/types" @@ -70,7 +71,7 @@ func (w engineWriter) WriteMerged(ctx context.Context, tenant, kb string, produc baseName := fmt.Sprintf("ragflow_%s", tenant) // Shard the rows and drive the inserts through the shared global pool // (docengine-bounded) instead of one monolithic InsertChunks call. - jobs := make([]compilerJob, 0, (len(products)+writeMergedBatchSize-1)/writeMergedBatchSize) + jobs := make([]CompilerJob, 0, (len(products)+writeMergedBatchSize-1)/writeMergedBatchSize) for start := 0; start < len(products); start += writeMergedBatchSize { end := start + writeMergedBatchSize if end > len(products) { @@ -108,6 +109,47 @@ func mergedChunkMap(tenant, kb string, p kccommon.Product) map[string]interface{ "source_doc_ids": srcDocIDs, "source_chunk_ids": srcChunkIDs, } + // Carry the wiki page metadata onto the merged row so the dataset-level + // products keep the fields the artifact API (ListArtifacts/ListWikiTopics) + // and page renderers read. Without this the merged rows lose page_type_kwd / + // topic_kwd / title_kwd and the compilation page would show no wiki pages + // even though per-document products carry them. + // + // slug_kwd follows the Python writer contract (api/db/db_models.py): it is + // stored as the full "/" form so GetWikiPage's filter + // (page_type + "/" + slug) matches directly. + pageType := metaString(p.Meta, "page_type") + if slug := metaString(p.Meta, "slug"); slug != "" { + // Normalize to the full "/" form (Python writer + // contract). Idempotent: slugs that already carry the prefix are kept. + fullSlug := slug + if pageType != "" && !strings.Contains(slug, "/") { + fullSlug = pageType + "/" + slug + } + m["slug_kwd"] = fullSlug + m["artifact_slug_kwd"] = fullSlug + } + if v := metaString(p.Meta, "title"); v != "" { + m["title_kwd"] = v + } + if pageType != "" { + m["page_type_kwd"] = pageType + } + if v := metaString(p.Meta, "topic"); v != "" { + m["topic_kwd"] = v + } + if v := metaString(p.Meta, "summary"); v != "" { + m["summary_with_weight"] = v + } + if v := metaStringSlice(p.Meta, "entity_names"); len(v) > 0 { + m["entity_names_kwd"] = v + } + if v := metaStringSlice(p.Meta, "related_kb_pages"); len(v) > 0 { + m["related_kb_pages_kwd"] = v + } + if v := metaStringSlice(p.Meta, "outlinks"); len(v) > 0 { + m["outlinks_kwd"] = v + } // Persist the merged product's embedding under the dimension-suffixed column // used elsewhere in the index, so dataset-level rows remain vector-searchable // and the Reader can reconstruct them (otherwise the vector is silently @@ -166,7 +208,7 @@ func (w engineWriter) StripMergedSources(ctx context.Context, tenant, kb string, const batchSize = 2000 var toDeleteIDs []string - var jobs []compilerJob + var jobs []CompilerJob offset := 0 for { res, err := eng.Search(ctx, &types.SearchRequest{ @@ -267,6 +309,15 @@ func hashStr(s string) string { return hex.EncodeToString(sum[:]) } +// metaString extracts a string from a map value, tolerating a missing or +// non-string entry. +func metaString(m map[string]any, key string) string { + if v, ok := m[key].(string); ok { + return v + } + return "" +} + func metaStringSlice(m map[string]any, key string) []string { switch v := m[key].(type) { case []string: @@ -282,3 +333,23 @@ func metaStringSlice(m map[string]any, key string) []string { } return nil } + +// metaInt extracts an integer from a map value that may be boxed as float64 +// (JSON number), int64, string, or a typed int — the engine/JSON round-trip does +// not guarantee a single numeric type. +func metaInt(m map[string]any, key string) (int64, bool) { + switch v := m[key].(type) { + case int64: + return v, true + case int: + return int64(v), true + case float64: + return int64(v), true + case string: + var n int64 + if _, err := fmt.Sscanf(v, "%d", &n); err == nil { + return n, true + } + } + return 0, false +} diff --git a/internal/ingestion/knowledge_compile/writer_test.go b/internal/ingestion/knowledge_compile/writer_test.go new file mode 100644 index 0000000000..8f7e9bfe42 --- /dev/null +++ b/internal/ingestion/knowledge_compile/writer_test.go @@ -0,0 +1,98 @@ +package knowledge_compile + +import ( + "testing" + + kccommon "ragflow/internal/ingestion/component/knowledge_compiler/common" +) + +// TestMergedChunkMapKeepsWikiFields locks the fix for the merged-row metadata +// gap: the dataset-level merged row written by mergedChunkMap must carry the +// wiki page fields (page_type_kwd/topic_kwd/title_kwd/slug_kwd/...) that the +// artifact API (ListArtifacts/ListWikiTopics) and page renderers read. Without +// them the compilation page surfaces no wiki pages from the merged products. +func TestMergedChunkMapKeepsWikiFields(t *testing.T) { + p := kccommon.Product{ + ID: "merged-1", DocID: "kb1", TenantID: "t1", Variant: kccommon.VariantWiki, + Content: "# Alpha\n\nBody", + Vector: []float32{0.1, 0.2, 0.3}, + Meta: map[string]any{ + "slug": "entity/alpha", + "title": "Alpha", + "page_type": "entity", + "topic": "Alpha", + "summary": "A page about Alpha", + "entity_names": []string{"Alpha"}, + "related_kb_pages": []string{"entity/beta"}, + "source_doc_ids": []string{"d1"}, + "source_chunk_ids": []string{"c1"}, + }, + } + m := mergedChunkMap("t1", "kb1", p) + + cases := map[string]string{ + "slug_kwd": "entity/alpha", + "artifact_slug_kwd": "entity/alpha", + "title_kwd": "Alpha", + "page_type_kwd": "entity", + "topic_kwd": "Alpha", + "summary_with_weight": "A page about Alpha", + } + for k, want := range cases { + if got, _ := m[k].(string); got != want { + t.Errorf("merged row[%q] = %q, want %q", k, got, want) + } + } + if v, _ := m["entity_names_kwd"].([]string); len(v) != 1 || v[0] != "Alpha" { + t.Errorf("entity_names_kwd = %#v, want [Alpha]", m["entity_names_kwd"]) + } + if m["doc_id"] != "kb1" || m["kc_merged"] != 1 || m["available_int"] != 1 { + t.Errorf("merged flags wrong: doc_id=%v kc_merged=%v available_int=%v", m["doc_id"], m["kc_merged"], m["available_int"]) + } + if m["q_3_vec"] == nil { + t.Errorf("vector column missing") + } +} + +// TestProductFromChunkMapRestoresWikiFields locks the reader side: the wiki page +// columns must be reconstructed into the product Meta so the merge step can carry +// them onto the merged row. +func TestProductFromChunkMapRestoresWikiFields(t *testing.T) { + c := map[string]interface{}{ + "id": "wiki/1", + "doc_id": "d1", + "compile_kwd": "wiki_page", + "content_with_weight": "# Alpha", + "kc_payload": "# Alpha\n\nBody", + "slug_kwd": "entity/alpha", + "page_type_kwd": "entity", + "topic_kwd": "Alpha", + "title_kwd": "Alpha", + "summary_with_weight": "A page about Alpha", + "entity_names_kwd": []interface{}{"Alpha"}, + "related_kb_pages_kwd": []interface{}{"entity/beta"}, + "section_level_int": float64(2), + } + p, ok := productFromChunkMap(c, "t1") + if !ok { + t.Fatalf("productFromChunkMap returned not-ok") + } + want := map[string]string{ + "slug": "entity/alpha", "page_type": "entity", "topic": "Alpha", + "title": "Alpha", "summary": "A page about Alpha", + } + for k, v := range want { + if got, _ := p.Meta[k].(string); got != v { + t.Errorf("meta[%q] = %q, want %q", k, got, v) + } + } + if v, _ := p.Meta["section_level"].(int64); v != 2 { + t.Errorf("section_level = %v, want 2", p.Meta["section_level"]) + } + if v, _ := p.Meta["entity_names"].([]string); len(v) != 1 || v[0] != "Alpha" { + t.Errorf("entity_names = %#v, want [Alpha]", p.Meta["entity_names"]) + } + if p.Merged { + t.Errorf("per-doc row must not be marked merged") + } +} diff --git a/internal/ingestion/pipeline/pipeline_knowledge_compiler_dsl_test.go b/internal/ingestion/pipeline/pipeline_knowledge_compiler_dsl_test.go index 66b76f356d..7a23c9e57f 100644 --- a/internal/ingestion/pipeline/pipeline_knowledge_compiler_dsl_test.go +++ b/internal/ingestion/pipeline/pipeline_knowledge_compiler_dsl_test.go @@ -95,10 +95,12 @@ func TestKnowledgeCompilerDSL_FixtureDecodesAndBindsParams(t *testing.T) { t.Errorf("fixture unexpectedly sets variant; frontend Compiler DSL omits it") } - // Authored as a single string group id in the frontend form. - gid, ok := params["compilation_template_group_id"].(string) - if !ok || gid != "c3aa748c8b2111f191f3047c16ec874f" { - t.Fatalf("compilation_template_group_id = %v, want single string id", params["compilation_template_group_id"]) + // Authored as a single string group id in the frontend form. The shipped + // template leaves it empty by design (the user selects the template group at + // runtime), so we only assert the DSL shape is a plain string (and that the + // node carries no variant). + if _, ok := params["compilation_template_group_id"].(string); !ok { + t.Fatalf("compilation_template_group_id = %v, want a single string id", params["compilation_template_group_id"]) } } @@ -137,8 +139,11 @@ func TestKnowledgeCompilerDSL_FrontendDSLDecodesAndConstructs(t *testing.T) { t.Fatal("default runtime factory not installed") } - // The fixture params (single group id, no variant) construct fine. - comp, err := f("KnowledgeCompiler", params) + // The fixture params carry no variant and a string group id. The shipped + // template leaves the group id empty (selected at runtime), so give it a + // concrete id here to verify the DSL constructs once configured. + params["compilation_template_group_id"] = "tpl-group" + comp, err := f("Compiler", params) if err != nil { t.Fatalf("construct from fixture params: %v", err) } @@ -154,16 +159,15 @@ func TestKnowledgeCompilerDSL_FrontendDSLDecodesAndConstructs(t *testing.T) { } // TestKnowledgeCompilerDSL_RegisteredAndConstructible confirms the Go runtime -// registers the component under the canonical name "KnowledgeCompiler" and that -// the runtime factory can build a component instance from a DSL params map that -// carries either compilation_template_id or compilation_template_group_id (the -// variant is no longer part of the DSL surface). The frontend label "Compiler" -// (see the fixture tests) maps to this runtime name at the API/canvas layer, -// not inside the pipeline DSL decoder. +// registers the knowledge-compiler component under the unified name "Compiler" +// (matching the Python side rag/flow/compiler/compiler.py) and that the runtime +// factory can build a component instance from a DSL params map that carries +// either compilation_template_id or compilation_template_group_id (the variant +// is no longer part of the DSL surface). func TestKnowledgeCompilerDSL_RegisteredAndConstructible(t *testing.T) { runtime.InstallDefaultRegistryFactory() - if _, _, _, ok := runtime.DefaultRegistry.Lookup("KnowledgeCompiler"); !ok { - t.Fatal("KnowledgeCompiler not registered in the runtime factory") + if _, _, _, ok := runtime.DefaultRegistry.Lookup("Compiler"); !ok { + t.Fatal("Compiler not registered in the runtime factory") } f := runtime.DefaultFactory() if f == nil { @@ -177,7 +181,7 @@ func TestKnowledgeCompilerDSL_RegisteredAndConstructible(t *testing.T) { {"compilation_template_id": "t1", "compilation_template_group_id": "g1"}, } for i, params := range cases { - comp, err := f("KnowledgeCompiler", params) + comp, err := f("Compiler", params) if err != nil { t.Fatalf("case %d construct: %v", i, err) } @@ -187,7 +191,7 @@ func TestKnowledgeCompilerDSL_RegisteredAndConstructible(t *testing.T) { } // Param map with neither id resolves to a parse error. - if _, err := f("KnowledgeCompiler", map[string]any{}); err == nil { + if _, err := f("Compiler", map[string]any{}); err == nil { t.Fatal("construct with no template spec: expected error") } } @@ -203,6 +207,8 @@ func TestKnowledgeCompilerDSL_ParamBinding(t *testing.T) { "compilation_template_group_id": "g1", "llm_id": "llm-1", "embedding_model": "emb-1", + "tenant_id": "tenant-1", + "dataset_id": "kb-1", "language": "Chinese", "extra": map[string]any{"prompt": "summarize"}, }) @@ -218,6 +224,8 @@ func TestKnowledgeCompilerDSL_ParamBinding(t *testing.T) { if p.Variant != "" { t.Errorf("Variant should be empty after ParseParam (derived from kind later), got %q", p.Variant) } + // TenantID/DatasetID are injected at runtime by the component (not via the + // DSL/ParseParam), so they are not asserted here. if p.LLMID != "llm-1" || p.EmbeddingModel != "emb-1" || p.Language != "Chinese" { t.Errorf("scalar fields = %+v", p) } diff --git a/internal/ingestion/pipeline/pipeline_knowledge_compiler_test.go b/internal/ingestion/pipeline/pipeline_knowledge_compiler_test.go index 2c7fd5d797..a2422f837d 100644 --- a/internal/ingestion/pipeline/pipeline_knowledge_compiler_test.go +++ b/internal/ingestion/pipeline/pipeline_knowledge_compiler_test.go @@ -57,7 +57,7 @@ func TestKnowledgeCompilerTemplate_RegisteredAndDecodable(t *testing.T) { if runtime.DefaultFactory() == nil { t.Fatal("default runtime factory not installed") } - for _, name := range []string{"File", "Parser", "TokenChunker", "KnowledgeCompiler"} { + for _, name := range []string{"File", "Parser", "TokenChunker", "Compiler"} { if _, _, _, ok := runtime.DefaultRegistry.Lookup(name); !ok { t.Errorf("component %q referenced by template is not registered in the runtime factory", name) } diff --git a/internal/ingestion/pipeline/template/ingestion_pipeline_knowledge_compiler.json b/internal/ingestion/pipeline/template/ingestion_pipeline_knowledge_compiler.json index 6b9bf3f5da..364bc66b09 100644 --- a/internal/ingestion/pipeline/template/ingestion_pipeline_knowledge_compiler.json +++ b/internal/ingestion/pipeline/template/ingestion_pipeline_knowledge_compiler.json @@ -6,9 +6,9 @@ "zh": "知识编译" }, "description": { - "en": "Compiles parsed chunks into structured knowledge units (graph/wiki/raptor/mindmap/datasetnav) via the KnowledgeCompiler component, emitting them as chunks merged into the upstream chunk stream. Ideal for building a retrievable knowledge layer on top of chunked documents.", - "de": "Kompiliert geparste Chunks über die KnowledgeCompiler-Komponente in strukturierte Wissenseinheiten und gibt sie als Chunks im upstream-Chunk-Strom zurück.", - "zh": "通过 KnowledgeCompiler 组件将解析后的分块编译为结构化知识单元（图谱/百科/RAPTOR/思维导图/数据集导航），以 chunks 形式合并进上游分块流输出，适合在分块文档之上构建可检索的知识层。" + "en": "Compiles parsed chunks into structured knowledge units (graph/wiki/raptor/mindmap/datasetnav) via the Compiler component, emitting them as chunks merged into the upstream chunk stream. Ideal for building a retrievable knowledge layer on top of chunked documents.", + "de": "Kompiliert geparste Chunks über die Compiler-Komponente in strukturierte Wissenseinheiten und gibt sie als Chunks im upstream-Chunk-Strom zurück.", + "zh": "通过 Compiler 组件将解析后的分块编译为结构化知识单元（图谱/百科/RAPTOR/思维导图/数据集导航），以 chunks 形式合并进上游分块流输出，适合在分块文档之上构建可检索的知识层。" }, "canvas_type": "Ingestion Pipeline", "canvas_category": "dataflow_canvas", @@ -150,7 +150,7 @@ }, "TokenChunker:SixApplesFall": { "downstream": [ - "KnowledgeCompiler:KnownSwiftLions" + "Compiler:KnownSwiftLions" ], "obj": { "component_name": "TokenChunker", @@ -182,11 +182,12 @@ "Parser:HipSignsRhyme" ] }, - "KnowledgeCompiler:KnownSwiftLions": { + "Compiler:KnownSwiftLions": { "downstream": [], "obj": { - "component_name": "KnowledgeCompiler", + "component_name": "Compiler", "params": { + "llm_id": "", "variant": "structure", "language": "English" } @@ -216,10 +217,10 @@ "targetHandle": "end" }, { - "id": "xy-edge__TokenChunker:SixApplesFallstart-KnowledgeCompiler:KnownSwiftLionsend", + "id": "xy-edge__TokenChunker:SixApplesFallstart-Compiler:KnownSwiftLionsend", "source": "TokenChunker:SixApplesFall", "sourceHandle": "start", - "target": "KnowledgeCompiler:KnownSwiftLions", + "target": "Compiler:KnownSwiftLions", "targetHandle": "end" } ], @@ -280,10 +281,10 @@ }, { "data": { - "label": "KnowledgeCompiler", + "label": "Compiler", "name": "Knowledge Compiler_0" }, - "id": "KnowledgeCompiler:KnownSwiftLions", + "id": "Compiler:KnownSwiftLions", "measured": { "height": 74, "width": 200 diff --git a/internal/ingestion/pipeline/template_integration_test.go b/internal/ingestion/pipeline/template_integration_test.go index bee1d14632..e5e43a2a21 100644 --- a/internal/ingestion/pipeline/template_integration_test.go +++ b/internal/ingestion/pipeline/template_integration_test.go @@ -821,8 +821,8 @@ func TestPipelineRun_AllIngestionTemplates_RealComponentsSmoke(t *testing.T) { if templateUsesComponent(t, templateBytes, "TagChunker") { t.Skip("template uses TagChunker which requires tag-structured content and parser setups not available for generic .md input; covered separately") } - if templateUsesComponent(t, templateBytes, "KnowledgeCompiler") { - t.Skip("template uses KnowledgeCompiler which requires LLM/embedder/ES wiring not available in the headless smoke run; covered by the knowledge_compiler component E2E tests") + if templateUsesComponent(t, templateBytes, "Compiler") { + t.Skip("template uses Compiler which requires LLM/embedder/ES wiring not available in the headless smoke run; covered by the knowledge_compiler component E2E tests") } terminalIDs := terminalComponentIDsFromTemplate(t, templateBytes) if len(terminalIDs) != 1 { diff --git a/internal/ingestion/task/embedder.go b/internal/ingestion/task/embedder.go index 6e0c1ec97a..8bd4ed0a97 100644 --- a/internal/ingestion/task/embedder.go +++ b/internal/ingestion/task/embedder.go @@ -38,6 +38,13 @@ func (e *embedder) MaxTokens() int { return e.model.MaxTokens } +func (e *embedder) BatchSize() int { + if e == nil || e.model == nil { + return models.DefaultEmbeddingBatchSize + } + return e.model.ResolveBatchSize() +} + func (e *embedder) Encode(ctx context.Context, texts []string) ([]componentpkg.EmbeddingResult, error) { if e.model.ModelDriver == nil { return nil, fmt.Errorf("embedder: embedding model driver is nil for model %v", e.model.ModelName) diff --git a/internal/ingestion/task/knowledge_compiler_wiring.go b/internal/ingestion/task/knowledge_compiler_wiring.go index 311e4e1970..a67d7266d8 100644 --- a/internal/ingestion/task/knowledge_compiler_wiring.go +++ b/internal/ingestion/task/knowledge_compiler_wiring.go @@ -28,7 +28,9 @@ import ( enginetypes "ragflow/internal/engine/types" "ragflow/internal/entity" "ragflow/internal/entity/models" + _ "ragflow/internal/ingestion/component/knowledge_compiler" kc "ragflow/internal/ingestion/component/knowledge_compiler/common" + "ragflow/internal/ingestion/knowledge_compile" "ragflow/internal/service" "gorm.io/gorm" @@ -97,12 +99,15 @@ func newKnowledgeCompilerDepsResolver() kc.DepsResolver { } // Resolve the chat model's context window so RAPTOR can truncate each // cluster's texts to fit the LLM context (mirrors Python self._llm_model.max_length). + // This uses content_length (PR #17839) — the total context window — not + // max_output. max_output is only the generation cap; using it as the + // budget source would collapse per-chunk input quotas. llmMax := kc.DefaultLLMContextLength // Bound the model-config lookup so a stalled provider/instance DB read // cannot block document ingestion indefinitely. ctx, cancel := context.WithTimeout(context.Background(), 30*time.Second) defer cancel() - if _, _, _, ml, merr := svc.ResolveModelConfig(ctx, tenantID, entity.ModelTypeChat, llmID); merr == nil && ml > 0 { + if ml, merr := svc.ResolveModelContextLength(ctx, tenantID, llmID); merr == nil && ml > 0 { llmMax = ml } @@ -113,7 +118,7 @@ func newKnowledgeCompilerDepsResolver() kc.DepsResolver { // HistoricalKNN / Redis are optional (wiki historical dedup, // datasetnav lock). They are wired separately when the // surrounding pipeline supplies the backing services. - LLMMaxLength: llmMax, + ModelContextLen: llmMax, }, nil } } @@ -174,33 +179,93 @@ func (e *kcEmbedder) Encode(ctx context.Context, texts []string) ([][]float32, e if len(texts) == 0 { return nil, nil } - embdID := strings.TrimSpace(e.embdID) - if embdID == "" { - return nil, fmt.Errorf("knowledge_compiler: embedding_model is required for production embedding") - } - mdl, err := e.svc.GetEmbeddingModel(ctx, e.tenantID, embdID) + mdl, err := e.resolveModel(ctx) if err != nil { - return nil, fmt.Errorf("knowledge_compiler: resolve embedding model: %w", err) - } - if mdl == nil || mdl.ModelDriver == nil { - return nil, fmt.Errorf("knowledge_compiler: embedding model %q is unavailable", embdID) + return nil, err } config := &models.EmbeddingConfig{} - // Embed expects *string for the model name; nil ModelUsage (not tracked here). - embeds, err := mdl.ModelDriver.Embed(ctx, mdl.ModelName, texts, mdl.APIConfig, config, nil) - if err != nil { - return nil, fmt.Errorf("knowledge_compiler: embed: %w", err) + // Slice inputs into per-provider batches: providers cap the per-request input + // count and reject larger batches rather than chunking internally. The batch + // size is resolved from the model's capability (all_models.json batch_size, + // added by #17877/#17878) via EmbeddingModel.ResolveBatchSize, which falls + // back to a conservative default. Batches are fanned out on the shared compiler + // pool and concatenated back in input order. + batchSize := mdl.ResolveBatchSize() + numBatches := (len(texts) + batchSize - 1) / batchSize + slots := make([][][]float32, numBatches) // per-batch vector lists, distinct indices => no race + jobs := make([]knowledge_compile.CompilerJob, 0, numBatches) + for b := 0; b < numBatches; b++ { + b := b + start := b * batchSize + end := start + batchSize + if end > len(texts) { + end = len(texts) + } + batchTexts := texts[start:end] + jobs = append(jobs, func() error { + if err := ctx.Err(); err != nil { + return err + } + embeds, err := mdl.ModelDriver.Embed(ctx, mdl.ModelName, batchTexts, mdl.APIConfig, config, nil) + if err != nil { + return fmt.Errorf("knowledge_compiler: embed: %w", err) + } + vecs := make([][]float32, len(embeds)) + for i, v := range embeds { + vecs[i] = float64sToFloat32(v.Embedding) + } + slots[b] = vecs + return nil + }) } - out := make([][]float32, len(embeds)) - for i, v := range embeds { - out[i] = float64sToFloat32(v.Embedding) + if err := knowledge_compile.SubmitCompilerJobs(ctx, jobs); err != nil { + return nil, err } - if len(out) > 0 { - e.dim.CompareAndSwap(0, int64(len(out[0]))) + // Flatten in input order and derive the vector dimension from the first + // batch's first vector. + out := make([][]float32, 0, len(texts)) + var batchDim int + for _, slot := range slots { + for _, vec := range slot { + out = append(out, vec) + if batchDim == 0 { + batchDim = len(vec) + } + } + } + if batchDim > 0 { + e.dim.CompareAndSwap(0, int64(batchDim)) } return out, nil } +// resolveModel returns the embedding model to embed with. It prefers the +// explicitly configured embedding_model; when the caller left it unset, it falls +// back to the tenant's default embedding model (mirrors Python, which uses the +// KB/tenant's configured embedding model for wiki compilation). A clear error is +// returned only when neither is available, so a KB with no embedding model fails +// loudly instead of silently producing empty vectors. +func (e *kcEmbedder) resolveModel(ctx context.Context) (*models.EmbeddingModel, error) { + if embdID := strings.TrimSpace(e.embdID); embdID != "" { + mdl, err := e.svc.GetEmbeddingModel(ctx, e.tenantID, embdID) + if err != nil { + return nil, fmt.Errorf("knowledge_compiler: resolve embedding model: %w", err) + } + if mdl == nil || mdl.ModelDriver == nil { + return nil, fmt.Errorf("knowledge_compiler: embedding model %q is unavailable", embdID) + } + return mdl, nil + } + driver, name, apiConfig, _, err := e.svc.GetTenantDefaultModelByType(ctx, e.tenantID, entity.ModelTypeEmbedding) + if err != nil { + return nil, fmt.Errorf("knowledge_compiler: embedding_model is required and no tenant default embedding model is set: %w", err) + } + if driver == nil || name == "" { + return nil, fmt.Errorf("knowledge_compiler: embedding_model is required (tenant default embedding model unavailable)") + } + return &models.EmbeddingModel{ModelDriver: driver, ModelName: &name, APIConfig: apiConfig}, nil +} + func (e *kcEmbedder) Dimensions() int { return int(e.dim.Load()) } // float64sToFloat32 converts an embedding vector to the product schema's @@ -230,9 +295,12 @@ func (s *kcWikiPageStore) FindSimilarPages(ctx context.Context, tenantID, datase KbIDs: []string{datasetID}, Limit: k, SelectFields: []string{"id", "slug_kwd", "title_kwd", "page_type_kwd", "topic_kwd", "summary_with_weight", "content_with_weight", "entity_names_kwd", "related_kb_pages_kwd", "outlinks_kwd", "kc_content_md_raw", "_score"}, + // compile_kwd="wiki_page" is the schema-backed discriminator for wiki + // pages (sections carry compile_kwd="wiki_section"); there is no + // "kc_kind" column in the chunk schema, so filtering on it would return + // empty on Infinity. Filter: map[string]interface{}{ - "compile_kwd": "artifact_page", - "kc_kind": "page", + "compile_kwd": "wiki_page", }, MatchExprs: []interface{}{&enginetypes.MatchDenseExpr{ VectorColumnName: fmt.Sprintf("q_%d_vec", len(vec)), @@ -264,9 +332,8 @@ func (s *kcWikiPageStore) GetPageBySlug(ctx context.Context, tenantID, datasetID Limit: 1, SelectFields: []string{"id", "slug_kwd", "title_kwd", "page_type_kwd", "topic_kwd", "summary_with_weight", "content_with_weight", "entity_names_kwd", "related_kb_pages_kwd", "outlinks_kwd", "kc_content_md_raw", "_score"}, Filter: map[string]interface{}{ - "compile_kwd": "artifact_page", + "compile_kwd": "wiki_page", "slug_kwd": slug, - "kc_kind": "page", }, } res, err := s.docEngine.Search(ctx, req) diff --git a/internal/ingestion/task/knowledge_compiler_wiring_test.go b/internal/ingestion/task/knowledge_compiler_wiring_test.go new file mode 100644 index 0000000000..ab4a81abf4 --- /dev/null +++ b/internal/ingestion/task/knowledge_compiler_wiring_test.go @@ -0,0 +1,43 @@ +// +// Copyright 2026 The InfiniFlow Authors. All Rights Reserved. +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. +// + +package task + +import ( + "testing" + + "ragflow/internal/agent/runtime" +) + +// TestKnowledgeCompilerRegisteredByWiring locks the composition-root contract: +// the task package imports the knowledge_compiler root package (via blank +// import in knowledge_compiler_wiring.go) so its init() registers the +// knowledge-compilation component under the unified name "Compiler" in the +// production runtime registry. Without this blank import the component is never +// registered and an ingestor consuming a canvas with a Compiler node fails with +// "unknown component". +func TestKnowledgeCompilerRegisteredByWiring(t *testing.T) { + factory, category, _, ok := runtime.DefaultRegistry.Lookup("Compiler") + if !ok { + t.Fatal("knowledge-compiler component \"Compiler\" is not registered; the task package blank-import must be present for its init() to run") + } + if category != runtime.CategoryIngestion { + t.Fatalf("component \"Compiler\" category = %q, want %q", category, runtime.CategoryIngestion) + } + if factory == nil { + t.Fatal("component \"Compiler\" registered with a nil factory") + } +} diff --git a/internal/ingestion/task/pipeline_executor.go b/internal/ingestion/task/pipeline_executor.go index c0cbfbc994..4a77724867 100644 --- a/internal/ingestion/task/pipeline_executor.go +++ b/internal/ingestion/task/pipeline_executor.go @@ -418,14 +418,6 @@ func (s *PipelineExecutor) runPipelineWithDSL(ctx context.Context, dsl string) ( // injected in place below without a nil-map assignment panic. parserConfig = map[string]interface{}{} } - common.InjectExtractorLLMID(parserConfig, s.taskCtx.Tenant.LLMID) - // When the dataset enables auto-metadata, ensure the Extractor node(s) - // carry the enable_metadata mode + field schema so the LLM extraction fires - // (mirrors Python task_executor.py:519 enabling gen_metadata_task). The - // dataset flag is authoritative: a node that already has enable_metadata - // turned on keeps its own config, but a shipped DSL defaulting it to 0 is - // still overridden so auto-metadata can activate. - common.InjectExtractorEnableMetadata(parserConfig) // Surface component params whose cpnID is absent from the DSL. The // runtime merge (override_params) silently drops such entries; diff --git a/internal/ingestion/task/pipeline_executor_defaults_test.go b/internal/ingestion/task/pipeline_executor_defaults_test.go index b7cb5ba9dd..b1eca5024f 100644 --- a/internal/ingestion/task/pipeline_executor_defaults_test.go +++ b/internal/ingestion/task/pipeline_executor_defaults_test.go @@ -51,7 +51,7 @@ var builtinComponentParamsGolden = map[string]string{ "qa": "{\"File\": {}, \"Parser:HipSignsRhyme\": {\"docx\": {\"flatten_media_to_text\": false, \"output_format\": \"json\", \"preprocess\": [\"main_content\"], \"suffix\": [\"docx\"], \"vlm\": {}}, \"markdown\": {\"flatten_media_to_text\": false, \"output_format\": \"json\", \"preprocess\": [\"main_content\"], \"suffix\": [\"md\", \"markdown\", \"mdx\"], \"vlm\": {}}, \"pdf\": {\"flatten_media_to_text\": false, \"output_format\": \"json\", \"parse_method\": \"DeepDOC\", \"preprocess\": [\"main_content\"], \"suffix\": [\"pdf\"], \"vlm\": {}}, \"spreadsheet\": {\"flatten_media_to_text\": false, \"output_format\": \"html\", \"parse_method\": \"DeepDOC\", \"preprocess\": [\"main_content\"], \"suffix\": [\"xls\", \"xlsx\", \"csv\"], \"vlm\": {}}, \"text&code\": {\"output_format\": \"json\", \"preprocess\": [\"main_content\"], \"suffix\": [\"txt\"]}}, \"Tokenizer:ColdCloudsDream\": {\"fields\": \"text\", \"filename_embd_weight\": 0.1, \"search_method\": [\"embedding\", \"full_text\"]}, \"QAChunker:TidyCloudsThink\": {}}", "resume": "{\"Extractor:ThreeDrinksAct\": {\"field_name\": \"metadata\", \"frequencyPenaltyEnabled\": true, \"frequency_penalty\": 0.7, \"llm_id\": \"THUDM/GLM-4.1V-9B-Thinking@SILICONFLOW\", \"maxTokensEnabled\": false, \"max_tokens\": 256, \"presencePenaltyEnabled\": true, \"presence_penalty\": 0.4, \"prompts\": [{\"content\": \"Content: {TitleChunker:FlatMiceFix@chunks}\", \"role\": \"user\"}], \"sys_prompt\": \"Act as a precise resume metadata extractor. Extract stable, chunk-supported metadata from the provided resume content.\\n\\nRules:\\n1. Use only information explicitly stated in the content. Do not infer, guess, normalize, or add missing facts.\\n2. The input may be only one chunk of a resume. Extract only what this content directly supports.\\n3. Use only these field names:\\ncandidate_name, gender, phone, email, city, location, nationality, linkedin, github, website, highest_degree, degree_levels, school_names, majors, graduation_years, work_experience_years, current_job_title, job_titles, company_names, job_experience, industries, target_job_titles, target_locations, employment_types, skills, certificates, awards, summary_tags\\n4. Ignore detailed responsibilities, project descriptions, achievement narratives, self-evaluation, and other low-value local details.\\n5. Keep values in the same language as the source text whenever possible.\\n6. Remove duplicates and keep only concise, high-value metadata.\\n7. Return only fields that are explicitly supported by the content. Do not return empty or unsupported fields.\\n\\nField guidance:\\n- highest_degree: highest explicit degree level mentioned\\n- degree_levels: all explicit degree levels mentioned\\n- school_names: explicit school, college, or university names\\n- majors: explicit fields of study\\n- graduation_years: explicit graduation years only\\n- work_experience_years: only if explicitly stated\\n- current_job_title: only if explicitly current or most recent\\n- job_titles: explicit role titles\\n- company_names: explicit employer names\\n- job_experience: concise structured work entries explicitly supported by the content, preferably including title, company, and time information when available\\n- industries: explicit industry names only\\n- target_job_titles: explicit desired roles only\\n- target_locations: explicit desired work locations only\\n- skills: concise, core, search-useful skills explicitly mentioned\\n- certificates: explicit certificate names only\\n- awards: explicit award names only\\n- summary_tags: short, high-value tags strictly supported by the content\\n\\nReturn only the extracted metadata. Do not output explanatory text.\", \"temperature\": 0.1, \"temperatureEnabled\": true, \"tenant_llm_id\": 29, \"topPEnabled\": true, \"top_p\": 0.3, \"auto_keywords\": 0, \"auto_questions\": 0, \"auto_tags\": 0, \"enable_metadata\": 0, \"metadata\": [], \"tag_file_id\": \"\"}, \"File\": {}, \"Parser:HipSignsRhyme\": {\"docx\": {\"flatten_media_to_text\": true, \"output_format\": \"json\", \"preprocess\": [\"main_content\"], \"suffix\": [\"docx\"], \"vlm\": {}}, \"pdf\": {\"flatten_media_to_text\": true, \"output_format\": \"json\", \"parse_method\": \"DeepDOC\", \"preprocess\": [\"main_content\"], \"suffix\": [\"pdf\"], \"vlm\": {}}, \"text&code\": {\"output_format\": \"json\", \"preprocess\": [\"main_content\"], \"suffix\": [\"txt\"]}}, \"TitleChunker:FlatMiceFix\": {\"hierarchy\": 1, \"include_heading_content\": false, \"levels\": [[\"^\\\\s*(?i:(?:\\\\d+[\\\\.\\\\)]\\\\s*)?(?:EDUCATION|ACADEMIC\\\\s*BACKGROUND|ACADEMIC\\\\s*HISTORY|EDUCATIONAL\\\\s*BACKGROUND|RELEVANT\\\\s*COURSEWORK|COURSEWORK|EXPERIENCE|WORK\\\\s*EXPERIENCE|PROFESSIONAL\\\\s*EXPERIENCE|RELEVANT\\\\s*EXPERIENCE|EMPLOYMENT\\\\s*HISTORY|CAREER\\\\s*HISTORY|INTERNSHIP\\\\s*EXPERIENCE|PROJECTS|PROJECT\\\\s*EXPERIENCE|ACADEMIC\\\\s*PROJECTS|PROFESSIONAL\\\\s*PROJECTS|SKILLS|TECHNICAL\\\\s*SKILLS|CORE\\\\s*COMPETENCIES|COMPETENCIES|QUALIFICATIONS|SUMMARY\\\\s*OF\\\\s*QUALIFICATIONS|CERTIFICATIONS|LICENSES|CERTIFICATES|AWARDS|HONORS|HONOURS|ACHIEVEMENTS|PUBLICATIONS|RESEARCH|RESEARCH\\\\s*EXPERIENCE|LEADERSHIP|LEADERSHIP\\\\s*EXPERIENCE|ACTIVITIES|EXTRACURRICULAR\\\\s*ACTIVITIES|ACTIVITIES\\\\s*(?:&|AND)\\\\s*SKILLS|INVOLVEMENT|CAMPUS\\\\s*INVOLVEMENT|VOLUNTEER\\\\s*EXPERIENCE|VOLUNTEERING|COMMUNITY\\\\s*SERVICE|LANGUAGES|INTERESTS|HOBBIES|PROFILE|PROFESSIONAL\\\\s*PROFILE|SUMMARY|PROFESSIONAL\\\\s*SUMMARY|CAREER\\\\s*SUMMARY|OBJECTIVE|CAREER\\\\s*OBJECTIVE|PERSONAL\\\\s*INFORMATION|CONTACT\\\\s*INFORMATION|ADDITIONAL\\\\s*INFORMATION|TRAINING))\\\\s*[:：]?\\\\s*$\"], [\"^\\\\s*(?:\\\\d+[\\\\.、\\\\)]\\\\s*)?(?:教育背景|教育经历|学历背景|学术背景|技术背景|工作经历|工作经验|实习经历|项目经历|项目经验|科研经历|研究经历|校园经历|实践经历|专业经历|职业经历|技能|专业技能|技能特长|核心技能|技术栈|个人技能|工作技能|职业技能|技能与评价|技能与自我评价|工作技能与自我评价|职业技能与自我评价|证书|资格证书|职业资格|资质证书|获奖情况|获奖经历|荣誉|荣誉奖项|奖项|科研成果|论文发表|发表论文|领导经历|学生工作|校园活动|社团经历|活动经历|志愿经历|志愿服务|社会实践|语言能力|语言|自我评价|个人评价|自我总结|个人总结|个人优势|个人简介|个人信息|基本信息|联系方式|求职意向|应聘意向|职业目标|求职目标|兴趣爱好|兴趣特长|培训经历|其他信息|附加信息)\\\\s*[:：]?\\\\s*$\"]], \"method\": \"hierarchy\"}, \"Tokenizer:KindHandsWin\": {\"fields\": \"text\", \"filename_embd_weight\": 0.1, \"search_method\": [\"embedding\", \"full_text\"]}}", "table": "{\"File\": {}, \"Parser:HipSignsRhyme\": {\"spreadsheet\": {\"flatten_media_to_text\": false, \"output_format\": \"html\", \"parse_method\": \"DeepDOC\", \"preprocess\": [\"main_content\"], \"suffix\": [\"xls\", \"xlsx\", \"csv\"], \"vlm\": {}, \"column_mode\": \"auto\", \"column_roles\": {}, \"column_names\": []}, \"text&code\": {\"output_format\": \"json\", \"preprocess\": [\"main_content\"], \"suffix\": [\"txt\"]}}, \"TableChunker:FastFoxesJump\": {}, \"Tokenizer:DeepLakesShine\": {\"fields\": \"text\", \"filename_embd_weight\": 0.1, \"search_method\": [\"embedding\", \"full_text\"]}}", - "knowledge_compiler": "{\"File\": {}, \"KnowledgeCompiler:KnownSwiftLions\": {\"language\": \"English\", \"variant\": \"structure\"}, \"Parser:HipSignsRhyme\": {\"setups\": {\"doc\": {\"output_format\": \"json\", \"preprocess\": [\"main_content\"], \"suffix\": [\"doc\"]}, \"docx\": {\"flatten_media_to_text\": false, \"output_format\": \"json\", \"preprocess\": [\"main_content\"], \"suffix\": [\"docx\"], \"vlm\": {}}, \"html\": {\"output_format\": \"json\", \"preprocess\": [\"main_content\"], \"suffix\": [\"htm\", \"html\"]}, \"markdown\": {\"flatten_media_to_text\": false, \"output_format\": \"json\", \"preprocess\": [\"main_content\"], \"suffix\": [\"md\", \"markdown\", \"mdx\"], \"vlm\": {}}, \"pdf\": {\"flatten_media_to_text\": false, \"output_format\": \"json\", \"parse_method\": \"DeepDOC\", \"preprocess\": [\"main_content\"], \"suffix\": [\"pdf\"], \"vlm\": {}}, \"spreadsheet\": {\"flatten_media_to_text\": false, \"output_format\": \"html\", \"parse_method\": \"DeepDOC\", \"preprocess\": [\"main_content\"], \"suffix\": [\"xls\", \"xlsx\", \"csv\"], \"vlm\": {}}, \"text&code\": {\"output_format\": \"json\", \"preprocess\": [\"main_content\"], \"suffix\": [\"txt\", \"py\", \"js\", \"java\", \"c\", \"cpp\", \"h\", \"php\", \"go\", \"ts\", \"sh\", \"cs\", \"kt\", \"sql\"]}}}, \"TokenChunker:SixApplesFall\": {\"children_delimiters\": [], \"chunk_token_size\": 512, \"delimiter_mode\": \"token_size\", \"delimiters\": [\"\\n\", \"!\", \"?\", \"。\", \"；\", \"！\", \"？\"], \"image_context_size\": 0, \"overlapped_percent\": 0, \"table_context_size\": 0}}", + "knowledge_compiler": "{\"File\": {}, \"Compiler:KnownSwiftLions\": {\"language\": \"English\", \"llm_id\": \"\", \"variant\": \"structure\"}, \"Parser:HipSignsRhyme\": {\"setups\": {\"doc\": {\"output_format\": \"json\", \"preprocess\": [\"main_content\"], \"suffix\": [\"doc\"]}, \"docx\": {\"flatten_media_to_text\": false, \"output_format\": \"json\", \"preprocess\": [\"main_content\"], \"suffix\": [\"docx\"], \"vlm\": {}}, \"html\": {\"output_format\": \"json\", \"preprocess\": [\"main_content\"], \"suffix\": [\"htm\", \"html\"]}, \"markdown\": {\"flatten_media_to_text\": false, \"output_format\": \"json\", \"preprocess\": [\"main_content\"], \"suffix\": [\"md\", \"markdown\", \"mdx\"], \"vlm\": {}}, \"pdf\": {\"flatten_media_to_text\": false, \"output_format\": \"json\", \"parse_method\": \"DeepDOC\", \"preprocess\": [\"main_content\"], \"suffix\": [\"pdf\"], \"vlm\": {}}, \"spreadsheet\": {\"flatten_media_to_text\": false, \"output_format\": \"html\", \"parse_method\": \"DeepDOC\", \"preprocess\": [\"main_content\"], \"suffix\": [\"xls\", \"xlsx\", \"csv\"], \"vlm\": {}}, \"text&code\": {\"output_format\": \"json\", \"preprocess\": [\"main_content\"], \"suffix\": [\"txt\", \"py\", \"js\", \"java\", \"c\", \"cpp\", \"h\", \"php\", \"go\", \"ts\", \"sh\", \"cs\", \"kt\", \"sql\"]}}}, \"TokenChunker:SixApplesFall\": {\"children_delimiters\": [], \"chunk_token_size\": 512, \"delimiter_mode\": \"token_size\", \"delimiters\": [\"\\n\", \"!\", \"?\", \"。\", \"；\", \"！\", \"？\"], \"image_context_size\": 0, \"overlapped_percent\": 0, \"table_context_size\": 0}}", } // Per-template test methods. Each resolves default component params from a diff --git a/internal/ingestion/task/pipeline_executor_test.go b/internal/ingestion/task/pipeline_executor_test.go index b366e90466..e5ccd059ff 100644 --- a/internal/ingestion/task/pipeline_executor_test.go +++ b/internal/ingestion/task/pipeline_executor_test.go @@ -32,7 +32,7 @@ func TestMarkCompiledProductsHidden(t *testing.T) { chunks := []map[string]any{ {"id": "src-1", "content_with_weight": "ordinary source chunk"}, {"id": "struct-1", "compile_kwd": "structure", "content_with_weight": "entity A"}, - {"id": "wiki-1", "compile_kwd": "artifact_page", "content_with_weight": "page X"}, + {"id": "wiki-1", "compile_kwd": "wiki_page", "content_with_weight": "page X"}, {"id": "src-2", "content_with_weight": "another source chunk"}, } markCompiledProductsHidden(chunks) diff --git a/internal/router/router.go b/internal/router/router.go index 1d679c4a79..fd2dc699c2 100644 --- a/internal/router/router.go +++ b/internal/router/router.go @@ -351,9 +351,9 @@ func (r *Router) Setup(engine *gin.Engine) { datasets.DELETE("/:dataset_id/tags", r.datasetsHandler.RemoveTags) datasets.POST("/:dataset_id/embedding/check", r.datasetsHandler.CheckEmbedding) datasets.POST("/:dataset_id/documents/batch-update-status", r.documentHandler.BatchUpdateDocumentStatus) - datasets.GET("/:dataset_id/index", r.datasetsHandler.TraceIndex) - datasets.POST("/:dataset_id/index", r.datasetsHandler.RunIndex) - datasets.DELETE("/:dataset_id/index", r.datasetsHandler.DeleteIndex) + // Scheduler compile-status contract (API_PROXY_SCHEME=go/hybrid); + // replaces the retired RunIndex/TraceIndex/DeleteIndex /index routes. + datasets.GET("/:dataset_id/compilation/status", r.datasetsHandler.GetCompilationStatus) // Knowledge-compilation wiki artifacts datasets.HEAD("/:dataset_id/artifacts", r.datasetArtifactHandler.AnyArtifact) @@ -380,8 +380,6 @@ func (r *Router) Setup(engine *gin.Engine) { datasets.GET("/:dataset_id/skills/:skill_kwd", r.datasetArtifactHandler.GetSkillPage) datasets.DELETE("/:dataset_id/skills/:skill_kwd", r.datasetArtifactHandler.DeleteSkill) - datasets.DELETE("/:dataset_id/:index_type", r.datasetsHandler.DeleteIndex) - //datasets.DELETE("/:dataset_id/graph", r.datasetsHandler.DeleteKnowledgeGraph) datasets.POST("", r.datasetsHandler.CreateDataset) datasets.DELETE("", r.datasetsHandler.DeleteDatasets) datasets.POST("/search", r.datasetsHandler.SearchDatasets) diff --git a/internal/server/config/base.go b/internal/server/config/base.go index 7a769b1636..dfbb624240 100644 --- a/internal/server/config/base.go +++ b/internal/server/config/base.go @@ -30,8 +30,8 @@ type Config struct { admin AdminConfig apiServer APIServerConfig - ingestor IngestorConfig syncer SyncerConfig + ingestor IngestorConfig log LogConfig smtp common.SMTPConfig diff --git a/internal/server/config/ingestor_config.go b/internal/server/config/ingestor_config.go index 7ea6edb0e5..684e289e06 100644 --- a/internal/server/config/ingestor_config.go +++ b/internal/server/config/ingestor_config.go @@ -12,19 +12,26 @@ // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. // See the License for the specific language governing permissions and // limitations under the License. -// package config import "github.com/spf13/viper" type IngestorConfig struct { + // MaxConcurrentWorkers bounds how many ingestion tasks the ingestor runs in + // parallel (the task channel width and dataset-level compile worker count + // default to this value). 0/negative falls back to runtime.NumCPU(). MaxConcurrentWorkers int `mapstructure:"max_concurrent_workers"` + // CompilerPoolSize bounds the process-wide knowledge-compilation worker + // pool that drives the cross-doc KNN / LLM-merge / write stages. 0/negative + // falls back to runtime.NumCPU() (or KC_COMPILE_CONCURRENCY if set). + CompilerPoolSize int `mapstructure:"compiler_pool_size"` } func (c *Config) ParseIngestorConfig(v *viper.Viper) error { // Default Ingestor config c.ingestor.MaxConcurrentWorkers = 1 + c.ingestor.CompilerPoolSize = 0 if !v.IsSet("ingestor") { return nil @@ -38,5 +45,13 @@ func (c *Config) ParseIngestorConfig(v *viper.Viper) error { c.ingestor.MaxConcurrentWorkers = sub.GetInt("max_concurrent_workers") } + if sub.IsSet("compiler_pool_size") { + c.ingestor.CompilerPoolSize = sub.GetInt("compiler_pool_size") + } + return nil } + +func (c *Config) GetIngestorConfig() *IngestorConfig { + return &c.ingestor +} diff --git a/internal/service/component_scoped_parser_config.go b/internal/service/component_scoped_parser_config.go new file mode 100644 index 0000000000..55ca470129 --- /dev/null +++ b/internal/service/component_scoped_parser_config.go @@ -0,0 +1,118 @@ +package service + +import ( + "strings" + + "ragflow/internal/entity" +) + +// ApplyComponentScopedParserConfig fills dataset-scoped component params onto a +// parser_config map without reintroducing top-level flat fields. It mutates the +// provided map in place and returns it for convenience. +func ApplyComponentScopedParserConfig( + parserConfig entity.JSONMap, + llmID string, +) entity.JSONMap { + if parserConfig == nil { + parserConfig = entity.JSONMap{} + } + + enableMetadata := parserConfigTruthy(parserConfig["enable_metadata"]) + metadataFields := mergeMetadataFields(parserConfig) + hasMetadataConfig := hasMetadataConfigShape(parserConfig) + + for cpnID, raw := range parserConfig { + params, ok := raw.(map[string]any) + if !ok { + continue + } + + cpnLower := strings.ToLower(cpnID) + switch { + case strings.HasPrefix(cpnLower, "extractor:") || strings.HasPrefix(cpnLower, "extractor_"): + if value, _ := params["llm_id"].(string); strings.TrimSpace(value) == "" && strings.TrimSpace(llmID) != "" { + params["llm_id"] = llmID + } + if enableMetadata && len(metadataFields) > 0 { + params["enable_metadata"] = 1 + params["metadata"] = metadataFields + } else if hasMetadataConfig { + params["enable_metadata"] = 0 + params["metadata"] = []any{} + } + case strings.HasPrefix(cpnLower, "compiler:") || strings.HasPrefix(cpnLower, "compiler_"): + if value, _ := params["llm_id"].(string); strings.TrimSpace(value) == "" && strings.TrimSpace(llmID) != "" { + params["llm_id"] = llmID + } + } + } + + return parserConfig +} + +func mergeMetadataFields(parserConfig entity.JSONMap) []any { + var out []any + for _, key := range []string{"metadata", "built_in_metadata"} { + for _, item := range anySlice(parserConfig[key]) { + field, ok := item.(map[string]any) + if !ok { + continue + } + name, _ := field["key"].(string) + if strings.TrimSpace(name) == "" { + continue + } + out = append(out, field) + } + } + return out +} + +func anySlice(value any) []any { + switch typed := value.(type) { + case []any: + return typed + case []map[string]any: + out := make([]any, 0, len(typed)) + for _, item := range typed { + out = append(out, item) + } + return out + default: + return nil + } +} + +func hasMetadataConfigShape(parserConfig entity.JSONMap) bool { + if parserConfig == nil { + return false + } + if _, ok := parserConfig["enable_metadata"]; ok { + return true + } + for _, key := range []string{"metadata", "built_in_metadata"} { + if anySlice(parserConfig[key]) != nil { + return true + } + } + return false +} + +func parserConfigTruthy(value any) bool { + switch typed := value.(type) { + case bool: + return typed + case string: + switch typed { + case "true", "True", "TRUE", "1": + return true + } + case float64: + return typed > 0 + case int: + return typed > 0 + case int64: + return typed > 0 + } + return false +} diff --git a/internal/service/component_scoped_parser_config_test.go b/internal/service/component_scoped_parser_config_test.go new file mode 100644 index 0000000000..bce355d2be --- /dev/null +++ b/internal/service/component_scoped_parser_config_test.go @@ -0,0 +1,147 @@ +package service + +import ( + "reflect" + "testing" + + "ragflow/internal/entity" +) + +func TestApplyComponentScopedParserConfig_SyncsExtractorAndCompiler(t *testing.T) { + parserConfig := entity.JSONMap{ + "enable_metadata": true, + "metadata": []any{ + map[string]any{"key": "author", "type": "string"}, + }, + "built_in_metadata": []any{ + map[string]any{"key": "document_name", "type": "string"}, + }, + "Extractor:AutoExtractDefault": map[string]any{}, + "Compiler:KnownSwiftLions": map[string]any{}, + } + + got := ApplyComponentScopedParserConfig( + parserConfig, + "llm-default", + ) + + extractor := got["Extractor:AutoExtractDefault"].(map[string]any) + if extractor["llm_id"] != "llm-default" { + t.Fatalf("extractor llm_id = %#v, want llm-default", extractor["llm_id"]) + } + if extractor["enable_metadata"] != 1 { + t.Fatalf("extractor enable_metadata = %#v, want 1", extractor["enable_metadata"]) + } + wantFields := []any{ + map[string]any{"key": "author", "type": "string"}, + map[string]any{"key": "document_name", "type": "string"}, + } + if !reflect.DeepEqual(extractor["metadata"], wantFields) { + t.Fatalf("extractor metadata = %#v, want %#v", extractor["metadata"], wantFields) + } + + compiler := got["Compiler:KnownSwiftLions"].(map[string]any) + if compiler["llm_id"] != "llm-default" { + t.Fatalf("compiler llm_id = %#v, want llm-default", compiler["llm_id"]) + } + if _, ok := compiler["embedding_model"]; ok { + t.Fatalf("compiler embedding_model = %#v, want absent", compiler["embedding_model"]) + } + if _, ok := compiler["tenant_id"]; ok { + t.Fatalf("compiler tenant_id = %#v, want absent", compiler["tenant_id"]) + } + if _, ok := compiler["dataset_id"]; ok { + t.Fatalf("compiler dataset_id = %#v, want absent", compiler["dataset_id"]) + } +} + +func TestApplyComponentScopedParserConfig_PreservesExplicitExtractorLLMID(t *testing.T) { + parserConfig := entity.JSONMap{ + "Extractor:Custom": map[string]any{ + "llm_id": "custom-llm", + }, + } + + got := ApplyComponentScopedParserConfig(parserConfig, "tenant-llm") + extractor := got["Extractor:Custom"].(map[string]any) + if extractor["llm_id"] != "custom-llm" { + t.Fatalf("extractor llm_id = %#v, want custom-llm", extractor["llm_id"]) + } +} + +func TestApplyComponentScopedParserConfig_AcceptsTypedMetadataSlices(t *testing.T) { + parserConfig := entity.JSONMap{ + "enable_metadata": true, + "metadata": []map[string]interface{}{ + {"key": "author", "type": "string"}, + }, + "built_in_metadata": []map[string]interface{}{ + {"key": "document_name", "type": "string"}, + }, + "Extractor:AutoExtractDefault": map[string]any{}, + } + + got := ApplyComponentScopedParserConfig(parserConfig, "tenant-llm") + extractor := got["Extractor:AutoExtractDefault"].(map[string]any) + + wantFields := []any{ + map[string]interface{}{"key": "author", "type": "string"}, + map[string]interface{}{"key": "document_name", "type": "string"}, + } + if !reflect.DeepEqual(extractor["metadata"], wantFields) { + t.Fatalf("extractor metadata = %#v, want %#v", extractor["metadata"], wantFields) + } +} + +func TestApplyComponentScopedParserConfig_ClearsExtractorMetadataWhenDisabled(t *testing.T) { + parserConfig := entity.JSONMap{ + "enable_metadata": false, + "metadata": []map[string]interface{}{}, + "built_in_metadata": []map[string]interface{}{ + {"key": "document_name", "type": "string"}, + }, + "Extractor:AutoExtractDefault": map[string]any{ + "enable_metadata": 1, + "metadata": []any{ + map[string]any{"key": "stale", "type": "string"}, + }, + }, + } + + got := ApplyComponentScopedParserConfig(parserConfig, "tenant-llm") + extractor := got["Extractor:AutoExtractDefault"].(map[string]any) + + if extractor["enable_metadata"] != 0 { + t.Fatalf("extractor enable_metadata = %#v, want 0", extractor["enable_metadata"]) + } + if !reflect.DeepEqual(extractor["metadata"], []any{}) { + t.Fatalf("extractor metadata = %#v, want empty list", extractor["metadata"]) + } +} + +func TestApplyComponentScopedParserConfig_DoesNotTreatDocumentMetadataValuesAsSchema(t *testing.T) { + parserConfig := entity.JSONMap{ + "metadata": map[string]interface{}{ + "author": "Alice", + }, + "Extractor:AutoExtractDefault": map[string]any{ + "enable_metadata": 1, + "metadata": []any{ + map[string]any{"key": "author", "type": "string"}, + }, + }, + } + + got := ApplyComponentScopedParserConfig(parserConfig, "tenant-llm") + extractor := got["Extractor:AutoExtractDefault"].(map[string]any) + + if extractor["enable_metadata"] != 1 { + t.Fatalf("extractor enable_metadata = %#v, want 1", extractor["enable_metadata"]) + } + want := []any{ + map[string]any{"key": "author", "type": "string"}, + } + if !reflect.DeepEqual(extractor["metadata"], want) { + t.Fatalf("extractor metadata = %#v, want %#v", extractor["metadata"], want) + } +} diff --git a/internal/service/dataset/compilation_status.go b/internal/service/dataset/compilation_status.go new file mode 100644 index 0000000000..0c996d565b --- /dev/null +++ b/internal/service/dataset/compilation_status.go @@ -0,0 +1,84 @@ +package dataset + +import ( + "context" + "encoding/json" + "errors" + "time" + + "gorm.io/gorm" + + "ragflow/internal/common" + "ragflow/internal/dao" + "ragflow/internal/entity" +) + +// CompilationStatus is the dataset-level knowledge-compile lifecycle state +// surfaced by GET /datasets/:id/compilation/status. It is the Go scheduler +// contract that replaces the Python-era RunIndex/TraceIndex task progress for +// API_PROXY_SCHEME=go / hybrid. +// +// State only takes one of idle/pending/running/completed. Error is NOT a fifth +// state: it is a diagnostic attached to a pending/running batch left for retry, +// so the frontend should test `error != ""` on its own (and hide the counts) +// rather than treating it as a peer of state. +type CompilationStatus struct { + State string `json:"state"` // idle | pending | running | completed + Error string `json:"error,omitempty"` // most recent batch diagnostic (empty when none) + Inflight int `json:"inflight"` // entries currently claimed (in-flight batch) + Backlog int `json:"backlog"` // entries still waiting to be claimed + LastCompletedAt *time.Time `json:"last_completed_at,omitempty"` // last backlog drain + UpdatedAt time.Time `json:"updated_at"` // last scheduling-row activity +} + +// GetDatasetCompilationStatus returns the scheduling-row lifecycle state for a +// dataset after verifying the calling user owns it. When no row exists the +// dataset has never had any compile work, so the state is idle. +func (d *DatasetService) GetDatasetCompilationStatus(ctx context.Context, userID, datasetID string) (CompilationStatus, common.ErrorCode, error) { + if datasetID == "" { + return CompilationStatus{}, common.CodeDataError, errors.New("dataset_id is required") + } + if !d.kbDAO.Accessible(ctx, dao.DB, datasetID, userID) { + return CompilationStatus{}, common.CodeDataError, errors.New("no authorization") + } + st := CompilationStatus{State: entity.DatasetStateIdle} + db := dao.GetDB() + if db == nil { + return st, common.CodeSuccess, nil + } + var row entity.KnowledgeCompileDataset + err := db.WithContext(ctx). + Where("dataset_id = ?", datasetID). + First(&row).Error + if errors.Is(err, gorm.ErrRecordNotFound) { + return st, common.CodeSuccess, nil // never compiled -> idle + } + if err != nil { + return st, common.CodeServerError, err + } + st.State = row.State + if st.State == "" { + st.State = entity.DatasetStateIdle + } + st.Error = row.ErrorMsg + st.Inflight = jsonArrayLen(row.InflightDocIDs) + st.Backlog = jsonArrayLen(row.BacklogDocIDs) + st.LastCompletedAt = row.LastCompletedAt + st.UpdatedAt = row.UpdatedAt + return st, common.CodeSuccess, nil +} + +// jsonArrayLen counts the top-level elements of a JSON array stored as TEXT. +// The *_doc_ids columns hold a `[]BacklogEntry` array ({doc_id,event_type,seq}), +// so each element is one scheduling entry (NOT a deduplicated doc). Empty or +// malformed strings count as 0. +func jsonArrayLen(s string) int { + if s == "" { + return 0 + } + var arr []json.RawMessage + if err := json.Unmarshal([]byte(s), &arr); err != nil { + return 0 + } + return len(arr) +} diff --git a/internal/service/dataset/compilation_status_test.go b/internal/service/dataset/compilation_status_test.go new file mode 100644 index 0000000000..9b910fbc4b --- /dev/null +++ b/internal/service/dataset/compilation_status_test.go @@ -0,0 +1,197 @@ +// +// Copyright 2026 The InfiniFlow Authors. All Rights Reserved. +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +package dataset + +import ( + "context" + "testing" + "time" + + "gorm.io/gorm" + + "ragflow/internal/common" + "ragflow/internal/dao" + "ragflow/internal/entity" +) + +// TestJSONArrayLen locks the inflight/backlog count derivation (plan v4.1 +// §9.3): counts are BacklogEntry array lengths, not deduplicated doc counts. +func TestJSONArrayLen(t *testing.T) { + cases := []struct { + name string + in string + want int + }{ + {name: "empty", in: "", want: 0}, + {name: "empty array", in: "[]", want: 0}, + {name: "single entry", in: `[{"doc_id":"d1","event_type":"completed","seq":1}]`, want: 1}, + {name: "two entries same doc", in: `[{"doc_id":"d1","event_type":"completed","seq":1},{"doc_id":"d1","event_type":"deleted","seq":2}]`, want: 2}, + {name: "malformed", in: "{not json", want: 0}, + {name: "null", in: "null", want: 0}, + } + for _, tc := range cases { + t.Run(tc.name, func(t *testing.T) { + if got := jsonArrayLen(tc.in); got != tc.want { + t.Fatalf("jsonArrayLen(%q) = %d, want %d", tc.in, got, tc.want) + } + }) + } +} + +// setupCompilationStatusTestDB migrates the minimal schema for +// GetDatasetCompilationStatus (Knowledgebase for the Accessible check plus the +// KnowledgeCompileDataset scheduling row) and pushes it onto dao.DB. +func setupCompilationStatusTestDB(t *testing.T) *gorm.DB { + t.Helper() + db := setupServiceTestDB(t) + if err := db.AutoMigrate(&entity.KnowledgeCompileDataset{}); err != nil { + t.Fatalf("migrate knowledge_compile_docs: %v", err) + } + pushServiceDB(t, db) + return db +} + +// insertCompilationOwnerKB inserts a valid KB owned by userID (TenantID == +// userID, so Accessible returns true). +func insertCompilationOwnerKB(t *testing.T, kbID, userID string) { + t.Helper() + status := string(entity.StatusValid) + kb := &entity.Knowledgebase{ + ID: kbID, + TenantID: userID, + Name: "compile-status-kb", + EmbdID: "BAAI/bge-large-zh-v1.5@Builtin", + CreatedBy: userID, + Permission: string(entity.TenantPermissionMe), + Status: &status, + } + if err := dao.DB.Create(kb).Error; err != nil { + t.Fatalf("insert kb: %v", err) + } +} + +func testCompilationStatusService() *DatasetService { + return &DatasetService{kbDAO: dao.NewKnowledgebaseDAO()} +} + +// TestGetDatasetCompilationStatus_NoRowIsIdle verifies a dataset with no +// scheduling row reports the idle state with zero counts. +func TestGetDatasetCompilationStatus_NoRowIsIdle(t *testing.T) { + setupCompilationStatusTestDB(t) + insertCompilationOwnerKB(t, "kb-no-row", "user-1") + + st, code, err := testCompilationStatusService().GetDatasetCompilationStatus( + t.Context(), "user-1", "kb-no-row") + if err != nil { + t.Fatalf("unexpected error: %v", err) + } + if code != common.CodeSuccess { + t.Fatalf("code=%d want %d", code, common.CodeSuccess) + } + if st.State != entity.DatasetStateIdle { + t.Fatalf("state=%q want idle", st.State) + } + if st.Inflight != 0 || st.Backlog != 0 { + t.Fatalf("expected zero counts for idle, got inflight=%d backlog=%d", st.Inflight, st.Backlog) + } + if st.Error != "" { + t.Fatalf("expected empty error, got %q", st.Error) + } +} + +// TestGetDatasetCompilationStatus_FullOutput locks the complete response +// mapping from the MySQL row: state, inflight/backlog counts, error diagnostic +// and last_completed_at. +func TestGetDatasetCompilationStatus_FullOutput(t *testing.T) { + db := setupCompilationStatusTestDB(t) + insertCompilationOwnerKB(t, "kb-full", "user-1") + + // A running row with 2 inflight + 1 backlog entries and a retained error + // diagnostic (error is NOT a fifth state: state stays running). + lastDone := time.Now().Add(-time.Hour).UTC() + row := entity.KnowledgeCompileDataset{ + DatasetID: "kb-full", + TenantID: "user-1", + BacklogDocIDs: `[{"doc_id":"d3","event_type":"completed","seq":3}]`, + InflightDocIDs: `[{"doc_id":"d1","event_type":"completed","seq":1},{"doc_id":"d2","event_type":"completed","seq":2}]`, + State: entity.DatasetStateRunning, + ErrorMsg: "merge failed: boom", + LastCompletedAt: &lastDone, + } + if err := db.Create(&row).Error; err != nil { + t.Fatalf("insert scheduling row: %v", err) + } + + st, code, err := testCompilationStatusService().GetDatasetCompilationStatus( + t.Context(), "user-1", "kb-full") + if err != nil { + t.Fatalf("unexpected error: %v", err) + } + if code != common.CodeSuccess { + t.Fatalf("code=%d want %d", code, common.CodeSuccess) + } + if st.State != entity.DatasetStateRunning { + t.Fatalf("state=%q want running", st.State) + } + if st.Inflight != 2 || st.Backlog != 1 { + t.Fatalf("want inflight=2 backlog=1, got inflight=%d backlog=%d", st.Inflight, st.Backlog) + } + if st.Error != "merge failed: boom" { + t.Fatalf("error=%q want %q", st.Error, "merge failed: boom") + } + if st.LastCompletedAt == nil || !st.LastCompletedAt.Equal(lastDone) { + t.Fatalf("last_completed_at=%v want %v", st.LastCompletedAt, lastDone) + } +} + +// TestGetDatasetCompilationStatus_Unauthorized verifies a user who does not +// own the dataset is rejected before reading the scheduling row. +func TestGetDatasetCompilationStatus_Unauthorized(t *testing.T) { + setupCompilationStatusTestDB(t) + // Owner is user-1; a different user-2 must be denied. + insertCompilationOwnerKB(t, "kb-other", "user-1") + if err := dao.DB.Create(&entity.KnowledgeCompileDataset{ + DatasetID: "kb-other", + TenantID: "user-1", + BacklogDocIDs: "[]", + InflightDocIDs: "[]", + State: entity.DatasetStatePending, + }).Error; err != nil { + t.Fatalf("insert scheduling row: %v", err) + } + + st, code, err := testCompilationStatusService().GetDatasetCompilationStatus( + t.Context(), "user-2", "kb-other") + if err == nil { + t.Fatalf("expected authorization error, got nil (status=%+v)", st) + } + if code != common.CodeDataError { + t.Fatalf("code=%d want %d", code, common.CodeDataError) + } +} + +// TestGetDatasetCompilationStatus_EmptyID validates the required-field guard. +func TestGetDatasetCompilationStatus_EmptyID(t *testing.T) { + setupCompilationStatusTestDB(t) + _, code, err := testCompilationStatusService().GetDatasetCompilationStatus( + context.Background(), "user-1", "") + if err == nil { + t.Fatal("expected error for empty dataset_id") + } + if code != common.CodeDataError { + t.Fatalf("code=%d want %d", code, common.CodeDataError) + } +} diff --git a/internal/service/dataset/create_test.go b/internal/service/dataset/create_test.go index 6132d0631d..4305065bb2 100644 --- a/internal/service/dataset/create_test.go +++ b/internal/service/dataset/create_test.go @@ -85,6 +85,54 @@ func TestCreateDataset_ComponentParamsPopulated(t *testing.T) { if !ok || len(parserConfig) == 0 { t.Fatal("expected non-empty parser_config for general pipeline") } + extractor, ok := parserConfig["Extractor:AutoExtractDefault"].(map[string]interface{}) + if !ok { + t.Fatalf("expected extractor component params, got %#v", parserConfig["Extractor:AutoExtractDefault"]) + } + if extractor["llm_id"] != "llm-default" { + t.Fatalf("extractor llm_id = %#v, want llm-default", extractor["llm_id"]) + } +} + +func TestCreateDataset_KnowledgeCompilerParamsPopulated(t *testing.T) { + db := setupServiceTestDB(t) + pushServiceDB(t, db) + insertCreateDatasetTenant(t, "tenant-1") + ctx := t.Context() + + parserID := "knowledge_compiler" + parseType := 1 + result, code, err := testDatasetCreateService(t).CreateDataset(ctx, &service.CreateDatasetRequest{ + Name: "ds-kc-cp", + ParserID: &parserID, + ParseType: &parseType, + }, "tenant-1") + if err != nil { + t.Fatalf("CreateDataset failed: %v", err) + } + if code != common.CodeSuccess { + t.Fatalf("expected success code, got %d", code) + } + parserConfig, ok := result["parser_config"].(entity.JSONMap) + if !ok || len(parserConfig) == 0 { + t.Fatal("expected non-empty parser_config for knowledge_compiler pipeline") + } + compiler, ok := parserConfig["Compiler:KnownSwiftLions"].(map[string]interface{}) + if !ok { + t.Fatalf("expected compiler component params, got %#v", parserConfig["Compiler:KnownSwiftLions"]) + } + if compiler["llm_id"] != "llm-default" { + t.Fatalf("compiler llm_id = %#v, want llm-default", compiler["llm_id"]) + } + if _, ok := compiler["embedding_model"]; ok { + t.Fatalf("compiler embedding_model = %#v, want absent", compiler["embedding_model"]) + } + if _, ok := compiler["tenant_id"]; ok { + t.Fatalf("compiler tenant_id = %#v, want absent", compiler["tenant_id"]) + } + if _, ok := compiler["dataset_id"]; ok { + t.Fatalf("compiler dataset_id = %#v, want absent", compiler["dataset_id"]) + } } func TestCreateDataset_ParseTypeBuiltinClearsPipelineID(t *testing.T) { diff --git a/internal/service/dataset/crud.go b/internal/service/dataset/crud.go index af10883cac..79d1081ac8 100644 --- a/internal/service/dataset/crud.go +++ b/internal/service/dataset/crud.go @@ -121,6 +121,11 @@ func (d *DatasetService) CreateDataset(ctx context.Context, req *service.CreateD // unique within the tenant. name = d.dedupeDatasetName(ctx, name, tenantID) + parserConfig = service.ApplyComponentScopedParserConfig( + parserConfig, + tenant.LLMID, + ) + kb := &entity.Knowledgebase{ ID: kbID, Name: name, @@ -521,6 +526,21 @@ func stringPtrIfNotEmpty(s string) *string { } // extractDocIDs returns the document IDs from a slice of documents. +// datasetIndexTaskIDs returns the deduplicated set of dataset-level index task +// ids recorded on the KB (graphrag/raptor/mindmap legacy task fields). It is +// used by deleteDataset to clear residual entity.Task rows when a KB is deleted. +// Kept here because it belongs to the dataset delete lifecycle, not the retired +// RunIndex scheduling path. +func datasetIndexTaskIDs(kb *entity.Knowledgebase) []string { + taskIDs := make([]string, 0, 3) + for _, taskID := range []*string{kb.GraphragTaskID, kb.RaptorTaskID, kb.MindmapTaskID} { + if taskID != nil && *taskID != "" { + taskIDs = append(taskIDs, *taskID) + } + } + return common.Deduplicate(taskIDs) +} + func extractDocIDs(docs []entity.Document) []string { ids := make([]string, 0, len(docs)) for _, doc := range docs { diff --git a/internal/service/dataset/embedding.go b/internal/service/dataset/embedding.go new file mode 100644 index 0000000000..6149976640 --- /dev/null +++ b/internal/service/dataset/embedding.go @@ -0,0 +1,285 @@ +package dataset + +import ( + "context" + "errors" + "fmt" + "math/rand" + "sort" + "strings" + + "ragflow/internal/common" + "ragflow/internal/dao" + "ragflow/internal/entity" + "ragflow/internal/entity/models" + "ragflow/internal/service" + + enginetypes "ragflow/internal/engine/types" +) + +// embeddingCheckSample is one sampled chunk with its stored vector, used by the +// embedding availability check. +type embeddingCheckSample struct { + ChunkID string + KbID string + DocID string + DocName string + VectorField string + Vector []float64 + PageNum interface{} + Position interface{} + Top interface{} + ContentWithWeight string + QuestionKeywords []string +} + +// CheckEmbedding verifies that a candidate embedding model is compatible with a +// dataset's existing vectors (the standard "switch embedding model" validation). +// It is independent of the retired RunIndex/graph_rag_queue scheduling path. +func (d *DatasetService) CheckEmbedding(ctx context.Context, userID, datasetID string, req *service.CheckEmbeddingRequest) (*service.EmbeddingCheckResponse, common.ErrorCode, error) { + if datasetID == "" { + return nil, common.CodeDataError, errors.New(`lack of "Dataset ID"`) + } + if !d.kbDAO.Accessible(ctx, dao.DB, datasetID, userID) { + return nil, common.CodeDataError, errors.New("no authorization") + } + + kb, err := d.kbDAO.GetByID(ctx, dao.DB, datasetID) + if err != nil { + if dao.IsNotFoundErr(err) { + return nil, common.CodeDataError, errors.New("invalid Dataset ID") + } + return nil, common.CodeServerError, errors.New("internal server error") + } + + if req == nil || strings.TrimSpace(req.EmbeddingID) == "" { + return nil, common.CodeDataError, errors.New("`embd_id` is required") + } + embeddingID := strings.TrimSpace(req.EmbeddingID) + if d.docEngine == nil { + return nil, common.CodeServerError, errors.New("doc engine not initialized") + } + + driver, modelName, apiConfig, maxTokens, err := service.NewModelProviderService().ResolveModelConfig(ctx, kb.TenantID, entity.ModelTypeEmbedding, embeddingID) + if err != nil { + return nil, common.CodeDataError, err + } + embeddingModel := models.NewEmbeddingModel(driver, &modelName, apiConfig, maxTokens) + + checkNum := defaultEmbeddingCheckNum + if req.CheckNum != nil { + checkNum = *req.CheckNum + } + if checkNum <= 0 { + checkNum = defaultEmbeddingCheckNum + } + + samples, err := d.sampleRandomChunksWithVectors(ctx, kb.TenantID, datasetID, checkNum) + if err != nil { + return nil, common.CodeServerError, err + } + if len(samples) == 0 { + return &service.EmbeddingCheckResponse{ + Summary: datasetEmbeddingCheckSummary(datasetID, embeddingID, 0, nil, ""), + Results: nil, + }, common.CodeSuccess, nil + } + + results := make([]service.EmbeddingCheckResult, 0, len(samples)) + effectiveSimilarities := make([]float64, 0, len(samples)) + sawTitleAndContent := false + for _, sample := range samples { + if sample.Vector == nil || len(sample.Vector) == 0 { + continue + } + + rawChunk, err := d.docEngine.GetChunk(ctx, fmt.Sprintf("ragflow_%s", kb.TenantID), sample.ChunkID, []string{datasetID}) + if err != nil { + continue + } + chunkMap := datasetMap(rawChunk) + if len(chunkMap) == 0 { + continue + } + + title := datasetString(chunkMap["title_tks"]) + content := datasetString(chunkMap["content_ltks"]) + + var titleVector [][]float64 + if title != "" { + titleVector, err = datasetEncodeEmbedding(ctx, embeddingModel, []string{title}) + if err != nil { + return nil, common.CodeServerError, err + } + } + var contentVector [][]float64 + if content != "" { + contentVector, err = datasetEncodeEmbedding(ctx, embeddingModel, []string{content}) + if err != nil { + return nil, common.CodeServerError, err + } + } + + var vectors [][]float64 + if len(titleVector) > 0 && len(contentVector) > 0 { + vectors = [][]float64{titleVector[0], contentVector[0]} + sawTitleAndContent = true + } else if len(titleVector) > 0 { + vectors = titleVector + } else if len(contentVector) > 0 { + vectors = contentVector + } else { + continue + } + + if len(vectors[0]) != len(sample.Vector) { + return nil, common.CodeDataError, fmt.Errorf("Embedding failure. The dimension (%d) of given embedding model is different from the original (%d)", len(vectors[0]), len(sample.Vector)) + } + + var sim float64 + if len(vectors) == 2 { + simContent := datasetCosSim(vectors[1], sample.Vector) + simMix := datasetCosSim(datasetMixVectors(vectors[0], vectors[1], 0.1), sample.Vector) + sim = simContent + if simMix > sim { + sim = simMix + sawTitleAndContent = true + } + } else { + sim = datasetCosSim(vectors[0], sample.Vector) + } + sim = datasetRoundFloat(sim, 6) + + effectiveSimilarities = append(effectiveSimilarities, sim) + results = append(results, service.EmbeddingCheckResult{ + ChunkID: sample.ChunkID, + DocID: sample.DocID, + DocName: sample.DocName, + VectorField: sample.VectorField, + VectorDim: len(sample.Vector), + CosSim: sim, + }) + } + + // Aggregate the batch mode explicitly: title_and_content when any sample was + // matched against the title+content mix, content_only otherwise. + matchMode := "content_only" + if sawTitleAndContent { + matchMode = "title_and_content" + } + summary := datasetEmbeddingCheckSummary(datasetID, embeddingID, len(samples), effectiveSimilarities, matchMode) + response := &service.EmbeddingCheckResponse{Summary: summary, Results: results} + if len(effectiveSimilarities) == 0 { + return nil, common.CodeDataError, errors.New("No embedded chunks are available to compare.") + } + if summary.AvgCosSim >= 0.9 { + return response, common.CodeSuccess, nil + } + return response, common.CodeNotEffective, errors.New("Embedding model switch failed: the average similarity between old and new vectors is below 0.9, indicating incompatible vector spaces.") +} + +func (d *DatasetService) sampleRandomChunksWithVectors(ctx context.Context, tenantID, datasetID string, n int) ([]embeddingCheckSample, error) { + indexName := fmt.Sprintf("ragflow_%s", tenantID) + totalResult, err := d.docEngine.Search(ctx, &enginetypes.SearchRequest{ + IndexNames: []string{indexName}, + KbIDs: []string{datasetID}, + Offset: 0, + Limit: 1, + Filter: map[string]interface{}{ + "kb_id": datasetID, + "available_int": 1, + }, + }) + if err != nil { + return nil, err + } + if totalResult == nil || totalResult.Total <= 0 { + return []embeddingCheckSample{}, nil + } + + total := int(totalResult.Total) + // Each sampled offset costs an engine Search + GetChunk plus up to two + // provider calls, so bound the client-controlled sample count server-side. + const maxEmbeddingSamples = 32 + if n < 0 { + return nil, fmt.Errorf("invalid sample size: %d", n) + } + if n > maxEmbeddingSamples { + n = maxEmbeddingSamples + } + if n > total { + n = total + } + limit := total + if limit > 1000 { + limit = 1000 + } + if n > limit { + n = limit + } + offsets := rand.Perm(limit) + offsets = offsets[:n] + sort.Ints(offsets) + + baseFields := []string{"docnm_kwd", "doc_id", "content_with_weight", "page_num_int", "position_int", "top_int"} + samples := make([]embeddingCheckSample, 0, n) + for _, offset := range offsets { + searchResult, err := d.docEngine.Search(ctx, &enginetypes.SearchRequest{ + IndexNames: []string{indexName}, + KbIDs: []string{datasetID}, + Offset: offset, + Limit: 1, + SelectFields: baseFields, + Filter: map[string]interface{}{ + "kb_id": datasetID, + "available_int": 1, + }, + }) + if err != nil { + return nil, err + } + if searchResult == nil || len(searchResult.Chunks) == 0 { + continue + } + chunkID := datasetChunkID(searchResult.Chunks[0]) + if chunkID == "" { + continue + } + fullChunk, err := d.docEngine.GetChunk(ctx, indexName, chunkID, []string{datasetID}) + if err != nil { + return nil, err + } + chunkMap := datasetMap(fullChunk) + if len(chunkMap) == 0 { + continue + } + vectorField := datasetGuessVecField(chunkMap) + vector := datasetAsFloatVec(chunkMap[vectorField]) + samples = append(samples, embeddingCheckSample{ + ChunkID: chunkID, + KbID: datasetID, + DocID: datasetString(chunkMap["doc_id"]), + DocName: datasetString(chunkMap["docnm_kwd"]), + VectorField: vectorField, + Vector: vector, + PageNum: chunkMap["page_num_int"], + Position: chunkMap["position_int"], + Top: chunkMap["top_int"], + ContentWithWeight: datasetString(chunkMap["content_with_weight"]), + QuestionKeywords: datasetStringSlice(chunkMap["question_keywords"]), + }) + } + + if len(samples) == 0 { + return nil, errors.New("no valid chunks with vectors found") + } + return samples, nil +} + +func (d *DatasetService) verifyEmbeddingAvailability(ctx context.Context, embdID string, tenantID string) (bool, string) { + _, _, _, _, err := service.NewModelProviderService().ResolveModelConfig(ctx, tenantID, entity.ModelTypeEmbedding, embdID) + if err != nil { + return false, err.Error() + } + return true, "" +} diff --git a/internal/service/dataset/helpers.go b/internal/service/dataset/helpers.go index d1559b594d..1dee37048b 100644 --- a/internal/service/dataset/helpers.go +++ b/internal/service/dataset/helpers.go @@ -269,6 +269,49 @@ func preserveDatasetParserConfigMetadata(next, existing entity.JSONMap, incoming return next } +func parserConfigJSONMap(value interface{}) entity.JSONMap { + switch typed := value.(type) { + case nil: + return nil + case entity.JSONMap: + return typed + case map[string]interface{}: + return entity.JSONMap(typed) + default: + return nil + } +} + +func cloneJSONMap(source entity.JSONMap) entity.JSONMap { + if source == nil { + return nil + } + cloned := make(entity.JSONMap, len(source)) + for key, value := range source { + cloned[key] = cloneJSONValue(value) + } + return cloned +} + +func cloneJSONValue(value interface{}) interface{} { + switch typed := value.(type) { + case map[string]interface{}: + nested := make(map[string]interface{}, len(typed)) + for key, item := range typed { + nested[key] = cloneJSONValue(item) + } + return nested + case []interface{}: + nested := make([]interface{}, len(typed)) + for idx, item := range typed { + nested[idx] = cloneJSONValue(item) + } + return nested + default: + return typed + } +} + func normalizeDatasetUpdateExt(ext map[string]interface{}) map[string]interface{} { if ext == nil { return nil diff --git a/internal/service/dataset/index.go b/internal/service/dataset/index.go deleted file mode 100644 index 8f0eda7ca7..0000000000 --- a/internal/service/dataset/index.go +++ /dev/null @@ -1,743 +0,0 @@ -package dataset - -import ( - "context" - "encoding/json" - "errors" - "fmt" - "math/rand" - "sort" - "strings" - "time" - - "ragflow/internal/common" - "ragflow/internal/dao" - redisengine "ragflow/internal/engine/redis" - enginetypes "ragflow/internal/engine/types" - "ragflow/internal/entity" - modelModule "ragflow/internal/entity/models" - "ragflow/internal/service" - "ragflow/internal/utility" - - "github.com/cespare/xxhash/v2" - "go.uber.org/zap" - "gorm.io/gorm" - "gorm.io/gorm/clause" -) - -func checkType(indexType string) bool { - haveType := false - for _, t := range validIndexTypes { - if indexType == t { - haveType = true - } - } - return haveType -} - -func (d *DatasetService) newRaptorOrGraphRagTask(ctx context.Context, sampleDoc *entity.Document, taskType string, taskDocID string, queueDocID string, docIDs []string) (*entity.Task, map[string]interface{}, error) { - if docIDs == nil || len(docIDs) == 0 { - docIDs = make([]string, 0) - } - if !checkIndexTaskType(taskType) { - return nil, nil, errors.New("type should be graphrag, raptor or mindmap") - } - - chunkingConfig, err := d.documentDAO.GetChunkingConfig(ctx, dao.DB, sampleDoc.ID) - if err != nil { - return nil, nil, err - } - - hasher := xxhash.New() - keys := make([]string, 0, len(chunkingConfig)) - for key := range chunkingConfig { - keys = append(keys, key) - } - sort.Strings(keys) - for _, key := range keys { - _, _ = hasher.Write([]byte(key)) - _, _ = hasher.Write([]byte{0}) - v, mErr := json.Marshal(chunkingConfig[key]) - if mErr != nil { - return nil, nil, mErr - } - _, _ = hasher.Write(v) - _, _ = hasher.Write([]byte{0}) - } - - taskID := utility.GenerateUUID() - beginAt := time.Now().Truncate(time.Second) - progressMsg := beginAt.Format("15:04:05") + " created task " + taskType - - for _, field := range []interface{}{taskDocID, maximumTaskPageNumber, maximumTaskPageNumber, taskType} { - _, _ = hasher.Write([]byte(fmt.Sprint(field))) - } - digest := fmt.Sprintf("%016x", hasher.Sum64()) - task := &entity.Task{ - ID: taskID, - DocID: taskDocID, - FromPage: maximumTaskPageNumber, - ToPage: maximumTaskPageNumber, - TaskType: taskType, - ProgressMsg: &progressMsg, - BeginAt: &beginAt, - Digest: &digest, - } - - queueMessage := map[string]interface{}{ - "id": taskID, - "doc_id": queueDocID, - "from_page": maximumTaskPageNumber, - "to_page": maximumTaskPageNumber, - "task_type": taskType, - "progress_msg": progressMsg, - "begin_at": beginAt.Format("2006-01-02 15:04:05"), - "digest": digest, - "doc_ids": docIDs, - } - - return task, queueMessage, nil -} - -func createDatasetIndexTaskInTx(tx *gorm.DB, task *entity.Task, queueDocID string) (*entity.Document, error) { - if task == nil { - return nil, errors.New("task is required") - } - if err := tx.Create(task).Error; err != nil { - return nil, err - } - - if queueDocID == "" { - return nil, nil - } - - var document entity.Document - err := tx.Select("id", "progress_msg", "process_begin_at").Where("id = ?", queueDocID).First(&document).Error - if err != nil { - if errors.Is(err, gorm.ErrRecordNotFound) { - return nil, nil - } - return nil, err - } - - beginAt := time.Now().Truncate(time.Second) - if task.BeginAt != nil { - beginAt = *task.BeginAt - } - if err = tx.Model(&entity.Document{}).Where("id = ?", queueDocID).Updates(map[string]interface{}{ - "progress_msg": "Task is queued...", - "process_begin_at": beginAt, - }).Error; err != nil { - return nil, err - } - - return &document, nil -} - -func enqueueDatasetIndexTask(ctx context.Context, priority int, queueMessage map[string]interface{}) error { - redisClient := redisengine.Get() - if redisClient == nil || !redisClient.QueueProduct(ctx, datasetIndexQueueName(priority), queueMessage) { - return errors.New("can't access Redis. Please check the Redis' status") - } - return nil -} - -func cleanupFailedDatasetIndexTask(taskID string, updatedDocument *entity.Document, kbID string, indexType string) error { - return dao.DB.Transaction(func(tx *gorm.DB) error { - if err := tx.Unscoped().Where("id = ?", taskID).Delete(&entity.Task{}).Error; err != nil { - return fmt.Errorf("delete task %s: %w", taskID, err) - } - - if column := datasetIndexTaskIDColumn(indexType); kbID != "" && column != "" { - if err := tx.Model(&entity.Knowledgebase{}).Where("id = ? AND "+column+" = ?", kbID, taskID).Update(column, nil).Error; err != nil { - return fmt.Errorf("clear dataset task id %s: %w", taskID, err) - } - } - - if updatedDocument == nil { - return nil - } - - return tx.Model(&entity.Document{}).Where("id = ?", updatedDocument.ID).Updates(map[string]interface{}{ - "progress_msg": updatedDocument.ProgressMsg, - "process_begin_at": updatedDocument.ProcessBeginAt, - }).Error - }) -} - -func datasetIndexTaskIDColumn(indexType string) string { - switch indexType { - case "graph": - return "graphrag_task_id" - case "raptor": - return "raptor_task_id" - case "mindmap": - return "mindmap_task_id" - default: - return "" - } -} - -func datasetIndexTaskFinishAtColumn(indexType string) string { - switch indexType { - case "graph": - return "graphrag_task_finish_at" - case "raptor": - return "raptor_task_finish_at" - case "mindmap": - return "mindmap_task_finish_at" - default: - return "" - } -} - -func checkIndexTaskType(taskType string) bool { - switch taskType { - case "graphrag", "raptor", "mindmap": - return true - default: - return false - } -} - -func datasetIndexTaskID(kb *entity.Knowledgebase, indexType string) string { - if kb == nil { - return "" - } - switch indexType { - case "graph": - if kb.GraphragTaskID != nil { - return *kb.GraphragTaskID - } - case "raptor": - if kb.RaptorTaskID != nil { - return *kb.RaptorTaskID - } - case "mindmap": - if kb.MindmapTaskID != nil { - return *kb.MindmapTaskID - } - } - return "" -} - -func datasetIndexTaskIDUpdate(indexType, taskID string) map[string]interface{} { - switch indexType { - case "graph": - return map[string]interface{}{"graphrag_task_id": taskID} - case "raptor": - return map[string]interface{}{"raptor_task_id": taskID} - case "mindmap": - return map[string]interface{}{"mindmap_task_id": taskID} - default: - return map[string]interface{}{} - } -} - -func datasetIndexTaskIDs(kb *entity.Knowledgebase) []string { - if kb == nil { - return nil - } - taskIDs := make([]string, 0, 3) - for _, taskID := range []*string{kb.GraphragTaskID, kb.RaptorTaskID, kb.MindmapTaskID} { - if taskID != nil && *taskID != "" { - taskIDs = append(taskIDs, *taskID) - } - } - return common.Deduplicate(taskIDs) -} - -func datasetIndexQueueName(priority int) string { - return fmt.Sprintf("%s.%d.common", serverQueueNamePrefix, priority) -} - -func clearGraphPhaseMarkers(ctx context.Context, redisClient *redisengine.Client, datasetID string) { - if redisClient == nil || datasetID == "" { - return - } - for _, phase := range []string{graphPhaseResolutionDone, graphPhaseCommunityDone} { - if !redisClient.Delete(ctx, fmt.Sprintf("graphrag:phase:%s:%s", datasetID, phase)) { - common.Warn("Failed to clear GraphRAG phase marker", zap.String("dataset_id", datasetID), zap.String("phase", phase)) - } - } -} - -func (d *DatasetService) RunIndex(ctx context.Context, userID, datasetID, indexType string) (map[string]interface{}, common.ErrorCode, error) { - if !checkType(indexType) { - return nil, common.CodeDataError, fmt.Errorf("invalid index type '%s'. Must be one of %v", indexType, validIndexTypes) - } - - if datasetID == "" { - return nil, common.CodeDataError, errors.New(`lack of "Dataset ID"`) - } - if !d.kbDAO.Accessible(ctx, dao.DB, datasetID, userID) { - return nil, common.CodeDataError, errors.New("no authorization") - } - - kb, err := d.kbDAO.GetByID(ctx, dao.DB, datasetID) - if err != nil { - if dao.IsNotFoundErr(err) { - return nil, common.CodeDataError, errors.New("invalid Dataset ID") - } - return nil, common.CodeDataError, errors.New("internal server error") - } - - taskType := indexTypeToTaskType[indexType] - displayName := indexTypeToDisplayName[indexType] - - documents, code, err := d.getDocumentsByDatasetForIndex(ctx, datasetID) - if err != nil { - return nil, code, err - } - _ = documents - - sampleDocument := documents[0] - documentIDs := make([]string, len(documents)) - - for i, doc := range documents { - documentIDs[i] = doc.ID - } - - task, queueMessage, err := d.newRaptorOrGraphRagTask(ctx, sampleDocument, taskType, sampleDocument.ID, graphRaptorQueueDocID, documentIDs) - if err != nil { - common.Warn("Failed to build dataset index task", zap.String("dataset_id", datasetID), zap.String("task_type", taskType), zap.Error(err)) - return nil, common.CodeDataError, errors.New("internal server error") - } - - var updatedDocument *entity.Document - var dataErr error - err = dao.DB.WithContext(ctx).Transaction(func(tx *gorm.DB) error { - var lockedKB entity.Knowledgebase - if err = tx.Clauses(clause.Locking{Strength: "UPDATE"}). - Where("id = ? AND status = ?", kb.ID, string(entity.StatusValid)). - First(&lockedKB).Error; err != nil { - return err - } - - existingTaskID := datasetIndexTaskID(&lockedKB, indexType) - if existingTaskID != "" { - var existingTask entity.Task - taskErr := tx.Where("id = ?", existingTaskID).First(&existingTask).Error - if taskErr != nil { - if errors.Is(taskErr, gorm.ErrRecordNotFound) { - } else { - return taskErr - } - } else if existingTask.Progress != 1 && existingTask.Progress != -1 { - dataErr = fmt.Errorf("task %s in progress with status %v. A %s Task is already running", existingTaskID, existingTask.Progress, displayName) - return dataErr - } - } - - updatedDocument, err = createDatasetIndexTaskInTx(tx, task, graphRaptorQueueDocID) - if err != nil { - return err - } - return tx.Model(&entity.Knowledgebase{}).Where("id = ?", lockedKB.ID).Updates(datasetIndexTaskIDUpdate(indexType, task.ID)).Error - }) - if err != nil { - if dataErr != nil { - return nil, common.CodeDataError, dataErr - } - common.Warn("Failed to create dataset index task", zap.String("dataset_id", datasetID), zap.String("task_type", taskType), zap.Error(err)) - return nil, common.CodeDataError, errors.New("internal server error") - } - - if err = enqueueDatasetIndexTask(ctx, 0, queueMessage); err != nil { - if cleanupErr := cleanupFailedDatasetIndexTask(task.ID, updatedDocument, kb.ID, indexType); cleanupErr != nil { - err = errors.Join(err, cleanupErr) - } - common.Warn("Failed to queue dataset index task", zap.String("dataset_id", datasetID), zap.String("task_type", taskType), zap.Error(err)) - return nil, common.CodeDataError, errors.New("internal server error") - } - - return map[string]interface{}{"task_id": task.ID}, common.CodeSuccess, nil -} - -func (d *DatasetService) getDocumentsByDatasetForIndex(ctx context.Context, datasetID string) ([]*entity.Document, common.ErrorCode, error) { - documents, _, err := d.documentDAO.GetByKBID(ctx, dao.DB, datasetID) - if err != nil { - common.Warn("Failed to load dataset documents for index", zap.String("dataset_id", datasetID), zap.Error(err)) - return nil, common.CodeDataError, errors.New("internal server error") - } - if len(documents) == 0 { - return nil, common.CodeDataError, fmt.Errorf("no documents in Dataset %s", datasetID) - } - return documents, common.CodeSuccess, nil -} - -func (d *DatasetService) TraceIndex(ctx context.Context, datasetID, userID, indexType string) (*entity.Task, common.ErrorCode, error) { - if !checkType(indexType) { - return nil, common.CodeDataError, fmt.Errorf("invalid index type '%s'. Must be one of %v", indexType, validIndexTypes) - } - - if datasetID == "" { - return nil, common.CodeDataError, errors.New(`lack of "Dataset ID"`) - } - if !d.kbDAO.Accessible(ctx, dao.DB, datasetID, userID) { - return nil, common.CodeDataError, errors.New("no authorization") - } - - kb, err := d.kbDAO.GetByID(ctx, dao.DB, datasetID) - if err != nil { - if dao.IsNotFoundErr(err) { - return nil, common.CodeDataError, errors.New("invalid Dataset ID") - } - return nil, common.CodeDataError, errors.New("internal server error") - } - - taskID := datasetIndexTaskID(kb, indexType) - - var task *entity.Task - if taskID != "" { - task, err = d.taskDAO.GetByID(ctx, dao.DB, taskID) - if err != nil { - if dao.IsNotFoundErr(err) { - return nil, common.CodeSuccess, nil - } - return nil, common.CodeServerError, errors.New("internal server error") - } - if task == nil { - return nil, common.CodeSuccess, nil - } - } - - return task, common.CodeSuccess, nil -} - -type embeddingCheckSample struct { - ChunkID string - KbID string - DocID string - DocName string - VectorField string - Vector []float64 - PageNum interface{} - Position interface{} - Top interface{} - ContentWithWeight string - QuestionKeywords []string -} - -func (d *DatasetService) CheckEmbedding(ctx context.Context, userID, datasetID string, req *service.CheckEmbeddingRequest) (*service.EmbeddingCheckResponse, common.ErrorCode, error) { - if datasetID == "" { - return nil, common.CodeDataError, errors.New(`lack of "Dataset ID"`) - } - if !d.kbDAO.Accessible(ctx, dao.DB, datasetID, userID) { - return nil, common.CodeDataError, errors.New("no authorization") - } - - kb, err := d.kbDAO.GetByID(ctx, dao.DB, datasetID) - if err != nil { - if dao.IsNotFoundErr(err) { - return nil, common.CodeDataError, errors.New("invalid Dataset ID") - } - return nil, common.CodeServerError, errors.New("internal server error") - } - - if req == nil || strings.TrimSpace(req.EmbeddingID) == "" { - return nil, common.CodeDataError, errors.New("`embd_id` is required") - } - embeddingID := strings.TrimSpace(req.EmbeddingID) - if ok, message := d.verifyEmbeddingAvailability(ctx, embeddingID, kb.TenantID); !ok { - return nil, common.CodeDataError, errors.New(message) - } - if d.docEngine == nil { - return nil, common.CodeServerError, errors.New("doc engine not initialized") - } - - driver, modelName, apiConfig, maxTokens, err := service.NewModelProviderService().ResolveModelConfig(ctx, kb.TenantID, entity.ModelTypeEmbedding, embeddingID) - if err != nil { - return nil, common.CodeDataError, err - } - embeddingModel := modelModule.NewEmbeddingModel(driver, &modelName, apiConfig, maxTokens) - - checkNum := defaultEmbeddingCheckNum - if req.CheckNum != nil { - checkNum = *req.CheckNum - } - if checkNum <= 0 { - checkNum = defaultEmbeddingCheckNum - } - - samples, err := d.sampleRandomChunksWithVectors(ctx, kb.TenantID, datasetID, checkNum) - if err != nil { - return nil, common.CodeServerError, err - } - if len(samples) == 0 { - return &service.EmbeddingCheckResponse{ - Summary: datasetEmbeddingCheckSummary(datasetID, embeddingID, 0, nil, ""), - Results: nil, - }, common.CodeSuccess, nil - } - - results := make([]service.EmbeddingCheckResult, 0, len(samples)) - effectiveSimilarities := make([]float64, 0, len(samples)) - matchMode := "content_only" - for _, sample := range samples { - if sample.Vector == nil || len(sample.Vector) == 0 { - continue - } - - rawChunk, err := d.docEngine.GetChunk(ctx, fmt.Sprintf("ragflow_%s", kb.TenantID), sample.ChunkID, []string{datasetID}) - if err != nil { - continue - } - chunkMap := datasetMap(rawChunk) - if len(chunkMap) == 0 { - continue - } - - title := datasetString(chunkMap["title_tks"]) - content := datasetString(chunkMap["content_ltks"]) - - var titleVector [][]float64 - if title != "" { - titleVector, err = datasetEncodeEmbedding(ctx, embeddingModel, []string{title}) - if err != nil { - return nil, common.CodeServerError, err - } - } - var contentVector [][]float64 - if content != "" { - contentVector, err = datasetEncodeEmbedding(ctx, embeddingModel, []string{content}) - if err != nil { - return nil, common.CodeServerError, err - } - } - - var vectors [][]float64 - if len(titleVector) > 0 && len(contentVector) > 0 { - vectors = [][]float64{titleVector[0], contentVector[0]} - matchMode = "title_and_content" - } else if len(titleVector) > 0 { - vectors = titleVector - } else if len(contentVector) > 0 { - vectors = contentVector - } else { - continue - } - - if len(vectors[0]) != len(sample.Vector) { - return nil, common.CodeDataError, fmt.Errorf("Embedding failure. The dimension (%d) of given embedding model is different from the original (%d)", len(vectors[0]), len(sample.Vector)) - } - - var sim float64 - if len(vectors) == 2 { - simContent := datasetCosSim(vectors[1], sample.Vector) - simMix := datasetCosSim(datasetMixVectors(vectors[0], vectors[1], 0.1), sample.Vector) - sim = simContent - if simMix > sim { - sim = simMix - matchMode = "title+content" - } - } else { - sim = datasetCosSim(vectors[0], sample.Vector) - } - sim = datasetRoundFloat(sim, 6) - - effectiveSimilarities = append(effectiveSimilarities, sim) - results = append(results, service.EmbeddingCheckResult{ - ChunkID: sample.ChunkID, - DocID: sample.DocID, - DocName: sample.DocName, - VectorField: sample.VectorField, - VectorDim: len(sample.Vector), - CosSim: sim, - }) - } - - summary := datasetEmbeddingCheckSummary(datasetID, embeddingID, len(samples), effectiveSimilarities, matchMode) - response := &service.EmbeddingCheckResponse{Summary: summary, Results: results} - if len(effectiveSimilarities) == 0 { - return nil, common.CodeDataError, errors.New("No embedded chunks are available to compare.") - } - if summary.AvgCosSim >= 0.9 { - return response, common.CodeSuccess, nil - } - return response, common.CodeNotEffective, errors.New("Embedding model switch failed: the average similarity between old and new vectors is below 0.9, indicating incompatible vector spaces.") -} - -func (d *DatasetService) sampleRandomChunksWithVectors(ctx context.Context, tenantID, datasetID string, n int) ([]embeddingCheckSample, error) { - indexName := fmt.Sprintf("ragflow_%s", tenantID) - totalResult, err := d.docEngine.Search(ctx, &enginetypes.SearchRequest{ - IndexNames: []string{indexName}, - KbIDs: []string{datasetID}, - Offset: 0, - Limit: 1, - Filter: map[string]interface{}{ - "kb_id": datasetID, - "available_int": 1, - }, - }) - if err != nil { - return nil, err - } - if totalResult == nil || totalResult.Total <= 0 { - return []embeddingCheckSample{}, nil - } - - total := int(totalResult.Total) - const maxEmbeddingSamples = 1024 - if n < 0 { - return nil, fmt.Errorf("invalid sample size: %d", n) - } - if n > maxEmbeddingSamples { - n = maxEmbeddingSamples - } - if n > total { - n = total - } - limit := total - if limit > 1000 { - limit = 1000 - } - if n > limit { - n = limit - } - offsets := rand.Perm(limit) - offsets = offsets[:n] - sort.Ints(offsets) - - baseFields := []string{"docnm_kwd", "doc_id", "content_with_weight", "page_num_int", "position_int", "top_int"} - samples := make([]embeddingCheckSample, 0, n) - for _, offset := range offsets { - searchResult, err := d.docEngine.Search(ctx, &enginetypes.SearchRequest{ - IndexNames: []string{indexName}, - KbIDs: []string{datasetID}, - Offset: offset, - Limit: 1, - SelectFields: baseFields, - Filter: map[string]interface{}{ - "kb_id": datasetID, - "available_int": 1, - }, - }) - if err != nil { - return nil, err - } - if searchResult == nil || len(searchResult.Chunks) == 0 { - continue - } - chunkID := datasetChunkID(searchResult.Chunks[0]) - if chunkID == "" { - continue - } - fullChunk, err := d.docEngine.GetChunk(ctx, indexName, chunkID, []string{datasetID}) - if err != nil { - return nil, err - } - chunkMap := datasetMap(fullChunk) - if len(chunkMap) == 0 { - continue - } - vectorField := datasetGuessVecField(chunkMap) - vector := datasetAsFloatVec(chunkMap[vectorField]) - samples = append(samples, embeddingCheckSample{ - ChunkID: chunkID, - KbID: datasetID, - DocID: datasetString(chunkMap["doc_id"]), - DocName: datasetString(chunkMap["docnm_kwd"]), - VectorField: vectorField, - Vector: vector, - PageNum: chunkMap["page_num_int"], - Position: chunkMap["position_int"], - Top: chunkMap["top_int"], - ContentWithWeight: datasetString(chunkMap["content_with_weight"]), - QuestionKeywords: datasetStringSlice(chunkMap["question_keywords"]), - }) - } - - if len(samples) == 0 { - return nil, errors.New("no valid chunks with vectors found") - } - return samples, nil -} - -func (d *DatasetService) verifyEmbeddingAvailability(ctx context.Context, embdID string, tenantID string) (bool, string) { - _, _, _, _, err := service.NewModelProviderService().ResolveModelConfig(ctx, tenantID, entity.ModelTypeEmbedding, embdID) - if err != nil { - return false, err.Error() - } - return true, "" -} - -func (d *DatasetService) DeleteIndex(ctx context.Context, userID, datasetID, indexType string, wipe bool) (common.ErrorCode, error) { - if !checkType(indexType) { - return common.CodeArgumentError, fmt.Errorf("invalid index type '%s'", indexType) - } - - if datasetID == "" { - return common.CodeDataError, errors.New(`lack of "Dataset ID"`) - } - - if !d.kbDAO.Accessible(ctx, dao.DB, datasetID, userID) { - return common.CodeDataError, errors.New("no authorization") - } - - kb, err := d.kbDAO.GetByID(ctx, dao.DB, datasetID) - if err != nil { - if dao.IsNotFoundErr(err) { - return common.CodeDataError, errors.New("invalid Dataset ID") - } - return common.CodeDataError, errors.New("internal server error") - } - - taskFinishAtField := datasetIndexTaskFinishAtColumn(indexType) - taskID := datasetIndexTaskID(kb, indexType) - - common.Info("delete_index", zap.String("dataset_id", datasetID), zap.String("index_type", indexType), zap.Bool("wipe", wipe)) - - if taskID != "" { - redisClient := redisengine.Get() - if redisClient == nil || !redisClient.Set(ctx, fmt.Sprintf("%s-cancel", taskID), "x", time.Hour) { - common.Warn("Failed to set dataset index cancellation marker", zap.String("dataset_id", datasetID), zap.String("task_id", taskID)) - } - if err := dao.DB.Unscoped().Where("id = ?", taskID).Delete(&entity.Task{}).Error; err != nil { - common.Warn("Failed to delete dataset index task", zap.String("dataset_id", datasetID), zap.String("task_id", taskID), zap.Error(err)) - return common.CodeDataError, errors.New("internal server error") - } - } - - if wipe && indexType == "graph" { - if d.docEngine == nil { - return common.CodeServerError, errors.New("document engine is not initialized") - } - indexName := fmt.Sprintf("ragflow_%s", kb.TenantID) - _, err = d.docEngine.DeleteChunks(ctx, map[string]interface{}{ - "knowledge_graph_kwd": []interface{}{"graph", "subgraph", "entity", "relation", "community_report"}, - "kb_id": datasetID, - }, indexName, datasetID) - if err != nil { - common.Warn("Failed to delete GraphRAG artefacts", zap.String("dataset_id", datasetID), zap.Error(err)) - return common.CodeDataError, errors.New("internal server error") - } - clearGraphPhaseMarkers(ctx, redisengine.Get(), datasetID) - common.Info("delete_index: cleared GraphRAG artefacts and phase markers", zap.String("dataset_id", datasetID)) - } else if wipe && indexType == "raptor" { - if d.docEngine == nil { - return common.CodeServerError, errors.New("document engine is not initialized") - } - indexName := fmt.Sprintf("ragflow_%s", kb.TenantID) - _, err = d.docEngine.DeleteChunks(ctx, map[string]interface{}{ - "raptor_kwd": []interface{}{"raptor"}, - "kb_id": datasetID, - }, indexName, datasetID) - if err != nil { - common.Warn("Failed to delete RAPTOR artefacts", zap.String("dataset_id", datasetID), zap.Error(err)) - return common.CodeDataError, errors.New("internal server error") - } - } - - updates := datasetIndexTaskIDUpdate(indexType, "") - if taskFinishAtField != "" { - updates[taskFinishAtField] = nil - } - if len(updates) > 0 { - if err = d.kbDAO.UpdateByID(ctx, dao.DB, kb.ID, updates); err != nil { - common.Warn("Failed to clear KB index task refs", zap.String("dataset_id", datasetID), zap.Error(err)) - } - } - - return common.CodeSuccess, nil -} diff --git a/internal/service/dataset/index_delete_test.go b/internal/service/dataset/index_delete_test.go deleted file mode 100644 index a4bb1cbdae..0000000000 --- a/internal/service/dataset/index_delete_test.go +++ /dev/null @@ -1,272 +0,0 @@ -// -// Copyright 2026 The InfiniFlow Authors. All Rights Reserved. -// -// Licensed under the Apache License, Version 2.0 (the "License"); -// you may not use this file except in compliance with the License. -// You may obtain a copy of the License at -// -// http://www.apache.org/licenses/LICENSE-2.0 -// -// Unless required by applicable law or agreed to in writing, software -// distributed under the License is distributed on an "AS IS" BASIS, -// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -// See the License for the specific language governing permissions and -// limitations under the License. -// - -package dataset - -import ( - "context" - "errors" - "testing" - "time" - - "gorm.io/gorm" - - "ragflow/internal/common" - "ragflow/internal/dao" - "ragflow/internal/entity" -) - -type deleteIndexDocEngine struct { - fakeChatDocEngine - deleteCalls []deleteIndexDocEngineCall -} - -type deleteIndexDocEngineCall struct { - condition map[string]interface{} - indexName string - datasetID string -} - -func (e *deleteIndexDocEngine) DeleteChunks(_ context.Context, condition map[string]interface{}, indexName string, datasetID string) (int64, error) { - e.deleteCalls = append(e.deleteCalls, deleteIndexDocEngineCall{ - condition: condition, - indexName: indexName, - datasetID: datasetID, - }) - return 1, nil -} - -func testDatasetServiceForDeleteIndex(docEngine *deleteIndexDocEngine) *DatasetService { - return &DatasetService{ - kbDAO: dao.NewKnowledgebaseDAO(), - taskDAO: dao.NewTaskDAO(), - docEngine: docEngine, - } -} - -func insertDeleteIndexKB(t *testing.T, indexType string, taskID string) { - t.Helper() - - finishAt := time.Date(2026, 6, 23, 10, 0, 0, 0, time.UTC) - kb := &entity.Knowledgebase{ - ID: "kb-1", - TenantID: "user-1", - Name: "test-kb", - EmbdID: "embedding@OpenAI", - CreatedBy: "user-1", - Permission: string(entity.TenantPermissionMe), - ParserID: "naive", - ParserConfig: entity.JSONMap{}, - Status: sptr("1"), - } - - switch indexType { - case "graph": - kb.GraphragTaskID = &taskID - kb.GraphragTaskFinishAt = &finishAt - case "raptor": - kb.RaptorTaskID = &taskID - kb.RaptorTaskFinishAt = &finishAt - case "mindmap": - kb.MindmapTaskID = &taskID - kb.MindmapTaskFinishAt = &finishAt - } - - if err := dao.DB.Create(kb).Error; err != nil { - t.Fatalf("insert kb: %v", err) - } - if taskID != "" { - if err := dao.DB.Create(&entity.Task{ID: taskID, DocID: "doc-1", TaskType: indexTypeToTaskType[indexType]}).Error; err != nil { - t.Fatalf("insert task: %v", err) - } - } -} - -func TestDatasetServiceDeleteIndexGraphWipeFalseOnlyCancelsTask(t *testing.T) { - db := setupServiceTestDB(t) - pushServiceDB(t, db) - insertDeleteIndexKB(t, "graph", "graph-task") - ctx := t.Context() - - docEngine := &deleteIndexDocEngine{} - code, err := testDatasetServiceForDeleteIndex(docEngine).DeleteIndex(ctx, "user-1", "kb-1", "graph", false) - if err != nil { - t.Fatalf("DeleteIndex failed: %v", err) - } - if code != common.CodeSuccess { - t.Fatalf("expected success code, got %d", code) - } - if len(docEngine.deleteCalls) != 0 { - t.Fatalf("wipe=false should not delete doc-store artefacts, got %#v", docEngine.deleteCalls) - } - - assertDeleteIndexTaskDeleted(t, "graph-task") - kb := getDeleteIndexKB(t) - if kb.GraphragTaskID == nil || *kb.GraphragTaskID != "" { - t.Fatalf("expected graphrag_task_id to be cleared to empty string, got %#v", kb.GraphragTaskID) - } - if kb.GraphragTaskFinishAt != nil { - t.Fatalf("expected graphrag_task_finish_at to be cleared, got %#v", kb.GraphragTaskFinishAt) - } -} - -func TestDatasetServiceDeleteIndexGraphWipeTrueDeletesArtefacts(t *testing.T) { - db := setupServiceTestDB(t) - pushServiceDB(t, db) - insertDeleteIndexKB(t, "graph", "graph-task") - ctx := t.Context() - - docEngine := &deleteIndexDocEngine{} - code, err := testDatasetServiceForDeleteIndex(docEngine).DeleteIndex(ctx, "user-1", "kb-1", "graph", true) - if err != nil { - t.Fatalf("DeleteIndex failed: %v", err) - } - if code != common.CodeSuccess { - t.Fatalf("expected success code, got %d", code) - } - if len(docEngine.deleteCalls) != 1 { - t.Fatalf("expected one doc-store delete call, got %#v", docEngine.deleteCalls) - } - - call := docEngine.deleteCalls[0] - if call.indexName != "ragflow_user-1" || call.datasetID != "kb-1" { - t.Fatalf("unexpected delete target: %#v", call) - } - if call.condition["kb_id"] != "kb-1" { - t.Fatalf("delete condition must include kb_id, got %#v", call.condition) - } - assertStringSet(t, call.condition["knowledge_graph_kwd"], []string{"graph", "subgraph", "entity", "relation", "community_report"}) - assertDeleteIndexTaskDeleted(t, "graph-task") -} - -func TestDatasetServiceDeleteIndexRaptorWipeTrueDeletesRaptorArtefacts(t *testing.T) { - db := setupServiceTestDB(t) - pushServiceDB(t, db) - insertDeleteIndexKB(t, "raptor", "raptor-task") - ctx := t.Context() - - docEngine := &deleteIndexDocEngine{} - code, err := testDatasetServiceForDeleteIndex(docEngine).DeleteIndex(ctx, "user-1", "kb-1", "raptor", true) - if err != nil { - t.Fatalf("DeleteIndex failed: %v", err) - } - if code != common.CodeSuccess { - t.Fatalf("expected success code, got %d", code) - } - if len(docEngine.deleteCalls) != 1 { - t.Fatalf("expected one doc-store delete call, got %#v", docEngine.deleteCalls) - } - call := docEngine.deleteCalls[0] - if call.condition["kb_id"] != "kb-1" { - t.Fatalf("delete condition must include kb_id, got %#v", call.condition) - } - assertStringSet(t, call.condition["raptor_kwd"], []string{"raptor"}) - assertDeleteIndexTaskDeleted(t, "raptor-task") - - kb := getDeleteIndexKB(t) - if kb.RaptorTaskID == nil || *kb.RaptorTaskID != "" { - t.Fatalf("expected raptor_task_id to be cleared to empty string, got %#v", kb.RaptorTaskID) - } - if kb.RaptorTaskFinishAt != nil { - t.Fatalf("expected raptor_task_finish_at to be cleared, got %#v", kb.RaptorTaskFinishAt) - } -} - -func TestDatasetServiceDeleteIndexMindmapDoesNotDeleteDocStore(t *testing.T) { - db := setupServiceTestDB(t) - pushServiceDB(t, db) - insertDeleteIndexKB(t, "mindmap", "mindmap-task") - ctx := t.Context() - - docEngine := &deleteIndexDocEngine{} - code, err := testDatasetServiceForDeleteIndex(docEngine).DeleteIndex(ctx, "user-1", "kb-1", "mindmap", true) - if err != nil { - t.Fatalf("DeleteIndex failed: %v", err) - } - if code != common.CodeSuccess { - t.Fatalf("expected success code, got %d", code) - } - if len(docEngine.deleteCalls) != 0 { - t.Fatalf("mindmap delete should not delete doc-store artefacts, got %#v", docEngine.deleteCalls) - } - assertDeleteIndexTaskDeleted(t, "mindmap-task") - - kb := getDeleteIndexKB(t) - if kb.MindmapTaskID == nil || *kb.MindmapTaskID != "" { - t.Fatalf("expected mindmap_task_id to be cleared to empty string, got %#v", kb.MindmapTaskID) - } - if kb.MindmapTaskFinishAt != nil { - t.Fatalf("expected mindmap_task_finish_at to be cleared, got %#v", kb.MindmapTaskFinishAt) - } -} - -func TestDatasetServiceDeleteIndexRejectsInvalidType(t *testing.T) { - db := setupServiceTestDB(t) - pushServiceDB(t, db) - ctx := t.Context() - - code, err := testDatasetServiceForDeleteIndex(&deleteIndexDocEngine{}).DeleteIndex(ctx, "user-1", "kb-1", "invalid", true) - if err == nil { - t.Fatal("expected invalid index type error") - } - if code != common.CodeArgumentError { - t.Fatalf("expected argument error code, got %d", code) - } -} - -func assertDeleteIndexTaskDeleted(t *testing.T, taskID string) { - t.Helper() - var task entity.Task - err := dao.DB.Where("id = ?", taskID).First(&task).Error - if !errors.Is(err, gorm.ErrRecordNotFound) { - t.Fatalf("expected task %s to be deleted, got err=%v task=%#v", taskID, err, task) - } -} - -func getDeleteIndexKB(t *testing.T) entity.Knowledgebase { - t.Helper() - var kb entity.Knowledgebase - if err := dao.DB.Where("id = ?", "kb-1").First(&kb).Error; err != nil { - t.Fatalf("fetch kb: %v", err) - } - return kb -} - -func assertStringSet(t *testing.T, actual interface{}, expected []string) { - t.Helper() - - items, ok := actual.([]interface{}) - if !ok { - t.Fatalf("expected []interface{}, got %#v", actual) - } - if len(items) != len(expected) { - t.Fatalf("expected %d items, got %#v", len(expected), items) - } - - seen := make(map[string]bool, len(items)) - for _, item := range items { - value, ok := item.(string) - if !ok { - t.Fatalf("expected string item, got %#v", item) - } - seen[value] = true - } - for _, item := range expected { - if !seen[item] { - t.Fatalf("missing %q in %#v", item, items) - } - } -} diff --git a/internal/service/dataset/metadata.go b/internal/service/dataset/metadata.go index a5e7676cdd..49b25d0e0e 100644 --- a/internal/service/dataset/metadata.go +++ b/internal/service/dataset/metadata.go @@ -38,6 +38,14 @@ func (d *DatasetService) UpdateDocumentMetadataConfig(ctx context.Context, userI parserConfig = entity.JSONMap{} } parserConfig["metadata"] = metadata + if kb, kbErr := d.kbDAO.GetByID(ctx, dao.DB, datasetID); kbErr == nil && kb != nil { + if tenant, tenantErr := d.tenantDAO.GetByID(ctx, dao.DB, kb.TenantID); tenantErr == nil && tenant != nil { + parserConfig = service.ApplyComponentScopedParserConfig( + parserConfig, + tenant.LLMID, + ) + } + } if err = d.documentDAO.UpdateByID(ctx, dao.DB, doc.ID, map[string]interface{}{"parser_config": parserConfig}); err != nil { return nil, common.CodeServerError, errors.New("database operation failed") @@ -114,6 +122,12 @@ func (d *DatasetService) UpdateMetadataConfig(ctx context.Context, datasetID, te } parserConfig["metadata"] = metadata parserConfig["built_in_metadata"] = builtInMetadata + if tenant, tenantErr := d.tenantDAO.GetByID(ctx, dao.DB, kb.TenantID); tenantErr == nil && tenant != nil { + parserConfig = service.ApplyComponentScopedParserConfig( + parserConfig, + tenant.LLMID, + ) + } if err = d.kbDAO.UpdateByID(ctx, dao.DB, kb.ID, map[string]interface{}{"parser_config": parserConfig}); err != nil { return nil, common.CodeServerError, errors.New("update auto-metadata error.(Database error)") diff --git a/internal/service/dataset/metadata_config_test.go b/internal/service/dataset/metadata_config_test.go index 5553d28461..ddc0f3f248 100644 --- a/internal/service/dataset/metadata_config_test.go +++ b/internal/service/dataset/metadata_config_test.go @@ -22,13 +22,30 @@ import ( "ragflow/internal/common" "ragflow/internal/dao" "ragflow/internal/entity" + "ragflow/internal/service" ) +func metadataFlagInt(t *testing.T, value interface{}) int { + t.Helper() + switch typed := value.(type) { + case int: + return typed + case int64: + return int(typed) + case float64: + return int(typed) + default: + t.Fatalf("unexpected metadata flag type %T (%#v)", value, value) + return 0 + } +} + func testDatasetServiceForDocumentMetadataConfig(t *testing.T) *DatasetService { t.Helper() return &DatasetService{ kbDAO: dao.NewKnowledgebaseDAO(), documentDAO: dao.NewDocumentDAO(), + tenantDAO: dao.NewTenantDAO(), } } @@ -207,3 +224,116 @@ func TestDatasetServiceUpdateDocumentMetadataConfigAllowsTeamMember(t *testing.T t.Fatalf("metadata was not updated: %#v", doc.ParserConfig) } } + +func TestDatasetServiceUpdateMetadataConfigSyncsExtractorSchema(t *testing.T) { + db := setupServiceTestDB(t) + pushServiceDB(t, db) + insertCreateDatasetTenant(t, "tenant-1") + insertDatasetMetadataConfigKB(t, "kb-1", "tenant-1") + if err := dao.DB.Model(&entity.Knowledgebase{}). + Where("id = ?", "kb-1"). + Update("parser_config", entity.JSONMap{ + "enable_metadata": false, + "Extractor:AutoExtractDefault": map[string]any{ + "enable_metadata": 1, + "metadata": []any{ + map[string]any{"key": "stale", "type": "string"}, + }, + }, + }).Error; err != nil { + t.Fatalf("seed parser_config: %v", err) + } + + ctx := t.Context() + result, code, err := (&DatasetService{ + kbDAO: dao.NewKnowledgebaseDAO(), + tenantDAO: dao.NewTenantDAO(), + }).UpdateMetadataConfig(ctx, "kb-1", "tenant-1", &service.MetadataConfigRequest{ + Metadata: []service.MetadataConfigField{ + {Key: "author", Type: "string"}, + }, + BuiltInMetadata: []service.MetadataConfigField{ + {Key: "document_name", Type: "string"}, + }, + }) + if err != nil { + t.Fatalf("UpdateMetadataConfig failed: %v", err) + } + if code != common.CodeSuccess { + t.Fatalf("expected success code, got %d", code) + } + if result["metadata"] == nil { + t.Fatalf("metadata response missing: %#v", result) + } + + persisted, err := dao.NewKnowledgebaseDAO().GetByID(ctx, db, "kb-1") + if err != nil { + t.Fatalf("failed to fetch persisted dataset: %v", err) + } + extractor, ok := persisted.ParserConfig["Extractor:AutoExtractDefault"].(map[string]interface{}) + if !ok { + t.Fatalf("expected extractor component params, got %#v", persisted.ParserConfig["Extractor:AutoExtractDefault"]) + } + if got := metadataFlagInt(t, extractor["enable_metadata"]); got != 0 { + t.Fatalf("extractor enable_metadata = %#v, want 0 when top-level flag stays disabled", extractor["enable_metadata"]) + } + + if err := dao.DB.Model(&entity.Knowledgebase{}). + Where("id = ?", "kb-1"). + Update("parser_config", entity.JSONMap{ + "enable_metadata": true, + "Extractor:AutoExtractDefault": map[string]any{}, + }).Error; err != nil { + t.Fatalf("reset parser_config: %v", err) + } + + _, code, err = (&DatasetService{ + kbDAO: dao.NewKnowledgebaseDAO(), + tenantDAO: dao.NewTenantDAO(), + }).UpdateMetadataConfig(ctx, "kb-1", "tenant-1", &service.MetadataConfigRequest{ + Metadata: []service.MetadataConfigField{ + {Key: "author", Type: "string"}, + }, + BuiltInMetadata: []service.MetadataConfigField{ + {Key: "document_name", Type: "string"}, + }, + }) + if err != nil { + t.Fatalf("UpdateMetadataConfig with enabled metadata failed: %v", err) + } + if code != common.CodeSuccess { + t.Fatalf("expected success code, got %d", code) + } + + persisted, err = dao.NewKnowledgebaseDAO().GetByID(ctx, db, "kb-1") + if err != nil { + t.Fatalf("failed to fetch persisted dataset: %v", err) + } + extractor, ok = persisted.ParserConfig["Extractor:AutoExtractDefault"].(map[string]interface{}) + if !ok { + t.Fatalf("expected extractor component params, got %#v", persisted.ParserConfig["Extractor:AutoExtractDefault"]) + } + if got := metadataFlagInt(t, extractor["enable_metadata"]); got != 1 { + t.Fatalf("extractor enable_metadata = %#v, want 1", extractor["enable_metadata"]) + } + gotFields, ok := extractor["metadata"].([]interface{}) + if !ok { + t.Fatalf("extractor metadata = %#v, want []interface{}", extractor["metadata"]) + } + wantFields := []map[string]interface{}{ + {"key": "author", "type": "string"}, + {"key": "document_name", "type": "string"}, + } + if len(gotFields) != len(wantFields) { + t.Fatalf("extractor metadata len = %d, want %d (%#v)", len(gotFields), len(wantFields), gotFields) + } + for i, want := range wantFields { + field, ok := gotFields[i].(map[string]interface{}) + if !ok { + t.Fatalf("extractor metadata[%d] = %#v, want map[string]interface{}", i, gotFields[i]) + } + if field["key"] != want["key"] || field["type"] != want["type"] { + t.Fatalf("extractor metadata[%d] = %#v, want key/type %#v", i, field, want) + } + } +} diff --git a/internal/service/dataset/task_cleanup_test.go b/internal/service/dataset/task_cleanup_test.go deleted file mode 100644 index 5c40f5190a..0000000000 --- a/internal/service/dataset/task_cleanup_test.go +++ /dev/null @@ -1,111 +0,0 @@ -// -// Copyright 2026 The InfiniFlow Authors. All Rights Reserved. -// -// Licensed under the Apache License, Version 2.0 (the "License"); -// you may not use this file except in compliance with the License. -// You may obtain a copy of the License at -// -// http://www.apache.org/licenses/LICENSE-2.0 -// -// Unless required by applicable law or agreed to in writing, software -// distributed under the License is distributed on an "AS IS" BASIS, -// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -// See the License for the specific language governing permissions and -// limitations under the License. -// - -package dataset - -import ( - "errors" - "testing" - "time" - - "gorm.io/gorm" - - "ragflow/internal/dao" - "ragflow/internal/entity" -) - -func TestCleanupFailedDatasetIndexTaskDeletesTaskAndRestoresDocument(t *testing.T) { - db := setupServiceTestDB(t) - pushServiceDB(t, db) - - previousMsg := "previous progress" - previousBeginAt := time.Date(2026, 6, 18, 10, 0, 0, 0, time.UTC) - queuedMsg := "Task is queued..." - queuedBeginAt := previousBeginAt.Add(time.Hour) - taskID := "task-1" - - kb := &entity.Knowledgebase{ - ID: "kb-1", - TenantID: "user-1", - Name: "test-kb", - EmbdID: "embedding@OpenAI", - CreatedBy: "user-1", - Permission: string(entity.TenantPermissionMe), - ParserID: "naive", - ParserConfig: entity.JSONMap{}, - GraphragTaskID: &taskID, - Status: sptr("1"), - } - if err := dao.DB.Create(kb).Error; err != nil { - t.Fatalf("insert kb: %v", err) - } - - doc := &entity.Document{ - ID: "doc-1", - KbID: "kb-1", - ParserID: "naive", - ParserConfig: entity.JSONMap{}, - SourceType: "local", - Type: "pdf", - CreatedBy: "user-1", - Suffix: ".pdf", - ProgressMsg: &queuedMsg, - ProcessBeginAt: &queuedBeginAt, - } - if err := dao.DB.Create(doc).Error; err != nil { - t.Fatalf("insert document: %v", err) - } - - task := &entity.Task{ID: taskID, DocID: doc.ID, TaskType: "graphrag"} - if err := dao.DB.Create(task).Error; err != nil { - t.Fatalf("insert task: %v", err) - } - - snapshot := &entity.Document{ - ID: doc.ID, - ProgressMsg: &previousMsg, - ProcessBeginAt: &previousBeginAt, - } - if err := cleanupFailedDatasetIndexTask(task.ID, snapshot, kb.ID, "graph"); err != nil { - t.Fatalf("cleanup failed: %v", err) - } - - var persistedTask entity.Task - err := dao.DB.Where("id = ?", task.ID).First(&persistedTask).Error - if !errors.Is(err, gorm.ErrRecordNotFound) { - t.Fatalf("expected task to be deleted, got err=%v task=%#v", err, persistedTask) - } - - ctx := t.Context() - persistedDoc, err := dao.NewDocumentDAO().GetByID(ctx, db, doc.ID) - if err != nil { - t.Fatalf("fetch document: %v", err) - } - if persistedDoc.ProgressMsg == nil || *persistedDoc.ProgressMsg != previousMsg { - t.Fatalf("expected progress_msg %q, got %#v", previousMsg, persistedDoc.ProgressMsg) - } - if persistedDoc.ProcessBeginAt == nil || !persistedDoc.ProcessBeginAt.Equal(previousBeginAt) { - t.Fatalf("expected process_begin_at %v, got %#v", previousBeginAt, persistedDoc.ProcessBeginAt) - } - - var persistedKB entity.Knowledgebase - if err = dao.DB.Where("id = ?", kb.ID).First(&persistedKB).Error; err != nil { - t.Fatalf("fetch kb: %v", err) - } - if persistedKB.GraphragTaskID != nil { - t.Fatalf("expected graphrag_task_id to be cleared, got %#v", *persistedKB.GraphragTaskID) - } -} diff --git a/internal/service/dataset/update.go b/internal/service/dataset/update.go index a2fbef107b..388ccd58f5 100644 --- a/internal/service/dataset/update.go +++ b/internal/service/dataset/update.go @@ -287,6 +287,23 @@ func (d *DatasetService) UpdateDataset(ctx context.Context, datasetID, tenantID updates["parser_config"] = preserveDatasetParserConfigMetadata(cpDefaults, lockedKB.ParserConfig, req.ParserConfig) } } + + effectiveParserConfig := parserConfigJSONMap(updates["parser_config"]) + if effectiveParserConfig == nil && embdIDProvided { + effectiveParserConfig = cloneJSONMap(lockedKB.ParserConfig) + } + if effectiveParserConfig != nil { + llmID := "" + if ownerTenant, tenantErr := d.tenantDAO.GetByID(ctx, tx, lockedKB.TenantID); tenantErr == nil && ownerTenant != nil { + llmID = ownerTenant.LLMID + } + effectiveParserConfig = service.ApplyComponentScopedParserConfig( + effectiveParserConfig, + llmID, + ) + updates["parser_config"] = effectiveParserConfig + } + if len(updates) > 0 { if err = tx.Model(&entity.Knowledgebase{}).Where("id = ?", lockedKB.ID).Updates(updates).Error; err != nil { if dao.IsDuplicateKeyErr(err) { diff --git a/internal/service/dataset_artifact_service.go b/internal/service/dataset_artifact_service.go index 1e91220679..a094f4b2ae 100644 --- a/internal/service/dataset_artifact_service.go +++ b/internal/service/dataset_artifact_service.go @@ -18,6 +18,7 @@ import ( "encoding/json" "fmt" "sort" + "strings" "ragflow/internal/engine" "ragflow/internal/engine/types" @@ -165,7 +166,15 @@ type WikiPageItem struct { // ListWikiPages lists wiki pages for a dataset with optional page_type/topic // filters and pagination. func (s *DatasetArtifactService) ListWikiPages(ctx context.Context, tenantID, datasetID, pageType, topic string, page, pageSize int) ([]WikiPageItem, int64, error) { - filter := map[string]interface{}{"compile_kwd": []string{CompileKwdWikiPage}} + // Only surface the merged dataset-level pages. Each unique (page_type, slug) + // can also have a per-document source row (available_int=0); without this + // filter the same entity/concept would appear once per source doc. Python's + // list_wiki_pages has no such duplication because its writer emits one row + // per page, so mirror that by selecting the merged rows (available_int=1). + filter := map[string]interface{}{ + "compile_kwd": []string{CompileKwdWikiPage}, + "available_int": 1, // merged dataset-level rows only (see engine available_int handling) + } if pageType != "" { filter["page_type_kwd"] = []string{pageType} } @@ -181,10 +190,18 @@ func (s *DatasetArtifactService) ListWikiPages(ctx context.Context, tenantID, da } items := make([]WikiPageItem, 0, len(chunks)) for _, c := range chunks { + pageType := firstStringValue(c["page_type_kwd"]) + // slug_kwd is stored as the full "/" form (Python + // contract); expose the bare slug to the frontend so it can be placed in + // a single URL path segment (gin :slug does not match '/'). + bareSlug := firstStringValue(c["slug_kwd"]) + if pageType != "" { + bareSlug = strings.TrimPrefix(bareSlug, pageType+"/") + } items = append(items, WikiPageItem{ - Slug: firstStringValue(c["slug_kwd"]), + Slug: bareSlug, Title: firstStringValue(c["title_kwd"]), - PageType: firstStringValue(c["page_type_kwd"]), + PageType: pageType, Topic: firstStringValue(c["topic_kwd"]), Summary: firstStringValue(c["summary_with_weight"]), }) @@ -192,13 +209,15 @@ func (s *DatasetArtifactService) ListWikiPages(ctx context.Context, tenantID, da return items, total, nil } -// WikiPageDetail is the full wiki page payload. +// WikiPageDetail is the full wiki page payload. The content field is exposed as +// content_md_rendered to match the frontend IArtifactPage contract (and Python's +// get_wiki_page), which renders it directly. type WikiPageDetail struct { Slug string `json:"slug"` Title string `json:"title"` PageType string `json:"page_type"` Topic string `json:"topic"` - ContentMd string `json:"content_md"` + ContentMd string `json:"content_md_rendered"` Summary string `json:"summary"` EntityNames []string `json:"entity_names"` Outlinks []string `json:"outlinks"` @@ -209,16 +228,21 @@ type WikiPageDetail struct { // GetWikiPage returns a single wiki page by page_type and slug. func (s *DatasetArtifactService) GetWikiPage(ctx context.Context, tenantID, datasetID, pageType, slug string) (*WikiPageDetail, error) { + // Match Python's get_wiki_page contract (dataset_api_service.py): slug_kwd + // is stored as the full "/" form, and list_wiki_pages + // returns the bare slug. Reconstruct the full form deterministically so the + // filter matches the stored value exactly. slugKwd := pageType + "/" + slug filter := map[string]interface{}{ "compile_kwd": []string{CompileKwdWikiPage}, "page_type_kwd": []string{pageType}, "slug_kwd": []string{slugKwd}, + "available_int": 1, // merged dataset-level page, not the per-doc source row } chunks, _, err := s.searchCompiled(ctx, tenantID, datasetID, filter, - []string{"slug_kwd", "title_kwd", "page_type_kwd", "topic_kwd", "content_with_weight", - "summary_with_weight", "entity_names_kwd", "outlinks_kwd", "related_kb_pages_kwd", - "source_chunk_ids", "source_doc_ids"}, + []string{"slug_kwd", "title_kwd", "page_type_kwd", "topic_kwd", "md_with_weight", + "content_with_weight", "summary_with_weight", "entity_names_kwd", "outlinks_kwd", + "related_kb_pages_kwd", "source_chunk_ids", "source_doc_ids"}, 0, 1, nil) if err != nil { return nil, err @@ -227,12 +251,26 @@ func (s *DatasetArtifactService) GetWikiPage(ctx context.Context, tenantID, data return nil, nil } c := chunks[0] + // Python stores the page body in md_with_weight (incremental writer), falling + // back to content_with_weight for legacy rows; mirror that here. + content := firstStringValue(c["md_with_weight"]) + if content == "" { + content = firstStringValue(c["content_with_weight"]) + } + // slug_kwd is the full "/" form; expose the bare slug so a + // client can pass it straight back to GetWikiPage/UpdateWikiPage without the + // "/" prefix being doubled (matches ListWikiPages). + detailPageType := firstStringValue(c["page_type_kwd"]) + detailSlug := firstStringValue(c["slug_kwd"]) + if detailPageType != "" { + detailSlug = strings.TrimPrefix(detailSlug, detailPageType+"/") + } detail := &WikiPageDetail{ - Slug: firstStringValue(c["slug_kwd"]), + Slug: detailSlug, Title: firstStringValue(c["title_kwd"]), - PageType: firstStringValue(c["page_type_kwd"]), + PageType: detailPageType, Topic: firstStringValue(c["topic_kwd"]), - ContentMd: firstStringValue(c["content_with_weight"]), + ContentMd: content, Summary: firstStringValue(c["summary_with_weight"]), EntityNames: toStringSlice(c["entity_names_kwd"]), Outlinks: toStringSlice(c["outlinks_kwd"]), @@ -251,11 +289,14 @@ func (s *DatasetArtifactService) UpdateWikiPage(ctx context.Context, tenantID, d if docEngine == nil { return nil, fmt.Errorf("document engine is not initialized") } + // Python contract: slug_kwd is stored as "page_type/slug"; reconstruct it + // deterministically from the bare slug (see GetWikiPage). slugKwd := pageType + "/" + slug filter := map[string]interface{}{ "compile_kwd": []string{CompileKwdWikiPage}, "page_type_kwd": []string{pageType}, "slug_kwd": []string{slugKwd}, + "available_int": 1, // merged dataset-level page only } chunks, _, err := s.searchCompiled(ctx, tenantID, datasetID, filter, []string{"id"}, 0, 1, nil) if err != nil { @@ -270,6 +311,9 @@ func (s *DatasetArtifactService) UpdateWikiPage(ctx context.Context, tenantID, d } update := map[string]interface{}{} if contentMd != "" { + // GetWikiPage prefers md_with_weight and falls back to content_with_weight, + // so write both to keep the edit readable regardless of the row's writer. + update["md_with_weight"] = contentMd update["content_with_weight"] = contentMd } if title != "" { @@ -300,9 +344,10 @@ func (s *DatasetArtifactService) ListWikiTopics(ctx context.Context, tenantID, d filter := map[string]interface{}{ "compile_kwd": []string{CompileKwdWikiPage}, "page_type_kwd": []string{"concept", "entity"}, + "available_int": 1, // count only merged pages, not per-doc source rows } chunks, _, err := s.searchCompiled(ctx, tenantID, datasetID, filter, - []string{"topic_kwd", "title_kwd", "slug_kwd"}, 0, 1000, nil) + []string{"topic_kwd", "title_kwd", "slug_kwd", "page_type_kwd"}, 0, 1000, nil) if err != nil { return nil, 0, err } @@ -313,12 +358,17 @@ func (s *DatasetArtifactService) ListWikiTopics(ctx context.Context, tenantID, d if t == "" { continue } + pageType := firstStringValue(c["page_type_kwd"]) + bareSlug := firstStringValue(c["slug_kwd"]) + if pageType != "" { + bareSlug = strings.TrimPrefix(bareSlug, pageType+"/") + } counts[t]++ if _, ok := metas[t]; !ok { metas[t] = WikiTopicItem{ Topic: t, Title: firstStringValue(c["title_kwd"]), - Slug: firstStringValue(c["slug_kwd"]), + Slug: bareSlug, } } } diff --git a/internal/service/dataset_types.go b/internal/service/dataset_types.go index cbcc7364d7..0864d66aa1 100644 --- a/internal/service/dataset_types.go +++ b/internal/service/dataset_types.go @@ -1,10 +1,5 @@ package service -// TraceIndexRequest is the request structure for tracing an index task. -type TraceIndexRequest struct { - Type string `json:"type" binding:"required"` -} - // CheckEmbeddingRequest is the request structure for checking embedding compatibility. type CheckEmbeddingRequest struct { EmbeddingID string `json:"embd_id" binding:"required"` diff --git a/internal/service/document/document_dataset_update.go b/internal/service/document/document_dataset_update.go index 9d3cc1b0cf..40ab76b5c1 100644 --- a/internal/service/document/document_dataset_update.go +++ b/internal/service/document/document_dataset_update.go @@ -166,6 +166,13 @@ func (s *DocumentService) UpdateDatasetDocument(ctx context.Context, userID, dat } } else { cleaned := pipelinepkg.BuildParserConfig(dslJSON, req.ParserConfig) + tenant, tenantErr := dao.NewTenantDAO().GetByID(ctx, dao.DB, kb.TenantID) + if tenantErr == nil && tenant != nil { + cleaned = service.ApplyComponentScopedParserConfig( + cleaned, + tenant.LLMID, + ) + } if err = s.documentDAO.UpdateByID(ctx, dao.DB, doc.ID, map[string]interface{}{ "parser_config": cleaned, }); err != nil { diff --git a/internal/service/model_service.go b/internal/service/model_service.go index 37b35f52c1..4ff4b0ed0c 100644 --- a/internal/service/model_service.go +++ b/internal/service/model_service.go @@ -3533,6 +3533,66 @@ func (m *ModelProviderService) ResolveModelConfig(ctx context.Context, tenantID return m.GetModelConfigFromProviderInstance(ctx, tenantID, modelType, modelRef) } +// ResolveModelContextLength returns the chat model's context window +// (content_length) in tokens, or 0 when unknown. After the all_models.json +// migration (PR #17839) content_length is the total context window and +// max_output is the generation cap; the knowledge_compiler prompt-budget logic +// needs the context window, not the output cap. modelRef accepts either a +// tenant model UUID or a "model@instance@provider" composite name. +func (m *ModelProviderService) ResolveModelContextLength(ctx context.Context, tenantID string, modelRef string) (int, error) { + if strings.TrimSpace(modelRef) == "" { + return 0, fmt.Errorf("model ref is required") + } + if modelObj, err := m.modelDAO.GetByID(ctx, dao.DB, modelRef); err == nil { + return m.modelContextLengthByID(ctx, modelObj) + } else if !errors.Is(err, gorm.ErrRecordNotFound) { + return 0, err + } + pureName, _, providerName, err := parseModelName(modelRef) + if err != nil { + return 0, err + } + return m.modelContextLengthByName(providerName, pureName) +} + +// modelContextLengthByID reads content_length from the factory catalog for a +// tenant model row (by id). +func (m *ModelProviderService) modelContextLengthByID(ctx context.Context, modelObj *entity.TenantModel) (int, error) { + if modelObj.Status != "active" { + return 0, fmt.Errorf("tenant model id=%s is disabled", modelObj.ID) + } + provider, err := m.modelProviderDAO.GetByID(ctx, dao.DB, modelObj.ProviderID) + if err != nil { + return 0, err + } + if provider == nil { + return 0, fmt.Errorf("provider id=%s not found for model id=%s", modelObj.ProviderID, modelObj.ID) + } + if mi, _ := dao.GetModelProviderManager().GetModelByName(provider.ProviderName, modelObj.ModelName); mi != nil && mi.ContentLength != nil { + return *mi.ContentLength, nil + } + return 0, nil +} + +// modelContextLengthByName reads content_length from the factory catalog for a +// "model@provider" style reference. It is best-effort: an unknown provider or +// model returns 0 (caller falls back to a default context length). +func (m *ModelProviderService) modelContextLengthByName(providerName, pureName string) (int, error) { + targetProvider := dao.GetModelProviderManager().FindProvider(providerName) + if targetProvider == nil { + return 0, fmt.Errorf("model provider config not found: %s", providerName) + } + for i := range targetProvider.Models { + if strings.EqualFold(targetProvider.Models[i].Name, pureName) { + if targetProvider.Models[i].ContentLength != nil { + return *targetProvider.Models[i].ContentLength, nil + } + return 0, nil + } + } + return 0, nil +} + func (m *ModelProviderService) ResolveModelID(ctx context.Context, tenantID string, modelType entity.ModelType, modelName string) (string, error) { if modelObj, err := m.modelDAO.GetByID(ctx, dao.DB, modelName); err == nil { if modelObj.Status != "active" { diff --git a/internal/service/model_service_test.go b/internal/service/model_service_test.go index 86e81f0607..8059a63772 100644 --- a/internal/service/model_service_test.go +++ b/internal/service/model_service_test.go @@ -201,6 +201,63 @@ func TestModelProviderServiceGetModelConfigByID(t *testing.T) { } } +func TestModelProviderServiceResolveModelContextLength(t *testing.T) { + db := setupModelProviderServiceTestDB(t) + useModelProviderServiceTestDB(t, db) + // Seed a tenant chat model that maps to a real factory-catalog model + // (Anthropic / claude-opus-4-8 has content_length=1000000, max_output=128000). + activeStatus := "1" + rows := []interface{}{ + &entity.UserTenant{ID: "user-tenant-cl", UserID: "user-1", TenantID: "tenant-cl", Role: "owner", InvitedBy: "user-1", Status: &activeStatus}, + &entity.TenantModelProvider{ID: "provider-anthropic", TenantID: "tenant-cl", ProviderName: "Anthropic"}, + &entity.TenantModelInstance{ID: "instance-anthropic", ProviderID: "provider-anthropic", InstanceName: "default", APIKey: "sk-anthropic", Status: "active", Extra: "{}"}, + &entity.TenantModel{ID: "model-claude", ProviderID: "provider-anthropic", InstanceID: "instance-anthropic", ModelName: "claude-opus-4-8", ModelType: int(entity.ModelTypeChat), Status: "active"}, + } + for _, row := range rows { + if err := db.Create(row).Error; err != nil { + t.Fatalf("failed to seed %T: %v", row, err) + } + } + + svc := NewModelProviderService() + ctx := t.Context() + + // UUID path: resolves content_length (context window) from the factory + // catalog, NOT max_output. + got, err := svc.ResolveModelContextLength(ctx, "user-1", "model-claude") + if err != nil { + t.Fatalf("ResolveModelContextLength(uuid) error = %v", err) + } + if got != 1000000 { + t.Fatalf("uuid content_length = %d, want 1000000 (must be the context window, not max_output=128000)", got) + } + + // Composite "model@instance@provider" path resolves the same value. + got2, err := svc.ResolveModelContextLength(ctx, "user-1", "claude-opus-4-8@default@Anthropic") + if err != nil { + t.Fatalf("ResolveModelContextLength(composite) error = %v", err) + } + if got2 != 1000000 { + t.Fatalf("composite content_length = %d, want 1000000", got2) + } +} + +func TestModelProviderServiceResolveModelContextLengthUnknownModel(t *testing.T) { + db := setupModelProviderServiceTestDB(t) + useModelProviderServiceTestDB(t, db) + + // A model that does not exist in the factory catalog resolves to 0 so the + // caller falls back to its default context length instead of failing. + got, err := NewModelProviderService().ResolveModelContextLength( + t.Context(), "user-1", "gpt-no-such-model@default@OpenAI") + if err != nil { + t.Fatalf("ResolveModelContextLength(unknown) error = %v", err) + } + if got != 0 { + t.Fatalf("unknown model content_length = %d, want 0", got) + } +} + func TestModelProviderServiceAlterModelRejectsInvalidStatus(t *testing.T) { ctx := t.Context() code, err := NewModelProviderService().AlterModel(ctx, "OpenAI", "default", "", "user-1", "model-1", map[string]interface{}{"status": "disabled"}) diff --git a/internal/service/wikisearch/engine_service.go b/internal/service/wikisearch/engine_service.go new file mode 100644 index 0000000000..9eac92b799 --- /dev/null +++ b/internal/service/wikisearch/engine_service.go @@ -0,0 +1,222 @@ +package wikisearch + +import ( + "context" + "fmt" + "strings" + + "ragflow/internal/engine" + "ragflow/internal/engine/types" +) + +// compileKWDWikiPage is the canonical compile_kwd for compiled wiki pages. It +// must match Python's WIKI_PAGE_COMPILE_KWD ("wiki_page", wiki.py:1661) AND the +// Go compiler's variantCompileKWD[VariantWiki] (component.go), so both Python- +// and Go-produced wiki pages are surfaced by this service. +const compileKWDWikiPage = "wiki_page" + +// tenantIndexName returns the tenant-scoped chunk index name +// ("ragflow_"), matching how the rest of the stack derives the index +// (internal/handler/dataset.go). The dataset IDs are passed as KB filters, NOT +// used as index names. +func tenantIndexName(tenantID string) string { + return fmt.Sprintf("ragflow_%s", tenantID) +} + +// engineWikiService is the concrete compiled-wiki search service, backed +// directly by the document engine. +// +// - Every operation derives the chunk index from the tenantID +// (ragflow_) and passes dataset IDs only as KB filters, so scope +// can never cross tenants. +// - QueryPages issues an engine Search restricted to compile_kwd="wiki_page" +// (+ supported kinds) so ordinary source chunks are never relabeled as wiki +// pages; the raw rows also carry source_chunk_ids, which are emitted for +// P7/R3 evidence backfill. +// - BackfillChunks fetches original chunks by id via GetChunk, scoped to the +// tenant index + dataset IDs. +type engineWikiService struct { + engine engine.DocEngine // may be nil -> degrade +} + +// NewEngineService builds the concrete wiki search service from the document +// engine (engine.Get() in production; nil disables all operations gracefully). +func NewEngineService(docEngine engine.DocEngine) Service { + return &engineWikiService{engine: docEngine} +} + +func (s *engineWikiService) AvailableFor(ctx context.Context, tenantID string, datasetIDs []string) bool { + if s.engine == nil || tenantID == "" || len(datasetIDs) == 0 { + return false + } + // "Wiki available" must mean the bound KBs actually carry wiki pages, not + // merely that a chunk store exists (a non-wiki KB would otherwise trigger a + // needless empty wiki-query/fallback round trip). Do a bounded existence + // search (Limit=1) filtered to compile_kwd="wiki_page", returning true if + // any page row exists. + res, err := s.engine.Search(ctx, &types.SearchRequest{ + IndexNames: []string{tenantIndexName(tenantID)}, + KbIDs: datasetIDs, + Limit: 1, + Filter: map[string]interface{}{ + "compile_kwd": compileKWDWikiPage, + "available_int": 1, + }, + }) + if err != nil { + return false + } + return res != nil && len(res.Chunks) > 0 +} + +func (s *engineWikiService) QueryPages(ctx context.Context, tenantID string, datasetIDs []string, query, keywords string, topN int) (SearchResult, error) { + if s.engine == nil || tenantID == "" || len(datasetIDs) == 0 || strings.TrimSpace(query) == "" { + return SearchResult{Chunks: []map[string]interface{}{}, DocAggs: []map[string]interface{}{}}, nil + } + if topN <= 0 { + topN = 12 + } + text := query + if kw := strings.TrimSpace(keywords); kw != "" { + text = query + " " + kw + } + req := &types.SearchRequest{ + IndexNames: []string{tenantIndexName(tenantID)}, + KbIDs: datasetIDs, + Limit: topN, + // Select the exact projection we consume. Infinity's default projection + // omits slug_kwd and source_chunk_ids (chunk.go:736-748); without them + // the page slug is blank and P7 evidence backfill has no provenance. + SelectFields: []string{ + "id", "kb_id", "doc_id", "docnm_kwd", "content_with_weight", + "slug_kwd", "source_chunk_ids", + }, + // Discriminate wiki pages by compile_kwd="wiki_page". There is NO + // "kc_kind" column in the chunk schema (infinity_mapping.json:47-56), so + // it must not be used as a filter. Sections (page.go kind:"section") are + // stamped compile_kwd="wiki_section", so this filter returns pages only. + Filter: map[string]interface{}{ + "compile_kwd": compileKWDWikiPage, + "available_int": 1, + }, + MatchExprs: []interface{}{ + &types.MatchTextExpr{ + Fields: []string{"content_with_weight^2", "title_tks^5", "content_ltks"}, + MatchingText: text, + TopN: topN, + }, + }, + } + res, err := s.engine.Search(ctx, req) + if err != nil || res == nil || len(res.Chunks) == 0 { + return SearchResult{Chunks: []map[string]interface{}{}, DocAggs: []map[string]interface{}{}}, nil + } + out := SearchResult{Chunks: []map[string]interface{}{}, DocAggs: []map[string]interface{}{}} + seenDoc := map[string]bool{} + for _, row := range res.Chunks { + // Engine raw rows carry the chunk id under "id" (shimmed to _id) and the + // dataset under "kb_id". Normalize explicitly to the agent chunk shape + // (chunk_id/dataset_id) so a stable id and KB scope are never lost. + id := firstString(row["id"]) + datasetID := firstString(row["kb_id"]) + content := firstString(row["content_with_weight"]) + if id == "" && content == "" { + continue + } + docID := firstString(row["doc_id"]) + chunk := map[string]interface{}{ + "chunk_id": id, "content_with_weight": content, + "doc_id": docID, "docnm_kwd": firstString(row["docnm_kwd"]), + "dataset_id": datasetID, + "wiki_slug_kwd": firstString(row["slug_kwd"]), + } + if src := stringArray(row["source_chunk_ids"]); len(src) > 0 { + chunk["source_chunk_ids"] = src + } + out.Chunks = append(out.Chunks, chunk) + if docID != "" && !seenDoc[docID] { + seenDoc[docID] = true + out.DocAggs = append(out.DocAggs, map[string]interface{}{"doc_id": docID, "doc_name": firstString(row["docnm_kwd"])}) + } + } + return out, nil +} + +func (s *engineWikiService) BackfillChunks(ctx context.Context, tenantID string, datasetIDs []string, chunkIDs []string) ([]map[string]interface{}, error) { + if s.engine == nil || tenantID == "" || len(chunkIDs) == 0 { + return nil, nil + } + const maxBackfill = 16 + seen := map[string]bool{} + ids := make([]string, 0, len(chunkIDs)) + for _, id := range chunkIDs { + if id == "" || seen[id] { + continue + } + seen[id] = true + ids = append(ids, id) + if len(ids) >= maxBackfill { + break + } + } + if len(ids) == 0 { + return nil, nil + } + out := make([]map[string]interface{}, 0, len(ids)) + for _, id := range ids { + raw, err := s.engine.GetChunk(ctx, tenantIndexName(tenantID), id, datasetIDs) + if err != nil { + continue + } + m, ok := raw.(map[string]interface{}) + if !ok { + continue + } + // GetChunk rows also carry id/kb_id (ES shims id to _id; source["id"] + // set in chunk.go:1991). Normalize to the agent chunk shape. + out = append(out, map[string]interface{}{ + "chunk_id": firstString(m["id"]), "content_with_weight": firstString(m["content_with_weight"]), + "doc_id": firstString(m["doc_id"]), "docnm_kwd": firstString(m["docnm_kwd"]), + "dataset_id": firstString(m["kb_id"]), + }) + } + return out, nil +} + +// firstString returns the first string of a possibly array-shaped engine field +// value (the document engine surfaces keyword fields as arrays), matching the +// firstStringValue helper used by the artifact service. +func firstString(v interface{}) string { + switch t := v.(type) { + case string: + return t + case []string: + if len(t) > 0 { + return t[0] + } + case []interface{}: + if len(t) > 0 { + if s, ok := t[0].(string); ok { + return s + } + } + } + return "" +} + +func stringArray(v interface{}) []string { + if raw, ok := v.([]string); ok { + return raw + } + arr, ok := v.([]interface{}) + if !ok { + return nil + } + out := make([]string, 0, len(arr)) + for _, item := range arr { + if s, ok := item.(string); ok { + out = append(out, s) + } + } + return out +} diff --git a/internal/service/wikisearch/engine_service_test.go b/internal/service/wikisearch/engine_service_test.go new file mode 100644 index 0000000000..57b4939d2f --- /dev/null +++ b/internal/service/wikisearch/engine_service_test.go @@ -0,0 +1,265 @@ +package wikisearch + +import ( + "context" + "reflect" + "testing" + + "ragflow/internal/engine" + "ragflow/internal/engine/types" +) + +// fakeDocEngine embeds engine.DocEngine (all methods nil) and overrides only the +// chunk-store existence probe, by-id chunk fetch, and search used by the service. +// It records the index/dataset params so tests can assert tenant-scoped index and +// dataset-scoped filters. +type fakeDocEngine struct { + engine.DocEngine + // existsDatasets lists dataset ids whose table exists; empty means all probe + // results follow `exists` (when false, nothing exists). + exists bool + existsDatasets map[string]bool + chunks map[string]interface{} // chunk id -> raw row (GetChunk) + searchRows []map[string]interface{} + searchReq *types.SearchRequest + gotChunkArgs [][3]interface{} // [indexName, chunkID, datasetIDs] + existsArgs [][2]string // [indexName, datasetID] +} + +func (f *fakeDocEngine) ChunkStoreExists(_ context.Context, indexName, datasetID string) (bool, error) { + f.existsArgs = append(f.existsArgs, [2]string{indexName, datasetID}) + if f.existsDatasets != nil { + return f.existsDatasets[datasetID], nil + } + return f.exists, nil +} + +func (f *fakeDocEngine) GetChunk(_ context.Context, indexName, chunkID string, datasetIDs []string) (interface{}, error) { + f.gotChunkArgs = append(f.gotChunkArgs, [3]interface{}{indexName, chunkID, datasetIDs}) + return f.chunks[chunkID], nil +} + +func (f *fakeDocEngine) Search(_ context.Context, req *types.SearchRequest) (*types.SearchResult, error) { + f.searchReq = req + rows := f.searchRows + // Honor the compile_kwd filter so availability semantics are exercised + // faithfully (only rows carrying the requested compile_kwd match). + if kwd, ok := req.Filter["compile_kwd"].(string); ok { + filtered := make([]map[string]interface{}, 0, len(rows)) + for _, r := range rows { + if r["compile_kwd"] == kwd { + filtered = append(filtered, r) + } + } + rows = filtered + } + return &types.SearchResult{Chunks: rows, Total: int64(len(rows))}, nil +} + +// realEngineRow returns a wiki page row in the engine's actual shape: id +// (shimmed _id), kb_id (dataset), compile_kwd, doc_id, docnm_kwd, slug_kwd, +// source_chunk_ids. +func realEngineRow(id, kbID, docID, name, slug string, source []string) map[string]interface{} { + row := map[string]interface{}{ + "id": id, "kb_id": kbID, "compile_kwd": "wiki_page", + "doc_id": docID, "docnm_kwd": name, + "content_with_weight": "content of " + id, "slug_kwd": slug, + } + if len(source) > 0 { + row["source_chunk_ids"] = source + } + return row +} + +func TestEngineService_QueryPages_NormalizesEngineRowShape(t *testing.T) { + eng := &fakeDocEngine{searchRows: []map[string]interface{}{ + realEngineRow("wiki/alpha", "kb1", "d1", "Alpha", "entity/alpha", []string{"c1", "c2"}), + }} + svc := NewEngineService(eng) + res, err := svc.QueryPages(context.Background(), "t1", []string{"kb1"}, "alpha", "", 5) + if err != nil { + t.Fatalf("QueryPages err = %v", err) + } + if len(res.Chunks) != 1 { + t.Fatalf("chunks = %d, want 1", len(res.Chunks)) + } + // id -> chunk_id, kb_id -> dataset_id must be normalized (Finding 2). + if res.Chunks[0]["chunk_id"] != "wiki/alpha" { + t.Errorf("chunk_id = %v, want wiki/alpha (from engine id)", res.Chunks[0]["chunk_id"]) + } + if res.Chunks[0]["dataset_id"] != "kb1" { + t.Errorf("dataset_id = %v, want kb1 (from engine kb_id)", res.Chunks[0]["dataset_id"]) + } + if res.Chunks[0]["wiki_slug_kwd"] != "entity/alpha" { + t.Errorf("wiki_slug_kwd = %v, want entity/alpha", res.Chunks[0]["wiki_slug_kwd"]) + } + src, ok := res.Chunks[0]["source_chunk_ids"].([]string) + if !ok || len(src) != 2 { + t.Fatalf("source_chunk_ids missing: %#v", res.Chunks[0]) + } + // The engine Search must be scoped to the tenant index + dataset filter and + // filtered to compiled wiki pages. + if eng.searchReq == nil { + t.Fatalf("no SearchRequest issued") + } + if !reflect.DeepEqual(eng.searchReq.IndexNames, []string{"ragflow_t1"}) { + t.Errorf("IndexNames = %v, want [ragflow_t1]", eng.searchReq.IndexNames) + } + if !reflect.DeepEqual(eng.searchReq.KbIDs, []string{"kb1"}) { + t.Errorf("KbIDs = %v, want [kb1]", eng.searchReq.KbIDs) + } + if f, ok := eng.searchReq.Filter["compile_kwd"].(string); !ok || f != "wiki_page" { + t.Errorf("compile_kwd filter = %v, want wiki_page", eng.searchReq.Filter["compile_kwd"]) + } + // SelectFields must project slug_kwd + source_chunk_ids (Infinity's default + // projection omits them), otherwise the page slug and P7 provenance are lost. + projected := map[string]bool{} + for _, f := range eng.searchReq.SelectFields { + projected[f] = true + } + for _, f := range []string{"id", "kb_id", "doc_id", "docnm_kwd", "content_with_weight", "slug_kwd", "source_chunk_ids"} { + if !projected[f] { + t.Errorf("SelectFields missing %q: %v", f, eng.searchReq.SelectFields) + } + } +} + +// TestEngineService_QueryPages_ArrayShapedFields locks firstString handling of +// array-shaped engine keyword fields: the document engine returns slug_kwd / +// docnm_kwd as arrays, and those must not be dropped. +func TestEngineService_QueryPages_ArrayShapedFields(t *testing.T) { + eng := &fakeDocEngine{searchRows: []map[string]interface{}{ + { + "id": "wiki/alpha", "kb_id": "kb1", "compile_kwd": "wiki_page", + "doc_id": "d1", "docnm_kwd": []string{"Alpha"}, + "content_with_weight": "content of wiki/alpha", + "slug_kwd": []string{"entity/alpha"}, + "source_chunk_ids": []string{"c1"}, + }, + }} + svc := NewEngineService(eng) + res, err := svc.QueryPages(context.Background(), "t1", []string{"kb1"}, "alpha", "", 5) + if err != nil { + t.Fatalf("QueryPages err = %v", err) + } + if len(res.Chunks) != 1 { + t.Fatalf("chunks = %d, want 1", len(res.Chunks)) + } + if res.Chunks[0]["wiki_slug_kwd"] != "entity/alpha" { + t.Errorf("wiki_slug_kwd = %v, want entity/alpha (array-shaped slug_kwd)", res.Chunks[0]["wiki_slug_kwd"]) + } + if res.Chunks[0]["docnm_kwd"] != "Alpha" { + t.Errorf("docnm_kwd = %v, want Alpha (array-shaped docnm_kwd)", res.Chunks[0]["docnm_kwd"]) + } + if len(res.DocAggs) != 1 || res.DocAggs[0]["doc_name"] != "Alpha" { + t.Errorf("DocAggs doc_name = %v, want Alpha", res.DocAggs) + } +} + +func TestEngineService_QueryPages_DegradesEmpty(t *testing.T) { + svc := NewEngineService(nil) + res, err := svc.QueryPages(context.Background(), "t1", []string{"kb1"}, "alpha", "", 5) + if err != nil || len(res.Chunks) != 0 { + t.Fatalf("got chunks=%d err=%v, want empty/noerr (no engine)", len(res.Chunks), err) + } + svc2 := NewEngineService(&fakeDocEngine{searchRows: []map[string]interface{}{ + realEngineRow("x", "kb1", "d", "D", "s", nil), + }}) + res2, _ := svc2.QueryPages(context.Background(), "t1", []string{"kb1"}, " ", "", 5) + if len(res2.Chunks) != 0 { + t.Fatalf("chunks = %d, want 0 for blank query", len(res2.Chunks)) + } +} + +func TestEngineService_BackfillChunks_ByIDScoped(t *testing.T) { + eng := &fakeDocEngine{chunks: map[string]interface{}{ + "c1": map[string]interface{}{"id": "c1", "content_with_weight": "raw 1", "doc_id": "d1", "docnm_kwd": "D", "kb_id": "kb1"}, + "c2": map[string]interface{}{"id": "c2", "content_with_weight": "raw 2", "doc_id": "d1", "docnm_kwd": "D", "kb_id": "kb1"}, + }} + svc := NewEngineService(eng) + out, err := svc.BackfillChunks(context.Background(), "t1", []string{"kb1", "kb2"}, []string{"c1", "c2", "c1", "missing"}) + if err != nil { + t.Fatalf("BackfillChunks err = %v", err) + } + if len(out) != 2 { + t.Fatalf("backfill = %d, want 2 (deduped c1,c2; missing skipped)", len(out)) + } + if out[0]["chunk_id"] != "c1" || out[1]["chunk_id"] != "c2" { + t.Errorf("backfill ids = %#v, want c1,c2", out) + } + // kb_id -> dataset_id must be normalized on evidence rows too (Low finding). + if out[0]["dataset_id"] != "kb1" || out[1]["dataset_id"] != "kb1" { + t.Errorf("backfill dataset_id = %v / %v, want kb1 for both (from engine kb_id)", + out[0]["dataset_id"], out[1]["dataset_id"]) + } + // Every GetChunk must be against the tenant index and the dataset scope. + for _, args := range eng.gotChunkArgs { + if args[0] != "ragflow_t1" { + t.Errorf("GetChunk index = %v, want ragflow_t1", args[0]) + } + ds, _ := args[2].([]string) + if !reflect.DeepEqual(ds, []string{"kb1", "kb2"}) { + t.Errorf("GetChunk datasetIDs = %v, want [kb1 kb2]", ds) + } + } +} + +func TestEngineService_BackfillChunks_DegradesNoEngine(t *testing.T) { + svc := NewEngineService(nil) + out, err := svc.BackfillChunks(context.Background(), "t1", []string{"kb1"}, []string{"c1"}) + if err != nil || len(out) != 0 { + t.Fatalf("got %d err=%v, want empty/noerr (no engine)", len(out), err) + } +} + +func TestEngineService_AvailableFor_BoundedExistenceSearch(t *testing.T) { + // A KB carrying wiki pages => AvailableFor true, and the request must be a + // bounded (Limit=1) search on the tenant index, filtered to + // compile_kwd="wiki_page" and the dataset KBs. + eng := &fakeDocEngine{searchRows: []map[string]interface{}{ + realEngineRow("wiki/p1", "kb1", "d1", "P", "p1", nil), + }} + svc := NewEngineService(eng) + if !svc.AvailableFor(context.Background(), "t1", []string{"kb1"}) { + t.Errorf("AvailableFor should be true when a wiki page row exists") + } + if eng.searchReq == nil { + t.Fatalf("no existence Search issued") + } + if !reflect.DeepEqual(eng.searchReq.IndexNames, []string{"ragflow_t1"}) { + t.Errorf("existence IndexNames = %v, want [ragflow_t1]", eng.searchReq.IndexNames) + } + if eng.searchReq.Limit != 1 { + t.Errorf("existence Limit = %d, want 1 (bounded)", eng.searchReq.Limit) + } + if f, ok := eng.searchReq.Filter["compile_kwd"].(string); !ok || f != "wiki_page" { + t.Errorf("existence compile_kwd filter = %v, want wiki_page", eng.searchReq.Filter["compile_kwd"]) + } + if !reflect.DeepEqual(eng.searchReq.KbIDs, []string{"kb1"}) { + t.Errorf("existence KbIDs = %v, want [kb1]", eng.searchReq.KbIDs) + } + // No engine / empty tenant / no datasets -> false. + if NewEngineService(nil).AvailableFor(context.Background(), "t1", []string{"kb1"}) { + t.Errorf("AvailableFor should be false with no engine") + } + if NewEngineService(eng).AvailableFor(context.Background(), "", []string{"kb1"}) { + t.Errorf("AvailableFor should be false with empty tenant") + } + if NewEngineService(eng).AvailableFor(context.Background(), "t1", nil) { + t.Errorf("AvailableFor should be false with no datasets") + } + // No wiki page rows (only ordinary chunks) -> false, even though a chunk + // store exists. This is the "wiki pages exist, not just a table" gate. + svcNoWiki := NewEngineService(&fakeDocEngine{searchRows: []map[string]interface{}{ + {"id": "c1", "kb_id": "kb1", "content_with_weight": "plain chunk"}, // no compile_kwd + }}) + if svcNoWiki.AvailableFor(context.Background(), "t1", []string{"kb1"}) { + t.Errorf("AvailableFor should be false for a KB with no wiki page rows") + } +} + +func TestTenantIndexName(t *testing.T) { + if got := tenantIndexName("t1"); got != "ragflow_t1" { + t.Errorf("tenantIndexName(t1) = %q, want ragflow_t1", got) + } +} diff --git a/internal/service/wikisearch/registry.go b/internal/service/wikisearch/registry.go new file mode 100644 index 0000000000..d29d0bccb2 --- /dev/null +++ b/internal/service/wikisearch/registry.go @@ -0,0 +1,26 @@ +package wikisearch + +import "sync" + +var ( + svcMu sync.RWMutex + svcInst Service +) + +// SetService installs the (production or test) wiki-search service singleton. +// Passing nil reverts to "not configured" (GetService returns nil), which is the +// pre-wiring state: the wiki_query tool then returns an empty result so the agent +// falls back to hybrid search. +func SetService(s Service) { + svcMu.Lock() + defer svcMu.Unlock() + svcInst = s +} + +// GetService returns the installed wiki-search service. It may be nil until +// SetService is called during server bootstrap. +func GetService() Service { + svcMu.RLock() + defer svcMu.RUnlock() + return svcInst +} diff --git a/internal/service/wikisearch/wikisearch.go b/internal/service/wikisearch/wikisearch.go new file mode 100644 index 0000000000..b7e9a756ed --- /dev/null +++ b/internal/service/wikisearch/wikisearch.go @@ -0,0 +1,52 @@ +// Package wikisearch defines the compiled-wiki search service contract. It is a +// dependency-light leaf package so that both the agent tool layer +// (internal/agent/tool) and any concrete engine-backed implementation can depend +// on it without an import cycle — mirroring the internal/service/nav pattern. +// +// The interface abstracts two independent concerns: +// +// - QueryPages: hybrid search over compiled wiki/artifact pages for a query, +// returning rendered page content + slug/title + stable doc aggregation. +// - BackfillChunks: fetch the ORIGINAL source chunks a compiled page was built +// from, by chunk id, scoped to tenant + datasets. This is how the harness P7 +// evidence expansion gets the raw evidence (the general retrieval Search API +// cannot fetch by chunk id). +// - AvailableFor: whether a dataset actually has wiki artifacts the tool may +// search, so the production runner only selects the wiki path when the bound +// KBs carry the artifact (Python's compilation_available gate). +// +// The concrete implementation lives in the same package +// (engineWikiService, see engine_service.go) and is backed by the document +// engine. It is installed at server bootstrap via SetService. +package wikisearch + +import "context" + +// SearchResult is the normalized wiki_query return shape (mirrors Python's +// {"answer":"", "chunks":[...], "doc_aggs":[...]}). +// +// Each chunk is a map with at least: chunk_id, content_with_weight, doc_id, +// docnm_kwd, dataset_id, and (for a compiled page) wiki_slug_kwd. +type SearchResult struct { + Chunks []map[string]interface{} + DocAggs []map[string]interface{} +} + +// Service is the single read entrypoint the wiki_query agent tool and the +// production runner use. +type Service interface { + // AvailableFor reports whether any of the given datasets carries searchable + // wiki/artifact pages (Python's compilation_available gate). Returning false + // means the wiki tool should not be selected for those datasets. + AvailableFor(ctx context.Context, tenantID string, datasetIDs []string) bool + // QueryPages hybrid-searches compiled wiki/artifact pages across the given + // datasets and returns rendered page chunks. Returns an empty result (never + // an error) when nothing matches or the backend is unavailable, so the agent + // can fall back to hybrid search. + QueryPages(ctx context.Context, tenantID string, datasetIDs []string, query, keywords string, topN int) (SearchResult, error) + // BackfillChunks fetches original source chunks by their ids, scoped to the + // tenant and datasets. It returns only chunks it could resolve, in the given + // order, deduped by id; unresolvable ids are skipped (never an error). Empty + // when the backend is unavailable or none resolve. + BackfillChunks(ctx context.Context, tenantID string, datasetIDs []string, chunkIDs []string) ([]map[string]interface{}, error) +} diff --git a/rag/advanced_rag/knowlege_compile/wiki.py b/rag/advanced_rag/knowlege_compile/wiki.py index f354028f01..b18758dfbe 100644 --- a/rag/advanced_rag/knowlege_compile/wiki.py +++ b/rag/advanced_rag/knowlege_compile/wiki.py @@ -821,9 +821,10 @@ async def _wiki_extract_one_batch( language: str, llm_timeout: int, parser_config: Optional[dict] = None, -) -> dict: +) -> Optional[dict]: """Single LLM call for one packed batch. Returns the raw (label-tagged) - extract dict. + extract dict, or ``None`` on a transient LLM timeout/error so the caller + can avoid persisting a poisoned empty result. The entity / relation schemas and the extra rules sections of the prompt are rendered from ``parser_config`` when supplied (mirroring @@ -851,10 +852,10 @@ async def _wiki_extract_one_batch( ) except asyncio.TimeoutError: logging.warning("wiki_map: batch extraction timed out after %ds (%d chunks)", llm_timeout, len(packed)) - return _wiki_empty_extract() + return None except Exception: logging.exception("wiki_map: batch extraction failed (%d chunks)", len(packed)) - return _wiki_empty_extract() + return None _ = language # reserved for future localization return _wiki_unwrap_extract(res) @@ -881,6 +882,12 @@ async def _wiki_process_batch( the top of ``wiki_map_from_chunks``; threaded through so the persisted resume rows record the right hash and the next incremental run can compare cleanly. + + On a transient LLM failure/timeout (``_wiki_extract_one_batch`` returns + ``None``) the batch is NOT persisted with a resume hash. The next + incremental run then sees those chunks as ``new`` and retries, instead of + replaying a permanently cached empty extract. Only a genuine LLM response + (even one with zero items) is persisted. """ if not packed: return _wiki_empty_extract() @@ -896,6 +903,10 @@ async def _wiki_process_batch( llm_timeout, parser_config=parser_config, ) + if raw_extract is None: + # LLM call failed/timed out: leave no resume hash so the next run + # re-extracts these chunks instead of locking in an empty result. + return _wiki_empty_extract() merged, per_chunk = _wiki_resolve_chunk_ids(raw_extract, label_to_id) await _wiki_persist_extracts( per_chunk, diff --git a/test/testcases/restful_api/conftest.py b/test/testcases/restful_api/conftest.py index ad3030287f..36e2b73516 100644 --- a/test/testcases/restful_api/conftest.py +++ b/test/testcases/restful_api/conftest.py @@ -25,6 +25,13 @@ from utils import wait_for GO_ONLY_SKIPS = { "Go route is not implemented": { + # Dataset-level graph/raptor/mindmap indexing via POST/GET/DELETE + # /datasets/:id/index is a Python-era RunIndex/TraceIndex/DeleteIndex + # contract; the Go port schedules dataset compilation through the + # knowledge_compile scheduler instead and does not serve /index. + "test_dataset_index_endpoints", + "test_dataset_index_trace_and_delete_type_contract", + "test_dataset_index_run_with_document_creates_task", "test_document_download_by_id_invalid_id_contract", "test_llm_factories_live_auth_contract", "test_llm_list_live_auth_contract", diff --git a/uv.lock b/uv.lock index 9d25eb4487..9b8fcc6eda 100644 --- a/uv.lock +++ b/uv.lock @@ -8167,7 +8167,6 @@ dependencies = [ { name = "webdav4" }, { name = "webdriver-manager" }, { name = "wechatpy" }, - { name = "werkzeug" }, { name = "wikipedia" }, { name = "word2number" }, { name = "xgboost" }, @@ -8324,7 +8323,6 @@ requires-dist = [ { name = "webdav4", specifier = ">=0.10.0,<0.11.0" }, { name = "webdriver-manager", specifier = "==4.0.1" }, { name = "wechatpy", specifier = ">=1.8.18" }, - { name = "werkzeug", specifier = ">=3.1.7,<4" }, { name = "wikipedia", specifier = "==1.4.0" }, { name = "word2number", specifier = "==1.1" }, { name = "xgboost", specifier = "==1.6.0" }, @@ -9922,14 +9920,14 @@ wheels = [ [[package]] name = "werkzeug" -version = "3.1.8" +version = "3.1.5" source = { registry = "https://mirrors.aliyun.com/pypi/simple" } dependencies = [ { name = "markupsafe" }, ] -sdist = { url = "https://mirrors.aliyun.com/pypi/packages/dd/b2/381be8cfdee792dd117872481b6e378f85c957dd7c5bca38897b08f765fd/werkzeug-3.1.8.tar.gz", hash = "sha256:9bad61a4268dac112f1c5cd4630a56ede601b6ed420300677a869083d70a4c44" } +sdist = { url = "https://mirrors.aliyun.com/pypi/packages/5a/70/1469ef1d3542ae7c2c7b72bd5e3a4e6ee69d7978fa8a3af05a38eca5becf/werkzeug-3.1.5.tar.gz", hash = "sha256:6a548b0e88955dd07ccb25539d7d0cc97417ee9e179677d22c7041c8f078ce67" } wheels = [ - { url = "https://mirrors.aliyun.com/pypi/packages/93/8c/2e650f2afeb7ee576912636c23ddb621c91ac6a98e66dc8d29c3c69446e1/werkzeug-3.1.8-py3-none-any.whl", hash = "sha256:63a77fb8892bf28ebc3178683445222aa500e48ebad5ec77b0ad80f8726b1f50" }, + { url = "https://mirrors.aliyun.com/pypi/packages/ad/e4/8d97cca767bcc1be76d16fb76951608305561c6e056811587f36cb1316a8/werkzeug-3.1.5-py3-none-any.whl", hash = "sha256:5111e36e91086ece91f93268bb39b4a35c1e6f1feac762c9c822ded0a4e322dc" }, ] [[package]] diff --git a/web/src/components/large-model-form-field.tsx b/web/src/components/large-model-form-field.tsx index 76da460131..f5040f84dc 100644 --- a/web/src/components/large-model-form-field.tsx +++ b/web/src/components/large-model-form-field.tsx @@ -39,9 +39,15 @@ export const LargeModelFilterFormSchema = { llm_filter: z.string().optional(), }; -type LargeModelFormFieldProps = Pick; +type LargeModelFormFieldProps = Pick< + NextInnerLLMSelectProps, + 'ownerTenantId' +> & { + name?: string; +}; export function LargeModelFormField({ ownerTenantId, + name = 'llm_id', }: LargeModelFormFieldProps) { const form = useFormContext(); const { t } = useTranslation(); @@ -51,7 +57,7 @@ export function LargeModelFormField({ <> ( diff --git a/web/src/components/pipeline-operator-tabs/pipeline-operator-form.tsx b/web/src/components/pipeline-operator-tabs/pipeline-operator-form.tsx index bb133e7a98..221a599fe8 100644 --- a/web/src/components/pipeline-operator-tabs/pipeline-operator-form.tsx +++ b/web/src/components/pipeline-operator-tabs/pipeline-operator-form.tsx @@ -1,5 +1,6 @@ import { Operator } from '@/constants/agent'; import { RAGFlowNodeType } from '@/interfaces/database/agent'; +import CompilationForm from '@/pages/agent/form/compilation-form'; import ExtractorForm from '@/pages/agent/form/extractor-form'; import ParserForm from '@/pages/agent/form/parser-form'; import TitleChunkerForm from '@/pages/agent/form/title-chunker-form'; @@ -61,6 +62,14 @@ const PipelineOperatorForm = ({ hideOutputs /> ); + case Operator.Compiler: + return ( + + ); case Operator.Tokenizer: return ( ({ + useParams: jest.fn(() => ({ id: 'kb1' })), +})); + +jest.mock('@/utils/api-proxy-scheme', () => ({ + isGoDatasetBackend: jest.fn(() => true), +})); + +jest.mock('@/services/knowledge-service', () => ({ + getDatasetCompilationStatus: jest.fn(), +})); + +// use-dataset-generate imports agent-service (and transitively register-server / +// next-request / locales config that touch import.meta.env). Mock it so the +// Go status path under test doesn't pull in that module graph. +jest.mock('@/services/agent-service', () => ({ + __esModule: true, + default: { cancelDataflow: jest.fn(), deletePipelineTask: jest.fn() }, +})); + +jest.mock('react-i18next', () => ({ + useTranslation: () => ({ t: (key: string) => key }), +})); + +import { getDatasetCompilationStatus } from '@/services/knowledge-service'; +import { isGoDatasetBackend } from '@/utils/api-proxy-scheme'; + +const mockStatus = jest.mocked(getDatasetCompilationStatus); + +function makeWrapper() { + const queryClient = new QueryClient({ + defaultOptions: { + queries: { retry: false }, + }, + }); + // esbuild-jest config here loads .ts with the "tsx" loader but .tsx with the + // plain "ts" loader, so JSX in this test file would not transform. Build the + // provider element with createElement instead. + // eslint-disable-next-line @typescript-eslint/no-explicit-any + const Wrapper = (props: { children: any }) => + React.createElement( + QueryClientProvider, + { client: queryClient }, + props.children, + ); + return Wrapper; +} + +describe('useTraceRunData (Go/hybrid compile-status contract)', () => { + beforeEach(() => { + jest.clearAllMocks(); + (isGoDatasetBackend as jest.Mock).mockReturnValue(true); + }); + + it('maps a successful status to the scheduler contract fields', async () => { + mockStatus.mockResolvedValue({ + data: { + code: 0, + data: { state: 'running', inflight: 2, backlog: 1, error: '' }, + }, + } as never); + + const { result } = renderHook( + () => useTraceRunData(GenerateType.Artifact), + { + wrapper: makeWrapper(), + }, + ); + + await waitFor(() => expect(result.current.isPending).toBe(false), { + timeout: 5000, + }); + + const info = result.current.data; + expect(info?.compilationState).toBe('running'); + expect(info?.inflight).toBe(2); + expect(info?.backlog).toBe(1); + expect(info?.compilationError).toBe(''); + }); + + it('rejects a non-zero business code instead of mapping to idle', async () => { + mockStatus.mockResolvedValue({ + data: { code: 1, message: 'no authorization', data: undefined }, + } as never); + + const { result } = renderHook( + () => useTraceRunData(GenerateType.Artifact), + { + wrapper: makeWrapper(), + }, + ); + + // The query configures retry: 3 with a 1s delay, so the terminal error state + // arrives after the retries drain. + await waitFor(() => expect(result.current.isError).toBe(true), { + timeout: 10000, + }); + + expect(result.current.error).toBeTruthy(); + // The error must not be silently treated as an idle/empty status. + expect(result.current.data).toBeUndefined(); + }); +}); diff --git a/web/src/hooks/use-dataset-generate.ts b/web/src/hooks/use-dataset-generate.ts index a0092f9ae4..8d568cef7c 100644 --- a/web/src/hooks/use-dataset-generate.ts +++ b/web/src/hooks/use-dataset-generate.ts @@ -9,9 +9,11 @@ import { import agentService from '@/services/agent-service'; import { deletePipelineTask, + getDatasetCompilationStatus, runIndex, traceIndex, } from '@/services/knowledge-service'; +import { isGoDatasetBackend } from '@/utils/api-proxy-scheme'; import { useMutation, useQuery, useQueryClient } from '@tanstack/react-query'; import { useMemo } from 'react'; import { useTranslation } from 'react-i18next'; @@ -48,6 +50,15 @@ export interface ITraceInfo { to_page: number; update_date: string; update_time: number; + // Go scheduler compile-status contract (API_PROXY_SCHEME=go/hybrid). These + // replace the legacy task percentage for the Go backend: state is the raw + // dataset-level lifecycle (idle/pending/running/completed), inflight/backlog + // are the MySQL scheduling-entry counts, and error is the batch diagnostic + // (NOT a peer state). Only populated by the Go/hybrid branch of useTraceQuery. + compilationState?: string; + inflight?: number; + backlog?: number; + compilationError?: string; } const useTraceQuery = ( @@ -61,6 +72,15 @@ const useTraceQuery = ( gcTime: 0, refetchInterval: (query) => { const progress = query.state.data?.progress; + // Go/hybrid: keep polling while the dataset compile is pending/running + // (a failed batch is left for retry, so the row stays running/pending and + // we keep polling until it drains to completed). + if (isGoDatasetBackend()) { + const state = query.state.data?.compilationState; + return state === 'pending' || state === 'running' + ? PollIntervalMs + : false; + } return progress != null && progress >= 0 && progress < 1 ? PollIntervalMs : false; @@ -69,6 +89,35 @@ const useTraceQuery = ( retryDelay: 1000, enabled: open && !!id, queryFn: async () => { + if (isGoDatasetBackend()) { + // Scheduler compile-status contract (dataset-level, variant-agnostic). + // The status is NOT a task percentage: we carry the raw state, the + // MySQL inflight/backlog entry counts and the error diagnostic, and + // derive the display status in useGenerateStatus. progress is only set + // so the shared refetch/status helpers keep their contract (idle->0, + // running/pending->0, completed->1, error->0). + const res = await getDatasetCompilationStatus(id!); + const data = res?.data; + // The handler returns HTTP 200 with a non-zero business code for + // authorization/business errors (e.g. "no authorization"). The request + // interceptor only shows a toast and does not reject, so without this + // explicit check a failed read would be mapped to a misleading idle + // state. Reject so the query surfaces the error instead. + if (!data || data.code !== 0) { + throw new Error(data?.message || 'Failed to read compilation status'); + } + const st = data.data ?? {}; + const state: string = st.state ?? 'idle'; + const error: string = st.error ?? ''; + return { + progress: state === 'completed' ? 1 : state === 'idle' ? 0 : 0, + progress_msg: error || state, + compilationState: state, + inflight: st.inflight ?? 0, + backlog: st.backlog ?? 0, + compilationError: error, + } as ITraceInfo; + } const { data } = await traceIndex(id!, traceType); return data?.data ?? {}; }, @@ -131,6 +180,14 @@ export const useDatasetGenerate = () => { } = useMutation({ mutationKey: [DatasetKey.generate], mutationFn: async ({ type }: { type: GenerateType }) => { + // Go/hybrid: dataset compilation is driven automatically by the scheduler + // on document completion; there is no manual RunIndex trigger, so a manual + // "generate" must NOT pretend to succeed. The UI hides/disables the + // control; if it is ever invoked, reject loudly so callers don't mistake + // it for a real run (plan v4.1 §4.2). + if (isGoDatasetBackend()) { + throw new Error(t('message.compileNotSupported')); + } const { data } = await runIndex(id!, TraceTypeMap[type]); if (data.code === 0) { message.success(t('message.operated')); @@ -151,6 +208,12 @@ export const useDatasetGenerate = () => { task_id: string; type: GenerateType; }) => { + // Go/hybrid: the scheduler has no task-level cancel; dataset compilation + // is auto-driven. There is no pause to perform, so reject rather than + // report success (the UI hides the pause control — plan v4.1 §4.2). + if (isGoDatasetBackend()) { + throw new Error(t('message.compileNotSupported')); + } const { data } = await agentService.cancelDataflow(task_id); // For GraphRAG, pause must preserve partial progress (subgraphs, @@ -177,6 +240,19 @@ export function useGenerateStatus(data?: ITraceInfo) { if (!data) { return GenerateStatus.Start; } + if (isGoDatasetBackend()) { + // Go/hybrid: derive from the scheduler contract, not a fake task + // percentage. Error diagnostic takes priority; otherwise map the raw + // dataset-level state (completed->Completed, idle->Start, running/pending + // ->Running). + if (data.compilationError) { + return GenerateStatus.Failed; + } + const st = data.compilationState; + if (st === 'completed') return GenerateStatus.Completed; + if (st === 'running' || st === 'pending') return GenerateStatus.Running; + return GenerateStatus.Start; + } if (data.progress >= 1) { return GenerateStatus.Completed; } else if (!data.progress && data.progress !== 0) { @@ -190,6 +266,13 @@ export function useGenerateStatus(data?: ITraceInfo) { }, [data]); const percent = useMemo(() => { + if (isGoDatasetBackend()) { + // No stable terminal state and no DocTotal/DocProcessed, so there is no + // meaningful percentage. Failures render as a full error marker; active + // runs show the inflight/backlog counts instead of a percent (see the + // UpdateRunProgress / EmptyState components). + return status === GenerateStatus.Failed ? 100 : 0; + } if (status === GenerateStatus.Failed) { return 100; } else if (status === GenerateStatus.Running) { diff --git a/web/src/hooks/use-knowledge-request.ts b/web/src/hooks/use-knowledge-request.ts index 138f05f19f..d33797d88f 100644 --- a/web/src/hooks/use-knowledge-request.ts +++ b/web/src/hooks/use-knowledge-request.ts @@ -1,5 +1,6 @@ import { useHandleFilterSubmit } from '@/components/list-filter-bar/use-handle-filter-submit'; import message from '@/components/ui/message'; +import { isGoDatasetBackend } from '@/utils/api-proxy-scheme'; import { GenerateType, ParseType } from '@/constants/knowledge'; import { ResponsePostType, ResponseType } from '@/interfaces/database/base'; import { @@ -964,6 +965,13 @@ export const useRunArtifactIndex = (kind: string) => { } = useMutation({ mutationKey: [KnowledgeApiAction.RunArtifactIndex], mutationFn: async () => { + // Go/hybrid: wiki compilation is auto-driven by the scheduler; there is no + // legacy RunIndex endpoint. Reject instead of reporting success so a wiki + // update can't be mistaken for a real re-merge (the UI hides/disables the + // update control — plan v4.1 §4.2). + if (isGoDatasetBackend()) { + throw new Error(i18n.t('message.compileNotSupported')); + } const { data } = await runIndex(knowledgeBaseId, 'artifact'); if (data?.code === 0) { message.success(i18n.t('message.operated')); diff --git a/web/src/locales/en.ts b/web/src/locales/en.ts index 38754fe3d3..dc02289284 100644 --- a/web/src/locales/en.ts +++ b/web/src/locales/en.ts @@ -476,6 +476,10 @@ Example: A 1 KB message with 1024-dim embedding uses ~9 KB. The 5 MB default lim log: 'Log', noSkills: 'No skills yet', generate: 'Generate', + compiling: 'Compiling…', + compilingCounts: '{{inflight}} processing / {{backlog}} queued', + autoCompiled: 'Compiled automatically when documents are parsed.', + raptor: 'RAPTOR', artifact: 'Artifact', toSkills: 'To skills', @@ -2412,6 +2416,10 @@ Example: Virtual Hosted Style`, noLangfuseConfigToDelete: 'No Langfuse configuration to delete', renamed: 'Renamed', operated: 'Operated', + compileAutoGenerated: + 'Compilation runs automatically when documents are parsed; no manual trigger is needed.', + compileNotSupported: + 'Manual compilation is not supported here; it runs automatically when documents are parsed.', updated: 'Updated', uploaded: 'Uploaded', 200: 'The server successfully returns the requested data.', diff --git a/web/src/locales/zh.ts b/web/src/locales/zh.ts index e8a93a1d3f..ae7a36cf5c 100644 --- a/web/src/locales/zh.ts +++ b/web/src/locales/zh.ts @@ -412,6 +412,9 @@ export default { generateToSkills: '从该数据集构建分层技能树，并存储生成的技能页面以供搜索和复用。', noWikiPages: '暂无 Wiki 页面', + compiling: '编译中…', + compilingCounts: '处理中 {{inflight}} / 待处理 {{backlog}}', + autoCompiled: '文档解析时自动编译。', clearWikiTitle: '清空 Wiki', clearWikiDescription: '确定要清空该数据集下的所有 Wiki 页面吗？此操作无法撤销。', @@ -2057,6 +2060,9 @@ NER：使用 spaCy NER 和基于规则的关键词提取来抽取实体和关系 noLangfuseConfigToDelete: '没有可删除的 Langfuse 配置', renamed: '重命名成功', operated: '操作成功', + compileAutoGenerated: '知识编译会在文档解析时自动进行，无需手动触发。', + compileNotSupported: + '此处不支持手动编译，知识编译会在文档解析时自动进行。', updated: '更新成功', uploaded: '上传成功', 200: '服务器成功返回请求的数据。', diff --git a/web/src/pages/agent/canvas/node/compilation-node.tsx b/web/src/pages/agent/canvas/node/compilation-node.tsx index b86659d51f..bb5f0aaf2e 100644 --- a/web/src/pages/agent/canvas/node/compilation-node.tsx +++ b/web/src/pages/agent/canvas/node/compilation-node.tsx @@ -2,7 +2,7 @@ import { useCompilationTemplateGroupOptions } from '@/hooks/use-compilation-temp import { IRagNode } from '@/interfaces/database/agent'; import { NodeProps } from '@xyflow/react'; import { get } from 'lodash'; -import { LabelCard } from './card'; +import { LabelCard, LLMLabelCard } from './card'; import { RagNode } from './index'; import { useTranslation } from 'react-i18next'; @@ -11,12 +11,14 @@ export function CompilationNode({ ...props }: NodeProps) { const { t } = useTranslation(); const options = useCompilationTemplateGroupOptions(); const groupId = get(data, 'form.compilation_template_group_id'); + const llmId = get(data, 'form.llm_id'); const groupName = options.find((option) => option.value === groupId)?.label ?? groupId; return (
+ {t('knowledgeConfiguration.compilationTemplate')} diff --git a/web/src/pages/agent/constant/pipeline.tsx b/web/src/pages/agent/constant/pipeline.tsx index f09cc15686..03dbb7e6f5 100644 --- a/web/src/pages/agent/constant/pipeline.tsx +++ b/web/src/pages/agent/constant/pipeline.tsx @@ -362,6 +362,7 @@ export const initialExtractorValues = { export const initialCompilationValues = { compilation_template_group_id: '', + llm_id: '', outputs: { chunks: { type: 'Array', value: [] }, }, diff --git a/web/src/pages/agent/form/compilation-form/index.tsx b/web/src/pages/agent/form/compilation-form/index.tsx index 3e2a7d3fdf..1cdbf0fc02 100644 --- a/web/src/pages/agent/form/compilation-form/index.tsx +++ b/web/src/pages/agent/form/compilation-form/index.tsx @@ -1,10 +1,13 @@ import { CompilationTemplateFormField } from '@/components/compilation-template-form-field'; +import { LargeModelFormField } from '@/components/large-model-form-field'; import { Form } from '@/components/ui/form'; import { zodResolver } from '@hookform/resolvers/zod'; import { memo } from 'react'; import { useForm } from 'react-hook-form'; import { z } from 'zod'; import { initialCompilationValues } from '../../constant/pipeline'; +import { useOwnerTenantId } from '../../context'; +import { useFormChangeCallback } from '../../hooks/use-form-change-callback'; import { useFormValues } from '../../hooks/use-form-values'; import { useWatchFormChange } from '../../hooks/use-watch-form-change'; import { INextOperatorForm } from '../../interface'; @@ -14,28 +17,44 @@ import { Output } from '../components/output'; export const FormSchema = z.object({ compilation_template_group_id: z.string().optional(), + llm_id: z.string().optional(), }); export type CompilationFormSchemaType = z.infer; const outputList = buildOutputList(initialCompilationValues.outputs); -const CompilationForm = ({ node }: INextOperatorForm) => { +const CompilationForm = ({ + node, + onValuesChange, + hideOutputs, +}: INextOperatorForm) => { const defaultValues = useFormValues(initialCompilationValues, node); + const ownerTenantId = useOwnerTenantId(); const form = useForm({ defaultValues, resolver: zodResolver(FormSchema), + mode: 'onChange', }); useWatchFormChange(node?.id, form); + useFormChangeCallback(form, onValuesChange); return (
- + + {!hideOutputs && ( +
+ +
+ )}
); }; diff --git a/web/src/pages/agent/hooks/use-add-node.ts b/web/src/pages/agent/hooks/use-add-node.ts index 82491da75e..1cbe828239 100644 --- a/web/src/pages/agent/hooks/use-add-node.ts +++ b/web/src/pages/agent/hooks/use-add-node.ts @@ -185,7 +185,7 @@ export const useInitializeOperatorParams = () => { sys_prompt: t('flow.prompts.system.summary'), prompts: t('flow.prompts.user.summary'), }, - [Operator.Compiler]: initialCompilationValues, + [Operator.Compiler]: { ...initialCompilationValues, llm_id: llmId }, [Operator.DataOperations]: initialDataOperationsValues, [Operator.ListOperations]: initialListOperationsValues, [Operator.VariableAssigner]: initialVariableAssignerValues, diff --git a/web/src/pages/dataset/compilation/empty-state.tsx b/web/src/pages/dataset/compilation/empty-state.tsx index 1c6556cf68..056cb08331 100644 --- a/web/src/pages/dataset/compilation/empty-state.tsx +++ b/web/src/pages/dataset/compilation/empty-state.tsx @@ -9,6 +9,7 @@ import { useDatasetGenerate, useGenerateStatus, } from '@/hooks/use-dataset-generate'; +import { isGoDatasetBackend } from '@/utils/api-proxy-scheme'; import { GenerableViewMode, @@ -58,29 +59,66 @@ export function CompilationEmptyState({ }, [pauseGenerate, data?.id, generateType]); const showProgress = status === 'running' || status === 'failed'; + const isGo = isGoDatasetBackend(); return (
{!showProgress ? (

{t(TitleKeyMap[type])}

- + {!isGo && ( + + )} + {isGo && ( +

+ {t('knowledgeDetails.autoCompiled')} +

+ )}
) : (
- + {isGo ? ( + // Go/hybrid: no stable percentage and no scheduler cancel, so show + // the MySQL inflight/backlog counts (or the error diagnostic). + status === 'failed' ? ( +
+ + + {data?.compilationError || t('message.operated')} + +
+ ) : ( +
+ + {t('knowledgeDetails.compiling', { + defaultValue: 'Compiling…', + })} + + + {t('knowledgeDetails.compilingCounts', { + inflight: data?.inflight ?? 0, + backlog: data?.backlog ?? 0, + defaultValue: + '{{inflight}} processing / {{backlog}} queued', + })} + +
+ ) + ) : ( + + )}
{t(ViewModeLabelKeyMap[type])} - {status === 'failed' && ( + {!isGo && status === 'failed' && ( )} - {status !== 'failed' && ( + {!isGo && status !== 'failed' && ( { @@ -33,6 +37,39 @@ export function UpdateRunProgress({ [pauseGenerate, data?.id, generateType], ); + // Go/hybrid: no scheduler task-level cancel, and no stable terminal state, so + // show the MySQL inflight/backlog entry counts instead of a percentage and no + // pause button. Error diagnostic takes priority (see plan v4.1 §4.2). + if (isGo && status === 'running') { + return ( + + + {data?.compilationError + ? data.compilationError + : t('knowledgeDetails.compiling', { + defaultValue: 'Compiling…', + })} + {!data?.compilationError && ( + + {t('knowledgeDetails.compilingCounts', { + inflight: data?.inflight ?? 0, + backlog: data?.backlog ?? 0, + defaultValue: '{{inflight}} processing / {{backlog}} queued', + })} + + )} + + ); + } + if (isGo && status === 'failed') { + return ( + + + {data?.compilationError || t('message.operated')} + + ); + } + return ( diff --git a/web/src/pages/dataset/compilation/utils/parse-wiki-link.ts b/web/src/pages/dataset/compilation/utils/parse-wiki-link.ts index 29e79d841c..7accc647b8 100644 --- a/web/src/pages/dataset/compilation/utils/parse-wiki-link.ts +++ b/web/src/pages/dataset/compilation/utils/parse-wiki-link.ts @@ -1,4 +1,4 @@ -export type WikiPageType = 'concept' | 'entity'; +export type WikiPageType = 'concept' | 'entity' | 'topic'; /** * Parse an internal wiki link href into pageType and slug. @@ -9,7 +9,7 @@ export type WikiPageType = 'concept' | 'entity'; * {pageType}/{slug} * /{pageType}/{slug} * - * Only entity/ and concept/ links are considered wiki navigation links. + * entity/, concept/ and topic/ links are all considered wiki navigation links. */ export function parseWikiLinkHref( href: string, @@ -19,7 +19,7 @@ export function parseWikiLinkHref( // Prefer the artifact/{datasetId}/{pageType}/{slug} form. const artifactMatch = normalized.match( - /(?:^|\/)artifact\/[^/]+\/(entity|concept)\/([^/\s"']+)/, + /(?:^|\/)artifact\/[^/]+\/(entity|concept|topic)\/([^/\s"']+)/, ); if (artifactMatch) { return { @@ -29,7 +29,9 @@ export function parseWikiLinkHref( } // Fallback to a plain {pageType}/{slug} form. - const simpleMatch = normalized.match(/(?:^|\/)(entity|concept)\/([^/\s"']+)/); + const simpleMatch = normalized.match( + /(?:^|\/)(entity|concept|topic)\/([^/\s"']+)/, + ); if (simpleMatch) { return { pageType: simpleMatch[1] as WikiPageType, diff --git a/web/src/pages/dataset/compilation/wiki-left-panel/hooks/use-wiki-navigation.ts b/web/src/pages/dataset/compilation/wiki-left-panel/hooks/use-wiki-navigation.ts index c37fa40208..8c3adf511b 100644 --- a/web/src/pages/dataset/compilation/wiki-left-panel/hooks/use-wiki-navigation.ts +++ b/web/src/pages/dataset/compilation/wiki-left-panel/hooks/use-wiki-navigation.ts @@ -6,7 +6,7 @@ import { IArtifactTopic } from '@/interfaces/database/dataset'; import { useDebounce } from 'ahooks'; import { useCallback, useMemo, useRef, useState } from 'react'; -export type WikiPageType = 'concept' | 'entity'; +export type WikiPageType = 'concept' | 'entity' | 'topic'; export function useWikiNavigation() { const scrollRef = useRef(null); diff --git a/web/src/pages/files/files-table.tsx b/web/src/pages/files/files-table.tsx index f25c814155..4c1268ccf5 100644 --- a/web/src/pages/files/files-table.tsx +++ b/web/src/pages/files/files-table.tsx @@ -52,8 +52,7 @@ import { LinkToDatasetDialog } from './link-to-dataset-dialog'; import { UseMoveDocumentShowType } from './use-move-file'; import { useNavigateToOtherFolder } from './use-navigate-to-folder'; import { isFolderType, isKnowledgeBaseType } from './util'; - -declare const __API_PROXY_SCHEME__: string; +import { isGoDatasetBackend } from '../../utils/api-proxy-scheme'; type FilesTableProps = Pick< ReturnType, @@ -104,13 +103,7 @@ export function FilesTable({ } = useRenameCurrentFile(); // Check if skills feature is enabled (only in hybrid or go mode) - const isSkillsEnabled = useMemo(() => { - const scheme = - typeof __API_PROXY_SCHEME__ !== 'undefined' - ? __API_PROXY_SCHEME__ - : 'python'; - return scheme === 'hybrid' || scheme === 'go'; - }, []); + const isSkillsEnabled = useMemo(() => isGoDatasetBackend(), []); // Sort files with skills folder first, then by time // Filter out skills folder if not in hybrid/go mode diff --git a/web/src/services/knowledge-service.ts b/web/src/services/knowledge-service.ts index 62dabd62ba..91dba1b196 100644 --- a/web/src/services/knowledge-service.ts +++ b/web/src/services/knowledge-service.ts @@ -10,7 +10,7 @@ import { } from '@/interfaces/request/knowledge'; import api from '@/utils/api'; import nextRequest from '@/utils/next-request'; -import registerServer from '@/utils/register-server'; +import registerServer, { registerNextServer } from '@/utils/register-server'; import request from '@/utils/request'; const { @@ -287,6 +287,20 @@ export const runIndex = (datasetId: string, indexType: string) => export const traceIndex = (datasetId: string, indexType: string) => request.get(api.traceIndex(datasetId, indexType)); +// getDatasetCompilationStatus reads the Go scheduler compile-status contract +// (GET /datasets/:id/compilation/status), used by API_PROXY_SCHEME=go/hybrid to +// replace the legacy traceIndex task-progress endpoint. Route it through the +// service-layer proxy (registerNextServer -> next-request) like the rest of the +// *-service.ts HTTP proxies. +const compilationStatusProxy = registerNextServer({ + getDatasetCompilationStatus: { + url: (datasetId: string) => api.compilationStatus(datasetId), + method: 'get', + }, +} as const); +export const getDatasetCompilationStatus = (datasetId: string) => + compilationStatusProxy.getDatasetCompilationStatus(datasetId); + // Using RESTful API: GET /api/v1/datasets/{dataset_id}/documents export const listDocument = ( params?: IFetchKnowledgeListRequestParams, diff --git a/web/src/utils/api-proxy-scheme.ts b/web/src/utils/api-proxy-scheme.ts new file mode 100644 index 0000000000..9f247da6a3 --- /dev/null +++ b/web/src/utils/api-proxy-scheme.ts @@ -0,0 +1,48 @@ +/** + * Shared helper for the Vite-injected `API_PROXY_SCHEME` runtime value. + * + * Vite (web/vite.config.ts) injects `__API_PROXY_SCHEME__` from + * `import.meta.env.API_PROXY_SCHEME`. It selects which backend serves the + * dataset compilation / index APIs: + * + * - `python`: all /api requests go to the Python backend (9380); the legacy + * `/datasets/:id/index` run/trace endpoints remain valid. + * - `go`: all /api requests go to the Go backend (9384); the legacy + * `/datasets/:id/index` routes are removed, so UI must use the scheduler + * compile-status contract instead. + * - `hybrid`: Vite routes `/api/v1/datasets` to the Go backend (9384); + * dataset compilation logic behaves like `go`. + * + * Default is `python` for safety (matches Vite's fallback). + * + * All pages must read the scheme through these predicates instead of scattering + * `typeof __API_PROXY_SCHEME__` checks (see files-table.tsx, which previously + * inlined this logic). + */ + +declare const __API_PROXY_SCHEME__: string; + +export type ApiProxyScheme = 'python' | 'go' | 'hybrid'; + +export function getApiProxyScheme(): ApiProxyScheme { + const raw = typeof __API_PROXY_SCHEME__ !== 'undefined' ? __API_PROXY_SCHEME__ : ''; + switch (raw) { + case 'go': + case 'hybrid': + case 'python': + return raw; + default: + return 'python'; + } +} + +/** True when the dataset compilation APIs are served by the Go backend. */ +export function isGoDatasetBackend(): boolean { + const s = getApiProxyScheme(); + return s === 'go' || s === 'hybrid'; +} + +/** True when the legacy Python run/trace `/index` endpoints are in use. */ +export function isPythonDatasetBackend(): boolean { + return getApiProxyScheme() === 'python'; +} diff --git a/web/src/utils/api.ts b/web/src/utils/api.ts index 83c8371b6b..b06f4bd1d6 100644 --- a/web/src/utils/api.ts +++ b/web/src/utils/api.ts @@ -215,6 +215,9 @@ export default { `${restAPIv1}/datasets/${datasetId}/index?type=${indexType.toLowerCase()}`, traceIndex: (datasetId: string, indexType: string) => `${restAPIv1}/datasets/${datasetId}/index?type=${indexType.toLowerCase()}`, + // Go scheduler compile-status contract (API_PROXY_SCHEME=go/hybrid). + compilationStatus: (datasetId: string) => + `${restAPIv1}/datasets/${datasetId}/compilation/status`, unbindPipelineTask: (datasetId: string, indexType: string, wipe?: boolean) => `${restAPIv1}/datasets/${datasetId}/${indexType.toLowerCase()}${wipe === false ? '?wipe=false' : ''}`, pipelineRerun: `${restAPIv1}/agents/rerun`,