diff --git a/AGENTS.md b/AGENTS.md index 37f9f28a51..733bd0ec42 100644 --- a/AGENTS.md +++ b/AGENTS.md @@ -77,7 +77,6 @@ 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 437524b766..f38526883c 100644 --- a/agent/templates/compiler.json +++ b/agent/templates/compiler.json @@ -5,8 +5,7 @@ "obj": { "component_name": "Compiler", "params": { - "compilation_template_group_id": "", - "llm_id": "", + "compilation_template_group_id": "c3aa748c8b2111f191f3047c16ec874f", "outputs": { "chunks": { "type": "Array", @@ -409,8 +408,7 @@ { "data": { "form": { - "compilation_template_group_id": "", - "llm_id": "", + "compilation_template_group_id": "c3aa748c8b2111f191f3047c16ec874f", "outputs": { "chunks": { "type": "Array", diff --git a/api/apps/restful_apis/agent_api.py b/api/apps/restful_apis/agent_api.py index 8d825a4a19..127b9f386a 100644 --- a/api/apps/restful_apis/agent_api.py +++ b/api/apps/restful_apis/agent_api.py @@ -682,7 +682,7 @@ _COMPILATION_TEMPLATE_GROUP_CATEGORY = "compilation_template_group" @add_tenant_id_to_kwargs def list_agents(tenant_id): keywords = request.args.get("keywords", "") - canvas_category = request.args.get("canvas_category") + canvas_category_list = [item for item in request.args.get("canvas_category", "").strip().split(",") if item] canvas_type = request.args.get("canvas_type") owner_ids = [item for item in request.args.get("owner_ids", "").strip().split(",") if item] tags = [item for item in request.args.get("tags", "").strip().split(",") if item] @@ -708,10 +708,11 @@ def list_agents(tenant_id): else: effective_owner_ids = list(authorized_owner_ids) - # Groups-only: an explicit ``compilation_template_group`` category returns - # just the caller's template groups (no agents) via list_saved, so the - # frontend can render a dedicated tab. list_saved paginates in Python. - if canvas_category == _COMPILATION_TEMPLATE_GROUP_CATEGORY: + # Groups-only: when ``compilation_template_group`` is the only selected + # category, return just the caller's template groups (no agents) via + # list_saved, so the frontend can render a dedicated tab. list_saved + # paginates in Python. + if canvas_category_list == [_COMPILATION_TEMPLATE_GROUP_CATEGORY]: from api.db.services.compilation_template_group_service import CompilationTemplateGroupService try: @@ -727,11 +728,17 @@ def list_agents(tenant_id): groups = groups[start : start + items_per_page] return get_json_result(data={"canvas": groups, "total": total}) + # Split selected categories: ``compilation_template_group`` is synthetic + # (resolves to template groups, not agents); everything else filters + # agents by canvas_category IN (...). + wants_groups = _COMPILATION_TEMPLATE_GROUP_CATEGORY in canvas_category_list + agent_categories = [c for c in canvas_category_list if c != _COMPILATION_TEMPLATE_GROUP_CATEGORY] + # Merge mode: with no ``canvas_category`` (and no agent-only filters), list # the caller's compilation template groups alongside agents, interleaved by # ``update_time``. ``canvas_type`` / ``tags`` are agent-only concepts, so # their presence keeps the response agent-only. - merge_groups = not canvas_category and not canvas_type and not tags + merge_groups = not canvas_category_list and not canvas_type and not tags if merge_groups: from api.db.services.compilation_template_group_service import CompilationTemplateGroupService @@ -777,6 +784,47 @@ def list_agents(tenant_id): return get_json_result(data={"canvas": items, "total": total}) + # Mixed mode: both template groups and agent categories are selected - fetch + # agents filtered by the agent categories and merge with template groups, + # interleaved by ``update_time`` (same merge strategy as merge mode). + if wants_groups and agent_categories: + from api.db.services.compilation_template_group_service import CompilationTemplateGroupService + + agents, _ = UserCanvasService.get_by_tenant_ids( + effective_owner_ids, + tenant_id, + 0, + 0, + order_by, + desc, + keywords, + agent_categories, + tags, + canvas_type, + ) + try: + groups = CompilationTemplateGroupService.list_saved(tenant_id, keywords, "", order_by, desc) + except Exception: + logging.exception("list_agents: compilation template group mixed failed for tenant=%s", tenant_id) + groups = [] + + items = [] + for agent in agents: + agent["type"] = "agent" + items.append(agent) + for group in groups: + group["type"] = _COMPILATION_TEMPLATE_GROUP_CATEGORY + group["title"] = group["name"] + items.append(group) + items.sort(key=lambda item: item.get("update_time") or 0, reverse=desc) + + total = len(items) + if page_number and items_per_page: + start = (page_number - 1) * items_per_page + items = items[start : start + items_per_page] + + return get_json_result(data={"canvas": items, "total": total}) + canvas, total = UserCanvasService.get_by_tenant_ids( effective_owner_ids, tenant_id, @@ -785,7 +833,7 @@ def list_agents(tenant_id): order_by, desc, keywords, - canvas_category, + agent_categories, tags, canvas_type, ) diff --git a/api/db/services/canvas_service.py b/api/db/services/canvas_service.py index e9b7d5750b..4bc3fd9d85 100644 --- a/api/db/services/canvas_service.py +++ b/api/db/services/canvas_service.py @@ -130,7 +130,7 @@ class UserCanvasService(CommonService): orderby, desc, keywords, - canvas_category=None, + canvas_category_list=None, tags=None, canvas_type=None, ): @@ -160,8 +160,8 @@ class UserCanvasService(CommonService): ) else: agents = cls.model.select(*fields).join(User, on=(cls.model.user_id == User.id)).where(owner_filter) - if canvas_category: - agents = agents.where(cls.model.canvas_category == canvas_category) + if canvas_category_list: + agents = agents.where(cls.model.canvas_category.in_(canvas_category_list)) if canvas_type: agents = agents.where(cls.model.canvas_type == canvas_type) if tags: diff --git a/cmd/ragflow_server.go b/cmd/ragflow_server.go index 7eab59a078..57266f78b6 100644 --- a/cmd/ragflow_server.go +++ b/cmd/ragflow_server.go @@ -41,7 +41,6 @@ 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" @@ -852,14 +851,6 @@ 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 383a518e88..25073a0fa3 100644 --- a/conf/service_conf.yaml +++ b/conf/service_conf.yaml @@ -68,6 +68,7 @@ otel: enable: false ingestor: mq_type: 'nats' + max_concurrent_workers: 1 file_syncer: max_concurrent_syncs: 1 sync_interval: 3 diff --git a/deepdoc/parser/figure_parser.py b/deepdoc/parser/figure_parser.py index b684df07ce..dc17df438b 100644 --- a/deepdoc/parser/figure_parser.py +++ b/deepdoc/parser/figure_parser.py @@ -28,7 +28,6 @@ from rag.nlp import append_context2table_image4pdf from rag.utils.lazy_image import ensure_pil_image, open_image_for_processing, is_image_like -# need to delete before pr def vision_figure_parser_figure_data_wrapper(figures_data_without_positions): if not figures_data_without_positions: return [] @@ -46,19 +45,29 @@ def vision_figure_parser_figure_data_wrapper(figures_data_without_positions): return res -def vision_figure_parser_docx_wrapper(sections, tbls, callback=None, **kwargs): +def _normalize_vision_language(lang): + return lang or "English" + + +def vision_figure_parser_docx_wrapper(sections, tbls, callback=None, lang="English", **kwargs): + lang = _normalize_vision_language(lang) if not sections: return tbls try: vision_model_config = get_tenant_default_model_by_type(kwargs["tenant_id"], LLMType.VISION) - vision_model = LLMBundle(kwargs["tenant_id"], vision_model_config) + vision_model = LLMBundle(kwargs["tenant_id"], vision_model_config, lang=lang) callback(0.7, "Visual model detected. Attempting to enhance figure extraction...") except Exception: vision_model = None if vision_model: figures_data = vision_figure_parser_figure_data_wrapper(sections) try: - docx_vision_parser = VisionFigureParser(vision_model=vision_model, figures_data=figures_data, **kwargs) + docx_vision_parser = VisionFigureParser( + vision_model=vision_model, + figures_data=figures_data, + lang=lang, + **kwargs, + ) boosted_figures = docx_vision_parser(callback=callback) tbls.extend(boosted_figures) except Exception as e: @@ -66,13 +75,14 @@ def vision_figure_parser_docx_wrapper(sections, tbls, callback=None, **kwargs): return tbls -def vision_figure_parser_figure_xlsx_wrapper(images, callback=None, **kwargs): +def vision_figure_parser_figure_xlsx_wrapper(images, callback=None, lang="English", **kwargs): + lang = _normalize_vision_language(lang) tbls = [] if not images: return [] try: vision_model_config = get_tenant_default_model_by_type(kwargs["tenant_id"], LLMType.VISION) - vision_model = LLMBundle(kwargs["tenant_id"], vision_model_config) + vision_model = LLMBundle(kwargs["tenant_id"], vision_model_config, lang=lang) callback(0.2, "Visual model detected. Attempting to enhance Excel image extraction...") except Exception: vision_model = None @@ -90,7 +100,12 @@ def vision_figure_parser_figure_xlsx_wrapper(images, callback=None, **kwargs): for img in images ] try: - parser = VisionFigureParser(vision_model=vision_model, figures_data=figures_data, **kwargs) + parser = VisionFigureParser( + vision_model=vision_model, + figures_data=figures_data, + lang=lang, + **kwargs, + ) callback(0.22, "Parsing images...") boosted_figures = parser(callback=callback) tbls.extend(boosted_figures) @@ -99,7 +114,8 @@ def vision_figure_parser_figure_xlsx_wrapper(images, callback=None, **kwargs): return tbls -def vision_figure_parser_pdf_wrapper(tbls, callback=None, **kwargs): +def vision_figure_parser_pdf_wrapper(tbls, callback=None, lang="English", **kwargs): + lang = _normalize_vision_language(lang) if not tbls: return [] sections = kwargs.get("sections") @@ -107,7 +123,7 @@ def vision_figure_parser_pdf_wrapper(tbls, callback=None, **kwargs): context_size = max(0, int(parser_config.get("image_context_size", 0) or 0)) try: vision_model_config = get_tenant_default_model_by_type(kwargs["tenant_id"], LLMType.VISION) - vision_model = LLMBundle(kwargs["tenant_id"], vision_model_config) + vision_model = LLMBundle(kwargs["tenant_id"], vision_model_config, lang=lang) callback(0.7, "Visual model detected. Attempting to enhance figure extraction...") except Exception: vision_model = None @@ -131,6 +147,7 @@ def vision_figure_parser_pdf_wrapper(tbls, callback=None, **kwargs): figures_data=figures_data, figure_contexts=figure_contexts, context_size=context_size, + lang=lang, **kwargs, ) boosted_figures = docx_vision_parser(callback=callback) @@ -141,12 +158,13 @@ def vision_figure_parser_pdf_wrapper(tbls, callback=None, **kwargs): return tbls -def vision_figure_parser_docx_wrapper_naive(chunks, idx_lst, callback=None, **kwargs): +def vision_figure_parser_docx_wrapper_naive(chunks, idx_lst, callback=None, lang="English", **kwargs): + lang = _normalize_vision_language(lang) if not chunks: return [] try: vision_model_config = get_tenant_default_model_by_type(kwargs["tenant_id"], LLMType.VISION) - vision_model = LLMBundle(kwargs["tenant_id"], vision_model_config) + vision_model = LLMBundle(kwargs["tenant_id"], vision_model_config, lang=lang) callback(0.7, "Visual model detected. Attempting to enhance figure extraction...") except Exception: vision_model = None @@ -164,12 +182,11 @@ def vision_figure_parser_docx_wrapper_naive(chunks, idx_lst, callback=None, **kw # context_above + caption if any context_above=ck.get("context_above") + ck.get("text", ""), context_below=ck.get("context_below"), + language=lang, ) logging.info(f"[VisionFigureParser] figure={idx} context_above_len={len(context_above)} context_below_len={len(context_below)} prompt=with_context") - logging.info(f"[VisionFigureParser] figure={idx} context_above_snippet={context_above[:512]}") - logging.info(f"[VisionFigureParser] figure={idx} context_below_snippet={context_below[:512]}") else: - prompt = vision_llm_figure_describe_prompt() + prompt = vision_llm_figure_describe_prompt(language=lang) logging.info(f"[VisionFigureParser] figure={idx} context_len=0 prompt=default") try: @@ -201,6 +218,7 @@ shared_executor = ThreadPoolExecutor(max_workers=10) class VisionFigureParser: def __init__(self, vision_model, figures_data, *args, **kwargs): self.vision_model = vision_model + self.language = kwargs.get("lang") or "English" self.figure_contexts = kwargs.get("figure_contexts") or [] self.context_size = max(0, int(kwargs.get("context_size", 0) or 0)) self._extract_figures_info(figures_data) @@ -261,14 +279,13 @@ class VisionFigureParser: prompt = vision_llm_figure_describe_prompt_with_context( context_above=context_above, context_below=context_below, + language=self.language, ) logging.info( f"[VisionFigureParser] figure={figure_idx} context_size={self.context_size} context_above_len={len(context_above)} context_below_len={len(context_below)} prompt=with_context" ) - logging.info(f"[VisionFigureParser] figure={figure_idx} context_above_snippet={context_above[:512]}") - logging.info(f"[VisionFigureParser] figure={figure_idx} context_below_snippet={context_below[:512]}") else: - prompt = vision_llm_figure_describe_prompt() + prompt = vision_llm_figure_describe_prompt(language=self.language) logging.info(f"[VisionFigureParser] figure={figure_idx} context_size={self.context_size} context_len=0 prompt=default") description_text = picture_vision_llm_chunk( binary=figure_binary, diff --git a/docker/service_conf.yaml.template b/docker/service_conf.yaml.template index f2103daaae..5b77399425 100644 --- a/docker/service_conf.yaml.template +++ b/docker/service_conf.yaml.template @@ -86,6 +86,7 @@ otel: enable: false ingestor: mq_type: 'nats' + max_concurrent_workers: 1 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 d843ff96ac..24b419199d 100644 --- a/internal/agent/harness/agentic_rag.go +++ b/internal/agent/harness/agentic_rag.go @@ -46,18 +46,7 @@ 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, @@ -67,8 +56,8 @@ func RunAgenticRAGWithRoute(ctx context.Context, db *gorm.DB, question, keywords return AnswerResult{FinalAnswer: emptyResultMessage, Empty: true} } - // ── route (pre-computed by the caller) ── - state.Route = route + // ── route ── + state.Route = RouteNode(ctx, db, state.Question, modeLabel) // ── pre_search (decomposition modes only) ── if state.Route.RequiresDecomposition { diff --git a/internal/agent/harness/production.go b/internal/agent/harness/production.go index 0b1f739585..2a6f3a9139 100644 --- a/internal/agent/harness/production.go +++ b/internal/agent/harness/production.go @@ -27,34 +27,24 @@ 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, 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. +// 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. 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(). 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). +// lazily via nav.GetNavService(). func NewProductionRunner(db *gorm.DB, tenantID string, datasetIDs []string) (*ProductionRunner, error) { searchBase, err := tool.BuildByName("hybrid_search", nil) if err != nil { @@ -64,11 +54,7 @@ func NewProductionRunner(db *gorm.DB, tenantID string, datasetIDs []string) (*Pr if !ok { return nil, fmt.Errorf("hybrid_search is not invokable") } - r := &ProductionRunner{db: db, tenantID: tenantID, datasetIDs: datasetIDs, searchTool: search} - if common.GetEnv(common.EnvTavilyApiKey) != "" { - r.webTool = tool.NewTavilyTool() - } - return r, nil + return &ProductionRunner{db: db, tenantID: tenantID, datasetIDs: datasetIDs, searchTool: search}, nil } // newProductionRunnerWithTools builds a ProductionRunner with an injected @@ -78,61 +64,19 @@ 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 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. +// Run executes the agentic-search graph with the real tools. It 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} } - 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 { + // The router tool returns a doc list; feed it as the search DocScope. 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 { @@ -141,55 +85,7 @@ func (r *ProductionRunner) hybridSearchFn(ctx context.Context, question, keyword } } } - 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) + return RunAgenticRAG(ctx, r.db, question, keywords, modeLabel, searchFn) } // search invokes the hybrid_search tool and normalizes its chunk output. @@ -212,204 +108,6 @@ 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 deleted file mode 100644 index 31c9344814..0000000000 --- a/internal/agent/harness/production_wiki_test.go +++ /dev/null @@ -1,276 +0,0 @@ -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 9e5ca228c1..5459cfd0d5 100644 --- a/internal/agent/harness/route.go +++ b/internal/agent/harness/route.go @@ -50,10 +50,9 @@ Output format (JSON): ` type routeResult struct { - QuestionType string `json:"question_type"` - RequiresDecomp *bool `json:"requires_decomposition"` - SuggestsCompilation string `json:"suggests_compilation"` - Reasoning string `json:"reasoning"` + QuestionType string `json:"question_type"` + RequiresDecomp *bool `json:"requires_decomposition"` + Reasoning string `json:"reasoning"` } // RouteNode mirrors Python route_node. It classifies the question into a @@ -105,28 +104,10 @@ 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 da9eec04a7..03d9170079 100644 --- a/internal/agent/harness/route_test.go +++ b/internal/agent/harness/route_test.go @@ -74,48 +74,6 @@ 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 2f3393d1c8..3b9c56d357 100644 --- a/internal/agent/harness/types.go +++ b/internal/agent/harness/types.go @@ -26,13 +26,7 @@ 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 - // 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 + Reasoning string } // ClaimTarget mirrors Python ClaimTarget. diff --git a/internal/agent/tool/registry.go b/internal/agent/tool/registry.go index f5143b9ec2..89fafeefd7 100644 --- a/internal/agent/tool/registry.go +++ b/internal/agent/tool/registry.go @@ -40,7 +40,6 @@ 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 deleted file mode 100644 index bd02c9a165..0000000000 --- a/internal/agent/tool/wiki_query.go +++ /dev/null @@ -1,115 +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 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 deleted file mode 100644 index 11e6911187..0000000000 --- a/internal/agent/tool/wiki_query_test.go +++ /dev/null @@ -1,125 +0,0 @@ -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/error_code.go b/internal/common/error_code.go index e3bd460755..9c7a8d60cf 100644 --- a/internal/common/error_code.go +++ b/internal/common/error_code.go @@ -43,6 +43,7 @@ const ( CodeLicenseTimeRollback ErrorCode = 324 CodeLicenseNotFound ErrorCode = 325 CodeLicenseUnexpectedError ErrorCode = 326 + CodeLicenseNotValidYet ErrorCode = 327 CodeBadRequest ErrorCode = 400 CodeUnauthorized ErrorCode = 401 CodeForbidden ErrorCode = 403 diff --git a/internal/common/parser_config.go b/internal/common/parser_config.go index df3c53588f..adc1793b96 100644 --- a/internal/common/parser_config.go +++ b/internal/common/parser_config.go @@ -1,5 +1,129 @@ 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 new file mode 100644 index 0000000000..e6d1195eba --- /dev/null +++ b/internal/common/parser_config_test.go @@ -0,0 +1,135 @@ +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 1c38685919..84f144b0c6 100644 --- a/internal/entity/knowledge_compile_doc.go +++ b/internal/entity/knowledge_compile_doc.go @@ -17,16 +17,6 @@ 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 @@ -52,20 +42,8 @@ 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"` - // 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"` + 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/anthropic.go b/internal/entity/models/anthropic.go index 94cff81ddb..2f00af7617 100644 --- a/internal/entity/models/anthropic.go +++ b/internal/entity/models/anthropic.go @@ -131,10 +131,11 @@ func applyAnthropicChatConfig(reqBody map[string]interface{}, chatModelConfig *C if chatModelConfig == nil { return } - // Deliberately do NOT set max_tokens: response length is controlled through - // the prompt, not a driver parameter. This matches the Python contract - // (rag/llm/chat_model.py strips max_tokens for claude models) and the - // "stop forwarding max token overrides" intent. + if chatModelConfig.MaxTokens != nil { + reqBody["max_tokens"] = *chatModelConfig.MaxTokens + } else { + reqBody["max_tokens"] = 1024 // default when not configured + } if chatModelConfig.Temperature != nil { reqBody["temperature"] = *chatModelConfig.Temperature } diff --git a/internal/entity/models/anthropic_test.go b/internal/entity/models/anthropic_test.go index ced72e8e76..8540f456af 100644 --- a/internal/entity/models/anthropic_test.go +++ b/internal/entity/models/anthropic_test.go @@ -66,10 +66,8 @@ func TestAnthropicChatHappyPath(t *testing.T) { if body["model"] != "claude-sonnet-4-5-20250929" { t.Errorf("model=%v", body["model"]) } - // max_tokens is deliberately NOT sent: response length is prompt-driven - // (matches the Python chat_model which strips it for claude models). - if _, ok := body["max_tokens"]; ok { - t.Errorf("max_tokens=%v, want absent", body["max_tokens"]) + if body["max_tokens"] != float64(1024) { + t.Errorf("max_tokens=%v want 1024", body["max_tokens"]) } msgs, ok := body["messages"].([]interface{}) if !ok || len(msgs) != 1 { @@ -117,9 +115,8 @@ func TestAnthropicChatMapsSystemConfigAndImages(t *testing.T) { if body["system"] != "be concise" { t.Errorf("system=%v, want be concise", body["system"]) } - // max_tokens is deliberately NOT sent (prompt-driven length control). - if _, ok := body["max_tokens"]; ok { - t.Errorf("max_tokens=%v, want absent", body["max_tokens"]) + if body["max_tokens"] != float64(64) { + t.Errorf("max_tokens=%v want 64", body["max_tokens"]) } if body["temperature"] != 0.25 { t.Errorf("temperature=%v want 0.25", body["temperature"]) diff --git a/internal/entity/models/astraflow_test.go b/internal/entity/models/astraflow_test.go index e538d279e2..c6c25eff7b 100644 --- a/internal/entity/models/astraflow_test.go +++ b/internal/entity/models/astraflow_test.go @@ -101,9 +101,6 @@ func TestAstraflowChatHappyPath(t *testing.T) { if body["stream"] != false { t.Errorf("stream=%v, want false", body["stream"]) } - if _, ok := body["max_tokens"]; ok { - t.Errorf("max_tokens should be omitted, got %v", body["max_tokens"]) - } if body["temperature"] != 0.3 { t.Errorf("temperature=%v, want 0.3", body["temperature"]) } diff --git a/internal/entity/models/avian_test.go b/internal/entity/models/avian_test.go index a8827bbcc6..0dac481ae1 100644 --- a/internal/entity/models/avian_test.go +++ b/internal/entity/models/avian_test.go @@ -91,9 +91,6 @@ func TestAvianChatHappyPath(t *testing.T) { if body["stream"] != false { t.Errorf("stream=%v, want false", body["stream"]) } - if _, ok := body["max_tokens"]; ok { - t.Errorf("max_tokens should be omitted, got %v", body["max_tokens"]) - } if body["temperature"] != 0.3 { t.Errorf("temperature=%v, want 0.3", body["temperature"]) } diff --git a/internal/entity/models/azure_openai.go b/internal/entity/models/azure_openai.go index aa317d0c8a..bebef259b3 100644 --- a/internal/entity/models/azure_openai.go +++ b/internal/entity/models/azure_openai.go @@ -94,6 +94,9 @@ func (a *AzureOpenAIModel) ChatWithMessages(ctx context.Context, modelName strin } if chatModelConfig != nil { + if chatModelConfig.MaxTokens != nil { + reqBody["max_tokens"] = *chatModelConfig.MaxTokens + } if chatModelConfig.Temperature != nil { reqBody["temperature"] = *chatModelConfig.Temperature } @@ -148,6 +151,9 @@ func (a *AzureOpenAIModel) ChatStreamlyWithSender(ctx context.Context, modelName if chatModelConfig.Stream != nil && !*chatModelConfig.Stream { return fmt.Errorf("stream must be true in ChatStreamlyWithSender") } + if chatModelConfig.MaxTokens != nil { + reqBody["max_tokens"] = *chatModelConfig.MaxTokens + } if chatModelConfig.Temperature != nil { reqBody["temperature"] = *chatModelConfig.Temperature } diff --git a/internal/entity/models/base_model_request_body_test.go b/internal/entity/models/base_model_request_body_test.go deleted file mode 100644 index 3d4376af9b..0000000000 --- a/internal/entity/models/base_model_request_body_test.go +++ /dev/null @@ -1,31 +0,0 @@ -package models - -import "testing" - -func TestBuildRequestBodyOmitsNonPositiveMaxTokens(t *testing.T) { - zero := 0 - body := buildRequestBody( - &ChatConfig{MaxTokens: &zero}, - "test-model", - []Message{{Role: "user", Content: "hello"}}, - false, - ) - - if _, ok := body["max_tokens"]; ok { - t.Fatalf("max_tokens should be omitted, got %#v", body["max_tokens"]) - } -} - -func TestBuildRequestBodyOmitsPositiveMaxTokens(t *testing.T) { - mt := 256 - body := buildRequestBody( - &ChatConfig{MaxTokens: &mt}, - "test-model", - []Message{{Role: "user", Content: "hello"}}, - false, - ) - - if _, ok := body["max_tokens"]; ok { - t.Fatalf("max_tokens should be omitted, got %#v", body["max_tokens"]) - } -} diff --git a/internal/entity/models/bedrock.go b/internal/entity/models/bedrock.go index 604e7c60eb..5d461cfb3e 100644 --- a/internal/entity/models/bedrock.go +++ b/internal/entity/models/bedrock.go @@ -499,6 +499,10 @@ 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 968636d200..c3a51cc860 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 { - t.Errorf("maxTokens should be omitted, got %v", inf.MaxTokens) + if inf.MaxTokens == nil || *inf.MaxTokens != 4096 { + t.Errorf("maxTokens=%v", inf.MaxTokens) } if inf.Temperature == nil || *inf.Temperature != 0.5 { t.Errorf("temperature=%v", inf.Temperature) diff --git a/internal/entity/models/cometapi_test.go b/internal/entity/models/cometapi_test.go index 4c15c608dd..a682e1b750 100644 --- a/internal/entity/models/cometapi_test.go +++ b/internal/entity/models/cometapi_test.go @@ -119,9 +119,6 @@ func TestCometAPIChatPropagatesConfig(t *testing.T) { withSSRFBypass(t) ctx := t.Context() srv := newCometAPIServer(t, "/v1/chat/completions", func(t *testing.T, body map[string]interface{}, w http.ResponseWriter) { - if _, ok := body["max_tokens"]; ok { - t.Errorf("max_tokens should be omitted, got %v", body["max_tokens"]) - } if body["temperature"] != 0.3 { t.Errorf("temperature=%v want 0.3", body["temperature"]) } diff --git a/internal/entity/models/google.go b/internal/entity/models/google.go index 09281e63cf..fb7d0230cb 100644 --- a/internal/entity/models/google.go +++ b/internal/entity/models/google.go @@ -20,6 +20,7 @@ import ( "context" "encoding/json" "fmt" + "math" "ragflow/internal/common" "ragflow/internal/entity" "strings" @@ -324,6 +325,12 @@ func googleGenerateContentConfig(chatModelConfig *ChatConfig, systemInstruction value := float32(*chatModelConfig.TopP) cfg.TopP = &value } + if chatModelConfig.MaxTokens != nil { + if *chatModelConfig.MaxTokens < 0 || *chatModelConfig.MaxTokens > math.MaxInt32 { + return nil, fmt.Errorf("gemini: max_tokens %d is out of range for int32", *chatModelConfig.MaxTokens) + } + cfg.MaxOutputTokens = int32(*chatModelConfig.MaxTokens) + } if chatModelConfig.Stop != nil { cfg.StopSequences = *chatModelConfig.Stop } diff --git a/internal/entity/models/google_test.go b/internal/entity/models/google_test.go index 52d6c09a72..d6a3a274a2 100644 --- a/internal/entity/models/google_test.go +++ b/internal/entity/models/google_test.go @@ -527,11 +527,11 @@ func TestGoogleGenerateContentConfigConvertsTools(t *testing.T) { func TestGoogleGenerateContentConfigRejectsMaxTokensOverflow(t *testing.T) { overflow := int(math.MaxInt32) + 1 cfg, err := googleGenerateContentConfig(&ChatConfig{MaxTokens: &overflow}, nil) - if err != nil { - t.Fatalf("googleGenerateContentConfig error = %v", err) + if err == nil { + t.Fatalf("expected an error for max_tokens overflowing int32, got cfg = %#v", cfg) } - if cfg != nil && cfg.MaxOutputTokens != 0 { - t.Fatalf("cfg.MaxOutputTokens = %#v, want 0", cfg.MaxOutputTokens) + if cfg != nil { + t.Fatalf("cfg = %#v, want nil on error", cfg) } maxInt32 := int(math.MaxInt32) @@ -539,8 +539,8 @@ func TestGoogleGenerateContentConfigRejectsMaxTokensOverflow(t *testing.T) { if err != nil { t.Fatalf("googleGenerateContentConfig error = %v", err) } - if cfg != nil && cfg.MaxOutputTokens != 0 { - t.Fatalf("cfg.MaxOutputTokens = %#v, want 0", cfg.MaxOutputTokens) + if cfg == nil || cfg.MaxOutputTokens != math.MaxInt32 { + t.Fatalf("cfg.MaxOutputTokens = %#v, want %d", cfg, int32(math.MaxInt32)) } } diff --git a/internal/entity/models/gpustack_test.go b/internal/entity/models/gpustack_test.go index 055dfb5dcf..84c2ad13fd 100644 --- a/internal/entity/models/gpustack_test.go +++ b/internal/entity/models/gpustack_test.go @@ -173,9 +173,6 @@ func TestGPUStackChatForwardsDocumentedFields(t *testing.T) { t.Errorf("documented field %q missing from request body", k) } } - if _, present := body["max_tokens"]; present { - t.Errorf("max_tokens should be omitted, got %v", body["max_tokens"]) - } _ = json.NewEncoder(w).Encode(map[string]interface{}{ "choices": []map[string]interface{}{{ "message": map[string]interface{}{"content": "ok"}, diff --git a/internal/entity/models/groq_test.go b/internal/entity/models/groq_test.go index 592da05eae..f1511b67dd 100644 --- a/internal/entity/models/groq_test.go +++ b/internal/entity/models/groq_test.go @@ -127,9 +127,6 @@ func TestGroqChatHappyPath(t *testing.T) { if body["stream"] != false { t.Errorf("stream=%v want false", body["stream"]) } - if _, ok := body["max_tokens"]; ok { - t.Errorf("max_tokens should be omitted, got %v", body["max_tokens"]) - } if body["temperature"] != 0.3 { t.Errorf("temperature=%v", body["temperature"]) } diff --git a/internal/entity/models/huaweicloud.go b/internal/entity/models/huaweicloud.go index d36514e4c6..dbbdd4c546 100644 --- a/internal/entity/models/huaweicloud.go +++ b/internal/entity/models/huaweicloud.go @@ -123,6 +123,9 @@ func huaweiCloudApplyChatConfig(req map[string]any, modelName string, chatModelC if chatModelConfig == nil { return } + if chatModelConfig.MaxTokens != nil { + req["max_tokens"] = *chatModelConfig.MaxTokens + } if chatModelConfig.Temperature != nil { req["temperature"] = *chatModelConfig.Temperature } diff --git a/internal/entity/models/hunyuan_test.go b/internal/entity/models/hunyuan_test.go index d7dbb4ac92..3fa6e7084b 100644 --- a/internal/entity/models/hunyuan_test.go +++ b/internal/entity/models/hunyuan_test.go @@ -104,9 +104,6 @@ func TestHunyuanChatHappyPath(t *testing.T) { if body["stream"] != false { t.Errorf("stream=%v, want false", body["stream"]) } - if _, ok := body["max_tokens"]; ok { - t.Errorf("max_tokens should be omitted, got %v", body["max_tokens"]) - } if body["temperature"] != 0.3 { t.Errorf("temperature=%v, want 0.3", body["temperature"]) } diff --git a/internal/entity/models/jina_test.go b/internal/entity/models/jina_test.go index f37a439c1c..27d4448e9f 100644 --- a/internal/entity/models/jina_test.go +++ b/internal/entity/models/jina_test.go @@ -134,9 +134,6 @@ func TestJinaChatPropagatesConfig(t *testing.T) { withSSRFBypass(t) ctx := t.Context() srv := newJinaServer(t, "/chat/completions", func(t *testing.T, body map[string]interface{}, w http.ResponseWriter) { - if _, ok := body["max_tokens"]; ok { - t.Errorf("max_tokens should be omitted, got %v", body["max_tokens"]) - } if body["temperature"] != 0.2 { t.Errorf("temperature=%v want 0.2", body["temperature"]) } diff --git a/internal/entity/models/longcat_test.go b/internal/entity/models/longcat_test.go index c2dc845bb3..37bbb7c0b0 100644 --- a/internal/entity/models/longcat_test.go +++ b/internal/entity/models/longcat_test.go @@ -289,9 +289,9 @@ func TestLongCatChatAcceptsReasoningOnlyResponse(t *testing.T) { // TestLongCatChatDropsUndocumentedFields guards against re-introducing // stop / reasoning_effort / response_format / tools etc. The LongCat -// docs only list model, messages, stream, max_tokens, temperature, -// top_p — anything else is undocumented and must not be sent, since -// the maintainer specifically flagged this on PR #14809. +// docs only list model, messages, stream, temperature, top_p — anything +// else is undocumented and must not be sent, since the maintainer +// specifically flagged this on PR #14809. func TestLongCatChatDropsUndocumentedFields(t *testing.T) { withSSRFBypass(t) ctx := t.Context() @@ -307,9 +307,6 @@ func TestLongCatChatDropsUndocumentedFields(t *testing.T) { t.Errorf("documented field %q missing from request body", k) } } - if _, present := body["max_tokens"]; present { - t.Errorf("max_tokens should be omitted, got %v", body["max_tokens"]) - } _ = json.NewEncoder(w).Encode(map[string]interface{}{ "choices": []map[string]interface{}{{ "message": map[string]interface{}{"content": "ok"}, diff --git a/internal/entity/models/mistral_test.go b/internal/entity/models/mistral_test.go index 06eb99a054..d4be9bf8e9 100644 --- a/internal/entity/models/mistral_test.go +++ b/internal/entity/models/mistral_test.go @@ -110,9 +110,6 @@ func TestMistralChatPropagatesConfig(t *testing.T) { withSSRFBypass(t) ctx := t.Context() srv := newMistralServer(t, "/chat/completions", func(t *testing.T, body map[string]interface{}, w http.ResponseWriter) { - if _, ok := body["max_tokens"]; ok { - t.Errorf("max_tokens should be omitted, got %v", body["max_tokens"]) - } if body["temperature"] != 0.3 { t.Errorf("temperature=%v want 0.3", body["temperature"]) } diff --git a/internal/entity/models/modelscope_test.go b/internal/entity/models/modelscope_test.go index 8d195794ee..3a963efb60 100644 --- a/internal/entity/models/modelscope_test.go +++ b/internal/entity/models/modelscope_test.go @@ -128,9 +128,6 @@ func TestModelScopeChatHappyPathNormalizesBaseURLAndOmitsEmptyAuth(t *testing.T) if seen["stream"] != false { t.Errorf("stream=%v, want false", seen["stream"]) } - if _, ok := seen["max_tokens"]; ok { - t.Errorf("max_tokens should be omitted, got %v", seen["max_tokens"]) - } if seen["temperature"] != 0.2 { t.Errorf("temperature=%v, want 0.2", seen["temperature"]) } diff --git a/internal/entity/models/ppio_test.go b/internal/entity/models/ppio_test.go index e4884e0a38..8e632d6f73 100644 --- a/internal/entity/models/ppio_test.go +++ b/internal/entity/models/ppio_test.go @@ -91,9 +91,6 @@ func TestPPIOChatHappyPath(t *testing.T) { if _, ok := body["reasoning_effort"]; ok { t.Errorf("reasoning_effort should not be sent: %v", body["reasoning_effort"]) } - if _, ok := body["max_tokens"]; ok { - t.Errorf("max_tokens should be omitted, got %v", body["max_tokens"]) - } if body["temperature"] != 0.3 { t.Errorf("temperature=%v", body["temperature"]) } diff --git a/internal/entity/models/replicate.go b/internal/entity/models/replicate.go index 2e2be6c913..c75a33cc55 100644 --- a/internal/entity/models/replicate.go +++ b/internal/entity/models/replicate.go @@ -164,6 +164,9 @@ 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 8ea68352b1..d4e80bbdbc 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 _, ok := input["max_new_tokens"]; ok { - t.Errorf("max_new_tokens should be omitted, got %v", input["max_new_tokens"]) + if input["max_new_tokens"] != float64(128) { + t.Errorf("max_new_tokens=%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/tokenpony_test.go b/internal/entity/models/tokenpony_test.go index 92330e1906..cb3f51e059 100644 --- a/internal/entity/models/tokenpony_test.go +++ b/internal/entity/models/tokenpony_test.go @@ -101,9 +101,6 @@ func TestTokenPonyChatHappyPath(t *testing.T) { if body["stream"] != false { t.Errorf("stream=%v, want false", body["stream"]) } - if _, ok := body["max_tokens"]; ok { - t.Errorf("max_tokens should be omitted, got %v", body["max_tokens"]) - } if body["temperature"] != 0.3 { t.Errorf("temperature=%v, want 0.3", body["temperature"]) } diff --git a/internal/entity/models/upstage_test.go b/internal/entity/models/upstage_test.go index a05b7d2b8e..56f5521fc8 100644 --- a/internal/entity/models/upstage_test.go +++ b/internal/entity/models/upstage_test.go @@ -218,9 +218,6 @@ func TestUpstageRequestBodyMatchesSolarAPIShape(t *testing.T) { if stopArr, ok := seen["stop"].([]interface{}); !ok || len(stopArr) != 1 || stopArr[0] != "END" { t.Errorf("body[stop]=%v want [END]", seen["stop"]) } - if _, ok := seen["max_tokens"]; ok { - t.Errorf("max_tokens should be omitted, got %v", seen["max_tokens"]) - } if _, ok := seen["messages"].([]interface{}); !ok { t.Errorf("body[messages] missing or wrong type") } diff --git a/internal/entity/models/xiaomi.go b/internal/entity/models/xiaomi.go index 9c8c3f3665..3065db7e50 100644 --- a/internal/entity/models/xiaomi.go +++ b/internal/entity/models/xiaomi.go @@ -77,9 +77,12 @@ func (x *XiaomiModel) ChatWithMessages(ctx context.Context, modelName string, me // Build request body reqBody := buildRequestBody(chatModelConfig, modelName, messages, false) - delete(reqBody, "max_tokens") if chatModelConfig != nil { + if chatModelConfig.MaxTokens != nil { + reqBody["max_completion_tokens"] = *chatModelConfig.MaxTokens + } + if chatModelConfig.Thinking != nil { if *chatModelConfig.Thinking { reqBody["thinking"] = map[string]interface{}{ @@ -121,12 +124,15 @@ func (x *XiaomiModel) ChatStreamlyWithSender(ctx context.Context, modelName stri // Build request body with streaming enabled reqBody := buildRequestBody(modelConfig, modelName, messages, true) - delete(reqBody, "max_tokens") reqBody["stream_options"] = map[string]interface{}{ "include_usage": true, } if modelConfig != nil { + if modelConfig.MaxTokens != nil { + reqBody["max_completion_tokens"] = *modelConfig.MaxTokens + } + if modelConfig.Thinking != nil { if *modelConfig.Thinking { reqBody["thinking"] = map[string]interface{}{ diff --git a/internal/entity/models/xiaomi_test.go b/internal/entity/models/xiaomi_test.go index 7f07842853..1fe90402b3 100644 --- a/internal/entity/models/xiaomi_test.go +++ b/internal/entity/models/xiaomi_test.go @@ -91,8 +91,8 @@ func TestXiaomiChatHappyPath(t *testing.T) { if body["max_tokens"] != nil { t.Errorf("max_tokens must not be sent: %v", body["max_tokens"]) } - if _, ok := body["max_completion_tokens"]; ok { - t.Errorf("max_completion_tokens should be omitted, got %v", body["max_completion_tokens"]) + if body["max_completion_tokens"] != float64(1024) { + t.Errorf("max_completion_tokens=%v", body["max_completion_tokens"]) } thinking, ok := body["thinking"].(map[string]interface{}) if !ok || thinking["type"] != "disabled" { diff --git a/internal/entity/models/xinference_test.go b/internal/entity/models/xinference_test.go index d04e94e61b..15f573070e 100644 --- a/internal/entity/models/xinference_test.go +++ b/internal/entity/models/xinference_test.go @@ -95,9 +95,6 @@ func TestXinferenceChatHappyPathNormalizesBaseURLAndOmitsEmptyAuth(t *testing.T) if seen["stream"] != false { t.Errorf("stream=%v, want false", seen["stream"]) } - if _, ok := seen["max_tokens"]; ok { - t.Errorf("max_tokens should be omitted, got %v", seen["max_tokens"]) - } if seen["temperature"] != 0.2 { t.Errorf("temperature=%v, want 0.2", seen["temperature"]) } diff --git a/internal/entity/models/zhipu-ai.go b/internal/entity/models/zhipu-ai.go index eb9151e72a..98920be0fe 100644 --- a/internal/entity/models/zhipu-ai.go +++ b/internal/entity/models/zhipu-ai.go @@ -21,7 +21,6 @@ import ( "context" "encoding/base64" "encoding/json" - "errors" "fmt" "io" "mime/multipart" @@ -30,8 +29,6 @@ import ( "path/filepath" "ragflow/internal/common" "strings" - - "go.uber.org/zap" ) // ZhipuAIModel implements ModelDriver for Zhipu AI @@ -58,43 +55,7 @@ func (z *ZhipuAIModel) Name() string { return "zhipu" } -type ZhipuChatResponse struct { - Choices []struct { - FinishReason string `json:"finish_reason"` - Index int `json:"index"` - Message struct { - Content string `json:"content"` - ReasoningContent string `json:"reasoning_content"` - Role string `json:"role"` - ToolCalls []map[string]any `json:"tool_calls"` - } `json:"message"` - } `json:"choices"` - Created int `json:"created"` - Id string `json:"id"` - Model string `json:"model"` - Object string `json:"object"` - RequestId string `json:"request_id"` - Usage struct { - CompletionTokens int `json:"completion_tokens"` - PromptTokens int `json:"prompt_tokens"` - PromptTokensDetails struct { - CachedTokens int `json:"cached_tokens"` - } `json:"prompt_tokens_details"` - TotalTokens int `json:"total_tokens"` - } `json:"usage"` -} - -// zhipuRetryable wraps errors that should be retried by the provider with -// exponential backoff: transient network failures, request timeouts, rate -// limiting (429) and server-side (5xx) errors. Request-level (4xx) and -// parse-level errors are not retried. -var zhipuRetryable = errors.New("zhipu: retryable transient error") - -// ChatWithMessages sends multiple messages with roles and returns response. -// The provider owns failure retry with exponential backoff so transient LLM -// outages (network blips, 429 rate limiting, 5xx, timeouts) do not fail the -// whole knowledge-compilation pipeline; the caller is still told about every -// failed attempt at error level. +// 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 @@ -112,115 +73,26 @@ func (z *ZhipuAIModel) ChatWithMessages(ctx context.Context, modelName string, m url := fmt.Sprintf("%s/%s", baseURL, z.baseModel.URLSuffix.Chat) reqBody := buildRequestBody(chatModelConfig, modelName, messages, false) - if chatModelConfig != nil && chatModelConfig.Thinking != nil { - if *chatModelConfig.Thinking { - reqBody["thinking"] = map[string]interface{}{ - "type": "enabled", - } - } else { - reqBody["thinking"] = map[string]interface{}{ - "type": "disabled", + if chatModelConfig != nil { + if chatModelConfig.Thinking != nil { + if *chatModelConfig.Thinking { + reqBody["thinking"] = map[string]interface{}{ + "type": "enabled", + } + } else { + reqBody["thinking"] = map[string]interface{}{ + "type": "disabled", + } } } } - jsonData, err := json.Marshal(reqBody) - if err != nil { - return nil, fmt.Errorf("failed to marshal request: %w", err) - } - - var resp *ChatResponse - err = common.RetryWithBackoff(ctx, common.DefaultRetryMax, common.DefaultRetryDelay, func() error { - r, err := z.doChatWithMessages(ctx, url, jsonData, apiConfig, chatModelConfig, modelUsage) - if err != nil { - // The caller wants every failed attempt surfaced at error level so - // slow/transient provider failures are observable, not silent. - common.Error("zhipu chat attempt failed", err, - zap.String("model", modelName), - zap.Any("max_tokens", reqBody["max_tokens"]), - zap.Any("thinking", reqBody["thinking"]), - ) - return err - } - resp = r - return nil - }, zhipuChatShouldRetry) + body, err := z.baseModel.doRequest(ctx, url, apiConfig, reqBody, nonStreamCallTimeout) if err != nil { return nil, err } - return resp, nil -} -// zhipuChatShouldRetry reports whether a chat attempt error is transient and -// worth another backoff round. -func zhipuChatShouldRetry(err error) bool { - if errors.Is(err, zhipuRetryable) || errors.Is(err, context.DeadlineExceeded) { - return true - } - return false -} - -// doChatWithMessages performs a single non-streaming chat round-trip under a -// fresh nonStreamCallTimeout window, returning a retryable-marker error for -// transient failures. -func (z *ZhipuAIModel) doChatWithMessages(ctx context.Context, url string, jsonData []byte, apiConfig *APIConfig, chatModelConfig *ChatConfig, modelUsage *common.ModelUsage) (*ChatResponse, error) { - ctx, cancel := context.WithTimeout(ctx, nonStreamCallTimeout) - defer cancel() - - req, err := http.NewRequestWithContext(ctx, "POST", url, bytes.NewBuffer(jsonData)) - if err != nil { - return nil, fmt.Errorf("failed to create request: %w", err) - } - - req.Header.Set("Content-Type", "application/json") - req.Header.Set("Authorization", fmt.Sprintf("Bearer %s", *apiConfig.ApiKey)) - - resp, err := z.baseModel.httpClient.Do(req) - if err != nil { - return nil, fmt.Errorf("%w: failed to send request: %v", zhipuRetryable, err) - } - defer resp.Body.Close() - - body, err := io.ReadAll(resp.Body) - if err != nil { - return nil, fmt.Errorf("%w: failed to read response: %v", zhipuRetryable, err) - } - - if resp.StatusCode == http.StatusTooManyRequests || resp.StatusCode >= http.StatusInternalServerError { - return nil, fmt.Errorf("%w: API request failed with status %d: %s", zhipuRetryable, resp.StatusCode, string(body)) - } - if resp.StatusCode != http.StatusOK { - return nil, fmt.Errorf("API request failed with status %d: %s", resp.StatusCode, string(body)) - } - - return parseChatCompletionResponse(body, chatModelConfig, modelUsage, func(body []byte, chatConfig *ChatConfig) (chatResponseParts, error) { - var result ZhipuChatResponse - if err := json.Unmarshal(body, &result); err != nil { - return chatResponseParts{}, fmt.Errorf("failed to parse response: %w", err) - } - - if len(result.Choices) == 0 { - return chatResponseParts{}, fmt.Errorf("empty response") - } - - choice := &result.Choices[0] - var reasonContent *string - if chatConfig != nil && chatConfig.Thinking != nil && *chatConfig.Thinking { - reasonContent = &choice.Message.ReasoningContent - } - - return chatResponseParts{ - RequestID: result.RequestId, - Content: &choice.Message.Content, - ReasonContent: reasonContent, - ToolCalls: choice.Message.ToolCalls, - Usage: &TokenUsage{ - PromptTokens: result.Usage.PromptTokens, - CompletionTokens: result.Usage.CompletionTokens, - TotalTokens: result.Usage.TotalTokens, - }, - }, nil - }) + return HandleNonStreamingResponse(body, modelUsage, chatModelConfig, OpenAIParserConfig) } // ChatStreamlyWithSender sends messages and streams response via sender function (best performance, no channel) diff --git a/internal/handler/compilation_status_test.go b/internal/handler/compilation_status_test.go deleted file mode 100644 index b76b490685..0000000000 --- a/internal/handler/compilation_status_test.go +++ /dev/null @@ -1,195 +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 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 1fefbb7650..59c9f2df47 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 -// (Compiler, Extractor, File, Parser, Tokenizer + 9 chunker variants). -// Names must be sorted ascending (plan §4 task 1 stable output). +// (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{ - "compiler", "extractor", "file", "grouptitlechunker", "hierarchytitlechunker", + "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{ - "compiler", "extractor", "file", "grouptitlechunker", "hierarchytitlechunker", + "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{ - "compiler", "extractor", "file", "grouptitlechunker", "hierarchytitlechunker", + "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 1fbaf8509e..37feaec5e4 100644 --- a/internal/handler/dataset.go +++ b/internal/handler/dataset.go @@ -1008,28 +1008,114 @@ func (h *DatasetsHandler) AggregateTags(c *gin.Context) { common.SuccessWithData(c, result, "success") } -// 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) { +// RunIndex Run an indexing task (graph/raptor/mindmap) for a dataset. +func (h *DatasetsHandler) RunIndex(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() - st, code, err := h.datasetsService.GetDatasetCompilationStatus(ctx, userID, datasetID) + indexType := strings.ToLower(strings.TrimSpace(c.Query("type"))) + data, code, err := h.datasetsService.RunIndex(ctx, userID, datasetID, indexType) if err != nil { common.ErrorWithCode(c, code, err.Error()) return } - common.SuccessWithData(c, st, "success") + + 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") } // ListMetadataFlattened handles GET /api/v1/datasets/metadata/flattened. diff --git a/internal/handler/dataset_artifact.go b/internal/handler/dataset_artifact.go index dcddfddda7..1c898f6469 100644 --- a/internal/handler/dataset_artifact.go +++ b/internal/handler/dataset_artifact.go @@ -119,9 +119,7 @@ func (h *DatasetArtifactHandler) ListArtifacts(c *gin.Context) { common.ErrorWithCode(c, common.CodeDataError, err.Error()) return } - // 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") + common.SuccessWithData(c, gin.H{"total": total, "pages": items}, "success") } // UpdateArtifact handles PUT /artifacts// — edit a wiki page. @@ -241,9 +239,7 @@ func (h *DatasetArtifactHandler) ListArtifactTopics(c *gin.Context) { common.ErrorWithCode(c, common.CodeDataError, err.Error()) return } - // 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") + common.SuccessWithData(c, gin.H{"total": total, "topics": items}, "success") } // GetArtifactAlteration handles GET /artifacts/alteration — wiki alteration summary. diff --git a/internal/ingestion/component/chunker/token.go b/internal/ingestion/component/chunker/token.go index e8a7fd060d..b74dd8de6a 100644 --- a/internal/ingestion/component/chunker/token.go +++ b/internal/ingestion/component/chunker/token.go @@ -547,10 +547,13 @@ func newChunkText(prevText, incoming string, target int, overlapPct float64, inc // mergeByTokenSize implements exact token-based chunk merging that mirrors // Python's naive_merge (rag/nlp/__init__.py) after the strict chunk_token_num // hard-cap fix. It uses tokenizeStr for precise token counting, treats the -// payload as a single section, splits oversized sections on production sentence -// delimiters, hard-caps atomic oversize units via splitOversizedUnit, and merges -// only when the projected total stays within chunk_token_size. Overlap is -// applied only when the resulting chunk still fits the budget. +// payload as a single section, and splits oversized sections on production +// sentence delimiters. An oversize unit (a single paragraph larger than the +// token budget) is kept whole as a standalone chunk — matching Python OVER_CAP, +// where the model layer truncates it later — instead of being atom-split. +// Sections are merged only when the projected total stays within +// chunk_token_size. Overlap is applied only when the resulting chunk still +// fits the budget. func (c *TokenChunkerComponent) mergeByTokenSize(text string, childrenPattern *regexp.Regexp) map[string]any { target := c.param.ChunkTokenSize overlapPct := c.param.OverlappedPercent @@ -608,18 +611,6 @@ func (c *TokenChunkerComponent) mergeByTokenSize(text string, childrenPattern *r } } - addUnit := func(unit string) { - if tokenizeStr(unit) <= target { - addChunk(unit) - return - } - slog.Debug("TokenChunker: splitting oversized unit via splitOversizedUnit", - "len", len(unit), "tokens", tokenizeStr(unit), "chunk_token_size", target) - for _, piece := range splitOversizedUnit(unit, target) { - addChunk(piece) - } - } - for _, sec := range sections { sec = strings.TrimSpace(sec) if sec == "" { @@ -630,8 +621,11 @@ func (c *TokenChunkerComponent) mergeByTokenSize(text string, childrenPattern *r addChunk(t) continue } - // Oversized section: split on production sentence delimiters, then - // hard-cap any unit that still exceeds the budget (unbroken atoms). + // Oversized section: split on production sentence delimiters into + // units. An oversize unit (still exceeds the budget) is passed through + // addChunk and kept whole — no atom-split, matching Python + // naive_merge. mergeDecision forces an oversize incoming unit to + // startNewChunk, so it stands alone as its own chunk. parts := sentenceDelimiter.Split(sec, -1) hadPart := false for _, part := range parts { @@ -651,10 +645,10 @@ func (c *TokenChunkerComponent) mergeByTokenSize(text string, childrenPattern *r continue } hadPart = true - addUnit("\n" + part) + addChunk("\n" + part) } if !hadPart { - addUnit(t) + addChunk(t) } } diff --git a/internal/ingestion/component/chunker/token_oversize_whole_test.go b/internal/ingestion/component/chunker/token_oversize_whole_test.go new file mode 100644 index 0000000000..a7065cfef7 --- /dev/null +++ b/internal/ingestion/component/chunker/token_oversize_whole_test.go @@ -0,0 +1,139 @@ +// 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 chunker + +import ( + "context" + "strings" + "testing" +) + +// TestTokenChunker_OversizeUnitKeptWhole pins the Python OVER_CAP contract for +// the text/markdown path: a single paragraph that exceeds chunk_token_size is +// kept as one standalone chunk. Python's naive_merge (_merge_paragraph_groups, +// rag/nlp/__init__.py) never atom-splits an oversize unit; it is kept whole and +// the model layer truncates it later. An unbroken input line with no delimiter +// forces the oversize path while isolating it from the delimiter-splitting logic. +func TestTokenChunker_OversizeUnitKeptWhole(t *testing.T) { + var longLine = strings.Repeat("word ", 400) // ~400 tokens, far above the 32 budget + + cases := []struct { + name string + conf map[string]any + input map[string]any + }{ + { + name: "text path", + conf: map[string]any{"chunk_token_size": 32, "delimiters": []string{"\n"}}, + input: map[string]any{ + "name": "t", "output_format": "text", "text": longLine, + }, + }, + { + name: "markdown path", + conf: map[string]any{"chunk_token_size": 32, "delimiters": []string{"\n"}}, + input: map[string]any{ + "name": "t", "output_format": "markdown", "markdown": longLine, + }, + }, + } + + for _, tc := range cases { + t.Run(tc.name, func(t *testing.T) { + c, err := NewTokenChunker(tc.conf) + if err != nil { + t.Fatalf("NewTokenChunker: %v", err) + } + out, err := c.Invoke(context.Background(), nil, tc.input) + if err != nil { + t.Fatalf("Invoke: %v", err) + } + chunks, ok := out["chunks"].([]map[string]any) + if !ok { + t.Fatalf("chunks missing or wrong type: %T", out["chunks"]) + } + if len(chunks) != 1 { + t.Fatalf("oversize unit: want 1 standalone chunk, got %d", len(chunks)) + } + if got, _ := chunks[0]["text"].(string); strings.TrimSpace(got) != strings.TrimSpace(longLine) { + t.Fatalf("oversize unit content not preserved: got %q", got) + } + }) + } +} + +// TestTokenChunker_OversizeUnitStandsAloneAfterInBudgetUnit exercises the +// mergeDecision oversize branch (incomingTokens > target -> startNewChunk), +// which TestTokenChunker_OversizeUnitKeptWhole never reaches because its lone +// oversize unit goes through the len(cks)==0 path of addChunk. A short +// in-budget sentence precedes the oversize paragraph; after the sentence +// delimiter split the oversize unit must stand alone as its own chunk +// (matching Python OVER_CAP), not be merged into or atom-split across the +// previous chunk. +func TestTokenChunker_OversizeUnitStandsAloneAfterInBudgetUnit(t *testing.T) { + var longLine = strings.Repeat("word ", 400) // ~400 tokens, far above the 32 budget + inBudget := "Hello world." // ASCII period is not a sentence delimiter; fits 32 + + cases := []struct { + name string + conf map[string]any + input map[string]any + }{ + { + name: "text path", + conf: map[string]any{"chunk_token_size": 32, "delimiters": []string{"\n"}}, + input: map[string]any{ + "name": "t", "output_format": "text", + "text": inBudget + "\n" + longLine, + }, + }, + { + name: "markdown path", + conf: map[string]any{"chunk_token_size": 32, "delimiters": []string{"\n"}}, + input: map[string]any{ + "name": "t", "output_format": "markdown", + "markdown": inBudget + "\n" + longLine, + }, + }, + } + + for _, tc := range cases { + t.Run(tc.name, func(t *testing.T) { + c, err := NewTokenChunker(tc.conf) + if err != nil { + t.Fatalf("NewTokenChunker: %v", err) + } + out, err := c.Invoke(context.Background(), nil, tc.input) + if err != nil { + t.Fatalf("Invoke: %v", err) + } + chunks, ok := out["chunks"].([]map[string]any) + if !ok { + t.Fatalf("chunks missing or wrong type: %T", out["chunks"]) + } + if len(chunks) != 2 { + t.Fatalf("oversize after in-budget: want 2 chunks (in-budget + standalone oversize), got %d", len(chunks)) + } + first, _ := chunks[0]["text"].(string) + if first != inBudget { + t.Fatalf("first chunk should be only the in-budget sentence %q, got %q", inBudget, first) + } + second, _ := chunks[1]["text"].(string) + if second != strings.TrimSpace(longLine) { + t.Fatalf("oversize chunk should be the whole long line kept whole, got %q", second) + } + }) + } +} diff --git a/internal/ingestion/component/chunker/token_strict_cap_test.go b/internal/ingestion/component/chunker/token_strict_cap_test.go index e5c4a3256e..209cfc5884 100644 --- a/internal/ingestion/component/chunker/token_strict_cap_test.go +++ b/internal/ingestion/component/chunker/token_strict_cap_test.go @@ -257,50 +257,6 @@ func TestMergeByTokenSize_UnderCapNoOverflow(t *testing.T) { } } -func TestMergeByTokenSize_UnbrokenAtomStrictCap(t *testing.T) { - // Unbroken dense string (no whitespace / sentence delim) must still - // hard-cap via the character-window fallback inside splitOversizedUnit. - const budget = 20 - // Use many distinct ASCII letters so cl100k does not collapse the whole - // run into a handful of tokens. - var b strings.Builder - for i := 0; i < 400; i++ { - b.WriteByte(byte('a' + i%26)) - } - text := b.String() - if tokenizeStr(text) <= budget { - t.Skipf("tokenizer collapsed unbroken atom to %d tokens (<= budget)", tokenizeStr(text)) - } - comp, err := NewTokenChunker(map[string]any{ - "delimiter_mode": "token_size", - "chunk_token_size": budget, - }) - if err != nil { - t.Fatalf("NewTokenChunker: %v", err) - } - tc := comp.(*TokenChunkerComponent) - out := tc.mergeByTokenSize(text, nil) - chunks, _ := out["chunks"].([]map[string]any) - if len(chunks) < 2 { - t.Fatalf("want multiple chunks for unbroken atom, got %d (total_tokens=%d)", len(chunks), tokenizeStr(text)) - } - var joined strings.Builder - for i, ck := range chunks { - s, _ := ck["text"].(string) - joined.WriteString(s) - // Sub-split pieces are <= budget+1; OVER_CAP merges at most two before - // closing, so a chunk can reach 2*(budget+1). - if n := tokenizeStr(s); n > 2*(budget+1) { - t.Errorf("chunk %d exceeds 2*(budget+1): tokens=%d text=%q", i, n, s) - } - } - // mergeByTokenSize prefixes "\n" on sections; stripping newlines recovers - // the original unbroken atom. - if strings.ReplaceAll(joined.String(), "\n", "") != text { - t.Errorf("content not preserved after stripping newlines: got %q", joined.String()) - } -} - func TestInvokeTextPayload_StrictCapEndToEnd(t *testing.T) { const budget = 32 unit := tokenizeStr(strings.TrimSpace(strings.Repeat("alpha ", 12))) diff --git a/internal/ingestion/component/knowledge_compiler/common/deps.go b/internal/ingestion/component/knowledge_compiler/common/deps.go index 2c4d00fefa..0f37d2c4be 100644 --- a/internal/ingestion/component/knowledge_compiler/common/deps.go +++ b/internal/ingestion/component/knowledge_compiler/common/deps.go @@ -106,11 +106,10 @@ type Deps struct { Redis RedisClient // optional (datasetnav) TenantID string DatasetID string - // 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 + // 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 } // 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 d5937d0a41..e6f868b91e 100644 --- a/internal/ingestion/component/knowledge_compiler/common/jsonchat.go +++ b/internal/ingestion/component/knowledge_compiler/common/jsonchat.go @@ -4,7 +4,6 @@ import ( "context" "encoding/json" "fmt" - "log" "regexp" "strings" ) @@ -35,13 +34,6 @@ 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)) } @@ -62,20 +54,15 @@ 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, fmt.Errorf("empty candidate") + return nil, false } var m map[string]any if err := json.Unmarshal([]byte(s), &m); err != nil { - return nil, err + return nil, false } - return m, nil + return m, true } 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 e1e1f25ac9..b88f762b93 100644 --- a/internal/ingestion/component/knowledge_compiler/component.go +++ b/internal/ingestion/component/knowledge_compiler/component.go @@ -9,11 +9,9 @@ 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" @@ -39,12 +37,7 @@ var chunkerOutputs = map[string]string{ "_ERROR": "Set only on validation failure.", } -// 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" +const componentNameKnowledgeCompiler = "KnowledgeCompiler" // KnowledgeCompilerComponent is the runtime.Component surface. Param is set at // construction from the DSL; per-call overrides flow through the inputs map. @@ -67,6 +60,10 @@ 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).", } } @@ -88,13 +85,14 @@ 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 - // 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", "") + 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 compilation template spec(s). Priority: // compilation_template_id > compilation_template_group_id. The variant is @@ -417,13 +415,10 @@ 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. 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"). +// routes retrieval-side filters (e.g. "compile_kwd": ["artifact_page"]). var variantCompileKWD = map[common.Variant]string{ common.VariantStructure: "structure", - common.VariantWiki: "wiki_page", + common.VariantWiki: "artifact_page", common.VariantTree: "tree", common.VariantMindmap: "mindmap", } @@ -475,14 +470,6 @@ 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) } @@ -587,25 +574,12 @@ 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. - // 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, "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 } } if v := metaString(p.Meta, "title"); v != "" { @@ -767,21 +741,7 @@ 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 { - // 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"]) - } + raw, _ := inputs["chunks"].([]any) merged := make([]any, 0, len(raw)+len(compiled)) for _, r := range raw { merged = append(merged, r) @@ -812,16 +772,6 @@ 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, @@ -831,45 +781,32 @@ func buildInputs(inputs map[string]any, param common.Param) (common.Inputs, erro if d, ok := inputs["doc_id"].(string); ok && d != "" { in.DocID = d } - // 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 raw, ok := inputs["chunks"].([]any); ok { + for _, r := range raw { + m, ok := r.(map[string]any) if !ok { continue } - 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 := 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) } - in.Chunks = append(in.Chunks, ch) } if hc, ok := inputs["historical_candidates"].([]common.Candidate); ok { in.HistoricalCandidates = hc @@ -887,17 +824,16 @@ func buildInputs(inputs map[string]any, param common.Param) (common.Inputs, erro } func init() { - // 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) + 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, + }) } diff --git a/internal/ingestion/component/knowledge_compiler/component_test.go b/internal/ingestion/component/knowledge_compiler/component_test.go index a9aaa5418a..c46a8847af 100644 --- a/internal/ingestion/component/knowledge_compiler/component_test.go +++ b/internal/ingestion/component/knowledge_compiler/component_test.go @@ -12,8 +12,6 @@ import ( "sync" "testing" - "ragflow/internal/agent/runtime" - "ragflow/internal/ingestion/component/globals" "ragflow/internal/ingestion/component/knowledge_compiler/common" "ragflow/internal/service/nav" @@ -175,7 +173,7 @@ func TestKnowledgeCompiler_Structure_EndToEnd(t *testing.T) { installMockDeps(t) installVariantTemplateResolver(t, "structure") - c, err := NewKnowledgeCompilerComponent("Compiler", map[string]any{ + c, err := NewKnowledgeCompilerComponent("KnowledgeCompiler", map[string]any{ "compilation_template_id": "tpl-structure", "llm_id": "llm1", "embedding_model": "emb1", }) if err != nil { @@ -234,7 +232,7 @@ func TestKnowledgeCompiler_Structure_EndToEnd(t *testing.T) { } func TestKnowledgeCompiler_UnknownVariant(t *testing.T) { - c, err := NewKnowledgeCompilerComponent("Compiler", map[string]any{"compilation_template_id": "nope"}) + c, err := NewKnowledgeCompilerComponent("KnowledgeCompiler", map[string]any{"compilation_template_id": "nope"}) if err != nil { t.Fatalf("NewKnowledgeCompilerComponent: %v", err) } @@ -380,7 +378,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("Compiler", map[string]any{ + c, err := NewKnowledgeCompilerComponent("KnowledgeCompiler", map[string]any{ "compilation_template_id": "mind_map", "llm_id": "llm1", "embedding_model": "emb1", }) if err != nil { @@ -410,7 +408,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("Compiler", params) + c, err := NewKnowledgeCompilerComponent("KnowledgeCompiler", params) if err != nil { t.Fatalf("NewKnowledgeCompilerComponent(%s): %v", variant, err) } @@ -533,7 +531,7 @@ func TestKnowledgeCompiler_EmitsChunks(t *testing.T) { installMockDeps(t) installVariantTemplateResolver(t, "structure") - c, err := NewKnowledgeCompilerComponent("Compiler", map[string]any{ + c, err := NewKnowledgeCompilerComponent("KnowledgeCompiler", map[string]any{ "compilation_template_id": "tpl-structure", "llm_id": "llm1", "embedding_model": "emb1", }) if err != nil { @@ -590,7 +588,7 @@ func TestKnowledgeCompiler_TemplateIDsAndProvenance(t *testing.T) { installMockDeps(t) installVariantTemplateResolver(t, "structure") - c, err := NewKnowledgeCompilerComponent("Compiler", map[string]any{ + c, err := NewKnowledgeCompilerComponent("KnowledgeCompiler", map[string]any{ "compilation_template_id": "tpl-structure", "llm_id": "llm1", "embedding_model": "emb1", @@ -674,7 +672,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("Compiler", map[string]any{ + c, err := NewKnowledgeCompilerComponent("KnowledgeCompiler", map[string]any{ "compilation_template_id": "tpl-tree", "llm_id": "llm1", "embedding_model": "emb1", "extra": map[string]any{"tree_order": 4}, }) @@ -729,7 +727,7 @@ func TestKnowledgeCompiler_Wiki_HistoricalDedupDropsDuplicates(t *testing.T) { t.Cleanup(func() { common.SetDepsResolver(nil) }) installVariantTemplateResolver(t, "wiki") - c, err := NewKnowledgeCompilerComponent("Compiler", map[string]any{ + c, err := NewKnowledgeCompilerComponent("KnowledgeCompiler", map[string]any{ "compilation_template_id": "tpl-wiki", "llm_id": "llm1", "embedding_model": "emb1", }) if err != nil { @@ -788,7 +786,7 @@ func TestKnowledgeCompiler_Wiki_UpdateMergesExistingPage(t *testing.T) { t.Cleanup(func() { common.SetDepsResolver(nil) }) installVariantTemplateResolver(t, "wiki") - c, err := NewKnowledgeCompilerComponent("Compiler", map[string]any{ + c, err := NewKnowledgeCompilerComponent("KnowledgeCompiler", map[string]any{ "compilation_template_id": "tpl-wiki", "llm_id": "llm1", "embedding_model": "emb1", }) if err != nil { @@ -815,7 +813,7 @@ func TestKnowledgeCompiler_Wiki_UpdateMergesExistingPage(t *testing.T) { if !ok { continue } - if cm["compile_kwd"] == "wiki_page" && cm["kc_kind"] == "page" && cm["slug_kwd"] == "entity/alpha" { + if cm["compile_kwd"] == "artifact_page" && cm["kc_kind"] == "page" && cm["slug_kwd"] == "entity/alpha" { page = cm break } @@ -869,7 +867,7 @@ func TestKnowledgeCompiler_Wiki_HistoricalDedupScopedByDataset(t *testing.T) { t.Cleanup(func() { common.SetDepsResolver(nil) }) installVariantTemplateResolver(t, "wiki") - c, err := NewKnowledgeCompilerComponent("Compiler", map[string]any{ + c, err := NewKnowledgeCompilerComponent("KnowledgeCompiler", map[string]any{ "compilation_template_id": "tpl-wiki", "llm_id": "llm1", "embedding_model": "emb1", "enable_historical_dedup": true, }) @@ -984,7 +982,7 @@ func TestKnowledgeCompiler_Structure_FencedJSONNotDropped(t *testing.T) { t.Cleanup(func() { common.SetDepsResolver(nil) }) installVariantTemplateResolver(t, "structure") - c, err := NewKnowledgeCompilerComponent("Compiler", map[string]any{ + c, err := NewKnowledgeCompilerComponent("KnowledgeCompiler", map[string]any{ "compilation_template_id": "tpl-structure", "llm_id": "llm1", "embedding_model": "emb1", }) if err != nil { @@ -1024,7 +1022,7 @@ func TestKnowledgeCompiler_Structure_MalformedJSONFailsLoud(t *testing.T) { t.Cleanup(func() { common.SetDepsResolver(nil) }) installVariantTemplateResolver(t, "structure") - c, err := NewKnowledgeCompilerComponent("Compiler", map[string]any{ + c, err := NewKnowledgeCompilerComponent("KnowledgeCompiler", map[string]any{ "compilation_template_id": "tpl-structure", "llm_id": "llm1", "embedding_model": "emb1", }) if err != nil { @@ -1050,7 +1048,7 @@ func TestKnowledgeCompiler_PassThroughEnvelope(t *testing.T) { installMockDeps(t) installVariantTemplateResolver(t, "structure") - c, err := NewKnowledgeCompilerComponent("Compiler", map[string]any{ + c, err := NewKnowledgeCompilerComponent("KnowledgeCompiler", map[string]any{ "compilation_template_id": "tpl-structure", "llm_id": "llm1", "embedding_model": "emb1", }) if err != nil { @@ -1118,7 +1116,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("Compiler", map[string]any{ + c, err := NewKnowledgeCompilerComponent("KnowledgeCompiler", map[string]any{ "compilation_template_group_id": "grp1", "llm_id": "llm1", "embedding_model": "emb1", @@ -1179,7 +1177,7 @@ func TestKnowledgeCompiler_GroupIDsWithoutResolverFailsLoud(t *testing.T) { installMockDeps(t) common.SetGroupResolver(nil) // ensure no resolver is installed - c, err := NewKnowledgeCompilerComponent("Compiler", map[string]any{ + c, err := NewKnowledgeCompilerComponent("KnowledgeCompiler", map[string]any{ "compilation_template_group_id": "grp1", "llm_id": "llm1", "embedding_model": "emb1", @@ -1234,148 +1232,6 @@ 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 7ffa3b16d2..09ed27dcc5 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("Compiler", params) + c, err := NewKnowledgeCompilerComponent("KnowledgeCompiler", 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 3d4a92b7f8..6eb4ec017a 100644 --- a/internal/ingestion/component/knowledge_compiler/pool_wiring.go +++ b/internal/ingestion/component/knowledge_compiler/pool_wiring.go @@ -20,7 +20,6 @@ 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" ) @@ -36,5 +35,4 @@ 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 5813e2e460..697b1a845e 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.ModelContextLen, maxToken) + content := buildClusterContent(texts, task.pointIdxs, deps.LLMMaxLength, 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.ModelContextLen, maxToken) + rootContent := buildClusterContent(topLevelTexts, allIndices(len(topLevelTexts)), deps.LLMMaxLength, 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, modelContextLen, maxToken int) string { +func buildClusterContent(texts []string, idxs []int, llmMaxLength, maxToken int) string { if len(idxs) == 0 { return "" } - if modelContextLen <= 0 { - modelContextLen = common.DefaultLLMContextLength + if llmMaxLength <= 0 { + llmMaxLength = common.DefaultLLMContextLength } - per := (modelContextLen - maxToken) / len(idxs) + per := (llmMaxLength - 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 e4dcab3216..5b8c563c7d 100644 --- a/internal/ingestion/component/knowledge_compiler/wiki/page_test.go +++ b/internal/ingestion/component/knowledge_compiler/wiki/page_test.go @@ -202,7 +202,6 @@ 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) @@ -212,26 +211,6 @@ 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 43f9cac2c9..6da9738c75 100644 --- a/internal/ingestion/component/knowledge_compiler/wiki/prompt.go +++ b/internal/ingestion/component/knowledge_compiler/wiki/prompt.go @@ -15,26 +15,6 @@ 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 @@ -154,16 +134,23 @@ 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 1224c0177f..2cb95570af 100644 --- a/internal/ingestion/component/knowledge_compiler/wiki/wiki.go +++ b/internal/ingestion/component/knowledge_compiler/wiki/wiki.go @@ -21,55 +21,6 @@ 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 @@ -84,12 +35,6 @@ 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 { @@ -155,33 +100,8 @@ 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 @@ -279,10 +199,6 @@ 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 @@ -297,66 +213,27 @@ func (p *wikiPipeline) run() error { } func (p *wikiPipeline) runMap() error { - // 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 + 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) } - 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 @@ -369,85 +246,31 @@ func (p *wikiPipeline) runPlan() (wikiPlan, error) { if len(batches) == 0 { batches = []wikiExtract{p.reduced} } - 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)) + plans := make([]wikiPlan, 0, len(batches)) for i, batch := range batches { - 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 + plan, err := p.runPlanBatch(batch, i+1, len(batches)) + if err != nil { + return wikiPlan{}, err } - 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 - }) + plans = append(plans, plan) } - if err := runBatches(p.ctx, jobs); err != nil { - return wikiPlan{}, err + if len(plans) == 1 { + plan := normalizeWikiPlan(plans[0], p.docID, p.reduced) + return p.reconcilePlan(plan) } - - 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) + plan, err := p.mergePlanCandidates(plans) if err != nil { return wikiPlan{}, err } - // 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 + return p.reconcilePlan(plan) } -// 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 == "" { @@ -455,155 +278,108 @@ 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, 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 + 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+"]]") } - res, err := p.runRefinePage(planItem, allPlanSlugs, pageTitles, slugToPageType, entityLookup, conceptLookup) + } + 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 err + return nil, err } - results[i] = res - return nil + } + 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, }) - } - if err := runBatches(p.ctx, jobs); err != nil { - return nil, err } return results, nil } -// 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) { +func (p *wikiPipeline) runPlanBatch(batch wikiExtract, batchIndex, batchTotal 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), @@ -621,52 +397,20 @@ func (p *wikiPipeline) runPlanBatch(batch wikiExtract, batchIndex, batchTotal, q return parseWikiPlan(raw, p.docID, batch), nil } -// 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) - } - } +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 } - merged = normalizeWikiPlan(merged, p.docID, reduced) - merged.Pages = normalizeWikiPlanPageLinks(merged.Pages) - return merged + return normalizeWikiPlan(parseWikiPlan(raw, p.docID, p.reduced), p.docID, p.reduced), nil } func (p *wikiPipeline) reconcilePlan(plan wikiPlan) (wikiPlan, error) { @@ -717,9 +461,6 @@ 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 @@ -1202,34 +943,6 @@ 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 == "" { @@ -1876,24 +1589,7 @@ var ( wikiArtifactMarkdownLink = regexp.MustCompile(`\[([^\]]+)\]\(([^)]+)\)`) ) -// 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) { +func transformWikiLinks(content, kbID string, pageTitles map[string]string) (string, []string) { kbID = strings.TrimSpace(kbID) seen := map[string]bool{} var outlinks []string @@ -1939,32 +1635,13 @@ func transformWikiLinks(content, kbID string, pageTitles, slugToPageType map[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 { - pageType := "" - bareSlug := slug - if idx := strings.Index(slug, "/"); idx >= 0 { - pageType = slug[:idx] - bareSlug = slug[idx+1:] - } else { - pageType = pageTypeOf(slug, slugToPageType) - } - 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 link(displayText(label, slug), slug) + return "[" + displayText(label, slug) + "](artifact/" + kbID + "/" + slug + ")" } out := wikiArtifactMarkdownLink.ReplaceAllStringFunc(content, func(match string) string { sub := wikiArtifactMarkdownLink.FindStringSubmatch(match) @@ -1980,7 +1657,7 @@ func transformWikiLinks(content, kbID string, pageTitles, slugToPageType map[str } slug := strings.TrimSpace(sub[1]) track(slug) - return link(sub[2], slug) + return "[" + sub[2] + "](artifact/" + kbID + "/" + slug + ")" }) out = wikiWikilinkSimpleRe.ReplaceAllStringFunc(out, func(match string) string { sub := wikiWikilinkSimpleRe.FindStringSubmatch(match) @@ -1989,7 +1666,7 @@ func transformWikiLinks(content, kbID string, pageTitles, slugToPageType map[str } slug := strings.TrimSpace(sub[1]) track(slug) - return link(displayText(slug, slug), slug) + return "[" + displayText(slug, slug) + "](artifact/" + kbID + "/" + slug + ")" }) return out, outlinks } @@ -2197,6 +1874,15 @@ 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 deleted file mode 100644 index c1807d3b4b..0000000000 --- a/internal/ingestion/component/knowledge_compiler/wiki/wiki_budget.go +++ /dev/null @@ -1,259 +0,0 @@ -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 deleted file mode 100644 index 2a26cb2c36..0000000000 --- a/internal/ingestion/component/knowledge_compiler/wiki/wiki_budget_test.go +++ /dev/null @@ -1,559 +0,0 @@ -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 deleted file mode 100644 index d506d630dd..0000000000 --- a/internal/ingestion/component/knowledge_compiler/wiki/wiki_reduce.go +++ /dev/null @@ -1,185 +0,0 @@ -package wiki - -import ( - "strings" - - "ragflow/internal/ingestion/component/knowledge_compiler/common" -) - -// This file implements the REDUCE-stage canonical-entity enhancement that -// narrows the gap with Python's wiki.py canonicalization: -// -// - entities with distinct names but high embedding similarity are treated as -// ambiguous and sent to an LLM merge decision (collapsing near-duplicates); -// - concepts keep exact-term dedup, matching Python's current semantic (any -// embedding/LLM dedup for concepts must be a separate new capability with -// its own quality bar, not an alignment claim). -// -// The exact-key merge in reduceExtracts stays the deterministic baseline; this -// step layers embedding + LLM disambiguation on top. When the embedder or chat -// seam is unavailable, entities pass through unchanged (degrade gracefully). - -// wikiEntityMergeThreshold is the embedding-cosine similarity at or above which -// two distinct-name entities are considered ambiguous and sent to the LLM merge -// decision. It is deliberately high so only genuinely similar candidates reach -// the LLM. -const wikiEntityMergeThreshold = 0.85 - -// wikiEntityMergeMaxCalls caps how many LLM disambiguation calls a single -// REDUCE run may make, bounding the cost on entity-dense documents. -const wikiEntityMergeMaxCalls = 16 - -// wikiEntityMergeMaxCandidates caps how many candidate partners one entity is -// checked against to keep the pairwise scan bounded. -const wikiEntityMergeMaxCandidates = 8 - -// dedupeEntities returns a copy of in with ambiguous near-duplicate entities -// collapsed via LLM disambiguation. It is a no-op when fewer than two entities -// are present or when deps.Embed / deps.Chat are unavailable. -func (p *wikiPipeline) dedupeEntities(in []wikiEntity) []wikiEntity { - if len(in) < 2 || p.deps.Embed == nil || p.deps.Chat == nil { - return in - } - names := make([]string, len(in)) - for i, e := range in { - names[i] = e.Name - } - vecs, err := p.deps.Embed.Encode(p.ctx, names) - if err != nil || len(vecs) != len(in) { - return in - } - - // Canonical entity per input index: which index owns the final entity. - canon := make([]int, len(in)) - for i := range canon { - canon[i] = i - } - llmCalls := 0 - - // Greedy best-partner scan in input order (already deterministic after - // reduceExtracts sorts by name). - for i := 0; i < len(in) && llmCalls < wikiEntityMergeMaxCalls; i++ { - if canon[i] != i { - // Already merged into another canonical entity. - continue - } - bestIdx, bestSim := -1, -1.0 - checked := 0 - for j := 0; j < len(in) && checked < wikiEntityMergeMaxCandidates; j++ { - if i == j { - continue - } - if canon[j] != j { - // Consumed by an earlier merge; never a standalone partner. - continue - } - if normKey(in[i].Name) == normKey(in[j].Name) { - // Exact-name duplicates are already merged by reduceExtracts; - // never treat them as a pair here. - continue - } - // Same-type-only candidate filtering (Python canonicalizes entities - // within the same type). Two entities with provably different types - // are never ambiguous regardless of embedding similarity. An empty - // type is treated as compatible (cannot prove a difference). - if in[i].Type != "" && in[j].Type != "" && !strings.EqualFold(in[i].Type, in[j].Type) { - continue - } - checked++ - sim := cosine32(vecs[i], vecs[j]) - if sim >= wikiEntityMergeThreshold && sim > bestSim { - bestSim = sim - bestIdx = j - } - } - if bestIdx < 0 { - continue - } - // i and bestIdx are guaranteed standalone by the loop guards above. - // Count the call BEFORE issuing it so a persistent failure cannot drive - // unbounded external requests: the llmCalls budget is consumed even when - // the request fails. The outer loop's `llmCalls < max` guard then stops - // further iterations once the budget is exhausted. - llmCalls++ - merge, err := p.llmMergeEntityDecision(in[i], in[bestIdx]) - if err != nil { - // A failed disambiguation call should not abort the whole REDUCE; - // keep the entities separate and move on. - continue - } - if !merge { - continue - } - // Merge j into i: i is canonical, j is consumed. - canon[bestIdx] = i - in[i].Aliases = mergeStrings(in[i].Aliases, in[bestIdx].Aliases) - if in[i].Name != in[bestIdx].Name { - in[i].Aliases = mergeStrings(in[i].Aliases, []string{in[bestIdx].Name}) - } - in[i].SourceChunkIDs = mergeStrings(in[i].SourceChunkIDs, in[bestIdx].SourceChunkIDs) - if in[i].Type == "" { - in[i].Type = in[bestIdx].Type - } - } - - out := make([]wikiEntity, 0, len(in)) - for i := range in { - if canon[i] != i { - continue - } - out = append(out, in[i]) - } - return out -} - -// llmMergeEntityDecision asks the chat seam whether two distinct-name entities -// refer to the same real-world concept. Returns true to merge. -func (p *wikiPipeline) llmMergeEntityDecision(a, b wikiEntity) (bool, error) { - 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 deleted file mode 100644 index 6f621535fe..0000000000 --- a/internal/ingestion/component/knowledge_compiler/wiki/wiki_reduce_test.go +++ /dev/null @@ -1,225 +0,0 @@ -package wiki - -import ( - "context" - "errors" - "testing" - - "ragflow/internal/ingestion/component/knowledge_compiler/common" -) - -// TestDedupeEntities_NoSeamIsNoop verifies dedupeEntities degrades to a no-op -// when the embedder or chat seam is nil (M1-style unit safety). -func TestDedupeEntities_NoSeamIsNoop(t *testing.T) { - p := &wikiPipeline{ctx: context.Background()} - in := []wikiEntity{ - {Name: "Alpha", SourceChunkIDs: []string{"c1"}}, - {Name: "Alpha Corp", SourceChunkIDs: []string{"c2"}}, - } - got := p.dedupeEntities(in) - if len(got) != 2 { - t.Fatalf("got %d entities, want 2 (no-op without seams)", len(got)) - } -} - -// TestDedupeEntities_LLMMergesAmbiguousPair verifies two distinct-name entities -// with high embedding similarity are collapsed into one canonical entity via the -// LLM merge decision, with aliases and provenance merged. -func TestDedupeEntities_LLMMergesAmbiguousPair(t *testing.T) { - p := &wikiPipeline{ - ctx: context.Background(), - llmID: "llm1", - deps: common.Deps{ - Chat: reconcileChatStub{resp: `{"merge":true,"reason":"same company"}`}, - Embed: mergeEmbedStub{}, - }, - } - in := []wikiEntity{ - {Name: "Alpha Inc", Type: "org", SourceChunkIDs: []string{"c1"}}, - {Name: "Alpha Incorporated", Type: "org", SourceChunkIDs: []string{"c2"}}, - } - got := p.dedupeEntities(in) - if len(got) != 1 { - t.Fatalf("got %d entities, want 1 (LLM merge)", len(got)) - } - if got[0].Name != "Alpha Inc" { - t.Fatalf("canonical name = %q, want Alpha Inc", got[0].Name) - } - if len(got[0].SourceChunkIDs) != 2 { - t.Fatalf("provenance = %#v, want 2 chunk ids", got[0].SourceChunkIDs) - } - if len(got[0].Aliases) == 0 { - t.Fatalf("aliases not merged: %#v", got[0].Aliases) - } -} - -// TestDedupeEntities_LLMRejectsDistinct verifies the LLM rejecting a merge keeps -// both entities distinct. -func TestDedupeEntities_LLMRejectsDistinct(t *testing.T) { - p := &wikiPipeline{ - ctx: context.Background(), - llmID: "llm1", - deps: common.Deps{ - Chat: reconcileChatStub{resp: `{"merge":false,"reason":"distinct products"}`}, - Embed: mergeEmbedStub{}, - }, - } - in := []wikiEntity{ - {Name: "Alpha", SourceChunkIDs: []string{"c1"}}, - {Name: "Beta", SourceChunkIDs: []string{"c2"}}, - } - got := p.dedupeEntities(in) - if len(got) != 2 { - t.Fatalf("got %d entities, want 2 (LLM rejected merge)", len(got)) - } -} - -// TestDedupeEntities_ExactNameIsNotAmbiguous verifies entities with identical -// normalized names (already collapsed by reduceExtracts before this stage) are -// not treated as ambiguous: they pass through untouched and no LLM call is made -// for the exact-name pair. -func TestDedupeEntities_ExactNameIsNotAmbiguous(t *testing.T) { - p := &wikiPipeline{ - ctx: context.Background(), - llmID: "llm1", - deps: common.Deps{ - Chat: reconcileChatStub{resp: `{"merge":true}`}, - Embed: mergeEmbedStub{}, - }, - } - in := []wikiEntity{ - {Name: "Alpha", SourceChunkIDs: []string{"c1"}}, - {Name: "Alpha", SourceChunkIDs: []string{"c2"}}, - } - got := p.dedupeEntities(in) - // Exact-name duplicates are out of scope for the embedding step; both are - // kept unchanged (they would already be one entity after reduceExtracts). - if len(got) != 2 { - t.Fatalf("got %d entities, want 2 (exact-name pairs are not ambiguous)", len(got)) - } -} - -// TestDedupeEntities_ConceptStaysExact validates the REDUCE boundary: concept -// dedup must remain exact (no embedding/LLM). This test guards that the entity -// enhancement never touches concepts. -func TestReduceExtracts_ConceptsStayExact(t *testing.T) { - reduced := reduceExtracts([]wikiExtract{ - {Concepts: []wikiConcept{{Term: "RAG", Definition: "d1", SourceChunkIDs: []string{"c1"}}}}, - {Concepts: []wikiConcept{{Term: "Retrieval Augmented Generation", Definition: "d2", SourceChunkIDs: []string{"c2"}}}}, - }) - if len(reduced.Concepts) != 2 { - t.Fatalf("concepts = %d, want 2 (exact-term dedup must not collapse distinct terms)", len(reduced.Concepts)) - } -} - -// TestDedupeEntities_FailingChatIsBudgeted locks F4: a chat seam that -// persistently fails must not drive unbounded external calls. The llmCalls -// budget is consumed before the request, so the loop stops after at most -// wikiEntityMergeMaxCalls attempts. -func TestDedupeEntities_FailingChatIsBudgeted(t *testing.T) { - calls := 0 - p := &wikiPipeline{ - ctx: context.Background(), - llmID: "llm1", - deps: common.Deps{ - Chat: chatFunc(func(_ context.Context, _ common.ChatRequest) (*common.ChatResponse, error) { - calls++ - return nil, errors.New("llm down") - }), - Embed: mergeEmbedStub{}, - }, - } - // 20 entities with distinct names but identical embeddings => every pair is - // ambiguous and would trigger an LLM call. - in := make([]wikiEntity, 0, 20) - for i := 0; i < 20; i++ { - in = append(in, wikiEntity{Name: "Entity " + itoa(i), Type: "person"}) - } - got := p.dedupeEntities(in) - if calls > wikiEntityMergeMaxCalls { - t.Fatalf("chat calls = %d, want <= %d despite persistent failures", calls, wikiEntityMergeMaxCalls) - } - // No merges happen because every call fails, so all 20 entities survive. - if len(got) != 20 { - t.Fatalf("entities = %d, want 20 (no merges on failure)", len(got)) - } -} - -// TestDedupeEntities_CrossTypeHighSimDoesNotCallLLM locks F5: entities with -// provably different types must never be treated as ambiguous, even with -// identical embeddings, so no LLM call is made for them. -func TestDedupeEntities_CrossTypeHighSimDoesNotCallLLM(t *testing.T) { - calls := 0 - p := &wikiPipeline{ - ctx: context.Background(), - llmID: "llm1", - deps: common.Deps{ - Chat: chatFunc(func(_ context.Context, _ common.ChatRequest) (*common.ChatResponse, error) { - calls++ - return &common.ChatResponse{Content: `{"merge":true}`}, nil - }), - Embed: mergeEmbedStub{}, - }, - } - // Both embed to [1,1,1] (identical vectors, cosine 1.0) but types differ. - in := []wikiEntity{ - {Name: "Alpha", Type: "person"}, - {Name: "Beta Corp", Type: "org"}, - } - got := p.dedupeEntities(in) - if calls != 0 { - t.Fatalf("chat calls = %d, want 0 (cross-type pairs must not reach the LLM)", calls) - } - if len(got) != 2 { - t.Fatalf("entities = %d, want 2 (cross-type entities must stay distinct)", len(got)) - } -} - -// mergeEmbedStub returns embeddings where identical names share a vector and -// distinct names are far apart (cosine ~0), so it can drive both the ambiguous -// and distinct test paths deterministically. -type mergeEmbedStub struct{} - -func (mergeEmbedStub) Encode(_ context.Context, texts []string) ([][]float32, error) { - out := make([][]float32, len(texts)) - for i, text := range texts { - // "Alpha Inc" and "Alpha Incorporated" both contain "alpha" -> same - // vector; "Beta" differs. - if containsFold(text, "beta") { - out[i] = []float32{1, 0, 0} - continue - } - out[i] = []float32{1, 1, 1} - } - return out, nil -} - -func (mergeEmbedStub) Dimensions() int { return 3 } - -func containsFold(s, sub string) bool { - return len(s) >= len(sub) && (len(sub) == 0 || indexFold(s, sub) >= 0) -} - -func indexFold(s, sub string) int { - if sub == "" { - return 0 - } - ls := toLowerASCII(s) - lsub := toLowerASCII(sub) - for i := 0; i+len(lsub) <= len(ls); i++ { - if ls[i:i+len(lsub)] == lsub { - return i - } - } - return -1 -} - -func toLowerASCII(s string) string { - b := []byte(s) - for i := range b { - if b[i] >= 'A' && b[i] <= 'Z' { - b[i] += 'a' - 'A' - } - } - return string(b) -} diff --git a/internal/ingestion/component/knowledge_compiler/wiki/wiki_refine_test.go b/internal/ingestion/component/knowledge_compiler/wiki/wiki_refine_test.go deleted file mode 100644 index a8bf19bedf..0000000000 --- a/internal/ingestion/component/knowledge_compiler/wiki/wiki_refine_test.go +++ /dev/null @@ -1,156 +0,0 @@ -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 2cf03c3b07..5f233ba648 100644 --- a/internal/ingestion/component/knowledge_compiler/wiki/wiki_test.go +++ b/internal/ingestion/component/knowledge_compiler/wiki/wiki_test.go @@ -3,9 +3,7 @@ package wiki import ( "context" "strings" - "sync" "testing" - "time" "ragflow/internal/ingestion/component/knowledge_compiler/common" ) @@ -46,77 +44,6 @@ 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"}, @@ -144,61 +71,6 @@ 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 } @@ -262,48 +134,6 @@ 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/knowledge_compile/consumer.go b/internal/ingestion/knowledge_compile/consumer.go index 2a663c7c87..1d8bb6c3df 100644 --- a/internal/ingestion/knowledge_compile/consumer.go +++ b/internal/ingestion/knowledge_compile/consumer.go @@ -21,11 +21,8 @@ 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 @@ -186,19 +183,10 @@ 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 { - common.Warn("knowledge_compile: ack failed", - zap.String("dataset_id", datasetID), zap.Error(err)) + _ = err } } @@ -207,10 +195,6 @@ 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{} @@ -371,7 +355,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 { @@ -453,10 +437,5 @@ 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 29f0cec042..2c3c9d9035 100644 --- a/internal/ingestion/knowledge_compile/consumer_test.go +++ b/internal/ingestion/knowledge_compile/consumer_test.go @@ -249,154 +249,3 @@ 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 6cc775769c..6c0cec3993 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 6b527dffd3..b6f5dd31a7 100644 --- a/internal/ingestion/knowledge_compile/pool.go +++ b/internal/ingestion/knowledge_compile/pool.go @@ -24,11 +24,10 @@ 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 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 +// 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 // compilerPool is the process-wide bounded worker pool that drives cross-doc // concurrency for every knowledge-compilation stage: the DocEngine KNN pass in @@ -44,10 +43,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 @@ -83,11 +82,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) @@ -122,7 +121,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 @@ -139,10 +138,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 a8deeee6ae..3b0332a85e 100644 --- a/internal/ingestion/knowledge_compile/reader.go +++ b/internal/ingestion/knowledge_compile/reader.go @@ -68,16 +68,6 @@ 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. @@ -102,7 +92,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: append(append([]string(nil), compiledSelectFields...), wikiSelectFields...), + SelectFields: compiledSelectFields, Limit: loadDocProductsLimit, Offset: offset, }) @@ -158,36 +148,8 @@ 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 } @@ -234,10 +196,9 @@ func (r engineReader) SearchSimilar(ctx context.Context, tenant, kb string, vari IndexNames: []string{fmt.Sprintf("ragflow_%s", tenant)}, KbIDs: []string{kb}, Limit: topN, - SelectFields: append([]string{"id", "doc_id", "kb_id", "content_with_weight", "kc_payload", + SelectFields: []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 03e9f15be4..323311877c 100644 --- a/internal/ingestion/knowledge_compile/scheduler.go +++ b/internal/ingestion/knowledge_compile/scheduler.go @@ -23,11 +23,9 @@ import ( "sync" "time" - "ragflow/internal/common" "ragflow/internal/engine" "ragflow/internal/entity" - "go.uber.org/zap" "gorm.io/gorm" "gorm.io/gorm/clause" ) @@ -43,15 +41,6 @@ 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 @@ -106,14 +95,6 @@ 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) @@ -216,23 +197,11 @@ 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) } @@ -269,17 +238,9 @@ 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 } @@ -325,61 +286,18 @@ 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 @@ -390,10 +308,6 @@ 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 } @@ -491,14 +405,9 @@ 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 @@ -511,13 +420,7 @@ func (s *mysqlScheduler) notify(ctx context.Context, datasetID string) error { return nil } payload, _ := json.Marshal(map[string]string{"dataset_id": datasetID}) - 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 + return s.mq.PublishKnowledgeCompile(notifySubject, payload) } func (s *mysqlScheduler) SubscribeNotify(ctx context.Context) (<-chan string, error) { @@ -558,8 +461,6 @@ 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 @@ -597,10 +498,6 @@ 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: @@ -631,8 +528,6 @@ 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 } @@ -650,12 +545,6 @@ 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 } @@ -671,39 +560,6 @@ 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) { @@ -717,7 +573,18 @@ func (f *FakeScheduler) TryClaim(ctx context.Context) (ClaimResult, bool, error) } } if readyID == "" { - expiredID = f.fakeReclaimExpired(now) + 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 + } + } } f.mu.Unlock() if readyID != "" { diff --git a/internal/ingestion/knowledge_compile/service.go b/internal/ingestion/knowledge_compile/service.go index 8141605b9c..7faedcd0f4 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.ModelContextLen), nil + return NewLLMDeduper(deps.Chat, deps.Embed, defaultLLMID, 0.99, deps.LLMMaxLength), nil } func generateHolder() string { diff --git a/internal/ingestion/knowledge_compile/writer.go b/internal/ingestion/knowledge_compile/writer.go index 45e6fded28..126e13d71c 100644 --- a/internal/ingestion/knowledge_compile/writer.go +++ b/internal/ingestion/knowledge_compile/writer.go @@ -20,7 +20,6 @@ import ( "crypto/sha256" "encoding/hex" "fmt" - "strings" "ragflow/internal/engine" "ragflow/internal/engine/types" @@ -71,7 +70,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) { @@ -109,47 +108,6 @@ 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 @@ -208,7 +166,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{ @@ -309,15 +267,6 @@ 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: @@ -333,23 +282,3 @@ 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 deleted file mode 100644 index 8f7e9bfe42..0000000000 --- a/internal/ingestion/knowledge_compile/writer_test.go +++ /dev/null @@ -1,98 +0,0 @@ -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 7a23c9e57f..66b76f356d 100644 --- a/internal/ingestion/pipeline/pipeline_knowledge_compiler_dsl_test.go +++ b/internal/ingestion/pipeline/pipeline_knowledge_compiler_dsl_test.go @@ -95,12 +95,10 @@ 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. 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"]) + // 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"]) } } @@ -139,11 +137,8 @@ func TestKnowledgeCompilerDSL_FrontendDSLDecodesAndConstructs(t *testing.T) { t.Fatal("default runtime factory not installed") } - // 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) + // The fixture params (single group id, no variant) construct fine. + comp, err := f("KnowledgeCompiler", params) if err != nil { t.Fatalf("construct from fixture params: %v", err) } @@ -159,15 +154,16 @@ func TestKnowledgeCompilerDSL_FrontendDSLDecodesAndConstructs(t *testing.T) { } // TestKnowledgeCompilerDSL_RegisteredAndConstructible confirms the Go runtime -// 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). +// 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. func TestKnowledgeCompilerDSL_RegisteredAndConstructible(t *testing.T) { runtime.InstallDefaultRegistryFactory() - if _, _, _, ok := runtime.DefaultRegistry.Lookup("Compiler"); !ok { - t.Fatal("Compiler not registered in the runtime factory") + if _, _, _, ok := runtime.DefaultRegistry.Lookup("KnowledgeCompiler"); !ok { + t.Fatal("KnowledgeCompiler not registered in the runtime factory") } f := runtime.DefaultFactory() if f == nil { @@ -181,7 +177,7 @@ func TestKnowledgeCompilerDSL_RegisteredAndConstructible(t *testing.T) { {"compilation_template_id": "t1", "compilation_template_group_id": "g1"}, } for i, params := range cases { - comp, err := f("Compiler", params) + comp, err := f("KnowledgeCompiler", params) if err != nil { t.Fatalf("case %d construct: %v", i, err) } @@ -191,7 +187,7 @@ func TestKnowledgeCompilerDSL_RegisteredAndConstructible(t *testing.T) { } // Param map with neither id resolves to a parse error. - if _, err := f("Compiler", map[string]any{}); err == nil { + if _, err := f("KnowledgeCompiler", map[string]any{}); err == nil { t.Fatal("construct with no template spec: expected error") } } @@ -207,8 +203,6 @@ 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"}, }) @@ -224,8 +218,6 @@ 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 a2422f837d..2c7fd5d797 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", "Compiler"} { + for _, name := range []string{"File", "Parser", "TokenChunker", "KnowledgeCompiler"} { 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 364bc66b09..6b9bf3f5da 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 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 形式合并进上游分块流输出,适合在分块文档之上构建可检索的知识层。" + "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 形式合并进上游分块流输出,适合在分块文档之上构建可检索的知识层。" }, "canvas_type": "Ingestion Pipeline", "canvas_category": "dataflow_canvas", @@ -150,7 +150,7 @@ }, "TokenChunker:SixApplesFall": { "downstream": [ - "Compiler:KnownSwiftLions" + "KnowledgeCompiler:KnownSwiftLions" ], "obj": { "component_name": "TokenChunker", @@ -182,12 +182,11 @@ "Parser:HipSignsRhyme" ] }, - "Compiler:KnownSwiftLions": { + "KnowledgeCompiler:KnownSwiftLions": { "downstream": [], "obj": { - "component_name": "Compiler", + "component_name": "KnowledgeCompiler", "params": { - "llm_id": "", "variant": "structure", "language": "English" } @@ -217,10 +216,10 @@ "targetHandle": "end" }, { - "id": "xy-edge__TokenChunker:SixApplesFallstart-Compiler:KnownSwiftLionsend", + "id": "xy-edge__TokenChunker:SixApplesFallstart-KnowledgeCompiler:KnownSwiftLionsend", "source": "TokenChunker:SixApplesFall", "sourceHandle": "start", - "target": "Compiler:KnownSwiftLions", + "target": "KnowledgeCompiler:KnownSwiftLions", "targetHandle": "end" } ], @@ -281,10 +280,10 @@ }, { "data": { - "label": "Compiler", + "label": "KnowledgeCompiler", "name": "Knowledge Compiler_0" }, - "id": "Compiler:KnownSwiftLions", + "id": "KnowledgeCompiler: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 e5e43a2a21..bee1d14632 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, "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") + 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") } terminalIDs := terminalComponentIDsFromTemplate(t, templateBytes) if len(terminalIDs) != 1 { diff --git a/internal/ingestion/task/knowledge_compiler_wiring.go b/internal/ingestion/task/knowledge_compiler_wiring.go index f1c46c73e5..311e4e1970 100644 --- a/internal/ingestion/task/knowledge_compiler_wiring.go +++ b/internal/ingestion/task/knowledge_compiler_wiring.go @@ -28,9 +28,7 @@ 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" @@ -99,15 +97,12 @@ 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.ResolveModelContextLength(ctx, tenantID, llmID); merr == nil && ml > 0 { + if _, _, _, ml, merr := svc.ResolveModelConfig(ctx, tenantID, entity.ModelTypeChat, llmID); merr == nil && ml > 0 { llmMax = ml } @@ -118,7 +113,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. - ModelContextLen: llmMax, + LLMMaxLength: llmMax, }, nil } } @@ -179,116 +174,33 @@ func (e *kcEmbedder) Encode(ctx context.Context, texts []string) ([][]float32, e if len(texts) == 0 { return nil, nil } - mdl, err := e.resolveModel(ctx) + 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) if err != nil { - return nil, err + 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) } config := &models.EmbeddingConfig{} - // Slice inputs into per-provider batches: providers cap the per-request input - // count (siliconflow allows at most 32, siliconflow.go:143) and reject larger - // batches rather than chunking internally. The batch size is derived from the - // resolved driver (see embeddingBatchSize). Batches are fanned out on the shared - // compiler pool and concatenated back in input order. - batchSize := embeddingBatchSize(mdl) - 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 - }) + // 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) } - if err := knowledge_compile.SubmitCompilerJobs(ctx, jobs); err != nil { - return nil, err + out := make([][]float32, len(embeds)) + for i, v := range embeds { + out[i] = float64sToFloat32(v.Embedding) } - // 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)) + if len(out) > 0 { + e.dim.CompareAndSwap(0, int64(len(out[0]))) } return out, nil } -// embeddingBatchSize returns the max texts per Embed request for the resolved -// model's driver. -// -// The ModelDriver interface exposes no embedding batch/input-limit capability -// (ListModels only reports MaxDimension/content-length; providers do not return -// a batch limit), so the limit cannot be queried at runtime. We therefore keep a -// provider-aware table keyed by the driver name: known strict providers get their -// documented cap, and unknown providers use a conservative default that is safe -// across OpenAI-compatible backends. See siliconflow.go:143 for the 32-input cap. -func embeddingBatchSize(mdl *models.EmbeddingModel) int { - // Driver Name() returns the provider's canonical (often upper-cased) name, - // e.g. SiliconflowModel.Name() == "SILICONFLOW"; compare case-insensitively. - name := strings.ToUpper(mdl.ModelDriver.Name()) - switch name { - case "SILICONFLOW": - return 32 - default: - // OpenAI-compatible providers generally accept large batches (often - // hundreds+); 64 is a safe conservative bound that still keeps a single - // embed round-trip small. - return 64 - } -} - -// 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 @@ -318,12 +230,9 @@ 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": "wiki_page", + "compile_kwd": "artifact_page", + "kc_kind": "page", }, MatchExprs: []interface{}{&enginetypes.MatchDenseExpr{ VectorColumnName: fmt.Sprintf("q_%d_vec", len(vec)), @@ -355,8 +264,9 @@ 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": "wiki_page", + "compile_kwd": "artifact_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 deleted file mode 100644 index ab4a81abf4..0000000000 --- a/internal/ingestion/task/knowledge_compiler_wiring_test.go +++ /dev/null @@ -1,43 +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 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 373f861b7a..8316db14f4 100644 --- a/internal/ingestion/task/pipeline_executor.go +++ b/internal/ingestion/task/pipeline_executor.go @@ -403,6 +403,14 @@ 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 b1eca5024f..b7cb5ba9dd 100644 --- a/internal/ingestion/task/pipeline_executor_defaults_test.go +++ b/internal/ingestion/task/pipeline_executor_defaults_test.go @@ -51,7 +51,7 @@ var builtinComponentParamsGolden = map[string]string{ "qa": "{\"File\": {}, \"Parser:HipSignsRhyme\": {\"docx\": {\"flatten_media_to_text\": false, \"output_format\": \"json\", \"preprocess\": [\"main_content\"], \"suffix\": [\"docx\"], \"vlm\": {}}, \"markdown\": {\"flatten_media_to_text\": false, \"output_format\": \"json\", \"preprocess\": [\"main_content\"], \"suffix\": [\"md\", \"markdown\", \"mdx\"], \"vlm\": {}}, \"pdf\": {\"flatten_media_to_text\": false, \"output_format\": \"json\", \"parse_method\": \"DeepDOC\", \"preprocess\": [\"main_content\"], \"suffix\": [\"pdf\"], \"vlm\": {}}, \"spreadsheet\": {\"flatten_media_to_text\": false, \"output_format\": \"html\", \"parse_method\": \"DeepDOC\", \"preprocess\": [\"main_content\"], \"suffix\": [\"xls\", \"xlsx\", \"csv\"], \"vlm\": {}}, \"text&code\": {\"output_format\": \"json\", \"preprocess\": [\"main_content\"], \"suffix\": [\"txt\"]}}, \"Tokenizer:ColdCloudsDream\": {\"fields\": \"text\", \"filename_embd_weight\": 0.1, \"search_method\": [\"embedding\", \"full_text\"]}, \"QAChunker:TidyCloudsThink\": {}}", "resume": "{\"Extractor:ThreeDrinksAct\": {\"field_name\": \"metadata\", \"frequencyPenaltyEnabled\": true, \"frequency_penalty\": 0.7, \"llm_id\": \"THUDM/GLM-4.1V-9B-Thinking@SILICONFLOW\", \"maxTokensEnabled\": false, \"max_tokens\": 256, \"presencePenaltyEnabled\": true, \"presence_penalty\": 0.4, \"prompts\": [{\"content\": \"Content: {TitleChunker:FlatMiceFix@chunks}\", \"role\": \"user\"}], \"sys_prompt\": \"Act as a precise resume metadata extractor. Extract stable, chunk-supported metadata from the provided resume content.\\n\\nRules:\\n1. Use only information explicitly stated in the content. Do not infer, guess, normalize, or add missing facts.\\n2. The input may be only one chunk of a resume. Extract only what this content directly supports.\\n3. Use only these field names:\\ncandidate_name, gender, phone, email, city, location, nationality, linkedin, github, website, highest_degree, degree_levels, school_names, majors, graduation_years, work_experience_years, current_job_title, job_titles, company_names, job_experience, industries, target_job_titles, target_locations, employment_types, skills, certificates, awards, summary_tags\\n4. Ignore detailed responsibilities, project descriptions, achievement narratives, self-evaluation, and other low-value local details.\\n5. Keep values in the same language as the source text whenever possible.\\n6. Remove duplicates and keep only concise, high-value metadata.\\n7. Return only fields that are explicitly supported by the content. Do not return empty or unsupported fields.\\n\\nField guidance:\\n- highest_degree: highest explicit degree level mentioned\\n- degree_levels: all explicit degree levels mentioned\\n- school_names: explicit school, college, or university names\\n- majors: explicit fields of study\\n- graduation_years: explicit graduation years only\\n- work_experience_years: only if explicitly stated\\n- current_job_title: only if explicitly current or most recent\\n- job_titles: explicit role titles\\n- company_names: explicit employer names\\n- job_experience: concise structured work entries explicitly supported by the content, preferably including title, company, and time information when available\\n- industries: explicit industry names only\\n- target_job_titles: explicit desired roles only\\n- target_locations: explicit desired work locations only\\n- skills: concise, core, search-useful skills explicitly mentioned\\n- certificates: explicit certificate names only\\n- awards: explicit award names only\\n- summary_tags: short, high-value tags strictly supported by the content\\n\\nReturn only the extracted metadata. Do not output explanatory text.\", \"temperature\": 0.1, \"temperatureEnabled\": true, \"tenant_llm_id\": 29, \"topPEnabled\": true, \"top_p\": 0.3, \"auto_keywords\": 0, \"auto_questions\": 0, \"auto_tags\": 0, \"enable_metadata\": 0, \"metadata\": [], \"tag_file_id\": \"\"}, \"File\": {}, \"Parser:HipSignsRhyme\": {\"docx\": {\"flatten_media_to_text\": true, \"output_format\": \"json\", \"preprocess\": [\"main_content\"], \"suffix\": [\"docx\"], \"vlm\": {}}, \"pdf\": {\"flatten_media_to_text\": true, \"output_format\": \"json\", \"parse_method\": \"DeepDOC\", \"preprocess\": [\"main_content\"], \"suffix\": [\"pdf\"], \"vlm\": {}}, \"text&code\": {\"output_format\": \"json\", \"preprocess\": [\"main_content\"], \"suffix\": [\"txt\"]}}, \"TitleChunker:FlatMiceFix\": {\"hierarchy\": 1, \"include_heading_content\": false, \"levels\": [[\"^\\\\s*(?i:(?:\\\\d+[\\\\.\\\\)]\\\\s*)?(?:EDUCATION|ACADEMIC\\\\s*BACKGROUND|ACADEMIC\\\\s*HISTORY|EDUCATIONAL\\\\s*BACKGROUND|RELEVANT\\\\s*COURSEWORK|COURSEWORK|EXPERIENCE|WORK\\\\s*EXPERIENCE|PROFESSIONAL\\\\s*EXPERIENCE|RELEVANT\\\\s*EXPERIENCE|EMPLOYMENT\\\\s*HISTORY|CAREER\\\\s*HISTORY|INTERNSHIP\\\\s*EXPERIENCE|PROJECTS|PROJECT\\\\s*EXPERIENCE|ACADEMIC\\\\s*PROJECTS|PROFESSIONAL\\\\s*PROJECTS|SKILLS|TECHNICAL\\\\s*SKILLS|CORE\\\\s*COMPETENCIES|COMPETENCIES|QUALIFICATIONS|SUMMARY\\\\s*OF\\\\s*QUALIFICATIONS|CERTIFICATIONS|LICENSES|CERTIFICATES|AWARDS|HONORS|HONOURS|ACHIEVEMENTS|PUBLICATIONS|RESEARCH|RESEARCH\\\\s*EXPERIENCE|LEADERSHIP|LEADERSHIP\\\\s*EXPERIENCE|ACTIVITIES|EXTRACURRICULAR\\\\s*ACTIVITIES|ACTIVITIES\\\\s*(?:&|AND)\\\\s*SKILLS|INVOLVEMENT|CAMPUS\\\\s*INVOLVEMENT|VOLUNTEER\\\\s*EXPERIENCE|VOLUNTEERING|COMMUNITY\\\\s*SERVICE|LANGUAGES|INTERESTS|HOBBIES|PROFILE|PROFESSIONAL\\\\s*PROFILE|SUMMARY|PROFESSIONAL\\\\s*SUMMARY|CAREER\\\\s*SUMMARY|OBJECTIVE|CAREER\\\\s*OBJECTIVE|PERSONAL\\\\s*INFORMATION|CONTACT\\\\s*INFORMATION|ADDITIONAL\\\\s*INFORMATION|TRAINING))\\\\s*[::]?\\\\s*$\"], [\"^\\\\s*(?:\\\\d+[\\\\.、\\\\)]\\\\s*)?(?:教育背景|教育经历|学历背景|学术背景|技术背景|工作经历|工作经验|实习经历|项目经历|项目经验|科研经历|研究经历|校园经历|实践经历|专业经历|职业经历|技能|专业技能|技能特长|核心技能|技术栈|个人技能|工作技能|职业技能|技能与评价|技能与自我评价|工作技能与自我评价|职业技能与自我评价|证书|资格证书|职业资格|资质证书|获奖情况|获奖经历|荣誉|荣誉奖项|奖项|科研成果|论文发表|发表论文|领导经历|学生工作|校园活动|社团经历|活动经历|志愿经历|志愿服务|社会实践|语言能力|语言|自我评价|个人评价|自我总结|个人总结|个人优势|个人简介|个人信息|基本信息|联系方式|求职意向|应聘意向|职业目标|求职目标|兴趣爱好|兴趣特长|培训经历|其他信息|附加信息)\\\\s*[::]?\\\\s*$\"]], \"method\": \"hierarchy\"}, \"Tokenizer:KindHandsWin\": {\"fields\": \"text\", \"filename_embd_weight\": 0.1, \"search_method\": [\"embedding\", \"full_text\"]}}", "table": "{\"File\": {}, \"Parser:HipSignsRhyme\": {\"spreadsheet\": {\"flatten_media_to_text\": false, \"output_format\": \"html\", \"parse_method\": \"DeepDOC\", \"preprocess\": [\"main_content\"], \"suffix\": [\"xls\", \"xlsx\", \"csv\"], \"vlm\": {}, \"column_mode\": \"auto\", \"column_roles\": {}, \"column_names\": []}, \"text&code\": {\"output_format\": \"json\", \"preprocess\": [\"main_content\"], \"suffix\": [\"txt\"]}}, \"TableChunker:FastFoxesJump\": {}, \"Tokenizer:DeepLakesShine\": {\"fields\": \"text\", \"filename_embd_weight\": 0.1, \"search_method\": [\"embedding\", \"full_text\"]}}", - "knowledge_compiler": "{\"File\": {}, \"Compiler:KnownSwiftLions\": {\"language\": \"English\", \"llm_id\": \"\", \"variant\": \"structure\"}, \"Parser:HipSignsRhyme\": {\"setups\": {\"doc\": {\"output_format\": \"json\", \"preprocess\": [\"main_content\"], \"suffix\": [\"doc\"]}, \"docx\": {\"flatten_media_to_text\": false, \"output_format\": \"json\", \"preprocess\": [\"main_content\"], \"suffix\": [\"docx\"], \"vlm\": {}}, \"html\": {\"output_format\": \"json\", \"preprocess\": [\"main_content\"], \"suffix\": [\"htm\", \"html\"]}, \"markdown\": {\"flatten_media_to_text\": false, \"output_format\": \"json\", \"preprocess\": [\"main_content\"], \"suffix\": [\"md\", \"markdown\", \"mdx\"], \"vlm\": {}}, \"pdf\": {\"flatten_media_to_text\": false, \"output_format\": \"json\", \"parse_method\": \"DeepDOC\", \"preprocess\": [\"main_content\"], \"suffix\": [\"pdf\"], \"vlm\": {}}, \"spreadsheet\": {\"flatten_media_to_text\": false, \"output_format\": \"html\", \"parse_method\": \"DeepDOC\", \"preprocess\": [\"main_content\"], \"suffix\": [\"xls\", \"xlsx\", \"csv\"], \"vlm\": {}}, \"text&code\": {\"output_format\": \"json\", \"preprocess\": [\"main_content\"], \"suffix\": [\"txt\", \"py\", \"js\", \"java\", \"c\", \"cpp\", \"h\", \"php\", \"go\", \"ts\", \"sh\", \"cs\", \"kt\", \"sql\"]}}}, \"TokenChunker:SixApplesFall\": {\"children_delimiters\": [], \"chunk_token_size\": 512, \"delimiter_mode\": \"token_size\", \"delimiters\": [\"\\n\", \"!\", \"?\", \"。\", \";\", \"!\", \"?\"], \"image_context_size\": 0, \"overlapped_percent\": 0, \"table_context_size\": 0}}", + "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}}", } // 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 e5ccd059ff..b366e90466 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": "wiki_page", "content_with_weight": "page X"}, + {"id": "wiki-1", "compile_kwd": "artifact_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 fd2dc699c2..1d679c4a79 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) - // 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) + datasets.GET("/:dataset_id/index", r.datasetsHandler.TraceIndex) + datasets.POST("/:dataset_id/index", r.datasetsHandler.RunIndex) + datasets.DELETE("/:dataset_id/index", r.datasetsHandler.DeleteIndex) // Knowledge-compilation wiki artifacts datasets.HEAD("/:dataset_id/artifacts", r.datasetArtifactHandler.AnyArtifact) @@ -380,6 +380,8 @@ 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.go b/internal/server/config.go index a0b95a47ad..ea25b22c42 100644 --- a/internal/server/config.go +++ b/internal/server/config.go @@ -121,6 +121,11 @@ func Init(configPath string) error { return fmt.Errorf("parse API server config error: %w", err) } + err = globalConfig.ParseIngestorConfig(v) + if err != nil { + return fmt.Errorf("parse ingestor config error: %w", err) + } + err = globalConfig.ParseSyncerConfig(v) if err != nil { return fmt.Errorf("parse syncer config error: %w", err) diff --git a/internal/server/config/base.go b/internal/server/config/base.go index 221126953d..7a769b1636 100644 --- a/internal/server/config/base.go +++ b/internal/server/config/base.go @@ -30,6 +30,7 @@ type Config struct { admin AdminConfig apiServer APIServerConfig + ingestor IngestorConfig syncer SyncerConfig log LogConfig diff --git a/internal/server/config/ingestor_config.go b/internal/server/config/ingestor_config.go new file mode 100644 index 0000000000..7ea6edb0e5 --- /dev/null +++ b/internal/server/config/ingestor_config.go @@ -0,0 +1,42 @@ +// +// 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 config + +import "github.com/spf13/viper" + +type IngestorConfig struct { + MaxConcurrentWorkers int `mapstructure:"max_concurrent_workers"` +} + +func (c *Config) ParseIngestorConfig(v *viper.Viper) error { + // Default Ingestor config + c.ingestor.MaxConcurrentWorkers = 1 + + if !v.IsSet("ingestor") { + return nil + } + sub := v.Sub("ingestor") + if sub == nil { + return nil + } + + if sub.IsSet("max_concurrent_workers") { + c.ingestor.MaxConcurrentWorkers = sub.GetInt("max_concurrent_workers") + } + + return nil +} diff --git a/internal/service/admin_client.go b/internal/service/admin_client.go index c13208f403..149048dbb0 100644 --- a/internal/service/admin_client.go +++ b/internal/service/admin_client.go @@ -20,6 +20,7 @@ import ( "encoding/json" "errors" "fmt" + "net/http" "ragflow/internal/common" "ragflow/internal/server" "ragflow/internal/utility" @@ -28,6 +29,7 @@ import ( "go.uber.org/zap" ) +var licenseStatusCode common.ErrorCode var AdminServiceClient *AdminClient // AdminClient is responsible for sending heartbeat reports to the admin server @@ -41,10 +43,12 @@ type AdminClient struct { version string lastSuccess bool attemptCount int + clusterInfo *utility.ClusterInfo } // NewAdminClient creates a new heartbeat service instance func NewAdminClient(logger *zap.Logger, serverType common.ServerType, serverName, host string, port int) *AdminClient { + licenseStatusCode = common.CodeSuccess return &AdminClient{ logger: logger, serverType: serverType, @@ -76,6 +80,11 @@ func (h *AdminClient) InitHTTPClient() error { zap.Int("admin_port", adminConfig.HTTPPort), ) + err := h.InitHTTPClientEE() + if err != nil { + h.logger.Fatal(fmt.Sprintf("Fail to init enterprise service: %v", err)) + } + return nil } @@ -110,6 +119,8 @@ func (h *AdminClient) SendHeartbeat() error { Ext: nil, } + message.Ext = h.clusterInfo + jsonData, err := json.Marshal(message) if err != nil { h.logger.Error("Failed to marshal heartbeat message", zap.Error(err)) @@ -122,22 +133,31 @@ func (h *AdminClient) SendHeartbeat() error { } defer resp.Body.Close() - if resp.StatusCode != 200 { - // extract the Code and Message field of the response - var responseBody map[string]interface{} - err = json.NewDecoder(resp.Body).Decode(&responseBody) - if err != nil { - return err - } - code, ok := responseBody["code"].(float64) - if !ok { - return fmt.Errorf("unexpected heartbeat response (status %d): missing or non-numeric \"code\" field", resp.StatusCode) - } - responseCode := common.ErrorCode(code) - if responseCode != common.CodeLicenseValid { - return errors.New(responseCode.Message()) - } + // extract the Code and Message field of the response + var responseBody map[string]interface{} + err = json.NewDecoder(resp.Body).Decode(&responseBody) + if err != nil { + return err } + code, ok := responseBody["code"].(float64) + if !ok { + return fmt.Errorf("unexpected heartbeat response (status %d): missing or non-numeric \"code\" field", resp.StatusCode) + } + + if resp.StatusCode != http.StatusOK { + return fmt.Errorf("error HTTP status code: %d", resp.StatusCode) + } + + responseCode := common.ErrorCode(code) + if responseCode != common.CodeLicenseValid { + if responseCode != licenseStatusCode { + licenseStatusCode = responseCode + h.logger.Warn(fmt.Sprintf("Heartbeat response error: %s, code: %d", responseCode.Message(), responseCode)) + } + + return errors.New(responseCode.Message()) + } + licenseStatusCode = responseCode h.logger.Debug("Heartbeat sent successfully", zap.String("server_id", h.serverName), diff --git a/internal/service/admin_client_ee.go b/internal/service/admin_client_ee.go new file mode 100644 index 0000000000..d2683d9bd8 --- /dev/null +++ b/internal/service/admin_client_ee.go @@ -0,0 +1,20 @@ +// 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 service + +func (h *AdminClient) InitHTTPClientEE() error { + return nil +} diff --git a/internal/service/component_scoped_parser_config.go b/internal/service/component_scoped_parser_config.go deleted file mode 100644 index 55ca470129..0000000000 --- a/internal/service/component_scoped_parser_config.go +++ /dev/null @@ -1,118 +0,0 @@ -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 deleted file mode 100644 index bce355d2be..0000000000 --- a/internal/service/component_scoped_parser_config_test.go +++ /dev/null @@ -1,147 +0,0 @@ -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 deleted file mode 100644 index 0c996d565b..0000000000 --- a/internal/service/dataset/compilation_status.go +++ /dev/null @@ -1,84 +0,0 @@ -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 deleted file mode 100644 index 9b910fbc4b..0000000000 --- a/internal/service/dataset/compilation_status_test.go +++ /dev/null @@ -1,197 +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" - "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 4305065bb2..6132d0631d 100644 --- a/internal/service/dataset/create_test.go +++ b/internal/service/dataset/create_test.go @@ -85,54 +85,6 @@ 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 79d1081ac8..af10883cac 100644 --- a/internal/service/dataset/crud.go +++ b/internal/service/dataset/crud.go @@ -121,11 +121,6 @@ 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, @@ -526,21 +521,6 @@ 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 deleted file mode 100644 index b1c926c899..0000000000 --- a/internal/service/dataset/embedding.go +++ /dev/null @@ -1,288 +0,0 @@ -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 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 := 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 1dee37048b..d1559b594d 100644 --- a/internal/service/dataset/helpers.go +++ b/internal/service/dataset/helpers.go @@ -269,49 +269,6 @@ 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 new file mode 100644 index 0000000000..8f0eda7ca7 --- /dev/null +++ b/internal/service/dataset/index.go @@ -0,0 +1,743 @@ +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 new file mode 100644 index 0000000000..a4bb1cbdae --- /dev/null +++ b/internal/service/dataset/index_delete_test.go @@ -0,0 +1,272 @@ +// +// 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 49b25d0e0e..a5e7676cdd 100644 --- a/internal/service/dataset/metadata.go +++ b/internal/service/dataset/metadata.go @@ -38,14 +38,6 @@ 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") @@ -122,12 +114,6 @@ 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 ddc0f3f248..5553d28461 100644 --- a/internal/service/dataset/metadata_config_test.go +++ b/internal/service/dataset/metadata_config_test.go @@ -22,30 +22,13 @@ 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(), } } @@ -224,116 +207,3 @@ 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 new file mode 100644 index 0000000000..5c40f5190a --- /dev/null +++ b/internal/service/dataset/task_cleanup_test.go @@ -0,0 +1,111 @@ +// +// 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 388ccd58f5..a2fbef107b 100644 --- a/internal/service/dataset/update.go +++ b/internal/service/dataset/update.go @@ -287,23 +287,6 @@ 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 a094f4b2ae..1e91220679 100644 --- a/internal/service/dataset_artifact_service.go +++ b/internal/service/dataset_artifact_service.go @@ -18,7 +18,6 @@ import ( "encoding/json" "fmt" "sort" - "strings" "ragflow/internal/engine" "ragflow/internal/engine/types" @@ -166,15 +165,7 @@ 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) { - // 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) - } + filter := map[string]interface{}{"compile_kwd": []string{CompileKwdWikiPage}} if pageType != "" { filter["page_type_kwd"] = []string{pageType} } @@ -190,18 +181,10 @@ 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: bareSlug, + Slug: firstStringValue(c["slug_kwd"]), Title: firstStringValue(c["title_kwd"]), - PageType: pageType, + PageType: firstStringValue(c["page_type_kwd"]), Topic: firstStringValue(c["topic_kwd"]), Summary: firstStringValue(c["summary_with_weight"]), }) @@ -209,15 +192,13 @@ func (s *DatasetArtifactService) ListWikiPages(ctx context.Context, tenantID, da return items, total, nil } -// 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. +// WikiPageDetail is the full wiki page payload. 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_rendered"` + ContentMd string `json:"content_md"` Summary string `json:"summary"` EntityNames []string `json:"entity_names"` Outlinks []string `json:"outlinks"` @@ -228,21 +209,16 @@ 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", "md_with_weight", - "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", "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 @@ -251,26 +227,12 @@ 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: detailSlug, + Slug: firstStringValue(c["slug_kwd"]), Title: firstStringValue(c["title_kwd"]), - PageType: detailPageType, + PageType: firstStringValue(c["page_type_kwd"]), Topic: firstStringValue(c["topic_kwd"]), - ContentMd: content, + ContentMd: firstStringValue(c["content_with_weight"]), Summary: firstStringValue(c["summary_with_weight"]), EntityNames: toStringSlice(c["entity_names_kwd"]), Outlinks: toStringSlice(c["outlinks_kwd"]), @@ -289,14 +251,11 @@ 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 { @@ -311,9 +270,6 @@ 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 != "" { @@ -344,10 +300,9 @@ 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", "page_type_kwd"}, 0, 1000, nil) + []string{"topic_kwd", "title_kwd", "slug_kwd"}, 0, 1000, nil) if err != nil { return nil, 0, err } @@ -358,17 +313,12 @@ 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: bareSlug, + Slug: firstStringValue(c["slug_kwd"]), } } } diff --git a/internal/service/dataset_types.go b/internal/service/dataset_types.go index 0864d66aa1..cbcc7364d7 100644 --- a/internal/service/dataset_types.go +++ b/internal/service/dataset_types.go @@ -1,5 +1,10 @@ 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 40ab76b5c1..9d3cc1b0cf 100644 --- a/internal/service/document/document_dataset_update.go +++ b/internal/service/document/document_dataset_update.go @@ -166,13 +166,6 @@ 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 4ff4b0ed0c..37b35f52c1 100644 --- a/internal/service/model_service.go +++ b/internal/service/model_service.go @@ -3533,66 +3533,6 @@ 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 8059a63772..86e81f0607 100644 --- a/internal/service/model_service_test.go +++ b/internal/service/model_service_test.go @@ -201,63 +201,6 @@ 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 deleted file mode 100644 index d1a45b6b07..0000000000 --- a/internal/service/wikisearch/engine_service.go +++ /dev/null @@ -1,220 +0,0 @@ -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, - }, - }) - 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, - }, - 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 deleted file mode 100644 index 57b4939d2f..0000000000 --- a/internal/service/wikisearch/engine_service_test.go +++ /dev/null @@ -1,265 +0,0 @@ -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 deleted file mode 100644 index d29d0bccb2..0000000000 --- a/internal/service/wikisearch/registry.go +++ /dev/null @@ -1,26 +0,0 @@ -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 deleted file mode 100644 index b7e9a756ed..0000000000 --- a/internal/service/wikisearch/wikisearch.go +++ /dev/null @@ -1,52 +0,0 @@ -// 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/internal/utility/fingerprint_ee.go b/internal/utility/fingerprint_ee.go new file mode 100644 index 0000000000..5c41b3124c --- /dev/null +++ b/internal/utility/fingerprint_ee.go @@ -0,0 +1,20 @@ +// +// 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 utility + +type ClusterInfo struct { +} diff --git a/pyproject.toml b/pyproject.toml index d9f0e1e214..a97ccfc132 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -100,6 +100,10 @@ dependencies = [ "qianfan==0.4.6", "quart-auth==0.11.0", "quart-cors==0.8.0", + # Werkzeug 3.1.5 can prepend \r\n to multipart file bodies when TCP chunks + # split at part headers (https://github.com/pallets/werkzeug/issues/3088); + # fixed in >=3.1.7. Pin so uv.lock cannot resolve back to the buggy release. + "werkzeug>=3.1.7,<4", "ranx==0.3.20", "readability-lxml>=0.8.4,<1.0.0", "replicate==0.31.0", diff --git a/rag/advanced_rag/knowlege_compile/wiki.py b/rag/advanced_rag/knowlege_compile/wiki.py index 5a2a350fd8..fc2437d3fb 100644 --- a/rag/advanced_rag/knowlege_compile/wiki.py +++ b/rag/advanced_rag/knowlege_compile/wiki.py @@ -821,10 +821,9 @@ async def _wiki_extract_one_batch( language: str, llm_timeout: int, parser_config: Optional[dict] = None, -) -> Optional[dict]: +) -> dict: """Single LLM call for one packed batch. Returns the raw (label-tagged) - extract dict, or ``None`` on a transient LLM timeout/error so the caller - can avoid persisting a poisoned empty result. + extract dict. The entity / relation schemas and the extra rules sections of the prompt are rendered from ``parser_config`` when supplied (mirroring @@ -852,10 +851,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 None + return _wiki_empty_extract() except Exception: logging.exception("wiki_map: batch extraction failed (%d chunks)", len(packed)) - return None + return _wiki_empty_extract() _ = language # reserved for future localization return _wiki_unwrap_extract(res) @@ -882,12 +881,6 @@ 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() @@ -903,10 +896,6 @@ 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/rag/app/book.py b/rag/app/book.py index 8ada2ec9b8..cb6132b5d5 100644 --- a/rag/app/book.py +++ b/rag/app/book.py @@ -86,7 +86,7 @@ def chunk(filename, binary=None, from_page=0, to_page=MAXIMUM_PAGE_NUMBER, lang= remove_contents_table(sections, eng=is_english(random_choices([t for t, _ in sections], k=200))) - tbls = vision_figure_parser_docx_wrapper(sections=sections, tbls=tbls, callback=callback, **kwargs) + tbls = vision_figure_parser_docx_wrapper(sections=sections, tbls=tbls, callback=callback, lang=lang, **kwargs) # tbls = [((None, lns), None) for lns in tbls] sections = [(item[0], item[1] if item[1] is not None else "") for item in sections if not isinstance(item[1], (Image.Image, LazyImage))] callback(0.8, "Finish parsing.") diff --git a/rag/app/manual.py b/rag/app/manual.py index 5f748276b3..50b614b033 100644 --- a/rag/app/manual.py +++ b/rag/app/manual.py @@ -268,6 +268,7 @@ def chunk(filename, binary=None, from_page=0, to_page=MAXIMUM_PAGE_NUMBER, lang= tbls=tbls, sections=sections, callback=callback, + lang=lang, **kwargs, ) res = tokenize_table(tbls, doc, eng, language=lang) @@ -283,7 +284,13 @@ def chunk(filename, binary=None, from_page=0, to_page=MAXIMUM_PAGE_NUMBER, lang= elif re.search(r"\.docx?$", filename, re.IGNORECASE): docx_parser = Docx() ti_list, tbls = docx_parser(filename, binary, from_page=0, to_page=MAXIMUM_PAGE_NUMBER, callback=callback) - tbls = vision_figure_parser_docx_wrapper(sections=ti_list, tbls=tbls, callback=callback, **kwargs) + tbls = vision_figure_parser_docx_wrapper( + sections=ti_list, + tbls=tbls, + callback=callback, + lang=lang, + **kwargs, + ) res = tokenize_table(tbls, doc, eng, language=lang) for text, image in ti_list: d = copy.deepcopy(doc) diff --git a/rag/app/naive.py b/rag/app/naive.py index 176755401b..99ab99713a 100644 --- a/rag/app/naive.py +++ b/rag/app/naive.py @@ -123,6 +123,7 @@ def by_deepdoc(filename, binary=None, from_page=0, to_page=MAXIMUM_PAGE_NUMBER, tbls=tables, sections=sections, callback=callback, + lang=lang, **kwargs, ) return sections, tables, pdf_parser @@ -1027,7 +1028,13 @@ def chunk(filename, binary=None, from_page=0, to_page=MAXIMUM_PAGE_NUMBER, lang= # images list - index of image chunk in chunks chunks, images = naive_merge_docx(sections, int(parser_config.get("chunk_token_num", 128)), parser_config.get("delimiter", "\n!?。;!?"), table_context_size, image_context_size) - vision_figure_parser_docx_wrapper_naive(chunks=chunks, idx_lst=images, callback=callback, **kwargs) + vision_figure_parser_docx_wrapper_naive( + chunks=chunks, + idx_lst=images, + callback=callback, + lang=lang, + **kwargs, + ) callback(0.8, "Finish parsing.") st = timer() @@ -1147,7 +1154,7 @@ def chunk(filename, binary=None, from_page=0, to_page=MAXIMUM_PAGE_NUMBER, lang= try: vision_model_config = get_tenant_default_model_by_type(kwargs["tenant_id"], LLMType.VISION) - vision_model = LLMBundle(kwargs["tenant_id"], vision_model_config) + vision_model = LLMBundle(kwargs["tenant_id"], vision_model_config, lang=lang) callback(0.2, "Visual model detected. Attempting to enhance figure extraction...") except Exception as e: logging.warning(f"Failed to detect figure extraction: {e}") @@ -1168,7 +1175,12 @@ def chunk(filename, binary=None, from_page=0, to_page=MAXIMUM_PAGE_NUMBER, lang= else: section_images = [None] * len(sections) section_images[idx] = combined_image - markdown_vision_parser = VisionFigureParser(vision_model=vision_model, figures_data=[((combined_image, ["markdown image"]), [(0, 0, 0, 0, 0)])], **kwargs) + markdown_vision_parser = VisionFigureParser( + vision_model=vision_model, + figures_data=[((combined_image, ["markdown image"]), [(0, 0, 0, 0, 0)])], + lang=lang, + **kwargs, + ) boosted_figures = markdown_vision_parser(callback=callback) sections[idx] = (section_text + "\n\n" + "\n\n".join([fig[0][1] for fig in boosted_figures]), sections[idx][1]) diff --git a/rag/app/one.py b/rag/app/one.py index 6e62fe6b3a..1a08d2c9cb 100644 --- a/rag/app/one.py +++ b/rag/app/one.py @@ -29,6 +29,9 @@ from common.constants import MAXIMUM_PAGE_NUMBER, MAXIMUM_TASK_PAGE_NUMBER from common.parser_config_utils import normalize_layout_recognizer +logger = logging.getLogger(__name__) + + class Pdf(PdfParser): def __call__(self, filename, binary=None, from_page=0, to_page=MAXIMUM_PAGE_NUMBER, zoomin=3, callback=None): from timeit import default_timer as timer @@ -83,7 +86,12 @@ def chunk(filename, binary=None, from_page=0, to_page=MAXIMUM_PAGE_NUMBER, lang= cks.append({"text": text, "image": image, "ck_type": ck_type}) - vision_figure_parser_docx_wrapper_naive(cks, image_idxs, callback, **kwargs) + logger.info( + "DOCX figure vision enhancement: language=%s image_count=%d", + lang or "English", + len(image_idxs), + ) + vision_figure_parser_docx_wrapper_naive(cks, image_idxs, callback, lang=lang, **kwargs) sections = [ck["text"] for ck in cks if ck.get("text")] callback(0.8, "Finish parsing.") diff --git a/rag/app/paper.py b/rag/app/paper.py index 89e735831c..19aeac46fa 100644 --- a/rag/app/paper.py +++ b/rag/app/paper.py @@ -184,6 +184,7 @@ def chunk(filename, binary=None, from_page=0, to_page=MAXIMUM_PAGE_NUMBER, lang= tbls=tbls, sections=sections, callback=callback, + lang=lang, **kwargs, ) paper["tables"] = tbls diff --git a/rag/app/table.py b/rag/app/table.py index feac9524b9..b04a6bc36f 100644 --- a/rag/app/table.py +++ b/rag/app/table.py @@ -63,7 +63,16 @@ def _deduplicate_column_names(columns): class Excel(ExcelParser): - def __call__(self, fnm, binary=None, from_page=0, to_page=MAXIMUM_TASK_PAGE_NUMBER, callback=None, **kwargs): + def __call__( + self, + fnm, + binary=None, + from_page=0, + to_page=MAXIMUM_TASK_PAGE_NUMBER, + callback=None, + lang="English", + **kwargs, + ): if not binary: wb = Excel._load_excel_to_workbook(fnm) else: @@ -80,7 +89,12 @@ class Excel(ExcelParser): images = Excel._extract_images_from_worksheet(ws, sheetname=sheet_name) pending_cell_images = [] if images: - image_descriptions = vision_figure_parser_figure_xlsx_wrapper(images=images, callback=callback, **kwargs) + image_descriptions = vision_figure_parser_figure_xlsx_wrapper( + images=images, + callback=callback, + lang=lang, + **kwargs, + ) if image_descriptions and len(image_descriptions) == len(images): for i, bf in enumerate(image_descriptions): desc = bf[0][1] @@ -406,7 +420,15 @@ def chunk(filename, binary=None, from_page=0, to_page=MAXIMUM_TASK_PAGE_NUMBER, if re.search(r"\.xlsx?$", filename, re.IGNORECASE): callback(0.1, "Start to parse.") excel_parser = Excel() - dfs, tbls = excel_parser(filename, binary, from_page=from_page, to_page=to_page, callback=callback, **kwargs) + dfs, tbls = excel_parser( + filename, + binary, + from_page=from_page, + to_page=to_page, + callback=callback, + lang=lang, + **kwargs, + ) elif re.search(r"\.txt$", filename, re.IGNORECASE): callback(0.1, "Start to parse.") txt = get_text(filename, binary) diff --git a/rag/flow/parser/parser.py b/rag/flow/parser/parser.py index 1d472ab4ff..4e80745813 100644 --- a/rag/flow/parser/parser.py +++ b/rag/flow/parser/parser.py @@ -774,6 +774,7 @@ class Parser(ProcessBase): self._canvas._tenant_id, conf.get("vlm"), callback=self.callback, + lang=getattr(self._canvas, "_language", None) or conf.get("lang") or "English", ) # Emit the requested final PDF output format. @@ -977,6 +978,7 @@ class Parser(ProcessBase): self._canvas._tenant_id, conf.get("vlm"), callback=self.callback, + lang=getattr(self._canvas, "_language", None) or conf.get("lang") or "English", ) self.set_output("json", sections) @@ -1113,6 +1115,7 @@ class Parser(ProcessBase): self._canvas._tenant_id, conf.get("vlm"), callback=self.callback, + lang=getattr(self._canvas, "_language", None) or conf.get("lang") or "English", ) self.set_output("json", json_results) else: diff --git a/rag/flow/parser/utils.py b/rag/flow/parser/utils.py index 76ca536180..db9eef0c24 100644 --- a/rag/flow/parser/utils.py +++ b/rag/flow/parser/utils.py @@ -164,16 +164,19 @@ def enhance_media_sections_with_vision( tenant_id, vlm_conf=None, callback=None, + lang="English", ): if not sections or not tenant_id: return sections + lang = lang or "English" + try: try: vision_model_config = resolve_model_config(tenant_id, LLMType.VISION, vlm_conf["llm_id"]) except Exception: vision_model_config = get_tenant_default_model_by_type(tenant_id, LLMType.VISION) - vision_model = LLMBundle(tenant_id, vision_model_config) + vision_model = LLMBundle(tenant_id, vision_model_config, lang=lang) except Exception: return sections @@ -189,6 +192,7 @@ def enhance_media_sections_with_vision( vision_model=vision_model, figures_data=[((item["image"], [""]), [(0, 0, 0, 0, 0)])], context_size=0, + lang=lang, )(callback=callback) except Exception: continue diff --git a/rag/prompts/generator.py b/rag/prompts/generator.py index 2ef35fdbdb..7867481a2d 100644 --- a/rag/prompts/generator.py +++ b/rag/prompts/generator.py @@ -360,14 +360,18 @@ def vision_llm_describe_prompt(page=None) -> str: return template.render(page=page) -def vision_llm_figure_describe_prompt() -> str: +def vision_llm_figure_describe_prompt(language: str = "English") -> str: template = PROMPT_JINJA_ENV.from_string(VISION_LLM_FIGURE_DESCRIBE_PROMPT) - return template.render() + return template.render(language=language) -def vision_llm_figure_describe_prompt_with_context(context_above: str, context_below: str) -> str: +def vision_llm_figure_describe_prompt_with_context(context_above: str, context_below: str, language: str = "English") -> str: template = PROMPT_JINJA_ENV.from_string(VISION_LLM_FIGURE_DESCRIBE_PROMPT_WITH_CONTEXT) - return template.render(context_above=context_above, context_below=context_below) + return template.render( + context_above=context_above, + context_below=context_below, + language=language, + ) def tool_schema(tools_description: list[dict], complete_task=False): diff --git a/rag/prompts/vision_llm_figure_describe_prompt.md b/rag/prompts/vision_llm_figure_describe_prompt.md index db17b44efe..1b6ddaa124 100644 --- a/rag/prompts/vision_llm_figure_describe_prompt.md +++ b/rag/prompts/vision_llm_figure_describe_prompt.md @@ -6,6 +6,12 @@ You are an expert visual data analyst. Analyze the image and produce a textual representation strictly based on what is visible in the image. +## OUTPUT LANGUAGE + +- Write all descriptions and field values in {{ language }}. +- Preserve all visible text verbatim in its original language; do not translate it. +- Keep the required output field names exactly as specified below. + ## DECISION RULE (CRITICAL) First, determine whether the image contains an explicit visual data representation with enumerable data units forming a coherent dataset. diff --git a/rag/prompts/vision_llm_figure_describe_prompt_with_context.md b/rag/prompts/vision_llm_figure_describe_prompt_with_context.md index 6843f7e7ef..07b13c51dd 100644 --- a/rag/prompts/vision_llm_figure_describe_prompt_with_context.md +++ b/rag/prompts/vision_llm_figure_describe_prompt_with_context.md @@ -7,6 +7,12 @@ You are an expert visual data analyst. Analyze the image and produce a textual representation strictly based on what is visible in the image. Surrounding context may be used only for minimal clarification or disambiguation of terms that appear in the image, not as a source of new information. +## OUTPUT LANGUAGE + +- Write all descriptions and field values in {{ language }}. +- Preserve all visible text verbatim in its original language; do not translate it. +- Keep the required output field names exactly as specified below. + ## CONTEXT (ABOVE) {{ context_above }} diff --git a/rag/svr/task_executor_refactor/dataset_skill_generator.py b/rag/svr/task_executor_refactor/dataset_skill_generator.py index c447d1fb03..77ce8e3fff 100644 --- a/rag/svr/task_executor_refactor/dataset_skill_generator.py +++ b/rag/svr/task_executor_refactor/dataset_skill_generator.py @@ -366,9 +366,10 @@ async def run_corpus2skill( llm_model=chat_mdl, embd_model=embedding_model, prompt="Please write a concise summary of the following texts:\n{cluster_content}", - max_token=256, - threshold=0.1, + max_token=512, max_errors=3, + clustering_threshold=0.3, + clustering_ratio=0.5, ) # ---- Phase 1: per-doc summaries. diff --git a/test/testcases/restful_api/conftest.py b/test/testcases/restful_api/conftest.py index 36e2b73516..ad3030287f 100644 --- a/test/testcases/restful_api/conftest.py +++ b/test/testcases/restful_api/conftest.py @@ -25,13 +25,6 @@ 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/test/testcases/test_web_api/test_agent_app/test_agents_webhook_unit.py b/test/testcases/test_web_api/test_agent_app/test_agents_webhook_unit.py index ae11b029bb..1b558fb4e7 100644 --- a/test/testcases/test_web_api/test_agent_app/test_agents_webhook_unit.py +++ b/test/testcases/test_web_api/test_agent_app/test_agents_webhook_unit.py @@ -521,7 +521,7 @@ def test_agents_crud_unit_branches(monkeypatch): captured = {} - def fake_get_by_tenant_ids(owner_ids, tenant_id, page, page_size, orderby, desc, keywords, canvas_category, tags): + def fake_get_by_tenant_ids(owner_ids, tenant_id, page, page_size, orderby, desc, keywords, canvas_category_list, tags, canvas_type=None): captured["owner_ids"] = owner_ids captured["tenant_id"] = tenant_id captured["page"] = page @@ -529,7 +529,7 @@ def test_agents_crud_unit_branches(monkeypatch): captured["orderby"] = orderby captured["desc"] = desc captured["keywords"] = keywords - captured["canvas_category"] = canvas_category + captured["canvas_category_list"] = canvas_category_list captured["tags"] = tags return [{"id": "agent-1"}], 1 diff --git a/test/unit_test/deepdoc/parser/test_figure_parser.py b/test/unit_test/deepdoc/parser/test_figure_parser.py new file mode 100644 index 0000000000..b92fa0e6e7 --- /dev/null +++ b/test/unit_test/deepdoc/parser/test_figure_parser.py @@ -0,0 +1,298 @@ +# +# 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. +# +import importlib.util +import sys +from pathlib import Path +from types import ModuleType, SimpleNamespace +from unittest.mock import Mock + +import pytest + + +def _package(monkeypatch, name): + package = ModuleType(name) + package.__path__ = [] + monkeypatch.setitem(sys.modules, name, package) + return package + + +def _module(monkeypatch, name, **attributes): + module = ModuleType(name) + for key, value in attributes.items(): + setattr(module, key, value) + monkeypatch.setitem(sys.modules, name, module) + return module + + +def _load_figure_parser(monkeypatch): + repo_root = Path(__file__).resolve().parents[4] + + for package_name in ( + "api", + "api.db", + "api.db.services", + "api.db.joint_services", + "common", + "rag", + "rag.app", + "rag.prompts", + "rag.utils", + ): + _package(monkeypatch, package_name) + + class FakeImage: + def close(self): + pass + + image_module = _module(monkeypatch, "PIL.Image", Image=FakeImage) + pil_module = _package(monkeypatch, "PIL") + pil_module.Image = image_module + + _module( + monkeypatch, + "common.constants", + LLMType=SimpleNamespace(VISION="vision"), + ) + _module( + monkeypatch, + "api.db.services.llm_service", + LLMBundle=Mock(), + ) + _module( + monkeypatch, + "api.db.joint_services.tenant_model_service", + get_tenant_default_model_by_type=Mock(), + ) + + def timeout(*_args, **_kwargs): + return lambda function: function + + _module(monkeypatch, "common.connection_utils", timeout=timeout) + _module( + monkeypatch, + "rag.app.picture", + vision_llm_chunk=Mock(return_value="description"), + ) + _module( + monkeypatch, + "rag.prompts.generator", + vision_llm_figure_describe_prompt=Mock(return_value="prompt"), + vision_llm_figure_describe_prompt_with_context=Mock(return_value="prompt"), + ) + _module( + monkeypatch, + "rag.nlp", + append_context2table_image4pdf=Mock(return_value=[]), + ) + _module( + monkeypatch, + "rag.utils.lazy_image", + ensure_pil_image=lambda image: image, + open_image_for_processing=lambda image, **_kwargs: (image, False), + is_image_like=lambda _image: True, + ) + + module_path = repo_root / "deepdoc" / "parser" / "figure_parser.py" + spec = importlib.util.spec_from_file_location( + "test_figure_parser_module", + module_path, + ) + module = importlib.util.module_from_spec(spec) + monkeypatch.setitem(sys.modules, spec.name, module) + spec.loader.exec_module(module) + return module, FakeImage + + +@pytest.mark.p1 +@pytest.mark.parametrize( + ("context_above", "context_below", "prompt_name", "expected_arguments"), + [ + ( + "", + "", + "vision_llm_figure_describe_prompt", + {}, + ), + ( + "Above ", + "Below", + "vision_llm_figure_describe_prompt_with_context", + { + "context_above": "Above Caption", + "context_below": "Below", + }, + ), + ], +) +@pytest.mark.parametrize( + ("language", "expected_language"), + [ + ("Chinese", "Chinese"), + ("", "English"), + ], +) +def test_docx_wrapper_passes_dataset_language_to_vision_model_and_prompt( + monkeypatch, + context_above, + context_below, + prompt_name, + expected_arguments, + language, + expected_language, +): + module, FakeImage = _load_figure_parser(monkeypatch) + model_config = {"llm_name": "vision-model"} + vision_model = object() + + module.get_tenant_default_model_by_type = Mock(return_value=model_config) + module.LLMBundle = Mock(return_value=vision_model) + module.picture_vision_llm_chunk = Mock(return_value="description") + + default_prompt = Mock(return_value="prompt") + contextual_prompt = Mock(return_value="prompt") + module.vision_llm_figure_describe_prompt = default_prompt + module.vision_llm_figure_describe_prompt_with_context = contextual_prompt + + chunks = [ + { + "image": FakeImage(), + "text": "Caption", + "context_above": context_above, + "context_below": context_below, + } + ] + + module.vision_figure_parser_docx_wrapper_naive( + chunks=chunks, + idx_lst=[0], + callback=lambda *_args, **_kwargs: None, + tenant_id="tenant-id", + lang=language, + ) + + module.LLMBundle.assert_called_once_with( + "tenant-id", + model_config, + lang=expected_language, + ) + + selected_prompt = getattr(module, prompt_name) + selected_prompt.assert_called_once_with( + **expected_arguments, + language=expected_language, + ) + assert chunks[0]["text"].endswith("description") + + +@pytest.mark.p1 +@pytest.mark.parametrize( + ("language", "expected_language"), + [ + ("Chinese", "Chinese"), + ("", "English"), + ], +) +def test_vision_figure_parser_passes_dataset_language_to_prompt( + monkeypatch, + language, + expected_language, +): + module, FakeImage = _load_figure_parser(monkeypatch) + prompt = Mock(return_value="prompt") + module.vision_llm_figure_describe_prompt = prompt + module.picture_vision_llm_chunk = Mock(return_value="description") + + parser = module.VisionFigureParser( + vision_model=object(), + figures_data=[(FakeImage(), ["caption"])], + lang=language, + ) + + parser(callback=lambda *_args, **_kwargs: None) + + prompt.assert_called_once_with(language=expected_language) + + +@pytest.mark.p1 +@pytest.mark.parametrize( + "wrapper_name", + [ + "vision_figure_parser_docx_wrapper", + "vision_figure_parser_figure_xlsx_wrapper", + "vision_figure_parser_pdf_wrapper", + ], +) +@pytest.mark.parametrize( + ("language", "expected_language"), + [ + ("Chinese", "Chinese"), + ("", "English"), + ], +) +def test_figure_wrappers_pass_dataset_language_to_model_and_parser( + monkeypatch, + wrapper_name, + language, + expected_language, +): + module, FakeImage = _load_figure_parser(monkeypatch) + model_config = {"llm_name": "vision-model"} + vision_model = object() + parser_instance = Mock(return_value=[]) + + module.get_tenant_default_model_by_type = Mock(return_value=model_config) + module.LLMBundle = Mock(return_value=vision_model) + module.VisionFigureParser = Mock(return_value=parser_instance) + + if wrapper_name == "vision_figure_parser_docx_wrapper": + arguments = { + "sections": [("caption", FakeImage())], + "tbls": [], + } + elif wrapper_name == "vision_figure_parser_figure_xlsx_wrapper": + arguments = { + "images": [ + { + "image": FakeImage(), + "image_description": "caption", + } + ], + } + else: + arguments = { + "tbls": [ + ( + (FakeImage(), ["caption"]), + [(0, 0, 0, 0, 0)], + ) + ], + "sections": [], + } + + getattr(module, wrapper_name)( + **arguments, + callback=lambda *_args, **_kwargs: None, + tenant_id="tenant-id", + lang=language, + ) + + module.LLMBundle.assert_called_once_with( + "tenant-id", + model_config, + lang=expected_language, + ) + assert module.VisionFigureParser.call_args.kwargs["lang"] == expected_language + parser_instance.assert_called_once() diff --git a/test/unit_test/rag/app/test_one.py b/test/unit_test/rag/app/test_one.py new file mode 100644 index 0000000000..8a2172986f --- /dev/null +++ b/test/unit_test/rag/app/test_one.py @@ -0,0 +1,50 @@ +# +# 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. +# +import logging +from unittest.mock import Mock + +import pytest + +from rag.app import one + + +@pytest.mark.p1 +def test_docx_chunk_forwards_language_to_vision_wrapper(monkeypatch, caplog): + docx_parser = Mock(return_value=[("caption", object(), None)]) + monkeypatch.setattr(one.naive, "Docx", Mock(return_value=docx_parser)) + + vision_wrapper = Mock() + monkeypatch.setattr(one, "vision_figure_parser_docx_wrapper_naive", vision_wrapper) + monkeypatch.setattr(one.rag_tokenizer, "tokenize", lambda text: text) + monkeypatch.setattr(one.rag_tokenizer, "fine_grained_tokenize", lambda text: text) + monkeypatch.setattr(one, "tokenize", Mock()) + + with caplog.at_level(logging.INFO, logger=one.__name__): + one.chunk( + "document.docx", + binary=b"docx", + lang="Japanese", + callback=lambda *_args, **_kwargs: None, + tenant_id="tenant-id", + ) + + vision_wrapper.assert_called_once() + args = vision_wrapper.call_args.args + kwargs = vision_wrapper.call_args.kwargs + assert args[1] == [0] + assert kwargs["lang"] == "Japanese" + assert kwargs["tenant_id"] == "tenant-id" + assert "DOCX figure vision enhancement: language=Japanese image_count=1" in caplog.messages diff --git a/test/unit_test/rag/app/test_vision_language_callers.py b/test/unit_test/rag/app/test_vision_language_callers.py new file mode 100644 index 0000000000..66a910dd5b --- /dev/null +++ b/test/unit_test/rag/app/test_vision_language_callers.py @@ -0,0 +1,100 @@ +# +# 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. +# +import ast +from pathlib import Path +from unittest.mock import Mock + +import pytest + +from rag.app import naive + +REPO_ROOT = Path(__file__).resolve().parents[4] + + +def _call_name(call): + if isinstance(call.func, ast.Name): + return call.func.id + if isinstance(call.func, ast.Attribute): + return call.func.attr + return None + + +@pytest.mark.p1 +@pytest.mark.parametrize( + ("relative_path", "expected_call_count"), + [ + ("rag/app/book.py", 1), + ("rag/app/manual.py", 2), + ("rag/app/naive.py", 2), + ("rag/app/one.py", 1), + ("rag/app/paper.py", 1), + ("rag/app/table.py", 1), + ], +) +def test_all_figure_wrapper_callers_forward_language(relative_path, expected_call_count): + tree = ast.parse((REPO_ROOT / relative_path).read_text()) + calls = [node for node in ast.walk(tree) if isinstance(node, ast.Call) and (_call_name(node) or "").startswith("vision_figure_parser_")] + + assert len(calls) == expected_call_count + for call in calls: + language = next((keyword.value for keyword in call.keywords if keyword.arg == "lang"), None) + assert isinstance(language, ast.Name), f"{relative_path}:{call.lineno} does not forward lang" + assert language.id == "lang" + + +@pytest.mark.p1 +def test_markdown_chunk_forwards_language_to_model_and_figure_parser(monkeypatch): + markdown_parser = Mock(return_value=([("section", "")], [], [object()])) + monkeypatch.setattr(naive, "Markdown", Mock(return_value=markdown_parser)) + monkeypatch.setattr(naive, "get_tenant_default_model_by_type", Mock(return_value={"llm_name": "vision-model"})) + + vision_model = object() + llm_bundle = Mock(return_value=vision_model) + monkeypatch.setattr(naive, "LLMBundle", llm_bundle) + + parser_instance = Mock(return_value=[((None, "description"), None)]) + parser_factory = Mock(return_value=parser_instance) + monkeypatch.setattr(naive, "VisionFigureParser", parser_factory) + + monkeypatch.setattr(naive.rag_tokenizer, "tokenize", lambda text: text) + monkeypatch.setattr(naive.rag_tokenizer, "fine_grained_tokenize", lambda text: text) + monkeypatch.setattr(naive, "num_tokens_from_string", lambda _text: 1) + monkeypatch.setattr(naive, "tokenize_table", Mock(return_value=[])) + monkeypatch.setattr(naive, "tokenize_chunks", Mock(return_value=[])) + monkeypatch.setattr(naive, "tokenize_chunks_with_images", Mock(return_value=[])) + + naive.chunk( + "document.md", + binary=b"markdown", + lang="Japanese", + callback=lambda *_args, **_kwargs: None, + tenant_id="tenant-id", + is_root=False, + parser_config={ + "chunk_token_num": 128, + "delimiter": "\n", + "analyze_hyperlink": False, + }, + ) + + llm_bundle.assert_called_once_with( + "tenant-id", + {"llm_name": "vision-model"}, + lang="Japanese", + ) + assert parser_factory.call_args.kwargs["vision_model"] is vision_model + assert parser_factory.call_args.kwargs["lang"] == "Japanese" + parser_instance.assert_called_once() diff --git a/test/unit_test/rag/flow/parser/test_vision_language.py b/test/unit_test/rag/flow/parser/test_vision_language.py new file mode 100644 index 0000000000..4c6622e4cd --- /dev/null +++ b/test/unit_test/rag/flow/parser/test_vision_language.py @@ -0,0 +1,122 @@ +# +# 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. +# +import ast +import importlib.util +import sys +from pathlib import Path +from types import ModuleType, SimpleNamespace +from unittest.mock import Mock + +import pytest + +REPO_ROOT = Path(__file__).resolve().parents[5] + + +def _package(monkeypatch, name): + package = ModuleType(name) + package.__path__ = [] + monkeypatch.setitem(sys.modules, name, package) + return package + + +def _module(monkeypatch, name, **attributes): + module = ModuleType(name) + for key, value in attributes.items(): + setattr(module, key, value) + monkeypatch.setitem(sys.modules, name, module) + return module + + +def _load_flow_utils(monkeypatch): + for package_name in ( + "api", + "api.db", + "api.db.services", + "api.db.joint_services", + "common", + "deepdoc", + "deepdoc.parser", + "rag", + ): + _package(monkeypatch, package_name) + + _module(monkeypatch, "api.db.services.llm_service", LLMBundle=Mock()) + _module( + monkeypatch, + "api.db.joint_services.tenant_model_service", + get_tenant_default_model_by_type=Mock(), + resolve_model_config=Mock(), + ) + _module(monkeypatch, "common.constants", LLMType=SimpleNamespace(VISION="vision")) + _module(monkeypatch, "deepdoc.parser.figure_parser", VisionFigureParser=Mock()) + _module( + monkeypatch, + "rag.nlp", + is_english=Mock(return_value=False), + random_choices=Mock(return_value=[]), + remove_contents_table=Mock(), + ) + + module_path = REPO_ROOT / "rag/flow/parser/utils.py" + spec = importlib.util.spec_from_file_location("test_flow_parser_utils_module", module_path) + module = importlib.util.module_from_spec(spec) + monkeypatch.setitem(sys.modules, spec.name, module) + spec.loader.exec_module(module) + return module + + +@pytest.mark.p1 +@pytest.mark.parametrize( + ("language", "expected_language"), + [ + ("Japanese", "Japanese"), + ("", "English"), + ], +) +def test_media_enhancement_forwards_language_to_model_and_parser(monkeypatch, language, expected_language): + utils = _load_flow_utils(monkeypatch) + model_config = {"llm_name": "vision-model"} + vision_model = object() + llm_bundle = Mock(return_value=vision_model) + parser_instance = Mock(return_value=[((None, "description"), None)]) + parser_factory = Mock(return_value=parser_instance) + + monkeypatch.setattr(utils, "resolve_model_config", Mock(return_value=model_config)) + monkeypatch.setattr(utils, "LLMBundle", llm_bundle) + monkeypatch.setattr(utils, "VisionFigureParser", parser_factory) + + sections = [{"text": "caption", "image": object(), "doc_type_kwd": "image"}] + result = utils.enhance_media_sections_with_vision( + sections, + "tenant-id", + {"llm_id": "vision-model"}, + lang=language, + ) + + llm_bundle.assert_called_once_with("tenant-id", model_config, lang=expected_language) + assert parser_factory.call_args.kwargs["vision_model"] is vision_model + assert parser_factory.call_args.kwargs["lang"] == expected_language + assert result[0]["text"] == "caption\ndescription" + + +@pytest.mark.p1 +def test_all_flow_media_enhancement_callers_forward_language(): + tree = ast.parse((REPO_ROOT / "rag/flow/parser/parser.py").read_text()) + calls = [node for node in ast.walk(tree) if isinstance(node, ast.Call) and isinstance(node.func, ast.Name) and node.func.id == "enhance_media_sections_with_vision"] + + assert len(calls) == 3 + for call in calls: + assert any(keyword.arg == "lang" for keyword in call.keywords), f"parser.py:{call.lineno} does not forward lang" diff --git a/test/unit_test/rag/prompts/test_vision_figure_prompt.py b/test/unit_test/rag/prompts/test_vision_figure_prompt.py new file mode 100644 index 0000000000..ca544d2636 --- /dev/null +++ b/test/unit_test/rag/prompts/test_vision_figure_prompt.py @@ -0,0 +1,117 @@ +# +# 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. +# +import importlib.util +import sys +from pathlib import Path +from types import ModuleType, SimpleNamespace + +import pytest + + +def _load_generator(monkeypatch): + repo_root = Path(__file__).resolve().parents[4] + + json_repair = ModuleType("json_repair") + json_repair.repair_json = lambda text, **_kwargs: text + monkeypatch.setitem(sys.modules, "json_repair", json_repair) + + common = ModuleType("common") + common.__path__ = [str(repo_root / "common")] + monkeypatch.setitem(sys.modules, "common", common) + + misc_utils = ModuleType("common.misc_utils") + misc_utils.hash_str2int = lambda value, _mod=500: 0 + monkeypatch.setitem(sys.modules, "common.misc_utils", misc_utils) + + constants = ModuleType("common.constants") + constants.TAG_FLD = "tag" + monkeypatch.setitem(sys.modules, "common.constants", constants) + + token_utils = ModuleType("common.token_utils") + token_utils.encoder = SimpleNamespace() + token_utils.num_tokens_from_string = len + monkeypatch.setitem(sys.modules, "common.token_utils", token_utils) + + rag = ModuleType("rag") + rag.__path__ = [str(repo_root / "rag")] + monkeypatch.setitem(sys.modules, "rag", rag) + + rag_nlp = ModuleType("rag.nlp") + rag_nlp.rag_tokenizer = SimpleNamespace() + monkeypatch.setitem(sys.modules, "rag.nlp", rag_nlp) + + prompts = ModuleType("rag.prompts") + prompts.__path__ = [str(repo_root / "rag" / "prompts")] + monkeypatch.setitem(sys.modules, "rag.prompts", prompts) + + template = ModuleType("rag.prompts.template") + template.load_prompt = lambda name: (repo_root / "rag" / "prompts" / f"{name}.md").read_text(encoding="utf-8").strip() + monkeypatch.setitem(sys.modules, "rag.prompts.template", template) + + module_path = repo_root / "rag" / "prompts" / "generator.py" + spec = importlib.util.spec_from_file_location( + "test_vision_figure_prompt_generator", + module_path, + ) + module = importlib.util.module_from_spec(spec) + monkeypatch.setitem(sys.modules, spec.name, module) + spec.loader.exec_module(module) + return module + + +@pytest.mark.p1 +@pytest.mark.parametrize( + ("function_name", "arguments", "expected_language"), + [ + ( + "vision_llm_figure_describe_prompt", + {}, + "English", + ), + ( + "vision_llm_figure_describe_prompt", + {"language": "Chinese"}, + "Chinese", + ), + ( + "vision_llm_figure_describe_prompt_with_context", + {"context_above": "Above", "context_below": "Below"}, + "English", + ), + ( + "vision_llm_figure_describe_prompt_with_context", + { + "context_above": "Above", + "context_below": "Below", + "language": "Chinese", + }, + "Chinese", + ), + ], +) +def test_figure_prompt_renders_output_language( + monkeypatch, + function_name, + arguments, + expected_language, +): + generator = _load_generator(monkeypatch) + + prompt = getattr(generator, function_name)(**arguments) + + assert f"Write all descriptions and field values in {expected_language}." in prompt + assert "Preserve all visible text verbatim in its original language" in prompt + assert "{{ language }}" not in prompt diff --git a/uv.lock b/uv.lock index 2b73825311..04de483b4d 100644 --- a/uv.lock +++ b/uv.lock @@ -8166,6 +8166,7 @@ dependencies = [ { name = "webdav4" }, { name = "webdriver-manager" }, { name = "wechatpy" }, + { name = "werkzeug" }, { name = "wikipedia" }, { name = "word2number" }, { name = "xgboost" }, @@ -8322,6 +8323,7 @@ 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" }, @@ -9919,14 +9921,14 @@ wheels = [ [[package]] name = "werkzeug" -version = "3.1.5" +version = "3.1.8" source = { registry = "https://mirrors.aliyun.com/pypi/simple" } dependencies = [ { name = "markupsafe" }, ] -sdist = { url = "https://mirrors.aliyun.com/pypi/packages/5a/70/1469ef1d3542ae7c2c7b72bd5e3a4e6ee69d7978fa8a3af05a38eca5becf/werkzeug-3.1.5.tar.gz", hash = "sha256:6a548b0e88955dd07ccb25539d7d0cc97417ee9e179677d22c7041c8f078ce67" } +sdist = { url = "https://mirrors.aliyun.com/pypi/packages/dd/b2/381be8cfdee792dd117872481b6e378f85c957dd7c5bca38897b08f765fd/werkzeug-3.1.8.tar.gz", hash = "sha256:9bad61a4268dac112f1c5cd4630a56ede601b6ed420300677a869083d70a4c44" } wheels = [ - { url = "https://mirrors.aliyun.com/pypi/packages/ad/e4/8d97cca767bcc1be76d16fb76951608305561c6e056811587f36cb1316a8/werkzeug-3.1.5-py3-none-any.whl", hash = "sha256:5111e36e91086ece91f93268bb39b4a35c1e6f1feac762c9c822ded0a4e322dc" }, + { url = "https://mirrors.aliyun.com/pypi/packages/93/8c/2e650f2afeb7ee576912636c23ddb621c91ac6a98e66dc8d29c3c69446e1/werkzeug-3.1.8-py3-none-any.whl", hash = "sha256:63a77fb8892bf28ebc3178683445222aa500e48ebad5ec77b0ad80f8726b1f50" }, ] [[package]] diff --git a/web/src/components/large-model-form-field.tsx b/web/src/components/large-model-form-field.tsx index f5040f84dc..76da460131 100644 --- a/web/src/components/large-model-form-field.tsx +++ b/web/src/components/large-model-form-field.tsx @@ -39,15 +39,9 @@ export const LargeModelFilterFormSchema = { llm_filter: z.string().optional(), }; -type LargeModelFormFieldProps = Pick< - NextInnerLLMSelectProps, - 'ownerTenantId' -> & { - name?: string; -}; +type LargeModelFormFieldProps = Pick; export function LargeModelFormField({ ownerTenantId, - name = 'llm_id', }: LargeModelFormFieldProps) { const form = useFormContext(); const { t } = useTranslation(); @@ -57,7 +51,7 @@ export function LargeModelFormField({ <> ( diff --git a/web/src/components/next-message-item/group-button.tsx b/web/src/components/next-message-item/group-button.tsx index 421309a9c4..a7bffa9c04 100644 --- a/web/src/components/next-message-item/group-button.tsx +++ b/web/src/components/next-message-item/group-button.tsx @@ -128,7 +128,7 @@ export const AssistantGroupButton = ({ )} - {!!attachment?.doc_id && !isShare && ( + {!!attachment?.doc_id && ( { 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 221a599fe8..bb133e7a98 100644 --- a/web/src/components/pipeline-operator-tabs/pipeline-operator-form.tsx +++ b/web/src/components/pipeline-operator-tabs/pipeline-operator-form.tsx @@ -1,6 +1,5 @@ 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'; @@ -62,14 +61,6 @@ 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-agent-request.ts b/web/src/hooks/use-agent-request.ts index 7903bea25d..91f79e4c86 100644 --- a/web/src/hooks/use-agent-request.ts +++ b/web/src/hooks/use-agent-request.ts @@ -102,14 +102,14 @@ const buildAgentListParams = ({ page, pageSize, keywords, - canvasCategory, + canvasCategoryIds, ownerIds, tags, }: { page: number; pageSize: number; keywords?: string; - canvasCategory?: string; + canvasCategoryIds?: string[]; ownerIds?: string[]; tags?: string[]; }) => { @@ -121,8 +121,8 @@ const buildAgentListParams = ({ if (keywords) { params.keywords = keywords; } - if (canvasCategory) { - params.canvas_category = canvasCategory; + if (Array.isArray(canvasCategoryIds) && canvasCategoryIds.length > 0) { + params.canvas_category = canvasCategoryIds.join(','); } if (Array.isArray(ownerIds) && ownerIds.length > 0) { params.owner_ids = ownerIds.join(','); @@ -139,8 +139,8 @@ export const useFetchAgentListByPage = () => { const { pagination, setPagination } = useGetPaginationWithRouter(); const debouncedSearchString = useDebounce(searchString, { wait: 500 }); const { filterValue, handleFilterSubmit } = useHandleFilterSubmit(); - const canvasCategory = Array.isArray(filterValue.canvasCategory) - ? (filterValue.canvasCategory[0] as string | undefined) + const canvasCategoryIds = Array.isArray(filterValue.canvasCategory) + ? (filterValue.canvasCategory as string[]) : undefined; const owner = filterValue.owner; const tags = Array.isArray(filterValue.tags) ? filterValue.tags : undefined; @@ -149,7 +149,7 @@ export const useFetchAgentListByPage = () => { page: pagination.current, pageSize: pagination.pageSize, keywords: debouncedSearchString, - canvasCategory, + canvasCategoryIds, ownerIds: Array.isArray(owner) ? owner : undefined, tags, }); diff --git a/web/src/hooks/use-dataset-generate.ts b/web/src/hooks/use-dataset-generate.ts index 8d568cef7c..a0092f9ae4 100644 --- a/web/src/hooks/use-dataset-generate.ts +++ b/web/src/hooks/use-dataset-generate.ts @@ -9,11 +9,9 @@ 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'; @@ -50,15 +48,6 @@ 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 = ( @@ -72,15 +61,6 @@ 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; @@ -89,35 +69,6 @@ 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 ?? {}; }, @@ -180,14 +131,6 @@ 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')); @@ -208,12 +151,6 @@ 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, @@ -240,19 +177,6 @@ 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) { @@ -266,13 +190,6 @@ 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 d33797d88f..138f05f19f 100644 --- a/web/src/hooks/use-knowledge-request.ts +++ b/web/src/hooks/use-knowledge-request.ts @@ -1,6 +1,5 @@ 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 { @@ -965,13 +964,6 @@ 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 f5bb947c39..2fc072f459 100644 --- a/web/src/locales/en.ts +++ b/web/src/locales/en.ts @@ -476,10 +476,6 @@ 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', @@ -2416,10 +2412,6 @@ 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 e8a7af27b9..dd7d21926f 100644 --- a/web/src/locales/zh.ts +++ b/web/src/locales/zh.ts @@ -412,9 +412,6 @@ export default { generateToSkills: '从该数据集构建分层技能树,并存储生成的技能页面以供搜索和复用。', noWikiPages: '暂无 Wiki 页面', - compiling: '编译中…', - compilingCounts: '处理中 {{inflight}} / 待处理 {{backlog}}', - autoCompiled: '文档解析时自动编译。', clearWikiTitle: '清空 Wiki', clearWikiDescription: '确定要清空该数据集下的所有 Wiki 页面吗?此操作无法撤销。', @@ -2060,9 +2057,6 @@ 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 bb5f0aaf2e..b86659d51f 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, LLMLabelCard } from './card'; +import { LabelCard } from './card'; import { RagNode } from './index'; import { useTranslation } from 'react-i18next'; @@ -11,14 +11,12 @@ 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 03dbb7e6f5..f09cc15686 100644 --- a/web/src/pages/agent/constant/pipeline.tsx +++ b/web/src/pages/agent/constant/pipeline.tsx @@ -362,7 +362,6 @@ 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 1cdbf0fc02..3e2a7d3fdf 100644 --- a/web/src/pages/agent/form/compilation-form/index.tsx +++ b/web/src/pages/agent/form/compilation-form/index.tsx @@ -1,13 +1,10 @@ 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'; @@ -17,44 +14,28 @@ 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, - onValuesChange, - hideOutputs, -}: INextOperatorForm) => { +const CompilationForm = ({ node }: 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 1cbe828239..82491da75e 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, llm_id: llmId }, + [Operator.Compiler]: initialCompilationValues, [Operator.DataOperations]: initialDataOperationsValues, [Operator.ListOperations]: initialListOperationsValues, [Operator.VariableAssigner]: initialVariableAssignerValues, diff --git a/web/src/pages/agents/agent-card.tsx b/web/src/pages/agents/agent-card.tsx index 23661bda9e..4f96472653 100644 --- a/web/src/pages/agents/agent-card.tsx +++ b/web/src/pages/agents/agent-card.tsx @@ -3,12 +3,18 @@ import { MoreButton } from '@/components/more-button'; import { SharedBadge } from '@/components/shared-badge'; import { Badge } from '@/components/ui/badge'; import { Button } from '@/components/ui/button'; +import { + Tooltip, + TooltipContent, + TooltipTrigger, +} from '@/components/ui/tooltip'; import { AgentCategory } from '@/constants/agent'; import { useNavigatePage } from '@/hooks/logic-hooks/navigate-hooks'; import { AgentListItemType, IFlow } from '@/interfaces/database/agent'; import { CanvasCategoryToFlowType, FlowType, FlowTypeConfig } from './constant'; import { AgentDropdown } from './agent-dropdown'; import { useRenameAgent } from './use-rename-agent'; +import { useRef, useState } from 'react'; import { Tag } from 'lucide-react'; export type DatasetCardProps = { @@ -45,20 +51,51 @@ function AgentTags({ tags }: { tags?: string }) { .split(',') .map((t) => t.trim()) .filter(Boolean); + const containerRef = useRef(null); + const [open, setOpen] = useState(false); + if (list.length === 0) return null; + + const handleOpenChange = (isOpen: boolean) => { + if (isOpen) { + const el = containerRef.current; + setOpen(el ? el.scrollHeight > el.clientHeight : false); + } else { + setOpen(false); + } + }; + return ( -
- {list.map((tag) => ( - - - {tag} - - ))} -
+ + +
+ {list.map((tag) => ( + + + {tag} + + ))} +
+
+ +
+ {list.map((tag) => ( + + + {tag} + + ))} +
+
+
); } diff --git a/web/src/pages/dataset/compilation/empty-state.tsx b/web/src/pages/dataset/compilation/empty-state.tsx index 056cb08331..1c6556cf68 100644 --- a/web/src/pages/dataset/compilation/empty-state.tsx +++ b/web/src/pages/dataset/compilation/empty-state.tsx @@ -9,7 +9,6 @@ import { useDatasetGenerate, useGenerateStatus, } from '@/hooks/use-dataset-generate'; -import { isGoDatasetBackend } from '@/utils/api-proxy-scheme'; import { GenerableViewMode, @@ -59,66 +58,29 @@ 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])} - {!isGo && status === 'failed' && ( + {status === 'failed' && ( )} - {!isGo && status !== 'failed' && ( + {status !== 'failed' && ( { @@ -37,39 +33,6 @@ 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 7accc647b8..29e79d841c 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' | 'topic'; +export type WikiPageType = 'concept' | 'entity'; /** * Parse an internal wiki link href into pageType and slug. @@ -9,7 +9,7 @@ export type WikiPageType = 'concept' | 'entity' | 'topic'; * {pageType}/{slug} * /{pageType}/{slug} * - * entity/, concept/ and topic/ links are all considered wiki navigation links. + * Only entity/ and concept/ links are 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|topic)\/([^/\s"']+)/, + /(?:^|\/)artifact\/[^/]+\/(entity|concept)\/([^/\s"']+)/, ); if (artifactMatch) { return { @@ -29,9 +29,7 @@ export function parseWikiLinkHref( } // Fallback to a plain {pageType}/{slug} form. - const simpleMatch = normalized.match( - /(?:^|\/)(entity|concept|topic)\/([^/\s"']+)/, - ); + const simpleMatch = normalized.match(/(?:^|\/)(entity|concept)\/([^/\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 8c3adf511b..c37fa40208 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' | 'topic'; +export type WikiPageType = 'concept' | 'entity'; 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 b59f778c9d..3f565c7e59 100644 --- a/web/src/pages/files/files-table.tsx +++ b/web/src/pages/files/files-table.tsx @@ -52,7 +52,8 @@ 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'; -import { isGoDatasetBackend } from '../../utils/api-proxy-scheme'; + +declare const __API_PROXY_SCHEME__: string; type FilesTableProps = Pick< ReturnType, @@ -103,7 +104,13 @@ export function FilesTable({ } = useRenameCurrentFile(); // Check if skills feature is enabled (only in hybrid or go mode) - const isSkillsEnabled = useMemo(() => isGoDatasetBackend(), []); + const isSkillsEnabled = useMemo(() => { + const scheme = + typeof __API_PROXY_SCHEME__ !== 'undefined' + ? __API_PROXY_SCHEME__ + : 'python'; + return scheme === 'hybrid' || scheme === 'go'; + }, []); // 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 91dba1b196..62dabd62ba 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, { registerNextServer } from '@/utils/register-server'; +import registerServer from '@/utils/register-server'; import request from '@/utils/request'; const { @@ -287,20 +287,6 @@ 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 deleted file mode 100644 index 9f247da6a3..0000000000 --- a/web/src/utils/api-proxy-scheme.ts +++ /dev/null @@ -1,48 +0,0 @@ -/** - * 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 b06f4bd1d6..83c8371b6b 100644 --- a/web/src/utils/api.ts +++ b/web/src/utils/api.ts @@ -215,9 +215,6 @@ 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`,