mtfyun/ragflow
1
0 Fork 0
mirror of https://github.com/infiniflow/ragflow.git synced 2026-11-10 21:32:43 +08:00
Code Issues Packages Projects Releases Wiki Activity

Revert "feat: Go knowledge compiler with scheduler-driven dataset compilation" (#17897)

Browse Source 
Reverts infiniflow/ragflow#17881
This commit is contained in:
Jin Hai authored and GitHub committed 2026-08-05 21:50:28 +08:00
1 parent eaf553320f
commit cf13082a1a
165 files changed
+3172 -6959
Download Patch File Download Diff File Expand all files Collapse all files

No files matched your search

AGENTS.md
-1
View File
@@ -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.
agent/templates/compiler.json
+2 -4
View File
@@ -5,8 +5,7 @@
"obj": {
"component_name": "Compiler",
"params": {
"compilation_template_group_id": "",
"llm_id": "",
"compilation_template_group_id": "c3aa748c8b2111f191f3047c16ec874f",
"outputs": {
"chunks": {
"type": "Array<Object>",
@@ -409,8 +408,7 @@
{
"data": {
"form": {
"compilation_template_group_id": "",
"llm_id": "",
"compilation_template_group_id": "c3aa748c8b2111f191f3047c16ec874f",
"outputs": {
"chunks": {
"type": "Array<Object>",
api/apps/restful_apis/agent_api.py
+55 -7
View File
@@ -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,
)
api/db/services/canvas_service.py
+3 -3
View File
@@ -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:
cmd/ragflow_server.go
-9
View File
@@ -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,
conf/service_conf.yaml
+1
View File
@@ -68,6 +68,7 @@ otel:
enable: false
ingestor:
mq_type: 'nats'
max_concurrent_workers: 1
file_syncer:
max_concurrent_syncs: 1
sync_interval: 3
deepdoc/parser/figure_parser.py
+34 -17
View File
@@ -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,
docker/service_conf.yaml.template
+1
View File
@@ -86,6 +86,7 @@ otel:
enable: false
ingestor:
mq_type: 'nats'
max_concurrent_workers: 1
file_syncer:
max_concurrent_syncs: 1
sync_interval: 3
internal/agent/harness/agentic_rag.go
+2 -13
View File
@@ -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 {
internal/agent/harness/production.go
+11 -313
View File
@@ -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
internal/agent/harness/production_wiki_test.go
-276
View File
@@ -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])
}
}
internal/agent/harness/route.go
+3 -22
View File
@@ -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",
internal/agent/harness/route_test.go
-42
View File
@@ -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) {
internal/agent/harness/types.go
+1 -7
View File
@@ -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.
internal/agent/tool/registry.go
-1
View File
@@ -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,
internal/agent/tool/wiki_query.go
-115
View File
@@ -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":[]}`
}
internal/agent/tool/wiki_query_test.go
-125
View File
@@ -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")
}
}
internal/common/error_code.go
+1
View File
@@ -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
internal/common/parser_config.go
+124
View File
@@ -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 {
internal/common/parser_config_test.go
+135
View File
@@ -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")
}
}
internal/entity/knowledge_compile_doc.go
+2 -24
View File
@@ -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.
internal/entity/models/anthropic.go
+5 -4
View File
@@ -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
}
internal/entity/models/anthropic_test.go
+4 -7
View File
@@ -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"])
internal/entity/models/astraflow_test.go
-3
View File
@@ -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"])
}
internal/entity/models/avian_test.go
-3
View File
@@ -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"])
}
internal/entity/models/azure_openai.go
+6
View File
@@ -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
}
internal/entity/models/base_model_request_body_test.go
-31
View File
@@ -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"])
}
}
internal/entity/models/bedrock.go
+4
View File
@@ -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
internal/entity/models/bedrock_test.go
+2 -2
View File
@@ -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)
internal/entity/models/cometapi_test.go
-3
View File
@@ -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"])
}
internal/entity/models/google.go
+7
View File
@@ -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
}
internal/entity/models/google_test.go
+6 -6
View File
@@ -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))
}
}
internal/entity/models/gpustack_test.go
-3
View File
@@ -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"},
internal/entity/models/groq_test.go
-3
View File
@@ -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"])
}
internal/entity/models/huaweicloud.go
+3
View File
@@ -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
}
internal/entity/models/hunyuan_test.go
-3
View File
@@ -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"])
}
internal/entity/models/jina_test.go
-3
View File
@@ -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"])
}
internal/entity/models/longcat_test.go
+3 -6
View File
@@ -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"},
internal/entity/models/mistral_test.go
-3
View File
@@ -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"])
}
internal/entity/models/modelscope_test.go
-3
View File
@@ -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"])
}
internal/entity/models/ppio_test.go
-3
View File
@@ -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"])
}
internal/entity/models/replicate.go
+3
View File
@@ -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
}
internal/entity/models/replicate_test.go
+2 -2
View File
@@ -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.
internal/entity/models/tokenpony_test.go
-3
View File
@@ -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"])
}
internal/entity/models/upstage_test.go
-3
View File
@@ -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")
}
internal/entity/models/xiaomi.go
+8 -2
View File
@@ -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{}{
internal/entity/models/xiaomi_test.go
+2 -2
View File
@@ -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" {
internal/entity/models/xinference_test.go
-3
View File
@@ -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"])
}
internal/entity/models/zhipu-ai.go
+13 -141
View File
@@ -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)
internal/handler/compilation_status_test.go
-195
View File
@@ -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")
}
}
internal/handler/components_testpkg/components_handler_test.go
+5 -5
View File
@@ -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",
}
internal/handler/dataset.go
+92 -6
View File
@@ -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.
internal/handler/dataset_artifact.go
+2 -6
View File
@@ -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/<page_type>/<slug> — 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.
internal/ingestion/component/chunker/token.go
+14 -20
View File
@@ -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)
}
}
internal/ingestion/component/chunker/token_oversize_whole_test.go
+139
View File
@@ -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)
}
})
}
}
internal/ingestion/component/chunker/token_strict_cap_test.go
-44
View File
@@ -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)))
internal/ingestion/component/knowledge_compiler/common/deps.go
+4 -5
View File
@@ -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.
internal/ingestion/component/knowledge_compiler/common/jsonchat.go
+3 -16
View File
@@ -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 {
internal/ingestion/component/knowledge_compiler/component.go
+55 -119
View File
@@ -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 "<page_type>/<slug>" 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 "<page_type>/<slug>" 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_<dim>_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_<dim>_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,
})
}
internal/ingestion/component/knowledge_compiler/component_test.go
+16 -160
View File
@@ -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)
internal/ingestion/component/knowledge_compiler/golden_test.go
+1 -1
View File
@@ -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)
}
internal/ingestion/component/knowledge_compiler/pool_wiring.go
-2
View File
@@ -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)
}
internal/ingestion/component/knowledge_compiler/tree/raptor.go
+6 -6
View File
@@ -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
}
internal/ingestion/component/knowledge_compiler/wiki/page_test.go
-21
View File
@@ -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"},
internal/ingestion/component/knowledge_compiler/wiki/prompt.go
+10 -23
View File
@@ -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
internal/ingestion/component/knowledge_compiler/wiki/wiki.go
+133 -447
View File
@@ -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/<kb>/<page_type>/<slug> 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/<kb>/<page_type>/<slug> 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 <page_type>/<name> 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/<kb>/<page_type>/<slug> so the frontend wiki-link
// parser (artifact/<kb>/<page_type>/<slug>, page_type in {entity,concept})
// can resolve and navigate to the target page. Slugs may already carry a
// <page_type>/<name> 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.
internal/ingestion/component/knowledge_compiler/wiki/wiki_budget.go
-259
View File
@@ -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
}
internal/ingestion/component/knowledge_compiler/wiki/wiki_budget_test.go
-559
View File
@@ -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")
}
}
internal/ingestion/component/knowledge_compiler/wiki/wiki_reduce.go
-185
View File
@@ -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)
}
internal/ingestion/component/knowledge_compiler/wiki/wiki_reduce_test.go
-225
View File
@@ -1,225 +0,0 @@
package wiki
import (
"context"
"errors"
"testing"
"ragflow/internal/ingestion/component/knowledge_compiler/common"
)
// TestDedupeEntities_NoSeamIsNoop verifies dedupeEntities degrades to a no-op
// when the embedder or chat seam is nil (M1-style unit safety).
func TestDedupeEntities_NoSeamIsNoop(t *testing.T) {
p := &wikiPipeline{ctx: context.Background()}
in := []wikiEntity{
{Name: "Alpha", SourceChunkIDs: []string{"c1"}},
{Name: "Alpha Corp", SourceChunkIDs: []string{"c2"}},
}
got := p.dedupeEntities(in)
if len(got) != 2 {
t.Fatalf("got %d entities, want 2 (no-op without seams)", len(got))
}
}
// TestDedupeEntities_LLMMergesAmbiguousPair verifies two distinct-name entities
// with high embedding similarity are collapsed into one canonical entity via the
// LLM merge decision, with aliases and provenance merged.
func TestDedupeEntities_LLMMergesAmbiguousPair(t *testing.T) {
p := &wikiPipeline{
ctx: context.Background(),
llmID: "llm1",
deps: common.Deps{
Chat: reconcileChatStub{resp: `{"merge":true,"reason":"same company"}`},
Embed: mergeEmbedStub{},
},
}
in := []wikiEntity{
{Name: "Alpha Inc", Type: "org", SourceChunkIDs: []string{"c1"}},
{Name: "Alpha Incorporated", Type: "org", SourceChunkIDs: []string{"c2"}},
}
got := p.dedupeEntities(in)
if len(got) != 1 {
t.Fatalf("got %d entities, want 1 (LLM merge)", len(got))
}
if got[0].Name != "Alpha Inc" {
t.Fatalf("canonical name = %q, want Alpha Inc", got[0].Name)
}
if len(got[0].SourceChunkIDs) != 2 {
t.Fatalf("provenance = %#v, want 2 chunk ids", got[0].SourceChunkIDs)
}
if len(got[0].Aliases) == 0 {
t.Fatalf("aliases not merged: %#v", got[0].Aliases)
}
}
// TestDedupeEntities_LLMRejectsDistinct verifies the LLM rejecting a merge keeps
// both entities distinct.
func TestDedupeEntities_LLMRejectsDistinct(t *testing.T) {
p := &wikiPipeline{
ctx: context.Background(),
llmID: "llm1",
deps: common.Deps{
Chat: reconcileChatStub{resp: `{"merge":false,"reason":"distinct products"}`},
Embed: mergeEmbedStub{},
},
}
in := []wikiEntity{
{Name: "Alpha", SourceChunkIDs: []string{"c1"}},
{Name: "Beta", SourceChunkIDs: []string{"c2"}},
}
got := p.dedupeEntities(in)
if len(got) != 2 {
t.Fatalf("got %d entities, want 2 (LLM rejected merge)", len(got))
}
}
// TestDedupeEntities_ExactNameIsNotAmbiguous verifies entities with identical
// normalized names (already collapsed by reduceExtracts before this stage) are
// not treated as ambiguous: they pass through untouched and no LLM call is made
// for the exact-name pair.
func TestDedupeEntities_ExactNameIsNotAmbiguous(t *testing.T) {
p := &wikiPipeline{
ctx: context.Background(),
llmID: "llm1",
deps: common.Deps{
Chat: reconcileChatStub{resp: `{"merge":true}`},
Embed: mergeEmbedStub{},
},
}
in := []wikiEntity{
{Name: "Alpha", SourceChunkIDs: []string{"c1"}},
{Name: "Alpha", SourceChunkIDs: []string{"c2"}},
}
got := p.dedupeEntities(in)
// Exact-name duplicates are out of scope for the embedding step; both are
// kept unchanged (they would already be one entity after reduceExtracts).
if len(got) != 2 {
t.Fatalf("got %d entities, want 2 (exact-name pairs are not ambiguous)", len(got))
}
}
// TestDedupeEntities_ConceptStaysExact validates the REDUCE boundary: concept
// dedup must remain exact (no embedding/LLM). This test guards that the entity
// enhancement never touches concepts.
func TestReduceExtracts_ConceptsStayExact(t *testing.T) {
reduced := reduceExtracts([]wikiExtract{
{Concepts: []wikiConcept{{Term: "RAG", Definition: "d1", SourceChunkIDs: []string{"c1"}}}},
{Concepts: []wikiConcept{{Term: "Retrieval Augmented Generation", Definition: "d2", SourceChunkIDs: []string{"c2"}}}},
})
if len(reduced.Concepts) != 2 {
t.Fatalf("concepts = %d, want 2 (exact-term dedup must not collapse distinct terms)", len(reduced.Concepts))
}
}
// TestDedupeEntities_FailingChatIsBudgeted locks F4: a chat seam that
// persistently fails must not drive unbounded external calls. The llmCalls
// budget is consumed before the request, so the loop stops after at most
// wikiEntityMergeMaxCalls attempts.
func TestDedupeEntities_FailingChatIsBudgeted(t *testing.T) {
calls := 0
p := &wikiPipeline{
ctx: context.Background(),
llmID: "llm1",
deps: common.Deps{
Chat: chatFunc(func(_ context.Context, _ common.ChatRequest) (*common.ChatResponse, error) {
calls++
return nil, errors.New("llm down")
}),
Embed: mergeEmbedStub{},
},
}
// 20 entities with distinct names but identical embeddings => every pair is
// ambiguous and would trigger an LLM call.
in := make([]wikiEntity, 0, 20)
for i := 0; i < 20; i++ {
in = append(in, wikiEntity{Name: "Entity " + itoa(i), Type: "person"})
}
got := p.dedupeEntities(in)
if calls > wikiEntityMergeMaxCalls {
t.Fatalf("chat calls = %d, want <= %d despite persistent failures", calls, wikiEntityMergeMaxCalls)
}
// No merges happen because every call fails, so all 20 entities survive.
if len(got) != 20 {
t.Fatalf("entities = %d, want 20 (no merges on failure)", len(got))
}
}
// TestDedupeEntities_CrossTypeHighSimDoesNotCallLLM locks F5: entities with
// provably different types must never be treated as ambiguous, even with
// identical embeddings, so no LLM call is made for them.
func TestDedupeEntities_CrossTypeHighSimDoesNotCallLLM(t *testing.T) {
calls := 0
p := &wikiPipeline{
ctx: context.Background(),
llmID: "llm1",
deps: common.Deps{
Chat: chatFunc(func(_ context.Context, _ common.ChatRequest) (*common.ChatResponse, error) {
calls++
return &common.ChatResponse{Content: `{"merge":true}`}, nil
}),
Embed: mergeEmbedStub{},
},
}
// Both embed to [1,1,1] (identical vectors, cosine 1.0) but types differ.
in := []wikiEntity{
{Name: "Alpha", Type: "person"},
{Name: "Beta Corp", Type: "org"},
}
got := p.dedupeEntities(in)
if calls != 0 {
t.Fatalf("chat calls = %d, want 0 (cross-type pairs must not reach the LLM)", calls)
}
if len(got) != 2 {
t.Fatalf("entities = %d, want 2 (cross-type entities must stay distinct)", len(got))
}
}
// mergeEmbedStub returns embeddings where identical names share a vector and
// distinct names are far apart (cosine ~0), so it can drive both the ambiguous
// and distinct test paths deterministically.
type mergeEmbedStub struct{}
func (mergeEmbedStub) Encode(_ context.Context, texts []string) ([][]float32, error) {
out := make([][]float32, len(texts))
for i, text := range texts {
// "Alpha Inc" and "Alpha Incorporated" both contain "alpha" -> same
// vector; "Beta" differs.
if containsFold(text, "beta") {
out[i] = []float32{1, 0, 0}
continue
}
out[i] = []float32{1, 1, 1}
}
return out, nil
}
func (mergeEmbedStub) Dimensions() int { return 3 }
func containsFold(s, sub string) bool {
return len(s) >= len(sub) && (len(sub) == 0 || indexFold(s, sub) >= 0)
}
func indexFold(s, sub string) int {
if sub == "" {
return 0
}
ls := toLowerASCII(s)
lsub := toLowerASCII(sub)
for i := 0; i+len(lsub) <= len(ls); i++ {
if ls[i:i+len(lsub)] == lsub {
return i
}
}
return -1
}
func toLowerASCII(s string) string {
b := []byte(s)
for i := range b {
if b[i] >= 'A' && b[i] <= 'Z' {
b[i] += 'a' - 'A'
}
}
return string(b)
}
internal/ingestion/component/knowledge_compiler/wiki/wiki_refine_test.go
-156
View File
@@ -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")
}
}
internal/ingestion/component/knowledge_compiler/wiki/wiki_test.go
-170
View File
@@ -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(),
internal/ingestion/knowledge_compile/consumer.go
+2 -23
View File
@@ -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
}
internal/ingestion/knowledge_compile/consumer_test.go
-151
View File
@@ -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)
}
}
internal/ingestion/knowledge_compile/dedup.go
+1 -1
View File
@@ -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}
}
internal/ingestion/knowledge_compile/pool.go
+11 -12
View File
@@ -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)
}
internal/ingestion/knowledge_compile/reader.go
+2 -41
View File
@@ -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 "<page_type>/<slug>" 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),
internal/ingestion/knowledge_compile/scheduler.go
+14 -147
View File
@@ -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 != "" {
internal/ingestion/knowledge_compile/service.go
+1 -1
View File
@@ -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 {
internal/ingestion/knowledge_compile/writer.go
+2 -73
View File
@@ -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 "<page_type>/<slug>" 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 "<page_type>/<slug>" 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
}
internal/ingestion/knowledge_compile/writer_test.go
-98
View File
@@ -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")
}
}
internal/ingestion/pipeline/pipeline_knowledge_compiler_dsl_test.go
+16 -24
View File
@@ -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)
}
internal/ingestion/pipeline/pipeline_knowledge_compiler_test.go
+1 -1
View File
@@ -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)
}
internal/ingestion/pipeline/template/ingestion_pipeline_knowledge_compiler.json
+10 -11
View File
@@ -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
internal/ingestion/pipeline/template_integration_test.go
+2 -2
View File
@@ -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 {
internal/ingestion/task/knowledge_compiler_wiring.go
+24 -114
View File
@@ -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)
internal/ingestion/task/knowledge_compiler_wiring_test.go
-43
View File
@@ -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")
}
}
internal/ingestion/task/pipeline_executor.go
+8
View File
@@ -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;
internal/ingestion/task/pipeline_executor_defaults_test.go
+1 -1
View File
@@ -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*)?(?:教育背景|教育经历|学历背景|学术背景|技术背景|工作经历|工作经验|实习经历|项目经历|项目经验|科研经历|研究经历|校园经历|实践经历|专业经历|职业经历|技能|专业技能|技能特长|核心技能|技术栈|个人技能|工作技能|职业技能|技能与评价|技能与自我评价|Line truncated
"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
internal/ingestion/task/pipeline_executor_test.go
+1 -1
View File
@@ -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)
internal/router/router.go
+5 -3
View File
@@ -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)
internal/server/config.go
+5
View File
@@ -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)
internal/server/config/base.go
+1
View File
@@ -30,6 +30,7 @@ type Config struct {
admin AdminConfig
apiServer APIServerConfig
ingestor IngestorConfig
syncer SyncerConfig
log LogConfig
internal/server/config/ingestor_config.go
+42
View File
@@ -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
}
internal/service/admin_client.go
+35 -15
View File
@@ -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),
internal/service/admin_client_ee.go
+20
View File
@@ -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
}
internal/service/component_scoped_parser_config.go
-118
View File
@@ -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
}
internal/service/component_scoped_parser_config_test.go
-147
View File
@@ -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)
}
}
internal/service/dataset/compilation_status.go
-84
View File
@@ -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)
}
internal/service/dataset/compilation_status_test.go
-197
View File
@@ -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)
}
}
Loaded 100 of 165 files, more files were not shown because too many files have changed in this diff. Show more