From 90f46b0b4dc79de8610daf3f16bff3c77a4fe154 Mon Sep 17 00:00:00 2001 From: Zhichang Yu Date: Wed, 29 Jul 2026 21:06:48 +0800 Subject: [PATCH] Go port: doc-level metadata extraction and knowledge compiler (#17536) Ports doc-level auto-metadata extraction to Go and adds the knowledge_compiler component with scheduler/routing. Fixes Extractor metadata injection type assertion and enable_metadata default-on. --- cmd/ragflow_server.go | 26 +- go.mod | 2 + go.sum | 1 + internal/common/metadata_utils.go | 134 ++ internal/common/metadata_utils_test.go | 86 + internal/common/parser_config.go | 98 + internal/common/parser_config_test.go | 82 + internal/common/task.go | 9 + internal/dao/compilation_template.go | 94 + internal/dao/database.go | 1 + internal/engine/engine.go | 19 + internal/engine/nats/knowledgecompile.go | 272 +++ internal/engine/nats/nats.go | 5 + internal/entity/knowledge_compile_doc.go | 46 + internal/ingestion/component/extractor.go | 469 +++- .../ingestion/component/extractor_test.go | 241 ++ .../component/knowledge_compiler/PORT_PLAN.md | 416 ++++ .../knowledge_compiler/common/batch.go | 23 + .../common/batch_packing.go | 46 + .../knowledge_compiler/common/common_test.go | 219 ++ .../knowledge_compiler/common/deps.go | 174 ++ .../component/knowledge_compiler/common/id.go | 45 + .../knowledge_compiler/common/id_test.go | 21 + .../knowledge_compiler/common/jsonchat.go | 74 + .../knowledge_compiler/common/memstore.go | 282 +++ .../knowledge_compiler/common/parserconfig.go | 98 + .../knowledge_compiler/common/tokenize.go | 20 + .../knowledge_compiler/common/types.go | 392 ++++ .../component/knowledge_compiler/component.go | 652 ++++++ .../knowledge_compiler/component_test.go | 1264 +++++++++++ .../datasetnav/datasetnav.go | 788 +++++++ .../datasetnav/datasetnav_test.go | 289 +++ .../knowledge_compiler/golden/corpus.go | 55 + .../golden/fixtures/raptor_baseline.json | 12 + .../fixtures/structure_wiki_baseline.json | 7 + .../knowledge_compiler/golden/golden_test.go | 140 ++ .../knowledge_compiler/golden/metrics.go | 113 + .../knowledge_compiler/golden_test.go | 287 +++ .../knowledge_compiler/mindmap/dictify.go | 417 ++++ .../knowledge_compiler/mindmap/mindmap.go | 250 +++ .../mindmap/mindmap_test.go | 239 ++ .../knowledge_compiler/mindmap/prompt.go | 88 + .../knowledge_compiler/raptor/clustering.go | 680 ++++++ .../raptor/clustering_test.go | 84 + .../knowledge_compiler/raptor/psi.go | 119 + .../knowledge_compiler/raptor/raptor.go | 341 +++ .../knowledge_compiler/structure/chain.go | 473 ++++ .../structure/chain_test.go | 210 ++ .../knowledge_compiler/structure/compile.go | 321 +++ .../knowledge_compiler/structure/graph.go | 156 ++ .../knowledge_compiler/structure/merge.go | 561 +++++ .../knowledge_compiler/structure/prompt.go | 384 ++++ .../knowledge_compiler/structure/structure.go | 205 ++ .../structure/structure_test.go | 613 +++++ .../component/knowledge_compiler/wiki/page.go | 312 +++ .../knowledge_compiler/wiki/page_test.go | 232 ++ .../knowledge_compiler/wiki/prompt.go | 459 ++++ .../component/knowledge_compiler/wiki/wiki.go | 1972 +++++++++++++++++ .../knowledge_compiler/wiki/wiki_test.go | 185 ++ .../ingestion/component/schema/extractor.go | 33 +- .../ingestion/knowledge_compile/consumer.go | 285 +++ .../knowledge_compile/consumer_test.go | 191 ++ internal/ingestion/knowledge_compile/dedup.go | 74 + internal/ingestion/knowledge_compile/event.go | 115 + .../ingestion/knowledge_compile/reader.go | 168 ++ .../ingestion/knowledge_compile/scheduler.go | 534 +++++ .../ingestion/knowledge_compile/service.go | 146 ++ .../ingestion/knowledge_compile/writer.go | 177 ++ .../pipeline_knowledge_compiler_test.go | 70 + .../template/ingestion_pipeline_audio.json | 4 +- .../template/ingestion_pipeline_book.json | 4 +- .../template/ingestion_pipeline_email.json | 4 +- .../template/ingestion_pipeline_general.json | 4 +- ...ingestion_pipeline_knowledge_compiler.json | 310 +++ .../template/ingestion_pipeline_laws.json | 4 +- .../template/ingestion_pipeline_manual.json | 4 +- .../template/ingestion_pipeline_one.json | 4 +- .../template/ingestion_pipeline_paper.json | 4 +- .../template/ingestion_pipeline_picture.json | 4 +- .../ingestion_pipeline_presentation.json | 4 +- .../template/ingestion_pipeline_resume.json | 4 +- .../pipeline/template_integration_test.go | 3 + internal/ingestion/service/doc_state.go | 22 +- internal/ingestion/service/doc_state_test.go | 50 + .../ingestion/service/ingestion_service.go | 139 +- .../task/knowledge_compiler_wiring.go | 291 +++ internal/ingestion/task/pipeline_executor.go | 45 +- .../task/pipeline_executor_defaults_test.go | 30 +- .../ingestion/task/pipeline_executor_test.go | 28 + internal/service/document/document_crud.go | 14 + internal/utility/metadata.go | 32 +- internal/utility/metadata_test.go | 46 +- ragflow_deps/download_go_deps.py | 3 +- 93 files changed, 17967 insertions(+), 186 deletions(-) create mode 100644 internal/dao/compilation_template.go create mode 100644 internal/engine/nats/knowledgecompile.go create mode 100644 internal/entity/knowledge_compile_doc.go create mode 100644 internal/ingestion/component/knowledge_compiler/PORT_PLAN.md create mode 100644 internal/ingestion/component/knowledge_compiler/common/batch.go create mode 100644 internal/ingestion/component/knowledge_compiler/common/batch_packing.go create mode 100644 internal/ingestion/component/knowledge_compiler/common/common_test.go create mode 100644 internal/ingestion/component/knowledge_compiler/common/deps.go create mode 100644 internal/ingestion/component/knowledge_compiler/common/id.go create mode 100644 internal/ingestion/component/knowledge_compiler/common/id_test.go create mode 100644 internal/ingestion/component/knowledge_compiler/common/jsonchat.go create mode 100644 internal/ingestion/component/knowledge_compiler/common/memstore.go create mode 100644 internal/ingestion/component/knowledge_compiler/common/parserconfig.go create mode 100644 internal/ingestion/component/knowledge_compiler/common/tokenize.go create mode 100644 internal/ingestion/component/knowledge_compiler/common/types.go create mode 100644 internal/ingestion/component/knowledge_compiler/component.go create mode 100644 internal/ingestion/component/knowledge_compiler/component_test.go create mode 100644 internal/ingestion/component/knowledge_compiler/datasetnav/datasetnav.go create mode 100644 internal/ingestion/component/knowledge_compiler/datasetnav/datasetnav_test.go create mode 100644 internal/ingestion/component/knowledge_compiler/golden/corpus.go create mode 100644 internal/ingestion/component/knowledge_compiler/golden/fixtures/raptor_baseline.json create mode 100644 internal/ingestion/component/knowledge_compiler/golden/fixtures/structure_wiki_baseline.json create mode 100644 internal/ingestion/component/knowledge_compiler/golden/golden_test.go create mode 100644 internal/ingestion/component/knowledge_compiler/golden/metrics.go create mode 100644 internal/ingestion/component/knowledge_compiler/golden_test.go create mode 100644 internal/ingestion/component/knowledge_compiler/mindmap/dictify.go create mode 100644 internal/ingestion/component/knowledge_compiler/mindmap/mindmap.go create mode 100644 internal/ingestion/component/knowledge_compiler/mindmap/mindmap_test.go create mode 100644 internal/ingestion/component/knowledge_compiler/mindmap/prompt.go create mode 100644 internal/ingestion/component/knowledge_compiler/raptor/clustering.go create mode 100644 internal/ingestion/component/knowledge_compiler/raptor/clustering_test.go create mode 100644 internal/ingestion/component/knowledge_compiler/raptor/psi.go create mode 100644 internal/ingestion/component/knowledge_compiler/raptor/raptor.go create mode 100644 internal/ingestion/component/knowledge_compiler/structure/chain.go create mode 100644 internal/ingestion/component/knowledge_compiler/structure/chain_test.go create mode 100644 internal/ingestion/component/knowledge_compiler/structure/compile.go create mode 100644 internal/ingestion/component/knowledge_compiler/structure/graph.go create mode 100644 internal/ingestion/component/knowledge_compiler/structure/merge.go create mode 100644 internal/ingestion/component/knowledge_compiler/structure/prompt.go create mode 100644 internal/ingestion/component/knowledge_compiler/structure/structure.go create mode 100644 internal/ingestion/component/knowledge_compiler/structure/structure_test.go create mode 100644 internal/ingestion/component/knowledge_compiler/wiki/page.go create mode 100644 internal/ingestion/component/knowledge_compiler/wiki/page_test.go create mode 100644 internal/ingestion/component/knowledge_compiler/wiki/prompt.go create mode 100644 internal/ingestion/component/knowledge_compiler/wiki/wiki.go create mode 100644 internal/ingestion/component/knowledge_compiler/wiki/wiki_test.go create mode 100644 internal/ingestion/knowledge_compile/consumer.go create mode 100644 internal/ingestion/knowledge_compile/consumer_test.go create mode 100644 internal/ingestion/knowledge_compile/dedup.go create mode 100644 internal/ingestion/knowledge_compile/event.go create mode 100644 internal/ingestion/knowledge_compile/reader.go create mode 100644 internal/ingestion/knowledge_compile/scheduler.go create mode 100644 internal/ingestion/knowledge_compile/service.go create mode 100644 internal/ingestion/knowledge_compile/writer.go create mode 100644 internal/ingestion/pipeline/pipeline_knowledge_compiler_test.go create mode 100644 internal/ingestion/pipeline/template/ingestion_pipeline_knowledge_compiler.json create mode 100644 internal/ingestion/task/knowledge_compiler_wiring.go diff --git a/cmd/ragflow_server.go b/cmd/ragflow_server.go index 46f6300716..043a7498dc 100644 --- a/cmd/ragflow_server.go +++ b/cmd/ragflow_server.go @@ -536,15 +536,27 @@ func runIngestor(ctx context.Context, cancel context.CancelFunc, args *serverArg } defer tokenizer.Close() + // The dataset-level post-processing consumer cluster (§11) is owned and run by + // the Ingestor: it is started inside ingestor.Start() and joined inside + // ingestor.Stop(), so its lifecycle matches the ingestor. The configured + // default LLM/embedding ids are passed so the LLM deduper is used (instead + // of the noop fallback that still emits merged products). Best-effort: a + // provisioning error is logged by the Ingestor and the pipeline still + // writes available_int=0 compiled chunks; they just won't be merged until + // the consumer is available. + cfg := server.GetConfig() ingestor := ingestion.NewIngestor(*args.name, 2, []string{"pdf", "docx", "txt"}) + ingestor.SetKnowledgeCompileModelConfig( + cfg.UserDefaultLLM.DefaultModels.ChatModel.Name, + cfg.UserDefaultLLM.DefaultModels.EmbeddingModel.Name, + ) - go func() { - err := ingestor.Start() - if err != nil { - common.Error("Failed to initialize ingestor", err) - return - } - }() + // Start returns immediately (it launches the owned consume/compile + // goroutines and joins them via Stop); a provisioning failure here is + // logged and the server keeps running without the ingestor. + if err := ingestor.Start(); err != nil { + common.Error("Failed to initialize ingestor", err) + } common.Info("\n ____ __ _\n" + " / _/___ ____ ____ _____/ /_(_)___ ____ ________ ______ _____ _____\n" + diff --git a/go.mod b/go.mod index 8ee7833934..7e788c96ba 100644 --- a/go.mod +++ b/go.mod @@ -51,6 +51,7 @@ require ( github.com/xuri/excelize/v2 v2.11.0 github.com/yfedoseev/office_oxide/go v0.1.8 github.com/yfedoseev/pdf_oxide/go v0.3.67 + github.com/yuin/goldmark v1.4.13 github.com/zeebo/xxh3 v1.0.2 go.opentelemetry.io/contrib/instrumentation/net/http/otelhttp v0.69.0 go.opentelemetry.io/otel v1.44.0 @@ -62,6 +63,7 @@ require ( golang.org/x/sync v0.21.0 golang.org/x/term v0.44.0 golang.org/x/text v0.38.0 + gonum.org/v1/gonum v0.17.0 google.golang.org/api v0.287.1 google.golang.org/genai v1.54.0 gopkg.in/natefinch/lumberjack.v2 v2.2.1 diff --git a/go.sum b/go.sum index b70624827f..69b65466dc 100644 --- a/go.sum +++ b/go.sum @@ -523,6 +523,7 @@ github.com/yfedoseev/pdf_oxide/go v0.3.67 h1:Fm1R/KtpmJPNbVmdT1fvYM/Yl41Uu2FdyT7 github.com/yfedoseev/pdf_oxide/go v0.3.67/go.mod h1:QbJ/nLbez0al2EnqEdEPIlGflFprWmiuUM4mo9rNNOI= github.com/yuin/goldmark v1.1.27/go.mod h1:3hX8gzYuyVAZsxl0MRgGTJEmQBFcNTphYh9decYSb74= github.com/yuin/goldmark v1.1.30/go.mod h1:3hX8gzYuyVAZsxl0MRgGTJEmQBFcNTphYh9decYSb74= +github.com/yuin/goldmark v1.4.13 h1:fVcFKWvrslecOb/tg+Cc05dkeYx540o0FuFt3nUVDoE= github.com/yuin/goldmark v1.4.13/go.mod h1:6yULJ656Px+3vBD8DxQVa3kxgyrAnzto9xy5taEt/CY= github.com/yuin/gopher-lua v1.1.1 h1:kYKnWBjvbNP4XLT3+bPEwAXJx262OhaHDWDVOPjL46M= github.com/yuin/gopher-lua v1.1.1/go.mod h1:GBR0iDaNXjAgGg9zfCvksxSRnQx76gclCIb7kdAd1Pw= diff --git a/internal/common/metadata_utils.go b/internal/common/metadata_utils.go index 1786e506f2..85e9e79da2 100644 --- a/internal/common/metadata_utils.go +++ b/internal/common/metadata_utils.go @@ -17,6 +17,7 @@ package common import ( + "regexp" "strconv" "strings" ) @@ -370,3 +371,136 @@ func toString(v interface{}) string { return "" } } + +// MetadataFieldDef describes one auto-metadata field, mirroring the +// {key, type, description, enum} shape stored in parser_config.metadata / +// built_in_metadata and the Python metadata_utils.py field contract. +type MetadataFieldDef struct { + Key string `json:"key"` + Type string `json:"type,omitempty"` + Description string `json:"description,omitempty"` + Enum []string `json:"enum,omitempty"` +} + +// Turn2JSONSchema converts a metadata field list into a JSON-Schema object, +// mirroring Python common/metadata_utils.py::turn2jsonschema / +// metadata_schema. The result is rendered into the auto-metadata extraction +// prompt (rag/prompts/meta_data.md equivalent) so the LLM knows the exact +// keys, descriptions and allowed enum values to extract. +func Turn2JSONSchema(fields []MetadataFieldDef) map[string]any { + properties := make(map[string]any, len(fields)) + for _, f := range fields { + if strings.TrimSpace(f.Key) == "" { + continue + } + prop := map[string]any{} + if f.Description != "" { + prop["description"] = f.Description + } + if len(f.Enum) > 0 { + prop["enum"] = f.Enum + prop["type"] = "string" + } else if f.Type != "" { + prop["type"] = f.Type + } + properties[f.Key] = prop + } + if len(properties) == 0 { + return map[string]any{} + } + return map[string]any{ + "type": "object", + "properties": properties, + "additionalProperties": false, + } +} + +// combinedValueDelim splits a single combined metadata value into its parts. +// Mirrors Python doc_metadata_service.py _split_combined_values regex +// r"[、,，;；|]+" (Chinese comma 、, ASCII comma ,, full-width comma ，, +// semicolon ;, full-width semicolon ；, pipe |). +var combinedValueDelim = regexp.MustCompile(`[、,，;；|]+`) + +// SplitCombinedMetadataValues post-processes a metadata map by splitting +// combined values written to the doc-metadata index, mirroring Python +// api/db/services/doc_metadata_service.py:_split_combined_values (applied at +// both insert_document_metadata:383 and update_document_metadata:468). +// +// Only list-typed values are split (each string element is split on the +// delimiter, trimmed, empties dropped, then flattened and de-duplicated); +// scalar string values are left untouched, matching the Python implementation. +// Non-string / non-list values pass through unchanged. +func SplitCombinedMetadataValues(meta map[string]any) map[string]any { + if len(meta) == 0 { + return meta + } + out := make(map[string]any, len(meta)) + for k, v := range meta { + switch val := v.(type) { + case []string: + out[k] = splitCombinedStringList(val) + case []any: + strs, allStr := toStringSlice(val) + if !allStr { + out[k] = v + continue + } + out[k] = splitCombinedStringList(strs) + default: + out[k] = v + } + } + return out +} + +// splitCombinedStringList splits each element on the combined-value delimiter, +// trims, drops empties, keeps single elements intact, and de-duplicates. +func splitCombinedStringList(items []string) []string { + out := make([]string, 0, len(items)) + for _, item := range items { + if item == "" { + continue + } + parts := combinedValueDelim.Split(strings.TrimSpace(item), -1) + got := false + for _, p := range parts { + p = strings.TrimSpace(p) + if p == "" { + continue + } + out = append(out, p) + got = true + } + if !got { + out = append(out, item) + } + } + return dedupeMetaStrings(out) +} + +// toStringSlice converts a []any to []string when every element is a string, +// reporting allStr=false otherwise. +func toStringSlice(val []any) ([]string, bool) { + strs := make([]string, 0, len(val)) + for _, e := range val { + s, ok := e.(string) + if !ok { + return nil, false + } + strs = append(strs, s) + } + return strs, true +} + +// dedupeMetaStrings removes duplicates while preserving order. +func dedupeMetaStrings(input []string) []string { + seen := make(map[string]struct{}, len(input)) + out := make([]string, 0, len(input)) + for _, s := range input { + if _, ok := seen[s]; !ok { + seen[s] = struct{}{} + out = append(out, s) + } + } + return out +} diff --git a/internal/common/metadata_utils_test.go b/internal/common/metadata_utils_test.go index 12bb87d36a..4a5b01101d 100644 --- a/internal/common/metadata_utils_test.go +++ b/internal/common/metadata_utils_test.go @@ -20,6 +20,92 @@ import ( "testing" ) +func TestTurn2JSONSchema(t *testing.T) { + fields := []MetadataFieldDef{ + {Key: "author", Description: "doc author", Enum: []string{"Alice", "Bob"}}, + {Key: "year", Type: "number"}, + {Key: "", Type: "string"}, // blank key must be dropped + } + schema := Turn2JSONSchema(fields) + if schema["type"] != "object" { + t.Fatalf("expected object type, got %v", schema["type"]) + } + if schema["additionalProperties"] != false { + t.Fatalf("expected additionalProperties false, got %v", schema["additionalProperties"]) + } + props, ok := schema["properties"].(map[string]any) + if !ok { + t.Fatal("expected properties map") + } + if len(props) != 2 { + t.Fatalf("expected 2 properties (blank key dropped), got %d", len(props)) + } + author, ok := props["author"].(map[string]any) + if !ok { + t.Fatal("author property missing") + } + if author["type"] != "string" { + t.Fatalf("enum field should be typed string, got %v", author["type"]) + } + if _, ok := author["enum"]; !ok { + t.Fatal("author should carry enum") + } + if _, ok := author["description"]; !ok { + t.Fatal("author should carry description") + } +} + +func TestTurn2JSONSchema_Empty(t *testing.T) { + if len(Turn2JSONSchema(nil)) != 0 { + t.Fatal("expected empty schema for nil fields") + } +} + +func TestSplitCombinedMetadataValues(t *testing.T) { + in := map[string]any{ + "people": []string{"关羽、孙权", "张辽", "刘备、关羽"}, + "level": "high", // scalar string is NOT split + "score": float64(0.9), // non-string passes through + "tags": []any{"a|b", "c"}, // []any of strings is split + } + out := SplitCombinedMetadataValues(in) + people, ok := out["people"].([]string) + if !ok { + t.Fatalf("people should be []string, got %T", out["people"]) + } + // 关羽、孙权 -> 关羽,孙权 ; 张辽 ; 刘备、关羽 -> 刘备,关羽 ; dedupe 关羽 + want := []string{"关羽", "孙权", "张辽", "刘备"} + if len(people) != len(want) { + t.Fatalf("people=%v, want %v", people, want) + } + for i := range want { + if people[i] != want[i] { + t.Fatalf("people[%d]=%q, want %q", i, people[i], want[i]) + } + } + if out["level"] != "high" { + t.Fatalf("scalar string must be preserved unchanged, got %v", out["level"]) + } + if out["score"] != float64(0.9) { + t.Fatalf("non-string passes through, got %v", out["score"]) + } + tags, ok := out["tags"].([]string) + if !ok || len(tags) != 3 { + t.Fatalf("tags=%v, want [a b c]", out["tags"]) + } +} + +func TestSplitCombinedMetadataValues_Noop(t *testing.T) { + var nilMap map[string]any + if SplitCombinedMetadataValues(nilMap) != nil { + t.Fatal("nil map must be returned as nil") + } + empty := map[string]any{} + if len(SplitCombinedMetadataValues(empty)) != 0 { + t.Fatal("empty map must stay empty") + } +} + func TestParseAndConvert_OperatorMapping(t *testing.T) { input := map[string]interface{}{ "conditions": []interface{}{ diff --git a/internal/common/parser_config.go b/internal/common/parser_config.go index b20be19618..adc1793b96 100644 --- a/internal/common/parser_config.go +++ b/internal/common/parser_config.go @@ -26,6 +26,104 @@ func InjectExtractorLLMID(parserConfig map[string]interface{}, llmID string) boo return updated } +// InjectExtractorEnableMetadata enables auto-metadata (enable_metadata) extraction +// on every Extractor node when the dataset has enable_metadata on and a +// non-empty field set (metadata and/or built_in_metadata). The dataset-level +// enable_metadata flag is authoritative (mirrors Python task_executor.py:519, +// which reads parser_config directly and never consults a per-node flag): a +// shipped DSL that defaults enable_metadata to 0 is still turned on. Only a +// node the user already turned ON (enable_metadata truthy) is left untouched, +// so an explicit per-node enablement keeps its own config. The field schema is +// taken from parserConfig["metadata"] and parserConfig["built_in_metadata"] +// (combined). Returns whether any entry was updated. +func InjectExtractorEnableMetadata(parserConfig map[string]interface{}) bool { + if parserConfig == nil { + return false + } + if !isTruthy(parserConfig["enable_metadata"]) { + return false + } + fields := metadataFieldDefs(parserConfig) + if len(fields) == 0 { + return false + } + updated := false + for cid, raw := range parserConfig { + compMap, ok := raw.(map[string]interface{}) + if !ok { + continue + } + cidLower := strings.ToLower(cid) + if !strings.HasPrefix(cidLower, "extractor:") && !strings.HasPrefix(cidLower, "extractor_") { + continue + } + // The dataset-level enable_metadata flag is authoritative (mirrors + // Python task_executor.py:519, which reads parser_config directly and + // never consults a per-node flag). Only a node the user already turned + // ON (truthy) is left alone; a shipped DSL that defaults the field to + // 0 must still be enabled by the dataset flag, otherwise auto-metadata + // could never turn on for any of the built-in pipelines. + if isTruthy(compMap["enable_metadata"]) { + continue + } + compMap["enable_metadata"] = 1 + compMap["metadata"] = fields + updated = true + } + return updated +} + +// metadataFieldDefs combines parserConfig["metadata"] and +// parserConfig["built_in_metadata"] into the field list injected as +// metadata (each entry keeps key/type/description/enum, mirroring the +// stored shape from dataset/helpers.go normalizeMetadataConfigFields). +// +// It returns []any (i.e. []interface{}) rather than []map[string]interface{} +// because the injected value is handed to NewExtractorComponent, which reads +// params["metadata"].([]any). A []map[string]interface{} value would fail that +// type assertion (Go slice types are not covariant) and the field schema would +// be silently dropped, so auto-metadata never reached ExtractorParam.Metadata. +func metadataFieldDefs(parserConfig map[string]interface{}) []any { + var out []any + for _, key := range []string{"metadata", "built_in_metadata"} { + raw, ok := parserConfig[key].([]interface{}) + if !ok { + continue + } + for _, item := range raw { + m, ok := item.(map[string]interface{}) + if !ok { + continue + } + k, _ := m["key"].(string) + if strings.TrimSpace(k) == "" { + continue + } + out = append(out, m) + } + } + return out +} + +// isTruthy reports whether a parserConfig flag (e.g. enable_metadata) is on. +// It tolerates bool, numeric >0 and the strings "true"/"1" so storage +// representation differences don't silently disable the feature. +func isTruthy(v interface{}) bool { + switch t := v.(type) { + case bool: + return t + case string: + return t == "true" || t == "1" || t == "True" || t == "TRUE" + case float64: + return t > 0 + case int: + return t > 0 + case int64: + return t > 0 + } + return false +} + // deepCopyMap duplicates a JSON-like map so later merges do not mutate shared defaults. func deepCopyMap(source map[string]interface{}) map[string]interface{} { if source == nil { diff --git a/internal/common/parser_config_test.go b/internal/common/parser_config_test.go index 77ac9946de..e6d1195eba 100644 --- a/internal/common/parser_config_test.go +++ b/internal/common/parser_config_test.go @@ -51,3 +51,85 @@ func TestInjectExtractorLLMID_NoExtractor(t *testing.T) { t.Fatal("expected Parser:X still present") } } + +func TestInjectExtractorEnableMetadata_Disabled(t *testing.T) { + pc := map[string]interface{}{ + "Extractor:A": map[string]interface{}{}, + } + if InjectExtractorEnableMetadata(pc) { + t.Fatal("expected no update when enable_metadata is off") + } + if _, ok := pc["Extractor:A"].(map[string]interface{})["enable_metadata"]; ok { + t.Fatal("enable_metadata should not be set") + } +} + +func TestInjectExtractorEnableMetadata_InjectsAndMerges(t *testing.T) { + pc := map[string]interface{}{ + "enable_metadata": true, + "metadata": []interface{}{ + map[string]interface{}{"key": "author", "type": "string", "description": "doc author", "enum": []interface{}{"Alice", "Bob"}}, + }, + "built_in_metadata": []interface{}{ + map[string]interface{}{"key": "year", "type": "number"}, + }, + "Extractor:A": map[string]interface{}{}, + "Parser:X": map[string]interface{}{}, + } + if !InjectExtractorEnableMetadata(pc) { + t.Fatal("expected update") + } + ext := pc["Extractor:A"].(map[string]interface{}) + if got, _ := ext["enable_metadata"].(int); got != 1 { + t.Fatalf("expected enable_metadata=1, got %v", ext["enable_metadata"]) + } + fields, ok := ext["metadata"].([]any) + if !ok || len(fields) != 2 { + t.Fatalf("expected 2 merged metadata, got %#v", ext["metadata"]) + } + if _, ok := pc["Parser:X"].(map[string]interface{})["enable_metadata"]; ok { + t.Fatal("Parser node must not be touched") + } +} + +func TestInjectExtractorEnableMetadata_RespectsExplicitEnabled(t *testing.T) { + // A node the user already turned ON keeps its own config (no clobber). + pc := map[string]interface{}{ + "enable_metadata": true, + "metadata": []interface{}{map[string]interface{}{"key": "author"}}, + "Extractor:A": map[string]interface{}{"enable_metadata": 1}, + } + if InjectExtractorEnableMetadata(pc) { + t.Fatal("expected no update when user explicitly enabled enable_metadata") + } +} + +func TestInjectExtractorEnableMetadata_OverridesDefaultZero(t *testing.T) { + // Shipped DSLs default enable_metadata to 0; the dataset flag must still + // turn them on (otherwise auto-metadata could never activate). + pc := map[string]interface{}{ + "enable_metadata": true, + "metadata": []interface{}{map[string]interface{}{"key": "author"}}, + "Extractor:A": map[string]interface{}{"enable_metadata": 0}, + } + if !InjectExtractorEnableMetadata(pc) { + t.Fatal("expected update: dataset flag must override default enable_metadata=0") + } + ext := pc["Extractor:A"].(map[string]interface{}) + if got, _ := ext["enable_metadata"].(int); got != 1 { + t.Fatalf("expected enable_metadata=1 after override, got %v", ext["enable_metadata"]) + } + if _, ok := ext["metadata"].([]any); !ok { + t.Fatalf("expected metadata field schema injected, got %#v", ext["metadata"]) + } +} + +func TestInjectExtractorEnableMetadata_NoFields(t *testing.T) { + pc := map[string]interface{}{ + "enable_metadata": true, + "Extractor:A": map[string]interface{}{}, + } + if InjectExtractorEnableMetadata(pc) { + t.Fatal("expected no update when no fields configured") + } +} diff --git a/internal/common/task.go b/internal/common/task.go index 52258eb4a0..fbca0b55cf 100644 --- a/internal/common/task.go +++ b/internal/common/task.go @@ -36,3 +36,12 @@ type TaskHandle interface { // periodically during long tasks to avoid in-flight redelivery. InProgress() error } + +// RawMessage is a broker message carrying opaque bytes (used by the dataset-level +// compile consumer, which publishes arbitrary JSON payloads rather than the +// TaskMessage shape). The NATS engine returns RawMessage from FetchKnowledgeCompileMessages. +type RawMessage interface { + Data() []byte + Ack() error + Nak() error +} diff --git a/internal/dao/compilation_template.go b/internal/dao/compilation_template.go new file mode 100644 index 0000000000..6791c7d002 --- /dev/null +++ b/internal/dao/compilation_template.go @@ -0,0 +1,94 @@ +// +// 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 dao + +import ( + "context" + "fmt" + "ragflow/internal/entity" +) + +// CompilationTemplateDAO is the data-access object for compilation templates and +// their groups. +type CompilationTemplateDAO struct{} + +// NewCompilationTemplateDAO creates a compilation template DAO. +func NewCompilationTemplateDAO() *CompilationTemplateDAO { + return &CompilationTemplateDAO{} +} + +// ResolveGroupTemplateIDs resolves compilation-template-group ids to their child +// compilation_template ids for the given tenant. It mirrors the Python +// CompilationTemplateGroupService.resolve_template_ids: only VALID templates +// whose owning group is VALID and belongs to the tenant are returned, ordered by +// create_time within each group and by the requested group order across groups +// (so the output is deterministic). +// +// An error is returned if any requested group id does not resolve to at least +// one valid template (group missing, wrong tenant, or empty). This surfaces the +// misconfiguration instead of silently dropping the compilation_template_ids +// stamp — the data-loss path the caller is guarding against. +func (dao *CompilationTemplateDAO) ResolveGroupTemplateIDs(ctx context.Context, tenantID string, groupIDs []string) ([]string, error) { + if len(groupIDs) == 0 { + return nil, nil + } + + db := DB.WithContext(ctx) + + // Verify the requested groups exist, are valid, and belong to the tenant. + var validGroups []entity.CompilationTemplateGroup + if err := db. + Where("id IN ? AND tenant_id = ? AND status = ?", groupIDs, tenantID, string(entity.StatusValid)). + Find(&validGroups).Error; err != nil { + return nil, fmt.Errorf("resolve compilation template groups: %w", err) + } + valid := make(map[string]struct{}, len(validGroups)) + for _, g := range validGroups { + valid[g.ID] = struct{}{} + } + for _, gid := range groupIDs { + if _, ok := valid[gid]; !ok { + return nil, fmt.Errorf("compilation_template_group %q not found for tenant %q", gid, tenantID) + } + } + + var templates []entity.CompilationTemplate + if err := db.Model(&entity.CompilationTemplate{}). + Where("group_id IN ? AND status = ?", groupIDs, string(entity.StatusValid)). + Order("create_time asc"). + Find(&templates).Error; err != nil { + return nil, fmt.Errorf("resolve compilation templates for groups: %w", err) + } + + byGroup := make(map[string][]string, len(groupIDs)) + for _, t := range templates { + if t.GroupID == nil { + continue + } + byGroup[*t.GroupID] = append(byGroup[*t.GroupID], t.ID) + } + + var out []string + for _, gid := range groupIDs { + ids := byGroup[gid] + if len(ids) == 0 { + return nil, fmt.Errorf("compilation_template_group %q resolved to no valid templates", gid) + } + out = append(out, ids...) + } + return out, nil +} diff --git a/internal/dao/database.go b/internal/dao/database.go index a40e4bf059..d11de3ce93 100644 --- a/internal/dao/database.go +++ b/internal/dao/database.go @@ -157,6 +157,7 @@ func InitDB(ctx context.Context, migrateDB bool) error { &entity.IngestionTaskLog{}, &entity.FileCommit{}, &entity.FileCommitItem{}, + &entity.KnowledgeCompileDoc{}, } if migrateDB { diff --git a/internal/engine/engine.go b/internal/engine/engine.go index 9040bbaa29..066d893e6d 100644 --- a/internal/engine/engine.go +++ b/internal/engine/engine.go @@ -18,6 +18,8 @@ package engine import ( "context" + "time" + "ragflow/internal/common" "ragflow/internal/engine/types" @@ -105,4 +107,21 @@ type MessageQueue interface { ListMessages(messageType string, pending bool) ([]map[string]string, error) ShowMessageQueue() (map[string]string, error) CheckStatus() string + + // dataset-level compile consumer (§11) surface. + InitKnowledgeCompileStream() error + InitKnowledgeCompileConsumer() error + PublishKnowledgeCompile(subject string, payload []byte) error + FetchKnowledgeCompileMessages(n int) ([]common.RawMessage, error) + InitKnowledgeCompileLeases() error + AcquireKnowledgeCompileLease(key, holder string, ttl time.Duration) (uint64, bool, error) + HeartbeatKnowledgeCompileLease(key, holder string, ttl time.Duration, revision uint64) (uint64, bool, error) + ReleaseKnowledgeCompileLease(key, holder string, revision uint64) error + + // SubscribeNotify returns a channel of dataset ids pushed by + // PublishKnowledgeCompile on the notify.kc.workers subject (Option E + // §11.4: NATS as a wake-up, MySQL as the scheduling system of record). + // Implementations return (nil, nil) when push wake-up is unavailable; + // callers must fall back to periodic polling in that case. + SubscribeNotify(ctx context.Context) (<-chan string, error) } diff --git a/internal/engine/nats/knowledgecompile.go b/internal/engine/nats/knowledgecompile.go new file mode 100644 index 0000000000..8e1fbbbbfc --- /dev/null +++ b/internal/engine/nats/knowledgecompile.go @@ -0,0 +1,272 @@ +// +// Copyright 2026 The InfiniFlow Authors. All Rights Reserved. +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +package nats + +import ( + "context" + "encoding/json" + "errors" + "fmt" + "strings" + "time" + + "ragflow/internal/common" + + "github.com/nats-io/nats.go" + "github.com/nats-io/nats.go/jetstream" +) + +// Knowledge-compile (§11) subjects, stream, queue group and KV bucket. The +// subject prefix is shared by the stream's Subjects filter and the consumer's +// FilterSubject, so every published event must sit under knowledge.compile.events.>. +const ( + knowledgeCompileStreamName = "RAGFLOW_KNOWLEDGE_COMPILE_EVENTS" + knowledgeCompileSubjectPrefix = "knowledge.compile.events.>" + knowledgeCompileQueueGroup = "knowledge_compile_events_q" + knowledgeCompileKVBucket = "knowledge_compile_leases" +) + +// knowledgeCompileLeaseValue is the CAS-protected payload stored under lock:. +type knowledgeCompileLeaseValue struct { + Holder string `json:"holder"` + Expiry int64 `json:"expiry"` // unix nanos; lease is stale when Expiry < now +} + +// InitKnowledgeCompileStream creates the dedicated RAGFLOW_KNOWLEDGE_COMPILE_EVENTS stream, +// isolated from the task queue (RAGFLOW_TASKS). +func (n *NatsEngine) InitKnowledgeCompileStream() error { + if n.jetStream == nil { + return fmt.Errorf("knowledgecompile: jetStream not initialized") + } + ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second) + defer cancel() + cfg := jetstream.StreamConfig{ + Name: knowledgeCompileStreamName, + Subjects: []string{knowledgeCompileSubjectPrefix}, + Retention: jetstream.WorkQueuePolicy, + Storage: jetstream.FileStorage, + MaxMsgs: 1024 * 1024, + MaxBytes: 1024 * 1024 * 64, + } + if _, err := n.jetStream.CreateStream(ctx, cfg); err != nil && !strings.Contains(err.Error(), "already exists") { + return fmt.Errorf("knowledgecompile: create stream: %w", err) + } + return nil +} + +// PublishKnowledgeCompile publishes a wake-up payload on the notify subject via +// core NATS. The subject (notify.kc.workers) is intentionally outside the +// knowledge.compile.events stream, so it must go through core NATS rather than +// the JetStream Publish API, which errors with "no stream matches subject". +func (n *NatsEngine) PublishKnowledgeCompile(subject string, payload []byte) error { + if n.nc == nil { + return fmt.Errorf("knowledgecompile: nats not initialized") + } + return n.nc.Publish(subject, payload) +} + +// InitKnowledgeCompileConsumer creates the competing-consumer (queue group) for the knowledge-compile stream. +func (n *NatsEngine) InitKnowledgeCompileConsumer() error { + if n.jetStream == nil { + return fmt.Errorf("knowledgecompile: jetStream not initialized") + } + ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second) + defer cancel() + stream, err := n.jetStream.Stream(ctx, knowledgeCompileStreamName) + if err != nil { + return fmt.Errorf("knowledgecompile: stream not found (call InitKnowledgeCompileStream first): %w", err) + } + n.knowledgeCompileStream = stream + cons, err := stream.CreateOrUpdateConsumer(ctx, jetstream.ConsumerConfig{ + Name: "KNOWLEDGE_COMPILE_CONSUMER", + Durable: "knowledge_compile_durable", + DeliverGroup: knowledgeCompileQueueGroup, + AckPolicy: jetstream.AckExplicitPolicy, + MaxDeliver: 16, + MaxAckPending: 1024 * 128, + FilterSubject: knowledgeCompileSubjectPrefix, + }) + if err != nil { + if strings.Contains(err.Error(), "max waiting can not be updated") { + cons, err = stream.Consumer(ctx, "KNOWLEDGE_COMPILE_CONSUMER") + if err != nil { + return fmt.Errorf("knowledgecompile: get existing consumer: %w", err) + } + } else { + return fmt.Errorf("knowledgecompile: create consumer: %w", err) + } + } + n.knowledgeCompileConsumer = cons + return nil +} + +// FetchKnowledgeCompileMessages pulls up to batchSize messages from the knowledge-compile consumer. Because +// the consumer filters by subject only (never payload), the returned batch is +// a mix of datasets; the caller keeps the KB it is processing and Naks the rest. +func (n *NatsEngine) FetchKnowledgeCompileMessages(batchSize int) ([]common.RawMessage, error) { + if n.knowledgeCompileConsumer == nil { + return nil, fmt.Errorf("knowledgecompile: consumer not initialized (call InitKnowledgeCompileConsumer first)") + } + messages, err := n.knowledgeCompileConsumer.Fetch(batchSize, jetstream.FetchMaxWait(1*time.Second)) + if err != nil { + // A max-wait timeout with zero messages is the normal "nothing to do" + // condition for a polling fetch, not an error: surface it as an empty + // batch so the caller can idle without logging or backoff-pacing. + if errors.Is(err, nats.ErrTimeout) { + return nil, nil + } + return nil, err + } + out := make([]common.RawMessage, 0, 8) + for msg := range messages.Messages() { + out = append(out, &natsRawHandle{msg: msg}) + } + return out, nil +} + +// natsRawHandle adapts a jetstream.Msg to common.RawMessage. +type natsRawHandle struct { + msg jetstream.Msg +} + +func (h *natsRawHandle) Data() []byte { return h.msg.Data() } +func (h *natsRawHandle) Ack() error { return h.msg.Ack() } +func (h *natsRawHandle) Nak() error { return h.msg.Nak() } + +// InitKnowledgeCompileLeases creates the KV bucket backing per-KB leases. +func (n *NatsEngine) InitKnowledgeCompileLeases() error { + if n.jetStream == nil { + return fmt.Errorf("knowledgecompile: jetStream not initialized") + } + ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second) + defer cancel() + kv, err := n.jetStream.CreateOrUpdateKeyValue(ctx, jetstream.KeyValueConfig{Bucket: knowledgeCompileKVBucket}) + if err != nil { + return fmt.Errorf("knowledgecompile: create kv: %w", err) + } + n.kv = kv + return nil +} + +// AcquireKnowledgeCompileLease attempts to take the lease for key (CAS). It succeeds when the +// key is absent or its previous holder's TTL has expired. Returns the new +// revision and acquired=true on success. +func (n *NatsEngine) AcquireKnowledgeCompileLease(key, holder string, ttl time.Duration) (uint64, bool, error) { + if n.kv == nil { + return 0, false, fmt.Errorf("knowledgecompile: kv not initialized (call InitKnowledgeCompileLeases first)") + } + ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second) + defer cancel() + val, _ := json.Marshal(knowledgeCompileLeaseValue{Holder: holder, Expiry: time.Now().Add(ttl).UnixNano()}) + + existing, err := n.kv.Get(ctx, key) + if err != nil { + if errors.Is(err, jetstream.ErrKeyNotFound) { + rev, cerr := n.kv.Create(ctx, key, val) + if cerr != nil { + return 0, false, nil // lost the race + } + return rev, true, nil + } + return 0, false, err + } + var lv knowledgeCompileLeaseValue + _ = json.Unmarshal(existing.Value(), &lv) + if lv.Expiry < time.Now().UnixNano() { + // Stale lease: CAS-overwrite with our revision. + nrev, uerr := n.kv.Update(ctx, key, val, existing.Revision()) + if uerr != nil { + return 0, false, nil + } + return nrev, true, nil + } + return 0, false, nil +} + +// HeartbeatKnowledgeCompileLease refreshes the TTL of a held lease. revision must match the +// current holder; a revision mismatch (returning ok=false) means the lease was +// taken over by another instance and the caller must abort. +func (n *NatsEngine) HeartbeatKnowledgeCompileLease(key, holder string, ttl time.Duration, revision uint64) (uint64, bool, error) { + if n.kv == nil { + return 0, false, fmt.Errorf("knowledgecompile: kv not initialized") + } + ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second) + defer cancel() + val, _ := json.Marshal(knowledgeCompileLeaseValue{Holder: holder, Expiry: time.Now().Add(ttl).UnixNano()}) + nrev, err := n.kv.Update(ctx, key, val, revision) + if err != nil { + return 0, false, nil + } + return nrev, true, nil +} + +// ReleaseKnowledgeCompileLease deletes the lease only if we still own it (revision match), +// so we never release a lease another instance has since acquired. +func (n *NatsEngine) ReleaseKnowledgeCompileLease(key, holder string, revision uint64) error { + if n.kv == nil { + return nil + } + ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second) + defer cancel() + err := n.kv.Delete(ctx, key, jetstream.LastRevision(revision)) + if err != nil && strings.Contains(err.Error(), "wrong") { + return nil // someone else owns it now; nothing to release + } + return err +} + +// knowledgeCompileNotifySubject is the Option E wake-up channel: workers +// subscribe here and publishers push a {dataset_id} payload after appending to +// a KB's MySQL backlog. It is intentionally outside the knowledge.compile.events +// stream because it carries no routing payload — MySQL is the scheduling truth. +const knowledgeCompileNotifySubject = "notify.kc.workers" + +// SubscribeNotify opens a core NATS subscription to the wake-up subject and +// streams dataset ids to the returned channel. The subscription is torn down +// when ctx is cancelled. A full channel is dropped (not blocked) so a slow +// worker can never stall publishers. +func (n *NatsEngine) SubscribeNotify(ctx context.Context) (<-chan string, error) { + if n.nc == nil { + return nil, fmt.Errorf("knowledgecompile: nats not initialized") + } + ch := make(chan string, 256) + // done guards the send so the callback never writes to a closed channel: + // Sub.Unsubscribe does not wait for an in-flight dispatcher callback, so we + // must not close(ch) while a callback may still run. The reader selects on + // ctx.Done as well, so leaving ch open is safe and avoids the panic. + done := make(chan struct{}) + sub, err := n.nc.Subscribe(knowledgeCompileNotifySubject, func(m *nats.Msg) { + var p struct { + DatasetID string `json:"dataset_id"` + } + if err := json.Unmarshal(m.Data, &p); err != nil || p.DatasetID == "" { + return + } + select { + case ch <- p.DatasetID: + case <-done: + } + }) + if err != nil { + return nil, fmt.Errorf("knowledgecompile: subscribe notify: %w", err) + } + go func() { + <-ctx.Done() + _ = sub.Unsubscribe() + close(done) + }() + return ch, nil +} diff --git a/internal/engine/nats/nats.go b/internal/engine/nats/nats.go index 3662a96f8b..273ec17bd7 100644 --- a/internal/engine/nats/nats.go +++ b/internal/engine/nats/nats.go @@ -37,6 +37,11 @@ type NatsEngine struct { jetStream jetstream.JetStream stream jetstream.Stream consumer jetstream.Consumer + + // dataset-level compile consumer (§11) state. + knowledgeCompileStream jetstream.Stream + knowledgeCompileConsumer jetstream.Consumer + kv jetstream.KeyValue } func NewNatsEngine(host string, port int) *NatsEngine { diff --git a/internal/entity/knowledge_compile_doc.go b/internal/entity/knowledge_compile_doc.go new file mode 100644 index 0000000000..80af64c9ab --- /dev/null +++ b/internal/entity/knowledge_compile_doc.go @@ -0,0 +1,46 @@ +// +// 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 entity + +import "time" + +// KnowledgeCompileDoc 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 +// doc entries for the KB, inflight_doc_ids the ones a worker has claimed (the +// closed batch), and the claim_* fields the owner/lease used for crash +// recovery. NATS notify is only a wake-up; same-KB serialization comes from +// these rows, not from the broker. +// +// The *_doc_ids columns store a JSON array of knowledge_compile.BacklogEntry +// (doc_id + event_type + seq) as TEXT so the consumer can re-apply the same +// out-of-order / tombstone guards as the broker-based design without re-reading +// the queue. +type KnowledgeCompileDoc struct { + DatasetID string `gorm:"primaryKey;column:dataset_id;size:64" json:"dataset_id"` + TenantID string `gorm:"column:tenant_id;size:64;not null;default:''" json:"tenant_id"` + BacklogDocIDs string `gorm:"column:backlog_doc_ids;type:text;not null;default:'[]'" json:"backlog_doc_ids"` + InflightDocIDs string `gorm:"column:inflight_doc_ids;type:text;not null;default:'[]'" json:"inflight_doc_ids"` + ClaimOwner string `gorm:"column:claim_owner;size:64;not null;default:''" json:"claim_owner"` + ClaimToken string `gorm:"column:claim_token;size:64;not null;default:''" json:"claim_token"` + ClaimExpiresAt *time.Time `gorm:"column:claim_expires_at;default:null" json:"claim_expires_at"` + Priority int `gorm:"column:priority;not null;default:0" json:"priority"` + CreatedAt time.Time `gorm:"column:created_at;autoCreateTime" json:"created_at"` + UpdatedAt time.Time `gorm:"column:updated_at;autoUpdateTime" json:"updated_at"` +} + +// TableName pins the scheduling table name. +func (KnowledgeCompileDoc) TableName() string { return "knowledge_compile_docs" } diff --git a/internal/ingestion/component/extractor.go b/internal/ingestion/component/extractor.go index 6d3b6a2ee9..2d33629441 100644 --- a/internal/ingestion/component/extractor.go +++ b/internal/ingestion/component/extractor.go @@ -65,22 +65,27 @@ import ( "encoding/json" "errors" "fmt" + "os" "regexp" "sort" + "strconv" "strings" "sync" "time" + "github.com/cespare/xxhash/v2" eschema "github.com/cloudwego/eino/schema" "go.uber.org/zap" "gorm.io/gorm" "ragflow/internal/agent/runtime" "ragflow/internal/common" + "ragflow/internal/engine/redis" "ragflow/internal/entity" "ragflow/internal/entity/models" "ragflow/internal/ingestion/component/schema" "ragflow/internal/tokenizer" + "ragflow/internal/utility" ) const componentNameExtractor = "Extractor" @@ -104,6 +109,51 @@ var ( // and question extraction calls (generator.py:230,245). const extractorTemperature = 0.2 +// defaultExtractorConcurrency bounds concurrent LLM extraction calls +// process-wide. It mirrors Python's chat_limiter bound +// (agent/component/base.py:353 asyncio.Semaphore(MAX_CONCURRENT_CHATS, +// default 10)) so the Go port rate-limits auto-keywords / auto-questions / +// auto-metadata extraction the same way Python does. +const defaultExtractorConcurrency = 10 + +// extractorJob is one chunk's auto-extraction unit of work. +type extractorJob func() error + +// extractorPool is the process-wide bounded worker pool that drives +// cross-chunk concurrency for auto-keywords / auto-questions / +// auto-metadata extraction. It mirrors the parser/structure/mindmap +// WorkerPool usage, but is held globally so every Extractor +// invocation shares one rate limiter instead of spinning up one pool +// per batch. The pool only bounds concurrency (it is never StopWait'd), +// so per-invocation completion is tracked by runAutoExtractions with its +// own sync.WaitGroup + first-error collection — concurrent Invoke calls +// therefore do not disturb each other. +var extractorPool = utility.NewWorkerPool[extractorJob, struct{}]( + extractorConcurrency(), + extractorConcurrency()*4, + func(_ context.Context, j extractorJob) (struct{}, error) { return struct{}{}, j() }, +) + +// extractorConcurrency resolves the pool size from MAX_CONCURRENT_CHATS +// (Python parity) with a safe default of defaultExtractorConcurrency. +func extractorConcurrency() int { + if v := os.Getenv("MAX_CONCURRENT_CHATS"); v != "" { + if n, err := strconv.Atoi(v); err == nil && n > 0 { + return n + } + } + return defaultExtractorConcurrency +} + +// SetExtractorConcurrency overrides the global extractor pool size at +// runtime (e.g. from service init or tests). Mirrors tuning Python's +// chat_limiter bound. +func SetExtractorConcurrency(n int) { + if n > 0 { + extractorPool.Resize(n) + } +} + const ( autoKeywordPrompt = `## Role You are a text analyzer. @@ -139,6 +189,21 @@ Propose questions about a given piece of text content. --- ## Text Content +%s` + + autoMetadataPrompt = `## Role: Metadata extraction expert. +## Rules: + - Strict Evidence Only: Extract a value ONLY if it is explicitly mentioned in the Content. + - Enum Filter: For any field with an 'enum' list, the list acts as a strict filter. If no element from the list (or its direct synonym) is found in the Content, you MUST NOT extract that field. + - No Meta-Inference: Do not infer values based on the document's nature, format, or category. If the text does not literally state the information, treat it as missing. + - Zero-Hallucination: Never invent information or pick a "likely" value from the enum to fill a field. + - Empty Result: If no matches are found for any field, or if the content is irrelevant, output ONLY {}. + - Output: ONLY a valid JSON string. No Markdown, no notes. + +## Schema for extraction: +%s + +## Content to analyze: %s` ) @@ -207,6 +272,38 @@ func NewExtractorComponent(params map[string]any) (runtime.Component, error) { if v, ok := params["auto_tags"]; ok { p.AutoTags = mapInt(v) } + if v, ok := params["enable_metadata"]; ok { + p.EnableMetadata = mapInt(v) + } + if v, ok := params["metadata"].([]any); ok { + fields := make([]common.MetadataFieldDef, 0, len(v)) + for _, f := range v { + m, ok := f.(map[string]any) + if !ok { + continue + } + key, _ := m["key"].(string) + if key = strings.TrimSpace(key); key == "" { + continue + } + def := common.MetadataFieldDef{Key: key} + if t, ok := m["type"].(string); ok { + def.Type = t + } + if d, ok := m["description"].(string); ok { + def.Description = d + } + if e, ok := m["enum"].([]any); ok { + for _, ev := range e { + if s, ok := ev.(string); ok { + def.Enum = append(def.Enum, s) + } + } + } + fields = append(fields, def) + } + p.Metadata = fields + } if v, ok := params["tag_file_id"].(string); ok { p.TagFileID = v } @@ -575,78 +672,48 @@ func (c *ExtractorComponent) Invoke(ctx context.Context, db *gorm.DB, inputs map in.chunks = []map[string]any{{in.fieldName: ans}} return nil } + // Auto keyword / question / metadata extraction runs with + // cross-chunk concurrency through the process-wide extractor + // pool (mirrors Python's chat_limiter). Each chunk job owns its + // own chunk map, so no per-chunk mutex is required; per-invocation + // completion is tracked inside runAutoExtractions. + if err := c.runAutoExtractions(timeoutCtx, db, in); err != nil { + return err + } + + // Primary field extraction stays sequential per chunk: each + // chunk's result lands under its own field_name key, preserving + // deterministic output ordering (see file header: -race tests). for i, ck := range in.chunks { + if in.fieldName == "" { + continue + } text, _ := ck["content_with_weight"].(string) if strings.TrimSpace(text) == "" { text, _ = ck["text"].(string) } - - if c.Param.AutoKeywords > 0 && c.Param.AutoQuestions > 0 { - // Run keyword and question extraction concurrently - // per chunk (mirrors Python's ThreadPoolExecutor: - // task_executor.py:444-448). The shared chunk map is - // guarded by mu inside runAutoKeywords/runAutoQuestions, - // which also short-circuit when the key already exists. - var kwErr, qErr error - var mu sync.Mutex - var wg sync.WaitGroup - wg.Add(2) - go func() { - defer wg.Done() - if err := c.runAutoKeywords(timeoutCtx, db, in, ck, text, &mu); err != nil { - kwErr = err - } - }() - go func() { - defer wg.Done() - if err := c.runAutoQuestions(timeoutCtx, db, in, ck, text, &mu); err != nil { - qErr = err - } - }() - wg.Wait() - if kwErr != nil { - return fmt.Errorf("chunk %d keywords: %w", i, kwErr) - } - if qErr != nil { - return fmt.Errorf("chunk %d questions: %w", i, qErr) - } - } else { - if c.Param.AutoKeywords > 0 { - if err := c.runAutoKeywords(timeoutCtx, db, in, ck, text, nil); err != nil { - return fmt.Errorf("chunk %d keywords: %w", i, err) - } - } - if c.Param.AutoQuestions > 0 { - if err := c.runAutoQuestions(timeoutCtx, db, in, ck, text, nil); err != nil { - return fmt.Errorf("chunk %d questions: %w", i, err) - } - } + // Substitute {field_name} placeholders with current + // chunk field values, mirroring Python's + // string_format at extractor.py:103. + callIn := in + callIn.prompt = substituteChunkPlaceholders(in.prompt, ck, text) + callIn.systemPrompt = substituteChunkPlaceholders(in.systemPrompt, ck, text) + // buildExtractorMessages appends chunkText to the user + // message unconditionally. When the chunk text was already + // embedded via a {text}/{chunks} placeholder above, passing + // it again would duplicate the content in the final prompt. + // Suppress the automatic append in that case (the keyword/ + // question helper paths do the same by passing ""). + callChunkText := text + if strings.Contains(in.prompt, "{text}") || strings.Contains(in.prompt, "{chunks}") || + strings.Contains(in.systemPrompt, "{text}") || strings.Contains(in.systemPrompt, "{chunks}") { + callChunkText = "" } - - if in.fieldName != "" { - // Substitute {field_name} placeholders with current - // chunk field values, mirroring Python's - // string_format at extractor.py:103. - callIn := in - callIn.prompt = substituteChunkPlaceholders(in.prompt, ck, text) - callIn.systemPrompt = substituteChunkPlaceholders(in.systemPrompt, ck, text) - // buildExtractorMessages appends chunkText to the user - // message unconditionally. When the chunk text was already - // embedded via a {text}/{chunks} placeholder above, passing - // it again would duplicate the content in the final prompt. - // Suppress the automatic append in that case (the keyword/ - // question helper paths do the same by passing ""). - callChunkText := text - if strings.Contains(in.prompt, "{text}") || strings.Contains(in.prompt, "{chunks}") || - strings.Contains(in.systemPrompt, "{text}") || strings.Contains(in.systemPrompt, "{chunks}") { - callChunkText = "" - } - ans, callErr := c.call(timeoutCtx, db, callIn, callChunkText) - if callErr != nil { - return fmt.Errorf("chunk %d: %w", i, callErr) - } - ck[in.fieldName] = ans + ans, callErr := c.call(timeoutCtx, db, callIn, callChunkText) + if callErr != nil { + return fmt.Errorf("chunk %d: %w", i, callErr) } + ck[in.fieldName] = ans } return nil }); err != nil { @@ -668,14 +735,8 @@ func (c *ExtractorComponent) Invoke(ctx context.Context, db *gorm.DB, inputs map // keyword/question goroutines stay race-free. The existence check and // the map writes are both guarded by mu. Keyword extraction pins // temperature to extractorTemperature (0.2) to mirror generator.py. -func (c *ExtractorComponent) runAutoKeywords(ctx context.Context, db *gorm.DB, in extractorInputs, ck map[string]any, chunkText string, mu *sync.Mutex) error { - if mu != nil { - mu.Lock() - } +func (c *ExtractorComponent) runAutoKeywords(ctx context.Context, db *gorm.DB, in extractorInputs, ck map[string]any, chunkText string) error { _, exists := ck["important_kwd"] - if mu != nil { - mu.Unlock() - } if exists { return nil } @@ -701,30 +762,17 @@ func (c *ExtractorComponent) runAutoKeywords(ctx context.Context, db *gorm.DB, i } tok := tokenizer.New(in.lang) tks, tkErr := tok.Tokenize(strings.Join(kwds, " ")) - if mu != nil { - mu.Lock() - } ck["important_kwd"] = kwds if tkErr == nil { ck["important_tks"] = tks } - if mu != nil { - mu.Unlock() - } return nil } // runAutoQuestions extracts questions for the current chunk and stores -// them on ck["question_kwd"]. See runAutoKeywords for the mu contract -// and the temperature pin. -func (c *ExtractorComponent) runAutoQuestions(ctx context.Context, db *gorm.DB, in extractorInputs, ck map[string]any, chunkText string, mu *sync.Mutex) error { - if mu != nil { - mu.Lock() - } +// them on ck["question_kwd"]. See runAutoKeywords for the temperature pin. +func (c *ExtractorComponent) runAutoQuestions(ctx context.Context, db *gorm.DB, in extractorInputs, ck map[string]any, chunkText string) error { _, exists := ck["question_kwd"] - if mu != nil { - mu.Unlock() - } if exists { return nil } @@ -758,19 +806,227 @@ func (c *ExtractorComponent) runAutoQuestions(ctx context.Context, db *gorm.DB, } tok := tokenizer.New(in.lang) tks, tkErr := tok.Tokenize(strings.Join(filtered, "\n")) - if mu != nil { - mu.Lock() - } ck["question_kwd"] = filtered if tkErr == nil { ck["question_tks"] = tks } - if mu != nil { - mu.Unlock() - } return nil } +// runAutoExtractions dispatches the auto keyword / question / metadata +// extraction for every chunk to the process-wide extractorPool so the +// work runs with bounded cross-chunk concurrency (mirrors Python's +// chat_limiter). The pool only bounds concurrency; per-invocation +// completion and first-error collection happen here with a local +// sync.WaitGroup, so concurrent Invoke calls do not disturb each other. +// Each chunk job owns its own chunk map, so no per-chunk mutex is needed. +func (c *ExtractorComponent) runAutoExtractions(ctx context.Context, db *gorm.DB, in extractorInputs) error { + if c.Param.AutoKeywords == 0 && c.Param.AutoQuestions == 0 && c.Param.EnableMetadata == 0 { + return nil + } + futs := make([]utility.WorkerPoolFuture[extractorJob, struct{}], 0, len(in.chunks)) + for i, ck := range in.chunks { + i, ck := i, ck + text, _ := ck["content_with_weight"].(string) + if strings.TrimSpace(text) == "" { + text, _ = ck["text"].(string) + } + fn := c.autoExtractionJob(ctx, db, in, i, ck, text) + f, err := extractorPool.Submit(ctx, fn) + if err != nil { + return err + } + futs = append(futs, f) + } + var firstErr error + var emu sync.Mutex + var wg sync.WaitGroup + for _, f := range futs { + wg.Add(1) + go func(f utility.WorkerPoolFuture[extractorJob, struct{}]) { + defer wg.Done() + res, werr := f.Wait(ctx) + if werr != nil { + emu.Lock() + if firstErr == nil { + firstErr = werr + } + emu.Unlock() + return + } + if res.Err != nil { + emu.Lock() + if firstErr == nil { + firstErr = res.Err + } + emu.Unlock() + } + }(f) + } + wg.Wait() + return firstErr +} + +// autoExtractionJob returns the unit of work for one chunk: it runs the +// enabled auto-extractions (keywords, questions, metadata) sequentially +// on the chunk's own map. Running them sequentially inside the job keeps +// the code free of per-chunk mutexes — cross-chunk parallelism is provided +// by the pool, not by intra-chunk goroutines. +func (c *ExtractorComponent) autoExtractionJob(ctx context.Context, db *gorm.DB, in extractorInputs, idx int, ck map[string]any, chunkText string) extractorJob { + return func() error { + if c.Param.AutoKeywords > 0 { + if err := c.runAutoKeywords(ctx, db, in, ck, chunkText); err != nil { + return fmt.Errorf("chunk %d keywords: %w", idx, err) + } + } + if c.Param.AutoQuestions > 0 { + if err := c.runAutoQuestions(ctx, db, in, ck, chunkText); err != nil { + return fmt.Errorf("chunk %d questions: %w", idx, err) + } + } + if c.Param.EnableMetadata > 0 { + if err := c.runEnableMetadata(ctx, db, in, ck, chunkText); err != nil { + return fmt.Errorf("chunk %d metadata: %w", idx, err) + } + } + return nil + } +} + +// runEnableMetadata extracts structured metadata for the current chunk and +// merges the parsed JSON object into ck["metadata"]. It mirrors the +// runAutoKeywords/runAutoQuestions shape but parses a JSON object and +// merges into the chunk's metadata map (which the existing +// mergeChunkMetadata → docState.apply aggregation path consumes). +func (c *ExtractorComponent) runEnableMetadata(ctx context.Context, db *gorm.DB, in extractorInputs, ck map[string]any, chunkText string) error { + if len(c.Param.Metadata) == 0 { + return nil + } + // Render the field schema into the prompt, mirroring Python's + // turn2jsonschema(metadata_conf) rendered into the META_DATA template. + schemaMap := common.Turn2JSONSchema(c.Param.Metadata) + if len(schemaMap) == 0 { + return nil + } + schemaJSON, err := json.Marshal(schemaMap) + if err != nil { + return err + } + schemaStr := string(schemaJSON) + + // LLM cache (mirrors Python get_llm_cache/set_llm_cache in + // task_executor.py:543/550 gen_metadata_task): identical (model + chunk + // text + schema) extractions are served from Redis within a 24h window so + // repeated runs / identical chunks don't re-pay the LLM call. + // Best-effort: a missing Redis client or any cache error falls through to + // a live call instead of failing the extraction. + var parsed map[string]any + if cached, hit := getMetadataLLMCache(in.llmID, schemaStr, chunkText); hit { + parsed = cached + } else { + metaTemp := extractorTemperature + metaIn := extractorInputs{ + llmID: in.llmID, + systemPrompt: fmt.Sprintf(autoMetadataPrompt, schemaStr, chunkText), + prompt: "Output: ", + temperature: &metaTemp, + } + result, err := c.call(ctx, db, metaIn, "") + if err != nil { + return err + } + // call JSON-parses a JSON object response into a map; plain-text + // (or -prefixed) responses come back as a string. Handle + // both so a valid ```json reply is never silently dropped. + var parsedObj map[string]any + switch r := result.(type) { + case map[string]any: + parsedObj = r + case string: + s := cleanExtractionResult(r) + if s == "" { + return nil + } + var ok bool + parsedObj, ok = tryParseJSONObject(s) + if !ok || len(parsedObj) == 0 { + return nil + } + default: + return nil + } + parsed = parsedObj + setMetadataLLMCache(in.llmID, schemaStr, chunkText, parsed) + } + // Merge into the chunk metadata map, preserving existing keys. + var meta map[string]any + if existing, ok := ck["metadata"].(map[string]any); ok && existing != nil { + meta = existing + } else { + meta = make(map[string]any, len(parsed)) + } + for k, v := range parsed { + if v == nil { + continue + } + if s, ok := v.(string); ok && strings.TrimSpace(s) == "" { + continue + } + meta[k] = v + } + ck["metadata"] = meta + return nil +} + +// metadataLLMCacheTTL mirrors Python get_llm_cache/set_llm_cache 24h TTL. +const metadataLLMCacheTTL = 24 * time.Hour + +// metadataLLMCacheKey builds a Redis key from (llm id, chunk text, "metadata", +// schema), mirroring Python get_llm_cache's xxh64(llmnm + txt + history + genconf). +func metadataLLMCacheKey(llmID, schemaJSON, chunkText string) string { + h := xxhash.New() + h.WriteString(llmID) + h.WriteString("\x00") + h.WriteString(chunkText) + h.WriteString("\x00") + h.WriteString("metadata") + h.WriteString("\x00") + h.WriteString(schemaJSON) + return fmt.Sprintf("kc:meta:%x", h.Sum64()) +} + +// getMetadataLLMCache returns a cached extraction for the given chunk, or +// (nil, false) on miss / Redis unavailable / decode error. Best-effort. +func getMetadataLLMCache(llmID, schemaJSON, chunkText string) (map[string]any, bool) { + client := redis.Get() + if client == nil { + return nil, false + } + data, err := client.Get(metadataLLMCacheKey(llmID, schemaJSON, chunkText)) + if err != nil || data == "" { + return nil, false + } + var parsed map[string]any + if err := json.Unmarshal([]byte(data), &parsed); err != nil { + return nil, false + } + return parsed, true +} + +// setMetadataLLMCache stores an extraction result for 24h. Best-effort: a +// missing Redis client or marshal error is silently ignored. +func setMetadataLLMCache(llmID, schemaJSON, chunkText string, parsed map[string]any) { + client := redis.Get() + if client == nil { + return + } + data, err := json.Marshal(parsed) + if err != nil { + return + } + client.Set(metadataLLMCacheKey(llmID, schemaJSON, chunkText), string(data), metadataLLMCacheTTL) +} + // cleanExtractionResult strips `` tags and rejects `**ERROR**` responses, // matching Python's keyword_extraction and question_proposal post-processing. func cleanExtractionResult(s string) string { @@ -1178,13 +1434,32 @@ func substituteChunkPlaceholders(prompt string, ck map[string]any, chunkText str // (```json ... ```) before parsing. func tryParseJSONObject(s string) (map[string]any, bool) { trimmed := strings.TrimSpace(s) - // Strip a single ``` fence pair if present. + // Strip a surrounding markdown code fence. Models commonly wrap JSON in + // ```json ... ``` (language tag on the opening line) but some emit the tag + // on its own line (```\njson\n{...}); Python's json_repair tolerates both, + // encoding/json does not, so we drop the fence and any bare leading + // "json"/"JSON" language line before parsing. if strings.HasPrefix(trimmed, "```") { - if idx := strings.Index(trimmed, "\n"); idx >= 0 { + // Drop the opening fence line (````` plus an optional inline language + // tag such as `json`) up to and including the first newline. + if idx := strings.IndexByte(trimmed, '\n'); idx >= 0 { trimmed = trimmed[idx+1:] + } else { + trimmed = trimmed[3:] } - if strings.HasSuffix(trimmed, "```") { - trimmed = trimmed[:len(trimmed)-3] + // Some models place the language tag on the next line + // (`````\njson\n{...}). Drop a bare leading `json`/`JSON` line so it is + // not mistaken for JSON content. + if firstNL := strings.IndexByte(trimmed, '\n'); firstNL >= 0 { + if firstLine := strings.TrimSpace(trimmed[:firstNL]); firstLine == "json" || firstLine == "JSON" { + trimmed = trimmed[firstNL+1:] + } + } else if firstLine := strings.TrimSpace(trimmed); firstLine == "json" || firstLine == "JSON" { + trimmed = "" + } + // Drop the closing fence if present. + if i := strings.LastIndex(trimmed, "```"); i >= 0 { + trimmed = trimmed[:i] } trimmed = strings.TrimSpace(trimmed) } diff --git a/internal/ingestion/component/extractor_test.go b/internal/ingestion/component/extractor_test.go index cff5bb206b..4cbfe5aa9b 100644 --- a/internal/ingestion/component/extractor_test.go +++ b/internal/ingestion/component/extractor_test.go @@ -29,7 +29,9 @@ import ( eschema "github.com/cloudwego/eino/schema" "ragflow/internal/agent/runtime" + "ragflow/internal/common" "ragflow/internal/ingestion/component/schema" + "ragflow/internal/utility" ) // stubExtractorChatInvoker is the test seam for the package-level @@ -582,6 +584,42 @@ func TestNewExtractorComponent_SysPromptAlias(t *testing.T) { } } +// TestNewExtractorComponent_MetadataAsAnySlice guards against the regression +// where InjectExtractorEnableMetadata injected the field schema as a +// []map[string]interface{} while NewExtractorComponent only accepted []any; +// the type assertion then failed and ExtractorParam.Metadata stayed empty, so +// auto-metadata never fired. The override_params path passes the injected +// value straight through (no JSON round-trip), so the slice element type must +// be []any for the assertion to succeed. +func TestNewExtractorComponent_MetadataAsAnySlice(t *testing.T) { + comp, err := NewExtractorComponent(map[string]any{ + "field_name": "out", + "enable_metadata": 1, + // This is exactly the dynamic type InjectExtractorEnableMetadata + // produces ([]any of map[string]any), NOT []map[string]interface{}. + "metadata": []any{ + map[string]any{"key": "author", "type": "string", "description": "doc author"}, + map[string]any{"key": "year", "type": "number", "enum": []any{"2020", "2021"}}, + }, + }) + if err != nil { + t.Fatalf("NewExtractorComponent: %v", err) + } + ec := comp.(*ExtractorComponent) + if ec.Param.EnableMetadata != 1 { + t.Fatalf("EnableMetadata = %d, want 1", ec.Param.EnableMetadata) + } + if len(ec.Param.Metadata) != 2 { + t.Fatalf("Metadata = %#v, want 2 fields", ec.Param.Metadata) + } + if ec.Param.Metadata[0].Key != "author" || ec.Param.Metadata[0].Type != "string" { + t.Errorf("Metadata[0] = %#v, want key=author type=string", ec.Param.Metadata[0]) + } + if len(ec.Param.Metadata[1].Enum) != 2 { + t.Errorf("Metadata[1].Enum = %#v, want 2 enum values", ec.Param.Metadata[1].Enum) + } +} + // TestNewExtractorComponent_PromptsArray verifies that the Python DSL // "prompts" array format is parsed into Param.Prompt. func TestNewExtractorComponent_PromptsArray(t *testing.T) { @@ -756,6 +794,13 @@ func TestTryParseJSONObject(t *testing.T) { {name: "object", in: `{"a":1}`, wantOK: true, wantKey: "a"}, {name: "object with fence", in: "```json\n{\"a\":1}\n```", wantOK: true, wantKey: "a"}, {name: "fence without json tag", in: "```\n{\"a\":1}\n```", wantOK: true, wantKey: "a"}, + // Language tag on its own line (```\njson\n{...}) — Python json_repair + // tolerates this, so encoding/json must not choke on the bare "json". + {name: "json tag on own line", in: "```\njson\n{\"a\":1}\n```", wantOK: true, wantKey: "a"}, + {name: "JSON tag on own line", in: "```\nJSON\n{\"a\":1}\n```", wantOK: true, wantKey: "a"}, + // Leading prose before the fence must not be stripped (only a real + // ``` fence prefix is handled). + {name: "leading prose no fence", in: "Here is the result: {\"a\":1}", wantOK: false}, {name: "plain string", in: "hello", wantOK: false}, {name: "array", in: `[1,2]`, wantOK: false}, {name: "empty object", in: `{}`, wantOK: false}, @@ -776,6 +821,202 @@ func TestTryParseJSONObject(t *testing.T) { } } +// TestCleanExtractionResult covers the chain-of-thought stripping +// and the **ERROR** guard that mirrors Python's metadata post-processing. +func TestCleanExtractionResult(t *testing.T) { + cases := []struct { + name string + in string + want string + }{ + {name: "plain", in: `{"a":1}`, want: `{"a":1}`}, + // Python re.sub(r"^.*", "", ans): everything up to and + // including the LAST is dropped. + {name: "thinks stripped", in: "let me thinkreasoning\n{\"a\":1}", want: `{"a":1}`}, + {name: "thinks no json", in: "thinkingno json here", want: "no json here"}, + // **ERROR** responses are rejected entirely. + {name: "error marker rejected", in: "**ERROR** could not extract", want: ""}, + {name: "error after think", in: "x**ERROR** boom", want: ""}, + {name: "whitespace trimmed", in: " {\"a\":1} ", want: `{"a":1}`}, + } + for _, tc := range cases { + t.Run(tc.name, func(t *testing.T) { + if got := cleanExtractionResult(tc.in); got != tc.want { + t.Errorf("cleanExtractionResult(%q) = %q, want %q", tc.in, got, tc.want) + } + }) + } +} + +// newMetadataExtractor returns an ExtractorComponent wired for doc-level +// metadata extraction with the given field definitions. +func newMetadataExtractor(fields ...common.MetadataFieldDef) *ExtractorComponent { + return &ExtractorComponent{Param: schema.ExtractorParam{ + EnableMetadata: 1, + Metadata: fields, + }} +} + +// TestExtractorComponent_runEnableMetadata_MergesIntoChunkMetadata verifies a +// JSON object from the LLM is parsed and merged into the chunk's metadata map, +// which mergeChunkMetadata then aggregates to the doc level. +func TestExtractorComponent_runEnableMetadata_MergesIntoChunkMetadata(t *testing.T) { + withStubChatInvoker(t, stubResponse{Content: `{"category":"finance","region":"east"}`}) + c := newMetadataExtractor( + common.MetadataFieldDef{Key: "category", Type: "string"}, + common.MetadataFieldDef{Key: "region", Type: "string"}, + ) + ck := map[string]any{} + if err := c.runEnableMetadata(t.Context(), nil, extractorInputs{llmID: "m"}, ck, "chunk text"); err != nil { + t.Fatalf("runEnableMetadata: %v", err) + } + meta, ok := ck["metadata"].(map[string]any) + if !ok { + t.Fatalf("ck[metadata] missing or wrong type: %T", ck["metadata"]) + } + if meta["category"] != "finance" || meta["region"] != "east" { + t.Errorf("metadata = %v, want category=finance region=east", meta) + } +} + +// TestExtractorComponent_runEnableMetadata_StripsJSONFence verifies the +// extraction path tolerates a fenced ```json response (the common model +// output) that would otherwise fail encoding/json parsing — mirroring Python +// json_repair. +func TestExtractorComponent_runEnableMetadata_StripsJSONFence(t *testing.T) { + withStubChatInvoker(t, stubResponse{Content: "```json\n{\"category\":\"law\"}\n```"}) + c := newMetadataExtractor(common.MetadataFieldDef{Key: "category", Type: "string"}) + ck := map[string]any{} + if err := c.runEnableMetadata(t.Context(), nil, extractorInputs{llmID: "m"}, ck, "chunk text"); err != nil { + t.Fatalf("runEnableMetadata: %v", err) + } + meta, ok := ck["metadata"].(map[string]any) + if !ok { + t.Fatalf("ck[metadata] missing: %T", ck["metadata"]) + } + if meta["category"] != "law" { + t.Errorf("metadata = %v, want category=law", meta) + } +} + +// TestExtractorComponent_runEnableMetadata_DegradesGracefully verifies that an +// empty / **ERROR** / unparseable / think-only LLM response does NOT block +// ingestion: the chunk metadata is left untouched and no error is returned +// (Python "no evidence → {}"). +func TestExtractorComponent_runEnableMetadata_DegradesGracefully(t *testing.T) { + cases := []struct { + name string + content string + }{ + {"empty", ""}, + {"error_marker", "**ERROR** something went wrong"}, + {"garbage", "I could not find any metadata in this text."}, + {"not_json", "{\"category\": } partial"}, + {"think_only", "let me think"}, + } + for _, tc := range cases { + t.Run(tc.name, func(t *testing.T) { + withStubChatInvoker(t, stubResponse{Content: tc.content}) + c := newMetadataExtractor(common.MetadataFieldDef{Key: "category", Type: "string"}) + ck := map[string]any{"metadata": map[string]any{"preexisting": "keep"}} + if err := c.runEnableMetadata(t.Context(), nil, extractorInputs{llmID: "m"}, ck, tc.name); err != nil { + t.Fatalf("runEnableMetadata returned error: %v", err) + } + meta, ok := ck["metadata"].(map[string]any) + if !ok { + t.Fatalf("ck[metadata] should remain a map, got %T", ck["metadata"]) + } + if meta["preexisting"] != "keep" { + t.Errorf("preexisting metadata must be preserved: %v", meta) + } + if _, has := meta["category"]; has { + t.Errorf("category should not be set on degraded response: %v", meta) + } + }) + } +} + +// TestExtractorComponent_runEnableMetadata_CrossChunkUnion simulates two chunks +// whose extraction returns overlapping list values for the same key. Aggregating +// the chunk metadata maps with utility.UpdateMetadataTo (as mergeChunkMetadata +// does) must produce a de-duplicated union, matching Python update_metadata_to. +func TestExtractorComponent_runEnableMetadata_CrossChunkUnion(t *testing.T) { + withStubChatInvoker(t, + stubResponse{Content: `{"people":["关羽","张辽"]}`}, + stubResponse{Content: `{"people":["张辽","刘备"]}`}, + ) + c := newMetadataExtractor(common.MetadataFieldDef{Key: "people", Type: "string"}) + ck1 := map[string]any{} + ck2 := map[string]any{} + if err := c.runEnableMetadata(t.Context(), nil, extractorInputs{llmID: "m"}, ck1, "chunk one"); err != nil { + t.Fatalf("ck1: %v", err) + } + if err := c.runEnableMetadata(t.Context(), nil, extractorInputs{llmID: "m"}, ck2, "chunk two"); err != nil { + t.Fatalf("ck2: %v", err) + } + // mirror mergeChunkMetadata: aggregate chunk metadata into doc metadata. + m1, ok := ck1["metadata"].(map[string]any) + if !ok { + t.Fatalf("ck1[metadata] missing: %T", ck1["metadata"]) + } + m2, ok := ck2["metadata"].(map[string]any) + if !ok { + t.Fatalf("ck2[metadata] missing: %T", ck2["metadata"]) + } + docMeta := map[string]any{} + docMeta = utility.UpdateMetadataTo(docMeta, m1) + docMeta = utility.UpdateMetadataTo(docMeta, m2) + people, ok := docMeta["people"].([]string) + if !ok { + t.Fatalf("people = %T, want []string", docMeta["people"]) + } + want := map[string]bool{"关羽": true, "张辽": true, "刘备": true} + if len(people) != len(want) { + t.Fatalf("people = %v, want union of %v", people, want) + } + for _, p := range people { + if !want[p] { + t.Errorf("unexpected person %q", p) + } + } +} + +// TestExtractorComponent_runEnableMetadata_CombinedValueSplit verifies a value +// the LLM combines with Chinese/comma delimiters is split when passed through +// common.SplitCombinedMetadataValues (as mergeDocMetadata does before writing), +// matching Python _split_combined_values (doc_metadata_service.py). +func TestExtractorComponent_runEnableMetadata_CombinedValueSplit(t *testing.T) { + withStubChatInvoker(t, stubResponse{Content: `{"people":["关羽、张辽、刘备"]}`}) + c := newMetadataExtractor(common.MetadataFieldDef{Key: "people", Type: "string"}) + ck := map[string]any{} + if err := c.runEnableMetadata(t.Context(), nil, extractorInputs{llmID: "m"}, ck, "chunk text"); err != nil { + t.Fatalf("runEnableMetadata: %v", err) + } + rawMeta, ok := ck["metadata"].(map[string]any) + if !ok { + t.Fatalf("ck[metadata] missing: %T", ck["metadata"]) + } + raw, ok := rawMeta["people"].([]any) + if !ok || len(raw) != 1 { + t.Fatalf("raw people = %v, want 1 combined element", rawMeta["people"]) + } + // mergeDocMetadata runs SplitCombinedMetadataValues before writing. + split := common.SplitCombinedMetadataValues(ck["metadata"].(map[string]any)) + people, ok := split["people"].([]string) + if !ok { + t.Fatalf("people = %T, want []string", split["people"]) + } + want := map[string]bool{"关羽": true, "张辽": true, "刘备": true} + if len(people) != len(want) { + t.Fatalf("people = %v, want 3 split elements", people) + } + for _, p := range people { + if !want[p] { + t.Errorf("unexpected %q", p) + } + } +} + // TestExtractorComponent_ConcurrentInvoke verifies the chat // invoker swap is safe under concurrent Invoke calls. This is // the canary for SetExtractorChatInvoker and the package-level diff --git a/internal/ingestion/component/knowledge_compiler/PORT_PLAN.md b/internal/ingestion/component/knowledge_compiler/PORT_PLAN.md new file mode 100644 index 0000000000..6ff18d7e26 --- /dev/null +++ b/internal/ingestion/component/knowledge_compiler/PORT_PLAN.md @@ -0,0 +1,416 @@ +# 将 `rag/advanced_rag/knowlege_compile` 移植到 Go 的计划 + +> 目标：在 `internal/ingestion/component` 下新增一个 pipeline 组件 `KnowledgeCompiler`， +> 把 Python 版「知识编译」全部变体（wiki / raptor / mind_map / structure / dataset_nav） +> 收敛到同一个组件，按 `variant` 参数分派。 + +**评审缺口索引**（code review 后补，正文以「缺口 X」标注）： +- 缺口 #1 历史去重决议（wiki 组件内 ES KNN + `HistoricalCandidates` 可选覆盖）→ §4 / 风险 #2 +- 缺口 A 容量护栏（防 OOM）→ `Param.Guardrails` / §4 Outputs / M1 单测 +- 缺口 C Raptor 验证门禁（NMI/ARI + 回退）→ M6 / 风险 #1 +- 缺口 D 输出合约前置冻结 → §4.1 +- 缺口 E golden 对齐必选门禁 → §6 +- 缺口 F 变体命名主名 + alias → §4 Inputs + +--- + +## 1. 范围与现状 + +### Python 源（`rag/advanced_rag/knowlege_compile/`） +| 文件 | 职责 | 关键导出 | +|------|------|----------| +| `_common.py` | 共享底座：stable id、embed 包装、tokenize、chunk 批处理引擎、bulk dedup（exact+embed+LLM）、ES I/O。
**注意**：Go 版删除「从 ES 拉源 chunk」与「组件内写 ES」，但保留 **wiki 变体历史去重的 ES KNN 读路径**（见 §1.1）；输入仍来自内存 inputs，产物以 chunks 形式合并进上游分块流返回 | `stable_row_id`, `encode`, `build_chunk_batches`, `run_chunked_pipeline`, `bulk_dedup_items`（内存去重 + wiki 历史候选 ES KNN） | +| `structure.py` | 文档级结构编译（list/set/hypergraph）+ 本地/ES 去重 + 紧凑图 JSON | `compile_structure_from_text`, `merge_compiled_structures`, `rebuild_structure_graph_json` | +| `wiki.py` | 工件管线 MAP→REDUCE→PLAN→REFINE（最大，~3600 行） | `wiki_map_from_chunks`, `wiki_reduce_from_extracts`, `wiki_plan_from_reduction`, `wiki_refine_from_plan` | +| `raptor.py` | RAPTOR 摘要树：classic + Psi 两种 builder，GMM/AHC/UMAP 聚类 | `RecursiveAbstractiveProcessing4TreeOrganizedRetrieval` | +| `mind_map_extractor.py` | 基于 `markdown_to_json` 的心智图抽取 | `MindMapExtractor` | +| `dataset_nav.py` | 数据集级导航的增量聚类（nav_cluster / nav_doc 树 + Redis 锁） | `upsert_dataset_nav_doc`, `remove_dataset_nav_doc` | +| `runner.py` | 非 tree 模板的批量编排（LLMCallPool、flush、synthesis 阶段） | `run_structure_compile_over_batches` | + +### 1.1 输入/输出/存储模型变更（关键，Go 版以内存为主 + wiki 历史去重读 ES） + +Python 版 `knowlege_compile` 在运行时**全程依赖 ES**：先连 ES 把已切好的源 chunks 拉出来 +（`_common.py` 的 `es_search` / `MatchDenseExpr` 取 `doc_id` 对应的 chunk 集合）， +编译过程中的去重/合并又**回查 ES 里已写入的编译产物**做 KNN 候选，最后把产物写回 ES。 + +Go 版是 **pipeline 组件，以内存为主**，并在 wiki 历史去重阶段读 ES KNN： + +1. **输入来自内存**：`KnowledgeCompiler` 的 `Inputs()` 直接消费上游 `chunker`/`parser` 输出的 + `chunks`（每项带 `id` + `text`/`content_with_weight` 等）。删去「按 `doc_id` 去 ES 拉源 chunk」。 +2. **编译产物驻留内存**：所有编译产物（structure 图 / wiki page / raptor summary / mindmap tree / + datasetnav 节点）在 `Invoke` 期间先存于内存结构中，**组件不负责写 ES**。 +3. **落库由 pipeline 调用者负责**：产物通过 `Outputs()` 返回给下游，由 pipeline 的 + 调用方（组合根 / 后续 writer 组件）统一决定何时、如何写入 ES。 +4. **去重分两层**： + - 运行内去重：在内存里已累积的编译产物集合中做向量相似检索（暴力余弦 / 小顶堆 top-k）； + - 历史去重（wiki 变体）：对 document 级 graph（entity/relationship）并行发起 ES KNN 候选检索， + 作用域受 `tenant_id` + `dataset_id` + `variant` 过滤约束，命中后在组件内直接 merge。 + +> 结论：`KnowledgeCompiler` 组件保留 `common/memstore.go`（运行内去重）， +> 并为 wiki 变体新增**只读 ES KNN 候选检索**依赖（不在组件内写 ES）。 +> `Deps` 包含 `Chat` / `Embed` / `Tokenizer`，以及 `HistoricalKNN`（或 `DocEngine.Search` 只读封装，见 §3/§4）。 + +### 移植原则（遵循 AGENTS.md） +- 收敛为**单一组件** `KnowledgeCompiler`，用 `variant` 参数分派，不保留 Python 的多文件分叉。 +- 复用现有 `runtime.Component` 注册模式（参考 `component/extractor.go`）。 +- 组件包 **不得** 反向 import `internal/service`；模型/引擎解析器在 `internal/ingestion/task`（组合根）注入，沿用 `embedder.go` 的 `DefaultEmbedderResolver` 模式。 +- 优先用 Go 已有能力，缺失的才新增；不引入与现有 runtime 重复的兼容层。 + +--- + +## 2. 模块映射（Python → Go） + +``` +rag/advanced_rag/knowlege_compile/ +├── _common.py → knowledge_compiler/common/ (id, embed, tokenize, batching, dedup engine, in-mem store) +├── structure.py → knowledge_compiler/structure/ (compile + merge + chain + graph json) +├── wiki.py → knowledge_compiler/wiki/ (map / reduce / plan / refine) +├── raptor.py → knowledge_compiler/raptor/ (classic + psi builders + clustering) +├── mind_map_extractor → knowledge_compiler/mindmap/ (markdown→json tree) +└── dataset_nav.py → knowledge_compiler/datasetnav/ (incremental nav clustering + redis lock) +``` + +`runner.py` 的编排逻辑（LLMCallPool / flush / synthesis）下沉为 +`knowledge_compiler/common` 里的并发原语（信号量 + 优先级队列 + 批累积器）， +由 `KnowledgeCompiler.Invoke` 直接驱动，不再单独保留一个 runner 文件。 + +--- + +## 3. Go 基础设施依赖：现状与缺口 + +### ✅ 已具备 +- **Embedding**：`componentpkg.Embedder` 接口 + `EmbedderResolver` + `DefaultEmbedderResolver` + （`internal/ingestion/task/embedder.go` 已在 `init()` 注入）。 +- **Tokenizer**：`internal/tokenizer`（C++ binding，含 `tokenize` / `fine_grained_tokenize` 等价能力）。 + 另需一个 `num_tokens_from_string` 等价函数（基于 tiktoken 或 RAG 分词器，放在 `common` 包）。 +- **DocEngine / 历史 KNN Reader（只读）**：组件不写 ES，但 wiki 变体会在历史去重阶段读取 ES KNN 候选 + （见 §1.1）。建议以 `HistoricalKNN` 接口封装 `Search` 能力，避免组件直接耦合全量 `DocEngine`。 + 输入 chunks 仍来自内存 inputs，写入 ES 与普通 chunk 走同一落库路径。 +- **Chat 工厂**：`models.NewModelFactory().CreateModelDriver(...)` + `models.NewEinoChatModel` + （`extractor.go:einoExtractorChatInvoker` 已示范）。 +- **xxhash**：`github.com/cespare/xxhash/v2`（已在 go.mod，被 `chunk_builder` 等使用）。 +- **组件注册**：`runtime.MustRegister(name, CategoryIngestion, ctor, Metadata{Inputs,Outputs})`。 + +### ⚠️ 缺口（需新建，按风险排序） +1. **JSON-mode chat helper（gen_json 等价物）** —— 最高频依赖。 + 封装 `einoExtractorChatInvoker`：请求时带 `response_format=json`（或 JSON schema）， + 返回解析后的 `map[string]any` / `[]any`，带 code-fence 剥离、重试、超时。 + 建议新增 `knowledge_compiler/common/jsonchat.go`，复用 `extractor` 的 driver 解析与 invoker 注入 seam。 +2. **内存产物存储 + 内存向量检索（`common/memstore.go`）** —— 承担运行内去重（见 §1.1）。提供： + - `Add(item)` / `Upsert(id, item)` / `Delete(id)`：维护本次 pipeline 内的编译产物集合； + - `TopK(vec []float32, k int, threshold float64) []Hit`：内存余弦 top-k，供去重/合并做候选查找。 + **实现建议（见下方①）**：内部把已存产物向量维护为一个 `[][]float32` + 预计算的 L2 norm + （embedder 若已归一化则可直接用点积）；`TopK` 对 query 做一次 matVec 取点积再阈值过滤 + + 部分 top-k。单次 pipeline 规模不大，朴素循环也够，但走矩阵形态（gonum `mat.Dense.Mul` + 或手写点积）更一致、更易与 RAPTOR Psi 的全对 `M·Mᵀ` 共用同一套向量运算。 + - `Snapshot() []Item`：`Invoke` 结束时把全部产物交给 `Outputs()`，由 pipeline 调用者落库。 +3. **历史去重 KNN 检索器（wiki 专用）** —— 新增 `common/historical_knn.go`（或等价位置）： + - `TopKHistory(ctx, qvec, k, threshold, scope) []Hit`，`scope` 至少含 `tenant_id/dataset_id/variant`； + - 仅用于读取历史候选；命中后在组件内执行 merge（仍不在组件内写 ES）。 +3. **聚类数值库（RAPTOR classic builder）** —— Python 用 `sklearn` 的 + `GaussianMixture`(BIC 选簇)、`AgglomerativeClustering`(ward + dendrogram gap)、`umap.UMAP`。 + **最终决策（M6 前已拍板）：纯 Go + `gonum`，不引入 `pa-m/sklearn`，不用 CGO/C++ 内核。** + + **放弃 `pa-m/sklearn`**：它只是 sklearn API 的 Go 模仿版，不成熟，尤其 GMM/UMAP 支持残缺， + 多一个需审计的第三方 ML 依赖得不偿失。 + **不用 CGO/C++ 内核（`libcluster.so` 之类）**：本仓库 `build.sh` 已在 juggling CGO 原生库 + （office_oxide / pdfium / pdf_oxide / lld），再加一个自维护 C++ 数值内核会显著加重构建与长期维护负担； + 而本场景数值热度根本不需要它（见下）。 + **只加 `gonum.org/v1/gonum`（纯 Go、无 CGO，不影响 `build.sh` 的 CGO 链路）到 go.mod。** + + **关键事实（已核对 raptor.py 源码）**： + - `clustering()` 在分支**之前**无条件跑 `umap.UMAP(metric="cosine")` 降到 ≤12 维（`raptor.py:326-333`）， + 所以**即便只支持 AHC 也仍需降维器**；Go 用 **PCA（gonum `mat.SVD`，SVDThin）** 替代 UMAP。 + 对文本 embedding（BGE/E5/...）PCA 通常保留大部分聚类结构，且比「cosine-UMAP→Ward」更契合 Ward 的欧氏几何。 + - **Psi builder 完全不聚类**（仅余弦矩阵 + union-find，`raptor.py:425-432`），零 GMM/AHC/UMAP 依赖。 + - **GMM 是 `covariance_type="diag"` + `reg_covar=1e-4`（`raptor.py:261,351`），不是 full 协方差**。 + 因此 EM 只需逐维均值/方差，无需 Cholesky / 全协方差行列式 / 逆矩阵 —— 实现量约 **200-300 行**， + 远低于 full 协方差的 500-800 行；BIC 参数 `p = K·d(均值) + K·d(对角协方差) + (K-1)(权重)`，可直接算。 + GMM 还用 `predict_proba` + `threshold`(默认 0.1) 软分配（取 prob>阈值 的首簇，否则 argmax）。 + - **Ward AHC**：维护 `{centroid, size, SSE}` 的 struct（cache locality 好），距离矩阵用扁平 `[]float64` + 连续内存（避免 `[][]float64` 反复分配）；典型 N=几千~几万、降维后 D≤12，O(n²) 距离矩阵完全可接受。 + 保存 merge history 即得到 linkage 矩阵 `Z=[left,right,distance,size]`，**dendrogram gap = 排序后 heights 的差分 argmax**。 + + **交付分解（M6）**： + - Psi builder：仅余弦矩阵（复用 ① 的 memstore 向量运算），零聚类依赖。 + - classic(AHC + PCA)：自写 Ward 链接 + gap 检测 + `_adjust_tree_nodes`(质心重分配，k-means 式)；PCA 走 gonum。 + - GMM/BIC + `predict_proba` 软分配：**留 `ClusteringMethod` 可插拔接口，后补**（纯 Go diagonal-GMM EM，已确认成本低）； + 不在 M6 阻塞项。 + + **接口骨架（现在就定义，避免将来改 `raptor.Run` 签名）** —— M6 落地前先把边界定死： + + ```go + // ClusteringMethod 标识聚类后端；M6 只交付 AHC，GMM 留 TODO。 + type ClusteringMethod string + + const ( + ClusteringAHC ClusteringMethod = "AHC" // 默认：Ward + gap 检测 + _adjust_tree_nodes + ClusteringGMM ClusteringMethod = "GMM" // TODO(M6+): diagonal-GMM EM + BIC + predict_proba 软分配 + ) + + // ClusterSpec 聚类配置，对应 Python RaptorConfig 的聚类相关字段。 + type ClusterSpec struct { + Method ClusteringMethod + Threshold float64 // AHC: dendrogram gap 阈值；GMM: predict_proba 软分配阈值(默认 0.1) + MinClusters int // _get_optimal_clusters 下界 + MaxClusters int // _get_optimal_clusters 上界(默认 20) + } + + // Clusterer 聚类后端接口；AHC/Psi 实现各自满足。 + // embeddings: 降维后矩阵 (n×d, d≤12)；返回每点簇标签 + 簇质心（供 _adjust_tree_nodes）。 + type Clusterer interface { + Fit(embeddings [][]float64, spec ClusterSpec) (labels []int, centroids [][]float64, err error) + } + + // Run 签名以 ClusterSpec 传入，后端按 Method 分派；GMM 未实现时返回明确 ErrNotImplemented。 + func Run(ctx context.Context, in *Inputs, spec ClusterSpec, deps Deps) (*Outputs, error) + ``` + + 注：Psi builder 不走 `Clusterer`（仅余弦矩阵 + union-find），但 `Run` 仍以 `ClusterSpec` 统一入口， + 内部按 `Method`/`builder` 分派到 Psi 或 classic 路径，保证签名单一、将来加 GMM 不破 `Run`。 +4. **`markdown_to_json`（mind_map）的 markdown 来源澄清** —— 注意它**不是文档级也不是 chunk 级 + 的源 markdown**，而是 **LLM 的回答 markdown**（`mind_map_extractor.py:168` + `markdown_to_json.dictify(response)`，`response` 是 `MIND_MAP_EXTRACTION_PROMPT` 产出的思维导图大纲； + 输入 `sections` 只是喂给 LLM 的源文本，`dictify` 作用在模型输出上）。 + 因此**不能复用 Go Parser 组件（`parser.go`/`parser_dispatch.go`）的输出**——那是源文档解析结果，与本题无关。 + Go 端用 `goldmark` 解析** LLM 回复的 markdown 大纲**，按 Python `dictify` 语义（标题层级→嵌套 dict、 + 列表项→数组）自实现轻量 AST 遍历转换即可；无需引入等价重库。计划措辞应改为「LLM 回复 markdown→tree」。 +5. **Redis 分布式锁** —— `dataset_nav` 用 `RedisDistributedLock`。 + `internal/engine/redis` 已有 `RedisClient`，新增一个 `spin_acquire/release`（SET NX + TTL）封装即可。 +6. **多语言 config localize** —— Python 的 `_struct_localize` 取 `cfg["en"]` 等。 + Go 端 `common` 加一个 `localize(value any, lang string) string` 处理 `string` / `[]string` / `map[string]string`。 +7. **`split_chunks` token budget（注意含义）** —— 它**不读 ES/磁盘、也不切分单个 chunk**， + 而是把**内存中**的 chunks **按 LLM 上下文窗口做 token-budget 打包成 batch** + （`generator.split_chunks:795` 注释「Do not split a single chunk, even if it exceeds max_length」， + 是 `build_chunk_batches`/`run_chunked_pipeline` 引擎里的预算装箱步骤）。既然 chunks 已在内存， + 仍需要它来尊重 `chat_mdl.max_length` 的窗口约束。**建议改名** `split_chunks.go` → `batch_packing.go` + （或并入 `batch.go`），明确其语义是「内存 chunks → LLM 批次」，避免误读成「从外部切分加载」。 + +--- + +## 4. 组件接口设计 + +```go +const componentNameKnowledgeCompiler = "KnowledgeCompiler" + +// Param 由 DSL params 构造；缺失键走 Defaults()。 +type Param struct { + Variant string // "structure" | "wiki" | "raptor" | "mindmap" | "datasetnav" + LLMID string + EmbeddingModel string + Language string // "en" / "zh" / ... + TemplateIDs []string // structure/wiki 用；group_ids 在 Invoke 时解析 + GroupIDs []string + SimilarityThreshold float64 // 去重阈值，默认 0.99(structure) / 0.9(artifact) + MaxWorkers int + // 变体特有：raptor 的 max_cluster / tree_builder / clustering_method； + // wiki 的 synthesis example；mindmap 的 prompt 覆盖等 + Extra map[string]any + + // 历史去重开关：wiki 变体在 document graph 阶段并行读取 ES KNN 候选。 + EnableHistoricalDedup bool + + // 容量护栏（缺口 A，防 OOM/GC 抖动）：限制单次 Invoke 产物规模，见 §4 Outputs。 + Guardrails CapacityGuardrails +} + +// CapacityGuardrails 单次 Invoke 产物容量上限；任一超限即触发对应处置（error / flush）。 +type CapacityGuardrails struct { + MaxItems int // 产物条目上限（默认按变体给安全值，如 structure 5e4、raptor 按树规模） + MaxVectorBytes int64 // 向量总字节上限（= Σ len(vec)*4） + MaxOutputBytes int64 // 全部产物序列化字节上限 + OnExceed ExceedAction // "error" | "flush" +} + +// ExceedAction 超限处置：error=返回 ErrCapacityExceeded；flush=经 ChunkedSink 分段交付。 +type ExceedAction string +``` + +- **变体命名口径（统一主名 + alias，缺口 F）**：组件内部 `Param.Variant` 与 `Run` 分派键一律用 + **`structure` / `wiki` / `raptor` / `mindmap` / `datasetnav`**（无下划线，Go 风格）。 + DSL 模板若填旧名 `mind_map` / `dataset_nav`，在 `Param` 构造处做 **alias 映射 + 弃用日志** + （`mind_map→mindmap`、`dataset_nav→datasetnav`），其余未知值返回明确 `ErrUnknownVariant`。 + 里程碑描述里的 `mind_map`/`dataset_nav` 仅为人读别名，代码与模板一律用主名。 + +- `Inputs()`：返回 `chunks`(**直接来自上游 chunker/parser 的输出**，每项带 `id`+`text`/`content_with_weight`， + 见 §1.1，**不再从 ES 读源 chunks**)、`llm_id`(可选覆盖)、`embedding_model`(可选覆盖)、变体特有键。 + - `HistoricalCandidates []Candidate`(**可选覆盖路径**)：用于测试或离线模式；若未提供且 + `EnableHistoricalDedup=true`，wiki 变体从 ES 并行检索历史候选（KNN）。 + 命中在组件内直接 merge（merge 统计仅留存于组件内部）。 + - `Sink ChunkedSink`(**可选**)：当 `CapacityGuardrails.OnExceed=="flush"` 时，组件边产边回调 + 分段交付，避免 `Outputs()` 一次性持有全部产物（缺口 A）。 +- `Outputs()`：**回传 `chunks`**（与上游 chunker 输出 schema 完全一致）—— + 各变体产物（structure 图行 / wiki page / raptor summary / mindmap tree / + datasetnav 节点）转换为 `schema.ChunkDoc`（对齐 `conf/infinity_mapping.json`，以 `compile_kwd` + 区分普通分块），与上游输入 `chunks` 合并后作为 `Outputs()` 的 `chunks` 返回；运行内去重/历史合并 + 统计不再外显到输出 surface（仅在组件内部 `common.Outputs` 中留存）。 + **容量护栏**：若 `Guardrails.OnExceed=="error"` 且超限，`Invoke` 返回 `ErrCapacityExceeded`； + 若 `=="flush"`，产物经 `ChunkedSink` 已分段交付。持久化与任何普通 chunk 一致，由调用方决定。 +- `Invoke(ctx, inputs)`：解析 `variant` → 调 `structure.Run` / `wiki.Run` / `raptor.Run` / + `mindmap.Run` / `datasetnav.Run`。每个子包暴露 `Run(ctx, deps, param, inputs) (map[string]any, error)`， + 共享 `common.Deps{ Chat ChatInvoker; Embed Embedder; Tokenizer; HistoricalKNN; ... }`。 + 各 `Run` 内部用 `common.MemStore` 做运行内去重；wiki 变体额外并行调用 `HistoricalKNN` 做历史候选标记， + 结束时 `Snapshot()` 进 outputs。 +- `init()`：`runtime.MustRegister(componentNameKnowledgeCompiler, runtime.CategoryIngestion, NewKnowledgeCompilerComponent, Metadata{...})`。 + +### 依赖注入 seam（沿用 embedder 模式） +`KnowledgeCompiler` 不直接构造模型，而是持有一个 `common.DepsResolver`： +在 `internal/ingestion/task` 的 `init()` 里把生产 resolver 注入 +（`GetChatModel(tenantID, llmID)`、`DefaultEmbedderResolver`、`HistoricalKNNResolver`； +datasetnav 若需 Redis 锁再单独注入 `RedisClient`，见 §5 M8）。 +测试时替换 resolver 注入 mock。组件仅执行历史候选读取，不在组件内写 ES；落库仍由 pipeline 调用者处理。 + +### 4.1 输出合约（chunk 对齐，缺口 D 演化） + +组件**不落库、也不再通过独立的 writer seam 写 ES**：编译产物在 `Invoke` 期间驻留内存， +结束时转换为 `schema.ChunkDoc`（对齐 `conf/infinity_mapping.json`，以 `compile_kwd` 区分普通 +分块），与上游输入 `chunks` 合并后作为 `Outputs()` 的 `chunks` 返回，schema 与上游 `chunker` +完全一致。行结构冻结项（M1 即与组件输出 schema 同步）： + +- **行结构 schema**：每条产物即一个 `schema.ChunkDoc`，含 `id`(xxhash 稳定)、`doc_id`/`tenant_id`、 + `text`/`content_with_weight`、`compile_kwd`(=`variant`)、`parent_kwd`(树形变体)、`q__vec`(向量)、 + 以及 `kc_*` 前缀保留原始 `meta`(kind/level/name/size)。 +- **幂等键**：`(tenant_id, doc_id, variant, id)` 唯一；若需落库，调用方以 upsert 语义写入，重复 id 覆盖不新增。 +- **历史去重职责**：wiki 变体在组件内执行**并行历史候选检索 + 直接 merge**（ES KNN 只读），merge 统计仅留存 + 于组件内部 `common.Outputs`，不再外显到输出 surface。 +- **错误语义**：容量超限(`ErrCapacityExceeded`)、未知变体(`ErrUnknownVariant`)、LLM/embed 失败均透传。 + +--- + +## 5. 分阶段实施 + +- **M1 脚手架 + common 底座** + `knowledge_compiler/common/`：`id.go`(xxhash 稳定 id)、`embed.go`(Embedder 包装)、 + `tokenize.go`(tokenizer + num_tokens)、`batch.go`(`build_chunk_batches` + `run_chunked_pipeline`)、 + `memstore.go`(内存产物存储 + 内存向量检索 `Add`/`Upsert`/`Delete`/`TopK`/`Snapshot`)、 + `localize.go`、`batch_packing.go`(内存 chunks → LLM token-budget 批次，原 `split_chunks`)。 + 单测：mock Embedder + `memstore` 的 `TopK` 精度/阈值断言； + **新增容量护栏单测**（缺口 A）：注入超大 chunks 集，断言 `CapacityGuardrails` 的 `error`/`flush` + 两条路径（`ErrCapacityExceeded` 与 `ChunkedSink` 分段回调次数），作为大文档内存压力的最小防护。 + +- **M2 JSON chat helper + 并发原语** + `common/jsonchat.go`（gen_json 等价：JSON 模式请求 + 解析 + 重试 + 超时）、 + `common/pool.go`（LLMCallPool 等价：优先级信号量）。 + 单测：注入 canned JSON 响应。 + +- **M3 structure 变体（编译 + 内存去重）** + `structure/compile.go`：`compile_structure_from_text`（list/set/hypergraph 提示词渲染、 + hypergraph 的 node→edge 两阶段、产物行构造）、`merge.go`：`merge_compiled_structures` + 的本地去重（余弦候选 + LLM merge-pair + 别名重写 relation 端点 + 紧凑图 JSON），全程走 `common.MemStore`。返回 docs 供断言。 + +- **M4 structure 内存合并闭环** + `merge.go` 接 `memstore.go`：**内存 `TopK` 候选查找**（替代 ES KNN）、分组 LLM 判定、 + 批量合并、`rebuild_structure_graph_json`，最终 `Snapshot()` 全部产物进 `Outputs()`。 + 单测用 `memstore` 双跑验证 inserted/updated/duplicates_dropped 计数（无需 DocEngine）。 + +- **M5 wiki 变体** + `wiki/`：`map_from_chunks` → `reduce_from_extracts` → `plan_from_reduction` → + `refine_from_plan`，含 5 个 `WIKI_*_COMPILE_KWD` 与 synthesis 阶段（example + compile_kwd 覆盖）。 + 在 document 级 graph（entity/relationship）产出后，启用并行历史去重：内存 `TopK` + ES KNN 候选， + 命中后组件内直接 merge 并写入统计（`historical_merged`）。 + +- **M6 raptor 变体** + `raptor/`：先交付 **Psi builder（仅余弦矩阵，零聚类依赖）** + **classic(AHC + PCA 降维)**； + AHC 仍需降维器（见缺口 #3：`clustering()` 无条件先 UMAP，Go 用 **gonum PCA** 替代）。 + 聚类后端：Ward AHC + gap 检测 + `_adjust_tree_nodes` 全部**纯 Go + gonum** 自实现； + GMM/BIC 留 `ClusteringMethod` 可插拔接口与 TODO（纯 Go diagonal-GMM，后补，不阻塞 M6）。 + `summarize_texts` 走 `jsonchat`/普通 chat + embed。 + **验证门禁（缺口 C，替代原「保真度可接受」弱约束）**：M6 落地前先跑 Python `raptor.py` 在固定样本集上 + 产出一个**基线快照**（簇标签、树高、叶覆盖、summary 文本），Go 版须对齐： + - 指标：与基线对比 **NMI/ARI**（簇一致性）、**平均/最大树高**、**叶节点覆盖率**(≥基线 95%)； + - **离线召回**：固定 query 集的 top-k hit 率回退不超过设定阈值（如 5%）； + - **门禁失败回退**：classic(AHC) 任一指标不达标即**禁用 classic、仅启用 Psi**（Psi 零聚类依赖、风险更低）， + 并告警，不阻塞发布。 + （开放问题 3：需确认是否存在可固化的 Python/Go 聚类一致性回归样本集；若无，先用合成 embedding 构造。） + +- **M7 mindmap 变体** + `mindmap/`：goldmark 解析 + `dictify` 语义转换 + 多段合并（`_merge`/`_be_children`）。 + +- **M8 datasetnav 变体** + `datasetnav/`：Redis 锁封装 + 分层 KNN 下降 + LLM merge/summary + 增量 split/rebalance。 + 单测用 miniredis 或内存锁。 + +- **M9 注册 + 模板 + 集成测试** + 在 `internal/ingestion/pipeline/template/` 新增一个使用 `KnowledgeCompiler` 的 + pipeline DSL（或一个参数化模板按 variant 切换），下游接普通 chunk 落库路径（组合根 / ES 写入） + （消费 `KnowledgeCompiler` 的 outputs 产物）；`builtin_registry` 自动加载； + 编写端到端用例：断言组件 outputs 产物正确，落库交由 pipeline 层验证。 + +--- + +## 6. 测试与验证 + +- 每个子包单测：`mockEmbedder`、`mockChatInvoker`（返回固定 JSON）、`common.MemStore`（真实内存实现）。 +- 包级运行（遵循 AGENTS.md，必须用 build.sh 以拿到 CGO/原生库）： + ```bash + bash build.sh --test ./internal/ingestion/component/knowledge_compiler/... + ``` +- 行为对齐（**structure + wiki 为 CI 必选门禁；raptor/mindmap/datasetnav 为可选 golden**）： + 用同一组 chunks 跑 Python 与 Go，对比 `compile_structure_from_text` 产物行数量与 `merge` 后 + inserted/updated 计数；wiki 额外对比 `reduce`/`refine` 关键阶段产物形状。未达标**阻断合并**。 +- 不默认跑全量 `go test`；只跑本组件子树。 + +--- + +## 7. 风险与决策 + +1. **聚类保真度（最高风险）**：sklearn GMM/UMAP 在 Go 无直接等价。 + 决策已定：**纯 Go + `gonum`，不引入 `pa-m/sklearn`、不用 CGO/C++ 内核**。 + - AHC 路径**仍依赖降维**（Python 分支前无条件跑 UMAP，`raptor.py:326`），Go 用 **gonum PCA** 替代； + PCA 输出的欧氏空间更契合 Ward，保真度可接受，需在 M6 单测固化聚类数/阈值。 + - GMM 是 `covariance_type="diag"`（`raptor.py:261,351`），diagonal-GMM EM+BIC 约 200-300 行， + 纯 Go 可行，留作可插拔后端后补，不阻塞 M6。 + - Psi builder 仅用余弦矩阵，不在此风险内。 + 结论：M6 先交付 **AHC + Psi（AHC 走 PCA 降维；Psi 仅依赖余弦）**，GMM/UMAP 作为可插拔聚类后端后续补。 + **验证闭环见缺口 C**：M6 须以 Python 基线快照做 NMI/ARI/树高/叶覆盖/离线召回门禁，失败则回退 Psi-only。 +2. **历史去重一致性**：Python 用 ES `MatchDenseExpr`（近似 KNN）在全量已存产物上查候选； + Go 版改为「运行内 memstore 精确余弦 + wiki 历史 ES KNN 并行候选」。关键点： + - 运行内去重：`MemStore.TopK` 精确余弦； + - 历史去重（wiki）：`HistoricalKNN.TopKHistory`（ES 近似 KNN）按 `tenant_id/dataset_id/variant` 过滤； + - 历史命中后组件内直接 merge，落库幂等兜底由调用方负责。 + 结论：不再将其定义为“仅本次运行去重”的 breaking change；行为更接近 Python，但仍保留 + `enable_historical_dedup` 开关用于灰度与回退。 +3. **依赖方向**：组件包不得 import `internal/service`；所有 resolver 在 `internal/ingestion/task` 注入。 +4. **并发语义**：Python 的 `asyncio.Semaphore` + 取消检查 → Go 用 `context.Context` + `errgroup`/信号量， + 取消用 `cancel_check` 包成 `ctx.Err()` 检查。 +5. **多语言/模板字段**：config 的 `en`/`zh` 多语言与 `entity`/`relation` 新模板形状需完整映射， + 避免只覆盖旧 `output.entities/relations` 形状。 + +--- + +## 8. 建议目录树 + +``` +internal/ingestion/component/knowledge_compiler/ +├── PORT_PLAN.md +├── component.go // KnowledgeCompilerComponent + Param + Invoke + init 注册 +├── deps.go // Deps / DepsResolver 接口与注入 seam +├── common/ +│ ├── id.go // stable_row_id (xxhash) +│ ├── embed.go // Embedder 包装 +│ ├── tokenize.go // tokenizer + num_tokens + localize +│ ├── batch.go // build_chunk_batches + run_chunked_pipeline +│ ├── batch_packing.go // 内存 chunks → LLM token-budget 批次 (原 split_chunks) +│ ├── memstore.go // 内存产物存储 + 内存向量检索 (Add/Upsert/Delete/TopK/Snapshot；不碰 ES) +│ ├── historical_knn.go // wiki 历史候选检索接口/实现（只读 ES KNN） +│ ├── jsonchat.go // gen_json 等价 +│ └── pool.go // LLMCallPool 等价 +├── structure/ (compile.go, merge.go, prompt.go, graph.go, chain.go) +├── wiki/ (map.go, reduce.go, plan.go, refine.go) +├── raptor/ (raptor.go, clustering.go, psi.go) +├── mindmap/ (mindmap.go, dictify.go) +└── datasetnav/(nav.go, lock.go, split.go) +``` + +--- + +## 9. 交付顺序建议（最小可用路径） + +`M1 → M2 → M3 → M4`（structure 变体端到端可用） +→ `M9` 注册 + 模板（让 pipeline 能跑 structure） +→ 之后再按 `M5/M6/M7/M8` 逐个补齐 wiki / raptor / mindmap / datasetnav， +每补一个变体就扩展 `Param.Variant` 与对应 `Run`，保持单一组件、单一注册入口。 + +**发布策略（回答开放问题 2）**：推荐**分阶段发布**——M9 先只上线 `structure` 变体并默认启用； +`wiki`/`raptor`/`mindmap`/`datasetnav` 经 feature flag 逐步放开，其中 `raptor` 的 classic(AHC) +受 M6 验证门禁约束（不达标则仅启用 Psi），避免低质量摘要树进入召回。 diff --git a/internal/ingestion/component/knowledge_compiler/common/batch.go b/internal/ingestion/component/knowledge_compiler/common/batch.go new file mode 100644 index 0000000000..0953621d3f --- /dev/null +++ b/internal/ingestion/component/knowledge_compiler/common/batch.go @@ -0,0 +1,23 @@ +package common + +import "context" + +// BatchHandler processes one batch of chunks. Returning an error aborts the +// pipeline (RunChunkedPipeline propagates it). +type BatchHandler func(ctx context.Context, batch []Chunk) error + +// RunChunkedPipeline packs chunks into token-budgeted batches and invokes h for +// each, in order, honouring ctx cancellation. It is the Go equivalent of +// Python's run_chunked_pipeline engine step. +func RunChunkedPipeline(ctx context.Context, chunks []Chunk, maxTokens int, tok Tokenizer, h BatchHandler) error { + batches := PackBatches(chunks, maxTokens, tok) + for _, b := range batches { + if err := ctx.Err(); err != nil { + return err + } + if err := h(ctx, b); err != nil { + return err + } + } + return nil +} diff --git a/internal/ingestion/component/knowledge_compiler/common/batch_packing.go b/internal/ingestion/component/knowledge_compiler/common/batch_packing.go new file mode 100644 index 0000000000..7c4b1e93ae --- /dev/null +++ b/internal/ingestion/component/knowledge_compiler/common/batch_packing.go @@ -0,0 +1,46 @@ +package common + +// PackBatches groups chunks into LLM-context-window-sized batches using greedy +// bin-packing by token budget. It mirrors Python's split_chunks / build_chunk_batches: +// a single chunk is NEVER split, even when it exceeds maxTokens — in that case it +// is placed alone in its own batch. Returns nil when there are no chunks. +func PackBatches(chunks []Chunk, maxTokens int, tok Tokenizer) [][]Chunk { + if len(chunks) == 0 { + return nil + } + if maxTokens <= 0 { + return [][]Chunk{chunks} + } + var batches [][]Chunk + var cur []Chunk + curTokens := 0 + + flush := func() { + if len(cur) > 0 { + batches = append(batches, cur) + cur = nil + curTokens = 0 + } + } + + for _, c := range chunks { + text := c.Text + if text == "" { + text = c.Content + } + t := numTokens(tok, text) + if t > maxTokens { + // Oversized chunk: emit current batch, then the chunk alone. + flush() + batches = append(batches, []Chunk{c}) + continue + } + if curTokens+t > maxTokens && len(cur) > 0 { + flush() + } + cur = append(cur, c) + curTokens += t + } + flush() + return batches +} diff --git a/internal/ingestion/component/knowledge_compiler/common/common_test.go b/internal/ingestion/component/knowledge_compiler/common/common_test.go new file mode 100644 index 0000000000..5076846fd8 --- /dev/null +++ b/internal/ingestion/component/knowledge_compiler/common/common_test.go @@ -0,0 +1,219 @@ +package common + +import ( + "context" + "testing" +) + +func vec(a, b, c float32) []float32 { return []float32{a, b, c} } + +func TestMemStoreTopK(t *testing.T) { + m := NewMemStore() + m.Add(Product{ID: "a", Vector: vec(1, 0, 0)}) + m.Add(Product{ID: "b", Vector: vec(0, 1, 0)}) + m.Add(Product{ID: "c", Vector: vec(0, 0, 1)}) + + // query aligned with "a" → exact match + hits := m.TopK(vec(1, 0, 0), 3, 0.0) + if len(hits) != 3 { + t.Fatalf("expected 3 hits, got %d", len(hits)) + } + if hits[0].ID != "a" || hits[0].Score < 0.999 { + t.Fatalf("expected top hit a with score ~1, got %s %.3f", hits[0].ID, hits[0].Score) + } + + // threshold filters out orthogonal vectors + hits = m.TopK(vec(1, 0, 0), 3, 0.5) + if len(hits) != 1 || hits[0].ID != "a" { + t.Fatalf("threshold=0.5 should keep only a, got %v", hits) + } + + // k caps results + hits = m.TopK(vec(0.6, 0.6, 0.6), 2, 0.0) + if len(hits) != 2 { + t.Fatalf("k=2 should cap to 2, got %d", len(hits)) + } +} + +func TestMemStoreUpsertDelete(t *testing.T) { + m := NewMemStore() + m.Add(Product{ID: "x", Vector: vec(1, 0, 0)}) + m.Add(Product{ID: "y", Vector: vec(0, 1, 0)}) + m.Upsert(Product{ID: "x", Vector: vec(0, 0, 1)}) + if got := m.TopK(vec(0, 0, 1), 1, 0.99); len(got) != 1 || got[0].ID != "x" { + t.Fatalf("upsert did not replace x vector: %v", got) + } + m.Delete("x") + if m.Len() != 1 { + t.Fatalf("delete failed, len=%d", m.Len()) + } + if got := m.Snapshot(); len(got) != 1 || got[0].ID != "y" { + t.Fatalf("snapshot wrong after delete: %v", got) + } +} + +func TestStableRowID(t *testing.T) { + a := StableRowID("t1", "d1", "structure", "hello") + b := StableRowID("t1", "d1", "structure", "hello") + c := StableRowID("t1", "d1", "structure", "world") + if a != b { + t.Fatalf("same parts must yield same id: %s vs %s", a, b) + } + if a == c { + t.Fatalf("different content must yield different id") + } + // part order must not collide + if StableRowID("ab", "c") == StableRowID("a", "bc") { + t.Fatalf("part-ordering collided") + } +} + +func TestLocalize(t *testing.T) { + if got := Localize("plain", "zh"); got != "plain" { + t.Fatalf("string passthrough: %q", got) + } + // Lists render as numbered lines (mirrors Python's _struct_localize). + if got := Localize([]string{"en-only"}, "zh"); got != "1. en-only" { + t.Fatalf("slice numbered join: %q", got) + } + if got := Localize([]any{"a", "b"}, "en"); got != "1. a\n2. b" { + t.Fatalf("any-slice numbered join: %q", got) + } + m := map[string]string{"en": "E", "zh": "中"} + if got := Localize(m, "zh"); got != "中" { + t.Fatalf("map lang pick: %q", got) + } + if got := Localize(m, "fr"); got != "E" { + t.Fatalf("map fallback en: %q", got) + } + // lang == "en" with no "en" key yields "" (no arbitrary-language fallback, + // matching Python). + if got := Localize(map[string]string{"zh": "中"}, "en"); got != "" { + t.Fatalf("no arbitrary fallback: %q", got) + } + // Nested map values localize recursively (list under lang key). + if got := Localize(map[string]any{"en": []any{"x", "y"}}, "zh"); got != "1. x\n2. y" { + t.Fatalf("nested en fallback: %q", got) + } +} + +func TestGuardrails(t *testing.T) { + g := CapacityGuardrails{MaxItems: 2, OnExceed: ExceedError} + o := Outputs{Products: make([]Product, 3)} + if !g.Exceeded(o) { + t.Fatalf("expected breach on MaxItems=2 with 3 products") + } + if err := checkGuardrails(g, o); err != ErrCapacityExceeded { + t.Fatalf("expected ErrCapacityExceeded, got %v", err) + } + + // flush path: a sink receives the buffer and it is reset + flushed := false + sink := sinkFunc(func(_ context.Context, items []Product) error { + flushed = true + if len(items) != 3 { + t.Fatalf("sink got %d items", len(items)) + } + return nil + }) + err := o.EnforceGuardrails(CapacityGuardrails{MaxItems: 2, OnExceed: ExceedFlush}, sink, context.Background()) + if err != nil || !o.Flushed || !flushed { + t.Fatalf("flush path failed: flushed=%v err=%v", flushed, err) + } + if len(o.Products) != 0 { + t.Fatalf("buffer should be reset after flush, got %d", len(o.Products)) + } +} + +func TestEnforceGuardrails(t *testing.T) { + // Error policy: returns ErrCapacityExceeded, leaves products untouched. + o := Outputs{Products: make([]Product, 5)} + err := o.EnforceGuardrails(CapacityGuardrails{MaxItems: 2, OnExceed: ExceedError}, nil, context.Background()) + if err != ErrCapacityExceeded { + t.Fatalf("EnforceGuardrails(error) = %v, want ErrCapacityExceeded", err) + } + if len(o.Products) != 5 { + t.Errorf("error policy must not drop products, got len=%d", len(o.Products)) + } + + // Flush policy with sink: emits and resets. + var captured []Product + sink := sinkFunc(func(_ context.Context, items []Product) error { + captured = append(captured, items...) + return nil + }) + o = Outputs{Products: make([]Product, 5)} + err = o.EnforceGuardrails(CapacityGuardrails{MaxItems: 2, OnExceed: ExceedFlush}, sink, context.Background()) + if err != nil { + t.Fatalf("EnforceGuardrails(flush+sink) err = %v", err) + } + if !o.Flushed { + t.Error("Flushed flag not set after flush") + } + if len(o.Products) != 0 { + t.Errorf("products not reset after flush, len=%d", len(o.Products)) + } + if len(captured) != 5 { + t.Errorf("sink got %d items, want 5", len(captured)) + } + + // Flush policy WITHOUT sink: degrades to no-op (preserves the buffer + // for downstream merging; the caller decides whether to continue). + o = Outputs{Products: make([]Product, 5)} + err = o.EnforceGuardrails(CapacityGuardrails{MaxItems: 2, OnExceed: ExceedFlush}, nil, context.Background()) + if err != nil { + t.Fatalf("EnforceGuardrails(flush+nil sink) err = %v", err) + } + if o.Flushed { + t.Error("Flushed should stay false when no sink is available") + } + if len(o.Products) != 5 { + t.Errorf("nil-sink flush should preserve products, got len=%d", len(o.Products)) + } + + // Under-limit: both error and flush policies must be no-ops. + for _, on := range []ExceedAction{ExceedError, ExceedFlush} { + o := Outputs{Products: make([]Product, 2)} + err := o.EnforceGuardrails(CapacityGuardrails{MaxItems: 5, OnExceed: on}, sink, context.Background()) + if err != nil { + t.Errorf("under-limit EnforceGuardrails(%s) err = %v", on, err) + } + if o.Flushed { + t.Errorf("under-limit Flushed should be false (on=%s)", on) + } + } +} + +func TestPackBatches(t *testing.T) { + chunks := []Chunk{ + {ID: "1", Text: "aaaa"}, + {ID: "2", Text: "bbbb"}, + {ID: "3", Text: "cccccccccccc"}, // oversized + {ID: "4", Text: "dddd"}, + } + // EstimateTokens ≈ len/4 → 1,1,3,1. budget 2 → [1,2],[3],[4] + batches := PackBatches(chunks, 2, nil) + if len(batches) != 3 { + t.Fatalf("expected 3 batches, got %d: %v", len(batches), batches) + } + if len(batches[1]) != 1 || batches[1][0].ID != "3" { + t.Fatalf("oversized chunk must be alone: %v", batches[1]) + } + + // zero budget → single batch + if got := PackBatches(chunks, 0, nil); len(got) != 1 || len(got[0]) != 4 { + t.Fatalf("zero budget should be one batch of all: %v", got) + } +} + +// checkGuardrails mirrors the error-path the variant Run applies before return. +func checkGuardrails(g CapacityGuardrails, o Outputs) error { + if g.OnExceed == ExceedError && g.Exceeded(o) { + return ErrCapacityExceeded + } + return nil +} + +type sinkFunc func(ctx context.Context, items []Product) error + +func (f sinkFunc) Emit(ctx context.Context, items []Product) error { return f(ctx, items) } diff --git a/internal/ingestion/component/knowledge_compiler/common/deps.go b/internal/ingestion/component/knowledge_compiler/common/deps.go new file mode 100644 index 0000000000..f02f4f7ce4 --- /dev/null +++ b/internal/ingestion/component/knowledge_compiler/common/deps.go @@ -0,0 +1,174 @@ +package common + +import ( + "context" + "fmt" + "sync" +) + +// ChatRequest is the minimal surface a variant needs to dispatch one LLM call. +type ChatRequest struct { + LLMID string + SystemPrompt string + UserPrompt string + JSONMode bool + // Temperature is optional; nil leaves the driver default. Python's + // knowledge compilation pins extraction at 0.1 and merge judging at 0.0, + // so variants set it per call site. + Temperature *float64 + APIKey string + BaseURL string +} + +// ChatResponse holds the LLM's text answer. +type ChatResponse struct { + Content string +} + +// ChatInvoker dispatches a single LLM chat call. The production path routes +// through internal/entity/models (same factory as other ingestion components); +// tests inject a canned responder. +type ChatInvoker interface { + Chat(ctx context.Context, req ChatRequest) (*ChatResponse, error) +} + +// Embedder wraps a text embedding model. Encode returns one vector per input. +type Embedder interface { + Encode(ctx context.Context, texts []string) ([][]float32, error) + Dimensions() int +} + +// Tokenizer estimates token counts for budget packing. +type Tokenizer interface { + NumTokens(text string) int +} + +// HistoricalHit is one cross-run KNN candidate returned by HistoricalKNN. +type HistoricalHit struct { + ID string + Score float64 +} + +// HistoricalKNN is the read-only ES KNN candidate lookup used by the wiki +// variant for cross-run dedup. nil in M1; wired in M5/M9. +type HistoricalKNN interface { + TopKHistory(ctx context.Context, tenantID, datasetID, variant string, vec []float32, k int, threshold float64) ([]HistoricalHit, error) +} + +// WikiPageCandidate is one existing wiki page returned from the backing store. +type WikiPageCandidate struct { + ID string + Slug string + Title string + PageType string + Topic string + Summary string + ContentMD string + ContentMDRaw string + EntityNames []string + RelatedKBPages []string + Outlinks []string + Score float64 +} + +// WikiPageStore is the read-only storage seam the wiki variant uses to +// reconcile against existing artifact_page rows and to load the current page +// body for UPDATE merges. +type WikiPageStore interface { + FindSimilarPages(ctx context.Context, tenantID, datasetID string, queryVec []float32, k int) ([]WikiPageCandidate, error) + GetPageBySlug(ctx context.Context, tenantID, datasetID, slug string) (*WikiPageCandidate, error) +} + +// RedisClient is injected only for the datasetnav variant (M8). +type RedisClient interface { + // Minimal lock surface; full API added in M8. +} + +// Deps bundles the runtime capabilities a variant needs. All fields are +// interfaces so tests inject stubs; production wiring lives in +// internal/ingestion/task (see PORT_PLAN.md §4 dependency injection seam). +type Deps struct { + Chat ChatInvoker + Embed Embedder + Tokenizer Tokenizer + HistoricalKNN HistoricalKNN // optional (wiki) + WikiPages WikiPageStore // optional (wiki) + Redis RedisClient // optional (datasetnav) + TenantID string + DatasetID string +} + +// DepsResolver resolves the per-run Deps from a tenant/llm/embedding triple. +type DepsResolver func(tenantID, llmID, embeddingModel string) (Deps, error) + +// DepsResolver / SetDepsResolver / ResolveDeps are the only injection seams the +// component exposes. The component is DB-independent: it compiles knowledge +// units into chunk-aligned docs (see conf/infinity_mapping.json) and returns +// them merged into the upstream chunk stream, so it owns the product schema but +// not the storage engine. Persistence (if any) is the caller's concern. + +var ( + depsResolverMu sync.RWMutex + depsResolver DepsResolver +) + +// SetDepsResolver installs the production (or test) DepsResolver. Tests call +// this to inject stubs; production wiring calls it from an init() in +// internal/ingestion/task. +func SetDepsResolver(r DepsResolver) { + depsResolverMu.Lock() + defer depsResolverMu.Unlock() + depsResolver = r +} + +// ResolveDeps resolves Deps via the installed resolver. +func ResolveDeps(tenantID, llmID, embeddingModel string) (Deps, error) { + depsResolverMu.RLock() + r := depsResolver + depsResolverMu.RUnlock() + if r == nil { + return Deps{}, fmt.Errorf("knowledge_compiler: no DepsResolver registered (call common.SetDepsResolver in production wiring)") + } + return r(tenantID, llmID, embeddingModel) +} + +// GroupResolver resolves compilation-template-group ids to the concrete +// compilation_template ids they contain. It is a DB-backed seam installed by +// production wiring in internal/ingestion/task; the knowledge_compiler package +// itself stays DB-independent. When a component is configured with +// compilation_template_group_ids but no GroupResolver is installed, the +// component fails loudly instead of silently emitting rows that miss +// compilation_template_ids (a data-loss path). +type GroupResolver func(ctx context.Context, tenantID string, groupIDs []string) ([]string, error) + +var ( + groupResolverMu sync.RWMutex + groupResolver GroupResolver +) + +// SetGroupResolver installs the production (or test) GroupResolver. Production +// wiring calls this from an init() in internal/ingestion/task; tests may inject +// a stub. +func SetGroupResolver(r GroupResolver) { + groupResolverMu.Lock() + defer groupResolverMu.Unlock() + groupResolver = r +} + +// ResolveGroupTemplateIDs resolves the given group ids to template ids via the +// installed resolver. Returns (nil, nil) when there are no group ids to +// resolve. Returns an error if group ids are requested but no resolver is +// installed, surfacing the misconfiguration instead of silently dropping the +// compilation_template_ids stamp. +func ResolveGroupTemplateIDs(ctx context.Context, tenantID string, groupIDs []string) ([]string, error) { + if len(groupIDs) == 0 { + return nil, nil + } + groupResolverMu.RLock() + r := groupResolver + groupResolverMu.RUnlock() + if r == nil { + return nil, fmt.Errorf("knowledge_compiler: compilation_template_group_ids provided but no GroupResolver installed (production wiring must call common.SetGroupResolver)") + } + return r(ctx, tenantID, groupIDs) +} diff --git a/internal/ingestion/component/knowledge_compiler/common/id.go b/internal/ingestion/component/knowledge_compiler/common/id.go new file mode 100644 index 0000000000..e5e1292c9f --- /dev/null +++ b/internal/ingestion/component/knowledge_compiler/common/id.go @@ -0,0 +1,45 @@ +package common + +import ( + "crypto/sha256" + "encoding/binary" + "encoding/hex" + "strings" +) + +// StableRowID returns a deterministic, collision-resistant hex id from the +// given parts (e.g. tenant_id, doc_id, variant, content hash). +// +// Encoding contract: each part is length-prefixed (a 4-byte big-endian +// uint32 byte-length followed by the raw part bytes) before hashing. Length +// prefixing — rather than a NUL/join delimiter — makes the encoding +// unambiguous, so distinct part sequences can never collide (e.g. ["a\x00b"] +// vs ["a","b"]). +// +// The digest is SHA-256 (not a 64-bit hash) so distinct products cannot +// accidentally share an id and be overwritten by Upsert/merge paths (M10). +// +// STABILITY: the encoding and hash algorithm form the row-id contract. This +// feature is pre-GA, so previously generated row ids are disposable and +// re-runs regenerate them deterministically; if the encoding or hash ever +// changes after GA, a backfill/migration of existing compiled rows is +// required before rollout. +func StableRowID(parts ...string) string { + h := sha256.New() + var lenBuf [4]byte + for _, p := range parts { + binary.BigEndian.PutUint32(lenBuf[:], uint32(len(p))) + h.Write(lenBuf[:]) + h.Write([]byte(p)) + } + return hex.EncodeToString(h.Sum(nil)) +} + +// ContentHash returns a stable hash of arbitrary text, used as a dedup key +// component for row ids and for in-run equality checks. SHA-256 for the same +// collision-resistance reason as StableRowID (M10). +func ContentHash(text string) string { + h := sha256.New() + _, _ = h.Write([]byte(strings.TrimSpace(text))) + return hex.EncodeToString(h.Sum(nil)) +} diff --git a/internal/ingestion/component/knowledge_compiler/common/id_test.go b/internal/ingestion/component/knowledge_compiler/common/id_test.go new file mode 100644 index 0000000000..33df9d63ea --- /dev/null +++ b/internal/ingestion/component/knowledge_compiler/common/id_test.go @@ -0,0 +1,21 @@ +package common + +import "testing" + +// TestStableRowID_NULCollision locks that distinct part sequences can never +// share an id even when a part contains a NUL byte (length-prefixed encoding). +func TestStableRowID_NULCollision(t *testing.T) { + a := StableRowID("a\x00b") + b := StableRowID("a", "b") + if a == b { + t.Fatalf("StableRowID collided for [\"a\\x00b\"] and [\"a\",\"b\"]: %s", a) + } + // Determinism: same inputs always yield the same id. + if StableRowID("a", "b") != b { + t.Fatal("StableRowID is not deterministic") + } + // Order matters. + if StableRowID("a", "b") == StableRowID("b", "a") { + t.Fatal("StableRowID ignored part ordering") + } +} diff --git a/internal/ingestion/component/knowledge_compiler/common/jsonchat.go b/internal/ingestion/component/knowledge_compiler/common/jsonchat.go new file mode 100644 index 0000000000..e6f868b91e --- /dev/null +++ b/internal/ingestion/component/knowledge_compiler/common/jsonchat.go @@ -0,0 +1,74 @@ +package common + +import ( + "context" + "encoding/json" + "fmt" + "regexp" + "strings" +) + +// fencedJSONRE matches a ```json ... ``` or ``` ... ``` fenced block. Models +// frequently wrap JSON in such fences even when JSONMode is requested, which +// previously slipped through as an unparseable "_raw" payload and was silently +// dropped by the extraction parser (a data-loss path). We strip the fence and +// retry before giving up. +var fencedJSONRE = regexp.MustCompile("(?s)```(?:json)?\\s*(.*?)\\s*```") + +// GenJSON dispatches a chat call in JSON mode and parses the response into a +// map. It first tries the raw content, then a fenced ```json ... ``` block the +// model may have wrapped around the JSON, and finally the outermost {...} span +// (handles "Here is the JSON: {...}" prose). A genuine parse failure is +// returned as an error (NOT a silent {"_raw": ...}) so the caller fails loudly +// or retries rather than silently dropping the extraction — the +// highest-frequency LLM integration must not lose knowledge units on a +// formatting hiccup. +func GenJSON(ctx context.Context, chat ChatInvoker, req ChatRequest) (map[string]any, error) { + req.JSONMode = true + resp, err := chat.Chat(ctx, req) + if err != nil { + return nil, err + } + for _, candidate := range jsonCandidates(resp.Content) { + if m, ok := tryUnmarshalJSON(candidate); ok { + return m, nil + } + } + return nil, fmt.Errorf("knowledge_compiler: LLM response is not parseable JSON: %q", truncate(resp.Content, 200)) +} + +// jsonCandidates yields progressively "cleaned" versions of an LLM reply that +// may contain JSON: the raw text, a fenced ```json ... ``` block, and the +// outermost {...} span. +func jsonCandidates(s string) []string { + cands := []string{s} + if loc := fencedJSONRE.FindStringSubmatch(s); len(loc) == 2 { + cands = append(cands, loc[1]) + } + if i := strings.Index(s, "{"); i >= 0 { + if j := strings.LastIndex(s, "}"); j > i { + cands = append(cands, s[i:j+1]) + } + } + return cands +} + +func tryUnmarshalJSON(s string) (map[string]any, bool) { + s = strings.TrimSpace(s) + if s == "" { + return nil, false + } + var m map[string]any + if err := json.Unmarshal([]byte(s), &m); err != nil { + return nil, false + } + return m, true +} + +func truncate(s string, n int) string { + r := []rune(s) + if len(r) <= n { + return s + } + return string(r[:n]) + "..." +} diff --git a/internal/ingestion/component/knowledge_compiler/common/memstore.go b/internal/ingestion/component/knowledge_compiler/common/memstore.go new file mode 100644 index 0000000000..220f246534 --- /dev/null +++ b/internal/ingestion/component/knowledge_compiler/common/memstore.go @@ -0,0 +1,282 @@ +package common + +import ( + "math" + "sort" + "sync" +) + +// Hit is one TopK match returned by MemStore. +type Hit struct { + Index int + ID string + Score float64 // cosine similarity in [−1, 1] +} + +// MemStore is an in-memory product store with exact-cosine TopK retrieval. +// It is the single source of truth for in-run dedup (replacing Python's ES +// KNN over the current run's products). Vectors are kept alongside precomputed +// L2 norms so TopK is a single matVec pass. +type MemStore struct { + mu sync.RWMutex + items []Product + vectors [][]float32 + norms []float64 + byID map[string]int +} + +// NewMemStore constructs an empty MemStore. +func NewMemStore() *MemStore { + return &MemStore{byID: make(map[string]int)} +} + +// KeepAction is the decision returned by a DedupCallback during DedupeAdd. +type KeepAction int + +const ( + // KeepAdd inserts the incoming product as a new entry. + KeepAdd KeepAction = iota + // KeepDrop discards the incoming product (it is a duplicate). + KeepDrop + // KeepMerge overwrites the best existing product with the incoming one + // (identity preserved via the existing id). + KeepMerge +) + +// DedupCallback decides, given the best existing match (score is the cosine +// similarity in [−1, 1]; existing is the zero Product when no candidate +// qualifies), whether to add, drop, or merge the incoming product. When it +// returns KeepMerge, the returned Product replaces the existing entry — its +// ID is forced to the existing entry's ID so merges preserve identity +// (mirrors Python's _struct_rebuild_doc_storage_doc preserve_id=True). +type DedupCallback func(existing Product, score float64) (KeepAction, Product, error) + +// Add appends a product and its vector to the store. +func (m *MemStore) Add(p Product) { + m.mu.Lock() + defer m.mu.Unlock() + m.addLocked(p) +} + +// DedupeAdd atomically checks the store for a near-duplicate of row and acts on +// the DedupCallback's decision. The callback is invoked OUTSIDE the store lock +// so it may perform slow work (e.g. an LLM merge judgment) without blocking +// concurrent callers. The check-then-act is race-free: the best candidate is +// selected under a read lock, the callback runs unlocked, then the add/merge is +// applied under a write lock using the previously resolved index. +func (m *MemStore) DedupeAdd(row Product, threshold float64, cb DedupCallback) (KeepAction, error) { + m.mu.RLock() + qn := l2Norm(row.Vector) + if qn == 0 { + qn = 1 + } + bestIdx, bestScore := -1, 0.0 + for i, v := range m.vectors { + vn := m.norms[i] + if vn == 0 { + vn = 1 + } + score := dotProduct(row.Vector, v) / (qn * vn) + if threshold <= 0 || score >= threshold { + if bestIdx == -1 || score > bestScore { + bestIdx, bestScore = i, score + } + } + } + var best Product + if bestIdx >= 0 && bestIdx < len(m.items) { + best = m.items[bestIdx] + } + m.mu.RUnlock() + + action, replacement, err := cb(best, bestScore) + if err != nil { + return 0, err + } + + m.mu.Lock() + defer m.mu.Unlock() + // Re-resolve the candidate by ID under the write lock. The unlocked + // callback may have allowed a concurrent Delete to delete or reindex the + // slice, so the index captured under the read lock (bestIdx) is stale and + // unsafe to index — it could point at a different product or be + // out-of-range (M11). + resolvedIdx := -1 + if best.ID != "" { + if idx, ok := m.byID[best.ID]; ok && idx < len(m.items) && m.items[idx].ID == best.ID { + resolvedIdx = idx + } + } + switch action { + case KeepDrop: + return KeepDrop, nil + case KeepMerge: + if resolvedIdx < 0 { + m.addLocked(row) + return KeepAdd, nil + } + // Merges preserve the existing entry's identity (Python preserve_id). + replacement.ID = m.items[resolvedIdx].ID + m.items[resolvedIdx] = replacement + m.vectors[resolvedIdx] = replacement.Vector + m.norms[resolvedIdx] = l2Norm(replacement.Vector) + return KeepMerge, nil + default: + m.addLocked(row) + return KeepAdd, nil + } +} + +func (m *MemStore) addLocked(p Product) { + m.items = append(m.items, p) + m.vectors = append(m.vectors, p.Vector) + m.norms = append(m.norms, l2Norm(p.Vector)) + if p.ID != "" { + m.byID[p.ID] = len(m.items) - 1 + } +} + +// Upsert replaces an existing product by ID, or appends if absent. +func (m *MemStore) Upsert(p Product) { + m.mu.Lock() + defer m.mu.Unlock() + if idx, ok := m.byID[p.ID]; ok { + m.items[idx] = p + m.vectors[idx] = p.Vector + m.norms[idx] = l2Norm(p.Vector) + return + } + m.items = append(m.items, p) + m.vectors = append(m.vectors, p.Vector) + m.norms = append(m.norms, l2Norm(p.Vector)) + m.byID[p.ID] = len(m.items) - 1 +} + +// Delete removes a product by ID. +func (m *MemStore) Delete(id string) { + m.mu.Lock() + defer m.mu.Unlock() + idx, ok := m.byID[id] + if !ok { + return + } + m.items = append(m.items[:idx], m.items[idx+1:]...) + m.vectors = append(m.vectors[:idx], m.vectors[idx+1:]...) + m.norms = append(m.norms[:idx], m.norms[idx+1:]...) + delete(m.byID, id) + m.reindexLocked() +} + +func (m *MemStore) reindexLocked() { + m.byID = make(map[string]int, len(m.items)) + for i, it := range m.items { + if it.ID != "" { + m.byID[it.ID] = i + } + } +} + +// Len returns the number of stored products. +func (m *MemStore) Len() int { + m.mu.RLock() + defer m.mu.RUnlock() + return len(m.items) +} + +// TopK returns up to k products whose cosine similarity to vec is >= threshold, +// sorted by descending similarity. A zero threshold returns the top-k by score +// regardless of absolute value. +func (m *MemStore) TopK(vec []float32, k int, threshold float64) []Hit { + m.mu.RLock() + defer m.mu.RUnlock() + qn := l2Norm(vec) + if qn == 0 { + qn = 1 + } + type cand struct { + idx int + score float64 + } + var cands []cand + for i, v := range m.vectors { + vn := m.norms[i] + if vn == 0 { + vn = 1 + } + score := dotProduct(vec, v) / (qn * vn) + if threshold <= 0 || score >= threshold { + cands = append(cands, cand{i, score}) + } + } + sort.Slice(cands, func(a, b int) bool { return cands[a].score > cands[b].score }) + if k > 0 && len(cands) > k { + cands = cands[:k] + } + hits := make([]Hit, 0, len(cands)) + for _, c := range cands { + hits = append(hits, Hit{Index: c.idx, ID: m.items[c.idx].ID, Score: c.score}) + } + return hits +} + +// Snapshot returns a shallow copy of the stored products. +func (m *MemStore) Snapshot() []Product { + m.mu.RLock() + defer m.mu.RUnlock() + out := make([]Product, len(m.items)) + copy(out, m.items) + return out +} + +// MergeSourceChunkIDs unions the given chunk ids into the source_chunk_ids +// list of the entry identified by id. Used by structure's dedup path to fold +// a dropped duplicate's provenance into the surviving canonical entity +// (mirrors Python's _struct_merge_graph_entities). No-op when id is absent. +func (m *MemStore) MergeSourceChunkIDs(id string, chunkIDs []string) { + if id == "" || len(chunkIDs) == 0 { + return + } + m.mu.Lock() + defer m.mu.Unlock() + idx, ok := m.byID[id] + if !ok { + return + } + p := m.items[idx] + if p.Meta == nil { + p.Meta = map[string]any{} + } + existing, _ := p.Meta["source_chunk_ids"].([]string) + seen := map[string]bool{} + for _, s := range existing { + seen[s] = true + } + for _, s := range chunkIDs { + if !seen[s] { + existing = append(existing, s) + seen[s] = true + } + } + p.Meta["source_chunk_ids"] = existing + m.items[idx] = p +} + +func dotProduct(a, b []float32) float64 { + n := len(a) + if len(b) < n { + n = len(b) + } + var s float64 + for i := 0; i < n; i++ { + s += float64(a[i]) * float64(b[i]) + } + return s +} + +func l2Norm(v []float32) float64 { + var s float64 + for _, x := range v { + s += float64(x) * float64(x) + } + return math.Sqrt(s) +} diff --git a/internal/ingestion/component/knowledge_compiler/common/parserconfig.go b/internal/ingestion/component/knowledge_compiler/common/parserconfig.go new file mode 100644 index 0000000000..374314079b --- /dev/null +++ b/internal/ingestion/component/knowledge_compiler/common/parserconfig.go @@ -0,0 +1,98 @@ +package common + +import ( + "fmt" + "strings" +) + +// This file holds the parser_config accessors shared by the structure and +// wiki variants. They mirror Python's structure.py helpers (_struct_get, +// _struct_localize), which wiki.py also imports — one code path in Go as well. + +// Get mirrors Python's _struct_get: case-insensitive lookup against the first +// matching key. Returns nil when no key is present. +func Get(cfg map[string]any, keys ...string) any { + if cfg == nil { + return nil + } + for _, k := range keys { + if v, ok := cfg[k]; ok { + return v + } + kl := strings.ToLower(k) + for ck, v := range cfg { + if strings.ToLower(ck) == kl { + return v + } + } + } + return nil +} + +// GetMap is Get narrowed to a map value (nil when absent or not a map). +func GetMap(cfg map[string]any, key string) map[string]any { + if m, ok := Get(cfg, key).(map[string]any); ok { + return m + } + return nil +} + +// Localize resolves a multilingual config value to the requested language. +// Mirrors Python's _struct_localize: +// - string → returned as-is +// - []string / []any → numbered lines ("1. item\n2. item") +// - map[string]string → value for lang, else "en" (only when lang != "en") +// - map[string]any → lang then "en", each recursively localized +// +// There is deliberately no arbitrary-first-value fallback: Python returns "" +// when neither the requested language nor "en" is present, and silently +// picking an unrelated language would drift the rendered prompts. +func Localize(value any, lang string) string { + switch v := value.(type) { + case string: + return v + case []string: + return numberedJoin(v) + case []any: + parts := make([]string, 0, len(v)) + for _, item := range v { + parts = append(parts, fmt.Sprint(item)) + } + return numberedJoin(parts) + case map[string]string: + if s, ok := v[lang]; ok { + return s + } + if lang != "en" { + return v["en"] + } + return "" + case map[string]any: + if raw, ok := v[lang]; ok { + if out := Localize(raw, lang); out != "" { + return out + } + } + if lang != "en" { + if raw, ok := v["en"]; ok { + return Localize(raw, "en") + } + } + return "" + default: + return "" + } +} + +// numberedJoin renders a list config value the way Python's _struct_localize +// does: "1. first\n2. second\n...". +func numberedJoin(items []string) string { + var b strings.Builder + for i, item := range items { + if i > 0 { + b.WriteByte('\n') + } + fmt.Fprintf(&b, "%d. %s", i+1, item) + } + return b.String() +} diff --git a/internal/ingestion/component/knowledge_compiler/common/tokenize.go b/internal/ingestion/component/knowledge_compiler/common/tokenize.go new file mode 100644 index 0000000000..548edf8f5a --- /dev/null +++ b/internal/ingestion/component/knowledge_compiler/common/tokenize.go @@ -0,0 +1,20 @@ +package common + +// EstimateTokens is a lightweight, dependency-free token estimate used when no +// real Tokenizer is injected (≈4 chars/token, matching common LLM tokenizers' +// order of magnitude). Production wiring should inject a real Tokenizer. +func EstimateTokens(s string) int { + n := (len(s) + 3) / 4 + if n == 0 { + return 0 + } + return n +} + +// numTokens returns the token count of s using tok when non-nil, else EstimateTokens. +func numTokens(tok Tokenizer, s string) int { + if tok != nil { + return tok.NumTokens(s) + } + return EstimateTokens(s) +} diff --git a/internal/ingestion/component/knowledge_compiler/common/types.go b/internal/ingestion/component/knowledge_compiler/common/types.go new file mode 100644 index 0000000000..1ff645d788 --- /dev/null +++ b/internal/ingestion/component/knowledge_compiler/common/types.go @@ -0,0 +1,392 @@ +// Package common holds the shared, dependency-light foundation for the +// KnowledgeCompiler component: parameter/IO types, capacity guardrails, +// the in-memory product store, tokenization/batching helpers, the LLM +// chat seam, and the concurrency pool. +// +// It deliberately imports only the standard library plus a stable-hash +// dependency, so it (and the M1 unit tests) build without the CGO native +// libraries that other ingestion paths require. +package common + +import ( + "context" + "errors" + "fmt" + "log" + "strings" +) + +// Variant identifies which knowledge-compile strategy to run. +type Variant string + +const ( + VariantStructure Variant = "structure" + VariantWiki Variant = "wiki" + VariantRaptor Variant = "raptor" + VariantMindmap Variant = "mindmap" + VariantDatasetnav Variant = "datasetnav" +) + +// Sentinel errors. +var ( + ErrCapacityExceeded = errors.New("knowledge_compiler: capacity exceeded") + ErrUnknownVariant = errors.New("knowledge_compiler: unknown variant") + ErrNotImplemented = errors.New("knowledge_compiler: variant not yet implemented") +) + +// Param is built from the DSL params map. Missing keys fall back to +// Defaults(); variant-name aliases are normalised (see ParseParam). +type Param struct { + Variant Variant + LLMID string + EmbeddingModel string + Language string + TemplateIDs []string + GroupIDs []string + SimilarityThreshold float64 + MaxWorkers int + EnableHistoricalDedup bool + Extra map[string]any + Guardrails CapacityGuardrails +} + +// Defaults returns a Param populated with safe production fallbacks. +func (p Param) Defaults() Param { + p.SimilarityThreshold = 0.99 + p.MaxWorkers = 4 + p.Extra = map[string]any{} + p.Guardrails = CapacityGuardrails{ + MaxItems: 50000, + MaxVectorBytes: 200 << 20, // 200 MiB + MaxOutputBytes: 400 << 20, // 400 MiB + OnExceed: ExceedError, + } + return p +} + +// Candidate is a historical product vector used for cross-run dedup. +type Candidate struct { + ID string + Vector []float32 +} + +// Chunk is one upstream chunk fed to the component. +type Chunk struct { + ID string + Text string // "text" field from upstream + Content string // "content_with_weight" field from upstream + Meta map[string]any +} + +// Inputs is the resolved input set passed to a variant Run. +type Inputs struct { + Chunks []Chunk + DocID string // owning document id (defaults to DatasetID when empty) + LLMID string + EmbeddingModel string + VariantSpecific map[string]any + HistoricalCandidates []Candidate // optional override path (test/offline) + Sink ChunkedSink // optional; used when guardrails flush +} + +// LogDeprecated emits a one-line deprecation notice for legacy param names. +func LogDeprecated(old, replacement string) { + log.Printf("[knowledge_compiler] deprecated variant name %q; use %q", old, replacement) +} + +// ChunkedSink receives products incrementally when capacity guardrails +// demand flushing instead of a single-shot Outputs(). +type ChunkedSink interface { + Emit(ctx context.Context, items []Product) error +} + +// Product is one compiled output row (schema_version=1). +type Product struct { + ID string + DocID string + TenantID string + Variant Variant + Content string + Vector []float32 + ParentID string + Meta map[string]any +} + +// Outputs is the result of a variant Run. +type Outputs struct { + Products []Product + DuplicatesDropped int + Items int + VectorBytes int64 + Flushed bool +} + +// SerializedBytes estimates the serialized footprint of the products. +func (o Outputs) SerializedBytes() int64 { + var n int64 + for _, p := range o.Products { + n += int64(len(p.Content)) + int64(len(p.Vector)*4) + 64 + } + return n +} + +// CapacityGuardrails bounds a single Invoke's product volume. +type CapacityGuardrails struct { + MaxItems int + MaxVectorBytes int64 + MaxOutputBytes int64 + OnExceed ExceedAction +} + +// ExceedAction controls behaviour when a guardrail is breached. +type ExceedAction string + +const ( + ExceedError ExceedAction = "error" + ExceedFlush ExceedAction = "flush" +) + +// Exceeded reports whether any guardrail is breached by o. +func (g CapacityGuardrails) Exceeded(o Outputs) bool { + if g.MaxItems > 0 && len(o.Products) > g.MaxItems { + return true + } + if g.MaxVectorBytes > 0 && o.VectorBytes > g.MaxVectorBytes { + return true + } + if g.MaxOutputBytes > 0 && o.SerializedBytes() > g.MaxOutputBytes { + return true + } + return false +} + +// EnforceGuardrails applies the variant's capacity policy uniformly: +// - OnExceed == ExceedError: returns ErrCapacityExceeded when any guardrail +// is breached (caller must propagate the error and drop the outputs). +// - OnExceed == ExceedFlush: emits the current products via sink and resets +// the buffer; subsequent variant Run iterations can keep producing +// without holding the full result set in memory. When sink is nil the +// policy degrades to "no-op" (the buffer is preserved; downstream +// callers may still see all products in the final merged chunks). +// +// Centralising the policy here keeps the five variants (structure/wiki/ +// raptor/mindmap/datasetnav) consistent and prevents the gap where only +// structure used to honour the flush path. +func (o *Outputs) EnforceGuardrails(g CapacityGuardrails, sink ChunkedSink, ctx context.Context) error { + if !g.Exceeded(*o) { + return nil + } + switch g.OnExceed { + case ExceedError: + return ErrCapacityExceeded + case ExceedFlush: + if sink == nil { + return nil + } + if err := sink.Emit(ctx, o.Products); err != nil { + return err + } + o.Flushed = true + // Hand off ownership to the sink (see FlushIfNeeded, M9) and reset all + // per-buffer accounting so a reused Outputs starts empty and does not + // re-trigger flushes on stale thresholds (M9 follow-up). + o.Products = nil + o.Items = 0 + o.VectorBytes = 0 + return nil + default: + return nil + } +} + +// ProductSink accumulates compiled products and flushes them to a ChunkedSink +// as the capacity guardrails are exceeded, bounding peak memory. A variant that +// may produce a large result set should accumulate through a ProductSink +// instead of appending to a plain slice and calling EnforceGuardrails only at +// the end: the end-of-run flush in EnforceGuardrails can only release memory +// AFTER the full result set has already been materialized, which defeats the +// anti-OOM goal (缺口 A). +// +// When OnExceed != ExceedFlush, or no sink is supplied, the sink behaves like a +// plain slice — every product stays in Products() until the caller decides what +// to do with it. This keeps the default (no-sink) pipeline behaviour unchanged: +// the full result set is returned in Outputs.Products for the component to merge +// into the upstream chunk stream. +type ProductSink struct { + products []Product + bytes int64 + total int + flushed bool + g CapacityGuardrails + sink ChunkedSink + ctx context.Context +} + +// NewProductSink constructs a sink bound to the given guardrails and (optional) +// flush target. +func NewProductSink(ctx context.Context, g CapacityGuardrails, sink ChunkedSink) *ProductSink { + return &ProductSink{g: g, sink: sink, ctx: ctx} +} + +// Add appends p to the buffer and, when the flush policy is active and the +// guardrails are breached, emits the current buffer to the sink and resets it. +// This is what keeps peak memory near the guardrail threshold rather than the +// full result-set size. +func (s *ProductSink) Add(p Product) error { + s.products = append(s.products, p) + s.bytes += int64(len(p.Vector) * 4) + s.total++ + if s.g.OnExceed == ExceedFlush && s.sink != nil { + o := Outputs{Products: s.products, VectorBytes: s.bytes, Items: len(s.products)} + if s.g.Exceeded(o) { + if err := s.sink.Emit(s.ctx, s.products); err != nil { + return err + } + s.flushed = true + // Hand off ownership of the backing array to the sink (M9). + s.products = nil + s.bytes = 0 + } + } + return nil +} + +// Products returns the not-yet-flushed buffer. +func (s *ProductSink) Products() []Product { return s.products } + +// Bytes returns the serialized-vector footprint of the not-yet-flushed buffer. +func (s *ProductSink) Bytes() int64 { return s.bytes } + +// TotalItems returns the cumulative number of products added (flushed + buffered), +// suitable for Outputs.Items. +func (s *ProductSink) TotalItems() int { return s.total } + +// Flushed reports whether at least one incremental flush has happened. +func (s *ProductSink) Flushed() bool { return s.flushed } + +// ParseParam builds a Param from the DSL params map, applying variant-name +// aliases (mind_map=>mindmap, dataset_nav=>datasetnav) and deprecation logs. +func ParseParam(m map[string]any) (Param, error) { + p := Param{}.Defaults() + if m == nil { + return p, fmt.Errorf("knowledge_compiler: params required (variant missing)") + } + if v, ok := m["variant"].(string); ok && v != "" { + p.Variant = normalizeVariant(v) + } + if v, ok := m["llm_id"].(string); ok { + p.LLMID = v + } + if v, ok := m["embedding_model"].(string); ok { + p.EmbeddingModel = v + } + if v, ok := m["language"].(string); ok { + p.Language = v + } + if v, ok := m["similarity_threshold"].(float64); ok { + p.SimilarityThreshold = v + } + if v, ok := m["enable_historical_dedup"].(bool); ok { + p.EnableHistoricalDedup = v + } + // compilation_template_ids / compilation_template_group_ids mirror the + // Python parser_config keys (see api/utils/validation_utils.py). Template + // ids are stamped onto every compiled row so the document-structure + // endpoint can group rows by template and the UI renders one tab per + // template. Group ids are resolved to concrete template ids by the caller + // (production wiring resolves them via the compilation-template-group + // service); when passed through raw they are carried as-is on GroupIDs. + p.TemplateIDs = parseStringList(m, "compilation_template_ids", "template_ids") + p.GroupIDs = parseStringList(m, "compilation_template_group_ids", "template_group_ids", "group_ids") + if raw, ok := m["extra"].(map[string]any); ok { + p.Extra = raw + } + if p.Variant == "" { + return p, fmt.Errorf("knowledge_compiler: 'variant' is required") + } + return p, nil +} + +// normalizeVariant maps a DSL variant name to the canonical Variant, +// logging a deprecation note for the legacy underscore forms. +func normalizeVariant(s string) Variant { + switch strings.ToLower(strings.TrimSpace(s)) { + case "structure": + return VariantStructure + case "wiki": + return VariantWiki + case "raptor": + return VariantRaptor + case "mindmap", "mind_map": + if strings.Contains(s, "_") { + LogDeprecated("mind_map", "mindmap") + } + return VariantMindmap + case "datasetnav", "dataset_nav": + if strings.Contains(s, "_") { + LogDeprecated("dataset_nav", "datasetnav") + } + return VariantDatasetnav + default: + return Variant(s) + } +} + +// FirstNonEmpty returns the first string that is non-empty after trimming. +func FirstNonEmpty(vals ...string) string { + for _, v := range vals { + if strings.TrimSpace(v) != "" { + return v + } + } + return "" +} + +// parseStringList reads a []string from the params map under any of the given +// keys (first hit wins). Accepts []string, []any (of strings), or a single +// whitespace/comma-separated string. Returns nil when no key is present. +func parseStringList(m map[string]any, keys ...string) []string { + for _, k := range keys { + raw, ok := m[k] + if !ok { + continue + } + switch v := raw.(type) { + case []string: + if len(v) > 0 { + return v + } + case []any: + out := make([]string, 0, len(v)) + for _, e := range v { + if s, ok := e.(string); ok { + s = strings.TrimSpace(s) + if s != "" { + out = append(out, s) + } + } + } + if len(out) > 0 { + return out + } + case string: + s := strings.TrimSpace(v) + if s == "" { + continue + } + // Tolerate comma- or whitespace-separated lists. + parts := strings.FieldsFunc(s, func(r rune) bool { return r == ',' || r == ' ' || r == '\t' || r == '\n' }) + var out []string + for _, p := range parts { + p = strings.TrimSpace(p) + if p != "" { + out = append(out, p) + } + } + if len(out) > 0 { + return out + } + } + } + return nil +} diff --git a/internal/ingestion/component/knowledge_compiler/component.go b/internal/ingestion/component/knowledge_compiler/component.go new file mode 100644 index 0000000000..a0b2130750 --- /dev/null +++ b/internal/ingestion/component/knowledge_compiler/component.go @@ -0,0 +1,652 @@ +// Package knowledge_compiler implements the KnowledgeCompiler ingestion +// component: a single runtime.Component that dispatches to one of the +// knowledge-compile variants (structure / wiki / raptor / mindmap / datasetnav) +// based on the `variant` param. See PORT_PLAN.md for the full design. +package knowledge_compiler + +import ( + "context" + "fmt" + "strings" + + "ragflow/internal/agent/runtime" + "ragflow/internal/ingestion/component/knowledge_compiler/common" + "ragflow/internal/ingestion/component/knowledge_compiler/datasetnav" + "ragflow/internal/ingestion/component/knowledge_compiler/mindmap" + "ragflow/internal/ingestion/component/knowledge_compiler/raptor" + "ragflow/internal/ingestion/component/knowledge_compiler/structure" + "ragflow/internal/ingestion/component/knowledge_compiler/wiki" + "ragflow/internal/ingestion/component/schema" + "ragflow/internal/tokenizer" + + "gorm.io/gorm" +) + +// chunkerOutputs mirrors chunker.ChunkerOutputs so this component's registered +// output schema is byte-for-byte identical to the upstream TokenChunker's. It is +// declared locally (rather than importing the chunker package) to keep the +// knowledge_compiler package free of the chunker's CGO-native dependencies. +var chunkerOutputs = map[string]string{ + "output_format": "Always \"chunks\" on success.", + "chunks": "list[object]: per-chunk map (text + optional meta keys).", + "name": "Source document name, carried forward from upstream (pass-through) when present — Tokenizer consumes it for title embedding.", + "tenant_id": "Carried forward from upstream (pass-through) when present — Tokenizer consumes it to resolve the embedding model.", + "kb_id": "Carried forward from upstream (pass-through) when present — Tokenizer consumes it to resolve the embedding model.", + "_ERROR": "Set only on validation failure.", +} + +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. +type KnowledgeCompilerComponent struct { + Param common.Param +} + +// NewKnowledgeCompilerComponent builds the component from a DSL params map. +// The name argument matches the runtime.ComponentFactory signature (ignored +// here; the component is registered under a single fixed name). +func NewKnowledgeCompilerComponent(name string, params map[string]any) (runtime.Component, error) { + p, err := common.ParseParam(params) + if err != nil { + return nil, err + } + return &KnowledgeCompilerComponent{Param: p}, nil +} + +// Inputs documents the component's input surface for the catalog. +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).", + } +} + +// Outputs documents the component's output surface. It is intentionally +// identical to the upstream chunker's output schema: the compiled knowledge +// units are expressed as chunks (schema-aligned to conf/infinity_mapping.json) +// and merged into the upstream input chunks, so downstream components (e.g. the +// Tokenizer) consume them exactly as they would normal chunks. +func (c *KnowledgeCompilerComponent) Outputs() map[string]string { + return chunkerOutputs +} + +// Invoke resolves deps, builds Inputs, and dispatches to the variant Run. The +// variant returns its compiled knowledge units as internal Product rows; those +// are converted to chunk-aligned docs (conf/infinity_mapping.json schema) and +// merged with the upstream input chunks. The result is the chunker-shaped map +// {output_format:"chunks", chunks:[...]}. +func (c *KnowledgeCompilerComponent) Invoke(ctx context.Context, db *gorm.DB, inputs map[string]any) (map[string]any, error) { + _ = db + param := c.Param + if v, ok := inputs["llm_id"].(string); ok && v != "" { + param.LLMID = v + } + if v, ok := inputs["embedding_model"].(string); ok && v != "" { + param.EmbeddingModel = v + } + tenantID, _ := inputs["tenant_id"].(string) + datasetID, _ := inputs["dataset_id"].(string) + + // Validate the variant before resolving deps so a bad variant fails fast + // with ErrUnknownVariant rather than a deps-resolution error. + switch param.Variant { + case common.VariantStructure, common.VariantWiki, common.VariantRaptor, + common.VariantMindmap, common.VariantDatasetnav: + // recognised; dispatch below + default: + return nil, fmt.Errorf("%w: %q", common.ErrUnknownVariant, param.Variant) + } + + deps, err := common.ResolveDeps(tenantID, param.LLMID, param.EmbeddingModel) + if err != nil { + return nil, err + } + deps.TenantID = tenantID + deps.DatasetID = datasetID + + // Resolve compilation-template-group ids to concrete template ids and merge + // them with any explicitly-passed template ids. Group-only configs that + // cannot be resolved fail loudly rather than silently emitting rows missing + // compilation_template_ids (a data-loss path). + resolvedGroups, err := common.ResolveGroupTemplateIDs(ctx, tenantID, param.GroupIDs) + if err != nil { + return nil, err + } + templateIDs := append(append([]string{}, param.TemplateIDs...), resolvedGroups...) + + in, err := buildInputs(inputs, param) + if err != nil { + return nil, err + } + + // Stamp the resolved template ids onto every product as it leaves the + // variant — including products streamed out through a sink under + // ExceedFlush — so they are not skipped by the post-Run stamping below + // (which only sees the buffered remainder in out.Products). The post-Run + // loop still covers that buffered remainder (M1). + // Only wrap a real sink. When no sink is supplied (in.Sink == nil) we must + // leave it nil so the product buffer is preserved under ExceedFlush; wrapping + // a nil sink would route flushed products into a delegate-less stamping sink + // that silently discards them (data loss). + if len(templateIDs) > 0 && in.Sink != nil { + in.Sink = &templateIDStampingSink{delegate: in.Sink, ids: templateIDs} + } + + var out common.Outputs + switch param.Variant { + case common.VariantStructure: + out, err = structure.Run(ctx, deps, param, in) + case common.VariantWiki: + out, err = wiki.Run(ctx, deps, param, in) + case common.VariantRaptor: + out, err = raptor.Run(ctx, deps, param, in) + case common.VariantMindmap: + out, err = mindmap.Run(ctx, deps, param, in) + case common.VariantDatasetnav: + out, err = datasetnav.Run(ctx, deps, param, in) + default: + return nil, fmt.Errorf("%w: %q", common.ErrUnknownVariant, param.Variant) + } + if err != nil { + return nil, err + } + + // Stamp the compilation template ids onto every product so the chunk + // converter can emit compilation_template_ids (Python stamps one + // template_id per row; here we carry the full resolved list since the Go + // component runs one variant per Invoke and the caller may pass multiple + // template ids from a compilation-template-group). The list is the union of + // explicitly-passed template ids and any resolved from group ids. + if len(templateIDs) > 0 { + for i := range out.Products { + if out.Products[i].Meta == nil { + out.Products[i].Meta = map[string]any{} + } + out.Products[i].Meta["compilation_template_ids"] = templateIDs + } + } + + // Convert the compiled products into chunk-aligned docs (matching + // conf/infinity_mapping.json) and merge them into the upstream input + // chunks. The component stays DB-independent and no longer routes through a + // separate writer seam: its output is plain chunks. + compiled, err := productsToChunkDocs(out.Products) + if err != nil { + return nil, err + } + return mergeChunks(inputs, compiled), nil +} + +// variantCompileKWD maps each Go variant to the compile_kwd discriminator value +// Python writes into ES (rag/advanced_rag/knowlege_compile). It is the primary +// key that distinguishes compiled knowledge units from ordinary chunks and +// routes retrieval-side filters (e.g. "compile_kwd": ["artifact_page"]). +var variantCompileKWD = map[common.Variant]string{ + common.VariantStructure: "structure", + common.VariantWiki: "artifact_page", + common.VariantRaptor: "raptor", + common.VariantMindmap: "mindmap", + common.VariantDatasetnav: "dataset_nav", +} + +// productsToChunkDocs converts the internal compiled Product rows into +// schema.ChunkDoc values aligned to conf/infinity_mapping.json (lines 1–77). +// The compile_kwd discriminator marks them as compiled knowledge units +// (distinct from plain chunks); variant-specific columns are populated from +// Product.Meta using stable keys (see each variant's build site for the +// contract). The original kind/level/name/size meta is also carried under +// "kc_"-prefixed Extra keys so no information is lost. +func productsToChunkDocs(products []common.Product) ([]schema.ChunkDoc, error) { + docs := make([]schema.ChunkDoc, 0, len(products)) + for _, p := range products { + doc := schema.ChunkDoc{ + Text: p.Content, + ContentWithWeight: p.Content, + } + // Populate content_ltks / content_sm_ltks the same way the chunker + // components do (see chunker/tag.go, chunker/qa.go): coarse tokenize + // the content, then fine-grained tokenize the coarse tokens. Errors + // are ignored (tokenizer pool may be uninitialised in no-CGo tests), + // leaving the fields empty — matching the chunker's graceful-degrade + // behaviour. + if ltks, err := tokenizer.Tokenize(p.Content); err == nil && ltks != "" { + doc.ContentLtks = ltks + if sm, err := tokenizer.FineGrainedTokenize(ltks); err == nil && sm != "" { + doc.ContentSmLtks = sm + } + } + // Common identity columns. + if err := doc.SetExtraValue("id", p.ID); err != nil { + return nil, err + } + if p.DocID != "" { + if err := doc.SetExtraValue("doc_id", p.DocID); err != nil { + return nil, err + } + } + if p.TenantID != "" { + if err := doc.SetExtraValue("tenant_id", p.TenantID); err != nil { + return nil, err + } + } + compileKWD := variantCompileKWD[p.Variant] + // A variant may pin a finer-grained compile_kwd per row via Meta + // (structure stamps the inferred compile kind — list/set/hypergraph — + // mirroring Python's per-row autotype stamp). + if v := metaString(p.Meta, "compile_kwd"); v != "" { + compileKWD = v + } + if compileKWD == "" { + compileKWD = string(p.Variant) + } + if err := doc.SetExtraValue("compile_kwd", compileKWD); err != nil { + return nil, err + } + if err := doc.SetExtraValue("compilation_template_kind_kwd", string(p.Variant)); err != nil { + return nil, err + } + if p.ParentID != "" { + if err := doc.SetExtraValue("parent_kwd", p.ParentID); err != nil { + return nil, err + } + } + if len(p.Vector) > 0 { + if err := doc.SetExtraValue(fmt.Sprintf("q_%d_vec", len(p.Vector)), p.Vector); err != nil { + return nil, err + } + } + // Provenance columns shared by all compiled rows. + if ids := metaStringSlice(p.Meta, "source_chunk_ids"); len(ids) > 0 { + if err := doc.SetExtraValue("source_chunk_ids", ids); err != nil { + return nil, err + } + } + if ids := metaStringSlice(p.Meta, "source_doc_ids"); len(ids) > 0 { + if err := doc.SetExtraValue("source_doc_ids", ids); err != nil { + return nil, err + } + } + // compilation_template_ids: the resolved template ids that produced + // this row (Python stamps one per row from the active template; the + // document-structure endpoint groups rows by this column). + if ids := metaStringSlice(p.Meta, "compilation_template_ids"); len(ids) > 0 { + if err := doc.SetExtraValue("compilation_template_ids", ids); err != nil { + return nil, err + } + } + // Per-variant fine-grained columns (conf/infinity_mapping.json §45–77). + if err := applyVariantColumns(&doc, p); err != nil { + return nil, err + } + // Preserve the raw Product.Meta under kc_* for round-trip fidelity. + for k, v := range p.Meta { + if err := doc.SetExtraValue("kc_"+k, v); err != nil { + return nil, err + } + } + docs = append(docs, doc) + } + return docs, nil +} + +// applyVariantColumns emits the compile-specific columns defined in +// conf/infinity_mapping.json lines 45–77, driven by Product.Meta keys that +// each variant's build site populates. Unknown/absent keys are skipped. +func applyVariantColumns(doc *schema.ChunkDoc, p common.Product) error { + kind := metaString(p.Meta, "kind") + + switch p.Variant { + case common.VariantStructure: + // knowledge_graph_kwd: "entity" | "relation" | "graph". + if kind != "" { + if err := doc.SetExtraValue("knowledge_graph_kwd", kind); err != nil { + return err + } + } + // Relations carry from/to entity endpoints (from_entity_kwd / to_entity_kwd). + if kind == "relation" { + if v := metaString(p.Meta, "from"); v != "" { + if err := doc.SetExtraValue("from_entity_kwd", v); err != nil { + return err + } + } + if v := metaString(p.Meta, "to"); v != "" { + if err := doc.SetExtraValue("to_entity_kwd", v); err != nil { + return err + } + } + } + // Entities carry their canonical name on name_kwd (lowercased, mirroring + // Python's _struct_to_doc_storage_doc; the structure-graph endpoints + // filter/sort on it) plus entity_type_kwd and mention_count_int. + if kind == "entity" { + if v := metaString(p.Meta, "name"); v != "" { + if err := doc.SetExtraValue("name_kwd", strings.ToLower(v)); err != nil { + return err + } + } + if v := metaString(p.Meta, "entity_type"); v != "" { + if err := doc.SetExtraValue("entity_type_kwd", v); err != nil { + return err + } + } + } + if v, ok := metaInt(p.Meta, "mention_count"); ok { + if err := doc.SetExtraValue("mention_count_int", v); err != nil { + return err + } + } + + case common.VariantWiki: + // One artifact_page row per wiki page; section rows reuse the same + // page-level columns so retrieval-side filters work uniformly. + if v := metaString(p.Meta, "slug"); v != "" { + if err := doc.SetExtraValue("slug_kwd", v); err != nil { + return err + } + if err := doc.SetExtraValue("artifact_slug_kwd", v); err != nil { + return err + } + } + if v := metaString(p.Meta, "title"); v != "" { + if err := doc.SetExtraValue("title_kwd", v); err != nil { + return err + } + setTitleTokens(doc, v) + } + if v := metaString(p.Meta, "page_type"); v != "" { + if err := doc.SetExtraValue("page_type_kwd", v); err != nil { + return err + } + } + if v := metaString(p.Meta, "topic"); v != "" { + if err := doc.SetExtraValue("topic_kwd", v); err != nil { + return err + } + } + if v := metaString(p.Meta, "summary"); v != "" { + if err := doc.SetExtraValue("summary_with_weight", v); err != nil { + return err + } + } + // Section rows also carry level/index so a retriever can scope to a + // sub-section of a wiki page. + if v, ok := metaInt(p.Meta, "section_level"); ok { + if err := doc.SetExtraValue("section_level_int", v); err != nil { + return err + } + if err := doc.SetExtraValue("depth_int", v); err != nil { + return err + } + } + if v, ok := metaInt(p.Meta, "section_index"); ok { + if err := doc.SetExtraValue("section_index_int", v); err != nil { + return err + } + } + if v := metaStringSlice(p.Meta, "entity_names"); len(v) > 0 { + if err := doc.SetExtraValue("entity_names_kwd", v); err != nil { + return err + } + } + if v := metaStringSlice(p.Meta, "outlinks"); len(v) > 0 { + if err := doc.SetExtraValue("outlinks_kwd", v); err != nil { + return err + } + if err := doc.SetExtraValue("outlinks_int", len(v)); err != nil { + return err + } + } + if v := metaStringSlice(p.Meta, "related_kb_pages"); len(v) > 0 { + if err := doc.SetExtraValue("related_kb_pages_kwd", v); err != nil { + return err + } + } + + case common.VariantRaptor: + // raptor_kwd tags summary/root nodes; raptor_layer_int records tree depth. + if kind != "" { + if err := doc.SetExtraValue("raptor_kwd", kind); err != nil { + return err + } + } + if v, ok := metaInt(p.Meta, "level"); ok { + if err := doc.SetExtraValue("raptor_layer_int", v); err != nil { + return err + } + if err := doc.SetExtraValue("depth_int", v); err != nil { + return err + } + } + if v := metaStringSlice(p.Meta, "children"); len(v) > 0 { + if err := doc.SetExtraValue("children_kwd", v); err != nil { + return err + } + } + + case common.VariantMindmap: + // Tree nodes: depth_int records the outline level. + if v, ok := metaInt(p.Meta, "level"); ok { + if err := doc.SetExtraValue("depth_int", v); err != nil { + return err + } + } + if v := metaString(p.Meta, "name"); v != "" { + if err := doc.SetExtraValue("title_kwd", v); err != nil { + return err + } + setTitleTokens(doc, v) + } + if v := metaStringSlice(p.Meta, "children"); len(v) > 0 { + if err := doc.SetExtraValue("children_kwd", v); err != nil { + return err + } + } + + case common.VariantDatasetnav: + // nav_cluster / nav_doc rows: type_kwd discriminates the row kind. + if v := metaString(p.Meta, "type"); v != "" { + if err := doc.SetExtraValue("type_kwd", v); err != nil { + return err + } + } + if v := metaString(p.Meta, "name"); v != "" { + if err := doc.SetExtraValue("title_kwd", v); err != nil { + return err + } + setTitleTokens(doc, v) + } + if v, ok := metaInt(p.Meta, "depth"); ok { + if err := doc.SetExtraValue("depth_int", v); err != nil { + return err + } + } + if v, ok := metaInt(p.Meta, "size"); ok { + if err := doc.SetExtraValue("doc_count_int", v); err != nil { + return err + } + } + if v := metaStringSlice(p.Meta, "doc_ids"); len(v) > 0 { + if err := doc.SetExtraValue("doc_ids_kwd", v); err != nil { + return err + } + } + } + + return nil +} + +// metaString reads a string-valued Product.Meta key. +func metaString(m map[string]any, key string) string { + v, _ := m[key].(string) + return v +} + +// setTitleTokens populates the ChunkDoc title_tks / title_sm_tks fields from a +// title string, mirroring how the chunker/tokenizer components tokenize titles +// (coarse → TitleTks, fine-grained → TitleSmTks). Errors are ignored: when the +// tokenizer pool is uninitialised (no-CGo test path) the fields stay empty, +// matching the chunker's graceful-degrade behaviour. +func setTitleTokens(doc *schema.ChunkDoc, title string) { + if title == "" { + return + } + if tks, err := tokenizer.Tokenize(title); err == nil && tks != "" { + doc.TitleTks = tks + if sm, err := tokenizer.FineGrainedTokenize(tks); err == nil && sm != "" { + doc.TitleSmTks = sm + } + } +} + +// metaInt reads an int-valued Product.Meta key (tolerant of float64 from JSON). +func metaInt(m map[string]any, key string) (int, bool) { + switch v := m[key].(type) { + case int: + return v, true + case float64: + return int(v), true + } + return 0, false +} + +// metaStringSlice reads a []string Product.Meta key (tolerant of []any). +func metaStringSlice(m map[string]any, key string) []string { + switch v := m[key].(type) { + case []string: + return v + case []any: + out := make([]string, 0, len(v)) + for _, e := range v { + if s, ok := e.(string); ok && s != "" { + out = append(out, s) + } + } + return out + } + return nil +} + +// mergeChunks returns the canonical chunker-shaped output: the upstream input +// chunks followed by the freshly compiled chunk docs, all under the "chunks" +// key with output_format "chunks". This makes KnowledgeCompiler's output schema +// byte-for-byte identical to the upstream chunker's, including the pass-through +// envelope (name / tenant_id / kb_id) the advertised contract promises. In a +// full pipeline those identity keys also live in CanvasState.Globals, but +// headless / manual chaining reads them from the component output map, so they +// must be forwarded when present. +func mergeChunks(inputs map[string]any, compiled []schema.ChunkDoc) map[string]any { + raw, _ := inputs["chunks"].([]any) + merged := make([]any, 0, len(raw)+len(compiled)) + for _, r := range raw { + merged = append(merged, r) + } + for _, c := range compiled { + merged = append(merged, c.ToMap()) + } + out := map[string]any{ + "output_format": "chunks", + "chunks": merged, + } + // Forward the chunker-shaped pass-through envelope so downstream + // components (e.g. Tokenizer) see the same top-level identity keys the + // upstream chunker would carry. Only present keys are forwarded. + for _, k := range []string{"name", "tenant_id", "kb_id"} { + if v, ok := inputs[k]; ok { + out[k] = v + } + } + return out +} + +// buildInputs converts the runtime inputs map into a typed common.Inputs. +func buildInputs(inputs map[string]any, param common.Param) (common.Inputs, error) { + in := common.Inputs{ + LLMID: param.LLMID, + EmbeddingModel: param.EmbeddingModel, + VariantSpecific: map[string]any{}, + } + if d, ok := inputs["doc_id"].(string); ok && d != "" { + in.DocID = d + } + if raw, ok := inputs["chunks"].([]any); ok { + for _, r := range raw { + m, ok := r.(map[string]any) + if !ok { + continue + } + ch := common.Chunk{Meta: m} + if id, ok := m["id"].(string); ok { + ch.ID = id + } + if t, ok := m["text"].(string); ok { + ch.Text = t + } + if cw, ok := m["content_with_weight"].(string); ok { + ch.Content = cw + } + in.Chunks = append(in.Chunks, ch) + } + } + if hc, ok := inputs["historical_candidates"].([]common.Candidate); ok { + in.HistoricalCandidates = hc + } + if sink, ok := inputs["sink"].(common.ChunkedSink); ok { + in.Sink = sink + } + known := map[string]bool{ + "doc_id": true, "chunks": true, "historical_candidates": true, "sink": true, + "llm_id": true, "embedding_model": true, "tenant_id": true, "dataset_id": true, + } + for k, v := range inputs { + if !known[k] { + in.VariantSpecific[k] = v + } + } + return in, nil +} + +// templateIDStampingSink wraps a ChunkedSink so every emitted product carries +// the resolved compilation template ids. It is used to stamp products that +// leave the variant through a streaming sink (ExceedFlush) — those never reach +// the buffered out.Products and would otherwise skip the template-id stamping +// (M1). +type templateIDStampingSink struct { + delegate common.ChunkedSink + ids []string +} + +// Emit stamps the template ids onto each product before delegating (a nil +// delegate is treated as a no-op drain). +func (s *templateIDStampingSink) Emit(ctx context.Context, items []common.Product) error { + for i := range items { + if items[i].Meta == nil { + items[i].Meta = map[string]any{} + } + items[i].Meta["compilation_template_ids"] = s.ids + } + if s.delegate == nil { + return nil + } + return s.delegate.Emit(ctx, items) +} + +func init() { + runtime.MustRegister(componentNameKnowledgeCompiler, runtime.CategoryIngestion, + NewKnowledgeCompilerComponent, runtime.Metadata{ + Version: "0.1.0", + Inputs: map[string]string{ + "chunks": "Upstream chunker/parser output chunks (id + text/content_with_weight).", + "llm_id": "Optional LLM id override.", + "embedding_model": "Optional embedding model override.", + "tenant_id": "Optional tenant scope.", + "dataset_id": "Optional dataset scope (wiki historical dedup).", + }, + Outputs: chunkerOutputs, + }) +} diff --git a/internal/ingestion/component/knowledge_compiler/component_test.go b/internal/ingestion/component/knowledge_compiler/component_test.go new file mode 100644 index 0000000000..582eb41f9c --- /dev/null +++ b/internal/ingestion/component/knowledge_compiler/component_test.go @@ -0,0 +1,1264 @@ +package knowledge_compiler + +import ( + "context" + "encoding/json" + "errors" + "fmt" + "hash/fnv" + "regexp" + "strings" + "sync" + "testing" + + "ragflow/internal/ingestion/component/knowledge_compiler/common" +) + +// mockChat answers the structure variant's three LLM call shapes under the +// Python-aligned items contract: node extraction (entity items), edge +// extraction (relation items constrained to Known Entities), and merge +// judging (Item A / Item B). It scans the prompt for the Greek-letter fixture +// universe (Alpha..Epsilon) so assertions hold regardless of batch packing. +type mockChat struct{} + +var greekNames = []string{"Alpha", "Beta", "Gamma", "Delta", "Epsilon"} +var greekPairs = [][2]string{{"Alpha", "Beta"}, {"Beta", "Gamma"}, {"Gamma", "Delta"}, {"Delta", "Epsilon"}, {"Alpha", "Delta"}, {"Gamma", "Epsilon"}} + +// graphParserConfig declares the entity/relation template shape so the +// structure variant compiles the hypergraph kind (mirrors a Python +// parser_config whose kind aliases "graph" -> hypergraph). +func graphParserConfig() map[string]any { + return map[string]any{ + "kind": "graph", + "entity": map[string]any{"fields": []any{ + map[string]any{"type": "project", "description": "a named project"}, + }}, + "relation": map[string]any{"fields": []any{ + map[string]any{"type": "linked", "description": "a link between projects"}, + }}, + } +} + +var chunkIDRe = regexp.MustCompile(`\[CHUNK_ID: ([^\]]+)\]`) + +func chunkIDsInPrompt(prompt string) []string { + var ids []string + for _, m := range chunkIDRe.FindAllStringSubmatch(prompt, -1) { + ids = append(ids, m[1]) + } + if len(ids) == 0 { + ids = []string{"c1"} + } + return ids +} + +func itemsJSON(items []map[string]any, chunkIDs []string) string { + for _, item := range items { + item["source_chunk_ids"] = chunkIDs + } + b, _ := json.Marshal(map[string]any{"items": items}) + return string(b) +} + +func parseMergePayloads(prompt string) (map[string]any, map[string]any) { + body := strings.TrimPrefix(prompt, "Item A (existing):\n") + parts := strings.SplitN(body, "\n\nItem B (incoming):\n", 2) + if len(parts) != 2 { + return nil, nil + } + var a, b map[string]any + if err := json.Unmarshal([]byte(parts[0]), &a); err != nil { + return nil, nil + } + if err := json.Unmarshal([]byte(parts[1]), &b); err != nil { + return nil, nil + } + return a, b +} + +func (mockChat) Chat(_ context.Context, req common.ChatRequest) (*common.ChatResponse, error) { + if !req.JSONMode { + // Non-JSON calls (mindmap outline, wiki prose paths) get plain text. + return &common.ChatResponse{Content: "Outline of: " + firstWords(req.UserPrompt, 24)}, nil + } + switch { + case strings.HasPrefix(req.UserPrompt, "Item A (existing):"): + a, b := parseMergePayloads(req.UserPrompt) + if a != nil && b != nil && a["name"] != nil && a["name"] == b["name"] { + merged, _ := json.Marshal(a) + return &common.ChatResponse{Content: fmt.Sprintf(`{"duplicated":true,"merged":%s}`, merged)}, nil + } + return &common.ChatResponse{Content: `{"duplicated":false,"merged":null}`}, nil + case strings.Contains(req.SystemPrompt, "## Known Entities:"): + var items []map[string]any + for _, pr := range greekPairs { + if strings.Contains(req.UserPrompt, pr[0]) && strings.Contains(req.UserPrompt, pr[1]) { + items = append(items, map[string]any{"type": "linked", "source": pr[0], "target": pr[1], "description": pr[0] + " linked to " + pr[1]}) + } + } + return &common.ChatResponse{Content: itemsJSON(items, chunkIDsInPrompt(req.UserPrompt))}, nil + default: + var items []map[string]any + for _, n := range greekNames { + if strings.Contains(req.UserPrompt, n) { + items = append(items, map[string]any{"type": "project", "name": n, "description": n + " project"}) + } + } + return &common.ChatResponse{Content: itemsJSON(items, chunkIDsInPrompt(req.UserPrompt))}, nil + } +} + +// proseChat returns generic, non-empty prose for the non-structure variants +// (wiki/raptor/mindmap/datasetnav), which all just need a summary/outline text. +type proseChat struct{} + +func (proseChat) Chat(_ context.Context, req common.ChatRequest) (*common.ChatResponse, error) { + if req.JSONMode { + return &common.ChatResponse{Content: `{"ok":true}`}, nil + } + // Echo a deterministic prose reply derived from the prompt. + return &common.ChatResponse{Content: "Summary of: " + firstWords(req.UserPrompt, 24)}, nil +} + +func firstWords(s string, n int) string { + fields := strings.Fields(s) + if len(fields) > n { + fields = fields[:n] + } + return strings.Join(fields, " ") +} + +func installProseDeps(t *testing.T) { + t.Helper() + common.SetDepsResolver(func(tenantID, llmID, embeddingModel string) (common.Deps, error) { + return common.Deps{Chat: proseChat{}, Embed: mockEmbedder{dim: 8}, TenantID: tenantID}, nil + }) + t.Cleanup(func() { common.SetDepsResolver(nil) }) +} + +type mockEmbedder struct{ dim int } + +func (m mockEmbedder) Dimensions() int { return m.dim } +func (m mockEmbedder) Encode(_ context.Context, texts []string) ([][]float32, error) { + out := make([][]float32, len(texts)) + for i, t := range texts { + out[i] = deterministicVec(t, m.dim) + } + return out, nil +} +func deterministicVec(s string, dim int) []float32 { + h := fnv.New32a() + _, _ = h.Write([]byte(s)) + seed := h.Sum32() + v := make([]float32, dim) + for j := 0; j < dim; j++ { + v[j] = float32((seed>>uint(j*5))&0xFF) / 255.0 + } + return v +} + +func installMockDeps(t *testing.T) { + t.Helper() + common.SetDepsResolver(func(tenantID, llmID, embeddingModel string) (common.Deps, error) { + return common.Deps{Chat: mockChat{}, Embed: mockEmbedder{dim: 8}, TenantID: tenantID}, nil + }) + t.Cleanup(func() { common.SetDepsResolver(nil) }) +} + +func TestKnowledgeCompiler_Structure_EndToEnd(t *testing.T) { + installMockDeps(t) + + c, err := NewKnowledgeCompilerComponent("KnowledgeCompiler", map[string]any{ + "variant": "structure", "llm_id": "llm1", "embedding_model": "emb1", + }) + if err != nil { + t.Fatalf("NewKnowledgeCompilerComponent: %v", err) + } + + out, err := c.Invoke(context.Background(), nil, map[string]any{ + "chunks": []any{ + map[string]any{"id": "c1", "text": "Alpha is a Beta"}, + map[string]any{"id": "c2", "text": "Beta related to Gamma"}, + }, + "doc_id": "d1", + "tenant_id": "t1", + "dataset_id": "ds1", + "parser_config": graphParserConfig(), + }) + if err != nil { + t.Fatalf("Invoke: %v", err) + } + + chunks, ok := out["chunks"].([]any) + if !ok { + t.Fatalf("chunks = %T, want []any", out["chunks"]) + } + // 2 input chunks + 3 entities + 2 relations + 1 graph = 8 total (the + // compiled knowledge units are merged into the upstream input chunks). + if len(chunks) != 8 { + t.Fatalf("len(chunks) = %d, want 8 (2 input + 3 entities + 2 relations + 1 graph)", len(chunks)) + } + + // Exactly one graph product, parseable to {entities, relations} (mirrors + // Python's _struct_rebuild_graph_json). + var graph map[string]any + for _, c := range chunks { + cm, ok := c.(map[string]any) + if !ok { + continue + } + if kind, _ := cm["kc_kind"].(string); kind == "graph" { + if err := json.Unmarshal([]byte(cm["text"].(string)), &graph); err != nil { + t.Fatalf("unmarshal graph: %v", err) + } + } + } + if graph == nil { + t.Fatal("no graph chunk emitted") + } + entities, _ := graph["entities"].([]any) + relations, _ := graph["relations"].([]any) + if len(entities) != 3 { + t.Fatalf("graph entities = %d, want 3 (Alpha/Beta/Gamma)", len(entities)) + } + if len(relations) != 2 { + t.Fatalf("graph relations = %d, want 2", len(relations)) + } +} + +func TestKnowledgeCompiler_UnknownVariant(t *testing.T) { + c, err := NewKnowledgeCompilerComponent("KnowledgeCompiler", map[string]any{"variant": "nope"}) + if err != nil { + t.Fatalf("NewKnowledgeCompilerComponent: %v", err) + } + _, err = c.Invoke(context.Background(), nil, map[string]any{"chunks": []any{}}) + if !errors.Is(err, common.ErrUnknownVariant) { + t.Fatalf("err = %v, want ErrUnknownVariant", err) + } +} + +func TestKnowledgeCompiler_Alias_Mindmap(t *testing.T) { + installMockDeps(t) + // "mind_map" is the deprecated alias for "mindmap"; both resolve to the + // implemented mindmap variant and must run (not ErrUnknownVariant / stub). + c, err := NewKnowledgeCompilerComponent("KnowledgeCompiler", map[string]any{ + "variant": "mind_map", "llm_id": "llm1", "embedding_model": "emb1", + }) + if err != nil { + t.Fatalf("NewKnowledgeCompilerComponent: %v", err) + } + out, err := c.Invoke(context.Background(), nil, map[string]any{ + "chunks": []any{map[string]any{"id": "c1", "text": "Alpha is a Beta"}}, + "doc_id": "d1", + "tenant_id": "t1", + }) + if err != nil { + t.Fatalf("Invoke mind_map: %v", err) + } + chunks, ok := out["chunks"].([]any) + if !ok || len(chunks) == 0 { + t.Fatalf("mind_map produced no chunks: %v", out["chunks"]) + } +} + +// runVariant is a shared harness that builds a KnowledgeCompiler component for +// the given variant and invokes it over a couple of chunks, returning the merged +// chunk list (upstream input chunks + compiled knowledge-unit chunks). +func runVariant(t *testing.T, variant string, extra map[string]any) []map[string]any { + t.Helper() + params := map[string]any{"variant": variant, "llm_id": "llm1", "embedding_model": "emb1"} + for k, v := range extra { + params[k] = v + } + c, err := NewKnowledgeCompilerComponent("KnowledgeCompiler", params) + if err != nil { + t.Fatalf("NewKnowledgeCompilerComponent(%s): %v", variant, err) + } + out, err := c.Invoke(context.Background(), nil, map[string]any{ + "chunks": []any{ + map[string]any{"id": "c1", "text": "The quick brown fox jumps over the lazy dog near the river bank."}, + map[string]any{"id": "c2", "text": "A red fox and a lazy dog rest beside a calm river at dawn."}, + map[string]any{"id": "c3", "text": "Rivers flow through valleys carrying water from the mountains."}, + }, + "doc_id": "d1", + "tenant_id": "t1", + }) + if err != nil { + t.Fatalf("Invoke(%s): %v", variant, err) + } + raw, ok := out["chunks"].([]any) + if !ok { + t.Fatalf("Invoke(%s): chunks = %T", variant, out["chunks"]) + } + if len(raw) == 0 { + t.Fatalf("Invoke(%s): produced no chunks", variant) + } + // Every emitted chunk must carry an id; every compiled knowledge unit must + // carry its embedding (q__vec) and a compile_kwd discriminator. + chunks := make([]map[string]any, 0, len(raw)) + seenCompiled := false + for _, r := range raw { + cm, ok := r.(map[string]any) + if !ok { + t.Fatalf("Invoke(%s): chunk not a map: %T", variant, r) + } + chunks = append(chunks, cm) + if cm["id"] == nil || cm["id"] == "" { + t.Fatalf("Invoke(%s): chunk missing id: %v", variant, cm) + } + if ck, _ := cm["compile_kwd"].(string); ck != "" { + seenCompiled = true + if _, ok := cm["q_8_vec"]; !ok { + t.Fatalf("Invoke(%s): compiled chunk missing vector: %v", variant, cm) + } + } + } + if !seenCompiled { + t.Fatalf("Invoke(%s): no compiled knowledge-unit chunks emitted", variant) + } + return chunks +} + +func TestKnowledgeCompiler_Wiki_EndToEnd(t *testing.T) { + installProseDeps(t) + chunks := runVariant(t, "wiki", nil) + // wiki produces a "page" chunk (kind stored as kc_kind). + foundPage := false + for _, c := range chunks { + if kind, _ := c["kc_kind"].(string); kind == "page" { + foundPage = true + } + } + if !foundPage { + t.Fatalf("wiki: no 'page' chunk; got %d chunks", len(chunks)) + } +} + +func TestKnowledgeCompiler_Raptor_EndToEnd(t *testing.T) { + installProseDeps(t) + // Psi builder (default): zero clustering dependency. + chunks := runVariant(t, "raptor", map[string]any{"extra": map[string]any{"clustering_method": "PSI"}}) + foundRoot := false + for _, c := range chunks { + if kind, _ := c["kc_kind"].(string); kind == "root" { + foundRoot = true + } + } + if !foundRoot { + t.Fatalf("raptor(PSI): no 'root' chunk; got %d chunks", len(chunks)) + } + + // AHC builder must also run (PCA + Ward). + chunksAHC := runVariant(t, "raptor", map[string]any{"extra": map[string]any{"clustering_method": "AHC"}}) + if len(chunksAHC) == 0 { + t.Fatalf("raptor(AHC): produced no chunks") + } + + // GMM backend (diagonal-GMM EM + BIC) is implemented and must run, + // producing a well-formed tree (root present, chunks non-empty). + chunksGMM := runVariant(t, "raptor", map[string]any{"extra": map[string]any{"clustering_method": "GMM"}}) + if len(chunksGMM) == 0 { + t.Fatalf("raptor(GMM): produced no chunks") + } + foundRootGMM := false + for _, c := range chunksGMM { + if kind, _ := c["kc_kind"].(string); kind == "root" { + foundRootGMM = true + } + } + if !foundRootGMM { + t.Fatalf("raptor(GMM): no 'root' chunk; got %d chunks", len(chunksGMM)) + } +} + +func TestKnowledgeCompiler_Mindmap_EndToEnd(t *testing.T) { + installProseDeps(t) + // The proseChat reply is flat text; parseOutline still yields a root + the + // reply as a child node, so chunks are non-empty and parent-linked. + chunks := runVariant(t, "mindmap", nil) + // Root chunk must have empty parent_kwd and kind "root". + var root map[string]any + for _, c := range chunks { + if kind, _ := c["kc_kind"].(string); kind == "root" { + root = c + } + } + if root == nil { + t.Fatalf("mindmap: no root chunk") + } + if pid, _ := root["parent_kwd"].(string); pid != "" { + t.Fatalf("mindmap: root parent_kwd = %q, want empty", pid) + } +} + +func TestKnowledgeCompiler_Datasetnav_EndToEnd(t *testing.T) { + installProseDeps(t) + chunks := runVariant(t, "datasetnav", nil) + foundRoot := false + for _, c := range chunks { + if kind, _ := c["kc_kind"].(string); kind == "root" { + foundRoot = true + } + } + if !foundRoot { + t.Fatalf("datasetnav: no 'root' chunk; got %d chunks", len(chunks)) + } +} + +// TestKnowledgeCompiler_EmitsChunks verifies that after compiling, the +// component returns the knowledge units merged into the chunk stream (no +// separate products/writer seam). The output shape is the chunker's +// {output_format:"chunks", chunks:[...]}. +func TestKnowledgeCompiler_EmitsChunks(t *testing.T) { + installMockDeps(t) + + c, err := NewKnowledgeCompilerComponent("KnowledgeCompiler", map[string]any{ + "variant": "structure", "llm_id": "llm1", "embedding_model": "emb1", + }) + if err != nil { + t.Fatalf("NewKnowledgeCompilerComponent: %v", err) + } + out, err := c.Invoke(context.Background(), nil, map[string]any{ + "chunks": []any{ + map[string]any{"id": "c1", "text": "Alpha is a Beta"}, + map[string]any{"id": "c2", "text": "Beta related to Gamma"}, + }, + "doc_id": "d1", + "tenant_id": "t1", + "dataset_id": "ds1", + "parser_config": graphParserConfig(), + }) + if err != nil { + t.Fatalf("Invoke: %v", err) + } + if got, _ := out["output_format"].(string); got != "chunks" { + t.Fatalf("output_format = %v, want chunks", out["output_format"]) + } + // No legacy products/written keys remain in the output surface. + if _, ok := out["products"]; ok { + t.Fatalf("legacy 'products' key still present in output") + } + if _, ok := out["written"]; ok { + t.Fatalf("legacy 'written' key still present in output") + } + compiled := 0 + for _, r := range out["chunks"].([]any) { + cm := r.(map[string]any) + // Structure rows carry the inferred compile kind (hypergraph here), + // mirroring Python's per-row autotype stamp. + if ck, _ := cm["compile_kwd"].(string); ck == "hypergraph" { + compiled++ + if cm["id"] == nil || cm["id"] == "" { + t.Fatal("compiled chunk missing id") + } + if _, ok := cm["q_8_vec"]; !ok { + t.Fatal("compiled chunk missing vector") + } + } + } + if compiled != 6 { + t.Fatalf("compiled structure chunks = %d, want 6 (3 entities + 2 relations + 1 graph)", compiled) + } +} + +// TestKnowledgeCompiler_TemplateIDsAndProvenance verifies that +// compilation_template_ids passed via the DSL params are stamped onto every +// compiled chunk, and that structure entity rows carry source_chunk_ids +// pointing back at their originating input chunks. +func TestKnowledgeCompiler_TemplateIDsAndProvenance(t *testing.T) { + installMockDeps(t) + + c, err := NewKnowledgeCompilerComponent("KnowledgeCompiler", map[string]any{ + "variant": "structure", + "llm_id": "llm1", + "embedding_model": "emb1", + "compilation_template_ids": []any{"tpl-A", "tpl-B"}, + }) + if err != nil { + t.Fatalf("NewKnowledgeCompilerComponent: %v", err) + } + out, err := c.Invoke(context.Background(), nil, map[string]any{ + "chunks": []any{ + map[string]any{"id": "c1", "text": "Alpha is a Beta"}, + map[string]any{"id": "c2", "text": "Beta related to Gamma"}, + }, + "doc_id": "d1", + "tenant_id": "t1", + "parser_config": graphParserConfig(), + }) + if err != nil { + t.Fatalf("Invoke: %v", err) + } + for _, r := range out["chunks"].([]any) { + cm := r.(map[string]any) + if ck, _ := cm["compile_kwd"].(string); ck != "hypergraph" { + continue + } + // Every compiled chunk must carry the resolved template ids. + var tidsCount int + switch tids := cm["compilation_template_ids"].(type) { + case []any: + tidsCount = len(tids) + case []string: + tidsCount = len(tids) + } + if tidsCount != 2 { + t.Fatalf("compiled chunk %v: compilation_template_ids = %v, want 2", cm["id"], cm["compilation_template_ids"]) + } + // Entity rows must carry source_chunk_ids. + if kg, _ := cm["knowledge_graph_kwd"].(string); kg == "entity" { + var idsCount int + switch ids := cm["source_chunk_ids"].(type) { + case []any: + idsCount = len(ids) + case []string: + idsCount = len(ids) + } + if idsCount == 0 { + t.Fatalf("entity chunk %v: missing source_chunk_ids", cm["id"]) + } + } + } + +} + +// constEmbedder returns one fixed vector for every input text. It is used to +// make historical-dedup deterministic: every compiled product shares the same +// vector, so a historical candidate carrying that vector is a guaranteed +// near-duplicate. +type constEmbedder struct { + dim int + vec []float32 +} + +func (m constEmbedder) Dimensions() int { return m.dim } +func (m constEmbedder) Encode(_ context.Context, texts []string) ([][]float32, error) { + out := make([][]float32, len(texts)) + for i := range out { + out[i] = m.vec + } + return out, nil +} + +// sinkRecorder is a ChunkedSink that tallies streamed products, used to assert +// the flush policy actually emits incrementally instead of holding the full set. +type sinkRecorder struct { + mu sync.Mutex + total int + batches int +} + +func (s *sinkRecorder) Emit(_ context.Context, items []common.Product) error { + s.mu.Lock() + defer s.mu.Unlock() + s.total += len(items) + s.batches++ + return nil +} + +// TestKnowledgeCompiler_Raptor_DegenerateNoInfiniteLoop is a regression for the +// High-1 bug: RAPTOR could recurse forever when a re-cluster returned a single +// label covering all points. The default mockEmbedder emits non-negative +// vectors, so under PSI (threshold=0) every pair has cosine >= 0 and the whole +// corpus collapses into one cluster; with more than RecursionMinClusterSize +// points the old code re-enqueued the identical work item at level+1 and hung. +// This test uses 6 chunks and asserts the run terminates with a well-formed root. +// (If the guard regresses, go test's timeout turns the hang into a failure.) +func TestKnowledgeCompiler_Raptor_DegenerateNoInfiniteLoop(t *testing.T) { + installProseDeps(t) + c, err := NewKnowledgeCompilerComponent("KnowledgeCompiler", map[string]any{ + "variant": "raptor", "llm_id": "llm1", "embedding_model": "emb1", + "extra": map[string]any{"clustering_method": "PSI"}, + }) + if err != nil { + t.Fatalf("NewKnowledgeCompilerComponent: %v", err) + } + out, err := c.Invoke(context.Background(), nil, map[string]any{ + "chunks": []any{ + map[string]any{"id": "c1", "text": "alpha beta gamma delta epsilon zeta"}, + map[string]any{"id": "c2", "text": "one two three four five six"}, + map[string]any{"id": "c3", "text": "red green blue yellow white black"}, + map[string]any{"id": "c4", "text": "cat dog bird fish frog snake"}, + map[string]any{"id": "c5", "text": "sun moon star planet comet meteor"}, + map[string]any{"id": "c6", "text": "king queen prince duke earl count"}, + }, + "doc_id": "d1", + "tenant_id": "t1", + }) + if err != nil { + t.Fatalf("Invoke hung or errored (RAPTOR infinite recursion?): %v", err) + } + raw, ok := out["chunks"].([]any) + if !ok { + t.Fatalf("chunks = %T", out["chunks"]) + } + foundRoot := false + for _, r := range raw { + if cm, ok := r.(map[string]any); ok { + if k, _ := cm["kc_kind"].(string); k == "root" { + foundRoot = true + } + } + } + if !foundRoot { + t.Fatalf("raptor(degenerate): no root chunk; got %d chunks", len(raw)) + } +} + +// TestKnowledgeCompiler_Wiki_HistoricalDedupDropsDuplicates is a regression for +// the High-2 bug: EnableHistoricalDedup / the HistoricalCandidates override used +// to be telemetry-only dead code. Here a constant embedder makes every wiki +// product share one vector, and we supply that vector as a historical candidate; +// the run must drop the near-duplicate products so none survive in the output. +func TestKnowledgeCompiler_Wiki_HistoricalDedupDropsDuplicates(t *testing.T) { + common.SetDepsResolver(func(tenantID, llmID, embeddingModel string) (common.Deps, error) { + return common.Deps{ + Chat: proseChat{}, + Embed: constEmbedder{dim: 8, vec: []float32{1, 0, 0, 0, 0, 0, 0, 0}}, + TenantID: tenantID, + }, nil + }) + t.Cleanup(func() { common.SetDepsResolver(nil) }) + + c, err := NewKnowledgeCompilerComponent("KnowledgeCompiler", map[string]any{ + "variant": "wiki", "llm_id": "llm1", "embedding_model": "emb1", + }) + if err != nil { + t.Fatalf("NewKnowledgeCompilerComponent: %v", err) + } + out, err := c.Invoke(context.Background(), nil, map[string]any{ + "chunks": []any{ + map[string]any{"id": "c1", "text": "The quick brown fox jumps over the lazy dog."}, + map[string]any{"id": "c2", "text": "A red fox and a lazy dog rest beside a calm river."}, + }, + "doc_id": "d1", + "tenant_id": "t1", + // Offline/test override: a historical candidate with the same vector as + // every compiled product -> all are near-duplicates and must be dropped. + "historical_candidates": []common.Candidate{ + {ID: "old-artifact", Vector: []float32{1, 0, 0, 0, 0, 0, 0, 0}}, + }, + }) + if err != nil { + t.Fatalf("Invoke: %v", err) + } + raw, ok := out["chunks"].([]any) + if !ok { + t.Fatalf("chunks = %T", out["chunks"]) + } + for _, r := range raw { + cm, ok := r.(map[string]any) + if !ok { + continue + } + if _, has := cm["compile_kwd"].(string); has { + t.Fatalf("wiki historical dedup failed: a compiled chunk survived despite a matching historical candidate: %v", cm) + } + } +} + +func TestKnowledgeCompiler_Wiki_UpdateMergesExistingPage(t *testing.T) { + common.SetDepsResolver(func(tenantID, llmID, embeddingModel string) (common.Deps, error) { + return common.Deps{ + Chat: wikiUpdateChat{}, + Embed: mockEmbedder{dim: 8}, + TenantID: tenantID, + WikiPages: wikiStoreTestStub{page: &common.WikiPageCandidate{ + ID: "existing-1", + Slug: "entity/alpha", + Title: "Alpha", + PageType: "entity", + Topic: "Alpha", + ContentMD: "# Alpha\n\nExisting fact.\n", + ContentMDRaw: "# Alpha\n\nExisting fact.\n", + EntityNames: []string{"Alpha"}, + Score: 0.85, + }}, + }, nil + }) + t.Cleanup(func() { common.SetDepsResolver(nil) }) + + c, err := NewKnowledgeCompilerComponent("KnowledgeCompiler", map[string]any{ + "variant": "wiki", "llm_id": "llm1", "embedding_model": "emb1", + }) + if err != nil { + t.Fatalf("NewKnowledgeCompilerComponent: %v", err) + } + out, err := c.Invoke(context.Background(), nil, map[string]any{ + "chunks": []any{ + map[string]any{"id": "c1", "text": "Alpha launched a new process in 2026."}, + }, + "doc_id": "d1", + "dataset_id": "ds1", + "tenant_id": "t1", + }) + if err != nil { + t.Fatalf("Invoke: %v", err) + } + raw, ok := out["chunks"].([]any) + if !ok { + t.Fatalf("chunks = %T", out["chunks"]) + } + var page map[string]any + for _, r := range raw { + cm, ok := r.(map[string]any) + if !ok { + continue + } + if cm["compile_kwd"] == "artifact_page" && cm["kc_kind"] == "page" && cm["slug_kwd"] == "entity/alpha" { + page = cm + break + } + } + if page == nil { + t.Fatalf("no merged wiki page found in output: %#v", raw) + } + text, _ := page["text"].(string) + if !strings.Contains(text, "Existing fact.") || !strings.Contains(text, "Alpha launched a new process in 2026.") { + t.Fatalf("merged page text = %q, want both existing and incoming facts", text) + } + if outlinks, ok := page["outlinks_kwd"].([]any); ok && len(outlinks) == 0 { + t.Fatalf("outlinks_kwd should include rewritten see-also link: %#v", page) + } +} + +// TestKnowledgeCompiler_GuardrailFlushStreams is a regression for the Medium-1 +// bug: the flush policy used to run only AFTER the full result set was +// materialised, so it could not cap peak memory. With the ProductSink, products +// are emitted to the sink as they are produced; here a tiny MaxItems forces +// streaming, and we assert the sink received every product while the returned +// chunk list holds only the upstream inputs (the compiled units went to the sink). +func TestKnowledgeCompiler_GuardrailFlushStreams(t *testing.T) { + installMockDeps(t) + + c, err := NewKnowledgeCompilerComponent("KnowledgeCompiler", map[string]any{ + "variant": "structure", "llm_id": "llm1", "embedding_model": "emb1", + }) + if err != nil { + t.Fatalf("NewKnowledgeCompilerComponent: %v", err) + } + // Override guardrails to a tiny cap with the flush policy. (Guardrails are + // not yet DSL-configurable; set directly on the exported Param.) + kc, ok := c.(*KnowledgeCompilerComponent) + if !ok { + t.Fatalf("component is not *KnowledgeCompilerComponent: %T", c) + } + kc.Param.Guardrails = common.CapacityGuardrails{ + MaxItems: 2, OnExceed: common.ExceedFlush, + } + + sink := &sinkRecorder{} + out, err := c.Invoke(context.Background(), nil, map[string]any{ + "chunks": []any{ + map[string]any{"id": "c1", "text": "Alpha is a Beta"}, + map[string]any{"id": "c2", "text": "Beta related to Gamma"}, + }, + "doc_id": "d1", + "tenant_id": "t1", + "parser_config": graphParserConfig(), + "sink": sink, + }) + if err != nil { + t.Fatalf("Invoke: %v", err) + } + + sink.mu.Lock() + emitted, batches := sink.total, sink.batches + sink.mu.Unlock() + if batches == 0 || emitted == 0 { + t.Fatalf("flush policy did not stream products to sink (batches=%d emitted=%d)", batches, emitted) + } + // The 6 compiled structure units were streamed to the sink; the returned + // chunk list must therefore hold only the 2 upstream inputs. + raw, ok := out["chunks"].([]any) + if !ok { + t.Fatalf("chunks = %T", out["chunks"]) + } + if len(raw) != 2 { + t.Fatalf("after flush, returned chunks = %d, want 2 (upstream only; compiled streamed to sink)", len(raw)) + } +} + +// fakeHistoricalKNN records the datasetID it was queried with and optionally +// returns a hit for every lookup, so a test can assert both the scope of the +// KNN query and that near-duplicate products are dropped. +type fakeHistoricalKNN struct { + mu sync.Mutex + lastDS string + hit bool +} + +func (f *fakeHistoricalKNN) TopKHistory(_ context.Context, _ string, datasetID, _ string, _ []float32, _ int, _ float64) ([]common.HistoricalHit, error) { + f.mu.Lock() + defer f.mu.Unlock() + f.lastDS = datasetID + if f.hit { + return []common.HistoricalHit{{ID: "old", Score: 1.0}}, nil + } + return nil, nil +} + +// TestKnowledgeCompiler_Wiki_HistoricalDedupScopedByDataset is a regression for +// the Medium bug: the historical KNN lookup used to be scoped by doc_id, so +// cross-document dedup within the same dataset did not work. This test supplies +// both a doc_id ("d1") and a dataset_id ("ds1") and asserts the lookup is +// scoped to the dataset, not the document. +func TestKnowledgeCompiler_Wiki_HistoricalDedupScopedByDataset(t *testing.T) { + knn := &fakeHistoricalKNN{hit: true} // every product is a near-dup -> dropped + common.SetDepsResolver(func(tenantID, llmID, embeddingModel string) (common.Deps, error) { + return common.Deps{ + Chat: proseChat{}, + Embed: constEmbedder{dim: 8, vec: []float32{1, 0, 0, 0, 0, 0, 0, 0}}, + TenantID: tenantID, + HistoricalKNN: knn, + }, nil + }) + t.Cleanup(func() { common.SetDepsResolver(nil) }) + + c, err := NewKnowledgeCompilerComponent("KnowledgeCompiler", map[string]any{ + "variant": "wiki", "llm_id": "llm1", "embedding_model": "emb1", + "enable_historical_dedup": true, + }) + if err != nil { + t.Fatalf("NewKnowledgeCompilerComponent: %v", err) + } + out, err := c.Invoke(context.Background(), nil, map[string]any{ + "chunks": []any{ + map[string]any{"id": "c1", "text": "The quick brown fox jumps over the lazy dog."}, + map[string]any{"id": "c2", "text": "A red fox and a lazy dog rest beside a calm river."}, + }, + "doc_id": "d1", + "dataset_id": "ds1", + "tenant_id": "t1", + }) + if err != nil { + t.Fatalf("Invoke: %v", err) + } + knn.mu.Lock() + ds := knn.lastDS + knn.mu.Unlock() + if ds != "ds1" { + t.Fatalf("historical KNN queried with datasetID=%q, want %q (lookup must be scoped to the dataset, not the document)", ds, "ds1") + } + // With hit=true the near-duplicate products must be dropped, so no compiled + // chunks survive. + raw, ok := out["chunks"].([]any) + if !ok { + t.Fatalf("chunks = %T", out["chunks"]) + } + for _, r := range raw { + if cm, ok := r.(map[string]any); ok { + if _, has := cm["compile_kwd"]; has { + t.Fatalf("wiki historical dedup failed: a compiled chunk survived despite a matching historical hit: %v", cm) + } + } + } +} + +// fakeDatasetnavLock records the key it was asked to acquire, so a test can +// assert the datasetnav rebuild lock is scoped to the dataset, not the document. +type fakeDatasetnavLock struct { + mu sync.Mutex + lastKey string +} + +func (f *fakeDatasetnavLock) Acquire(_ context.Context, key string) (bool, error) { + f.mu.Lock() + defer f.mu.Unlock() + f.lastKey = key + return true, nil +} +func (f *fakeDatasetnavLock) Release(_ context.Context, _ string) error { return nil } + +// TestKnowledgeCompiler_Datasetnav_LockKeyedByDataset is a regression for the +// Medium bug: the distributed rebuild lock used to be keyed by doc_id, so two +// runs against the same dataset (with different per-document doc_id values) +// took different locks and could interleave. This test supplies both doc_id +// ("d1") and dataset_id ("ds1") and asserts the lock key uses the dataset. +func TestKnowledgeCompiler_Datasetnav_LockKeyedByDataset(t *testing.T) { + lk := &fakeDatasetnavLock{} + common.SetDepsResolver(func(tenantID, llmID, embeddingModel string) (common.Deps, error) { + return common.Deps{ + Chat: proseChat{}, + Embed: mockEmbedder{dim: 8}, + TenantID: tenantID, + Redis: lk, + }, nil + }) + t.Cleanup(func() { common.SetDepsResolver(nil) }) + + c, err := NewKnowledgeCompilerComponent("KnowledgeCompiler", map[string]any{ + "variant": "datasetnav", "llm_id": "llm1", "embedding_model": "emb1", + }) + if err != nil { + t.Fatalf("NewKnowledgeCompilerComponent: %v", err) + } + if _, err := c.Invoke(context.Background(), nil, map[string]any{ + "chunks": []any{ + map[string]any{"id": "c1", "text": "The quick brown fox jumps over the lazy dog near the river bank."}, + map[string]any{"id": "c2", "text": "A red fox and a lazy dog rest beside a calm river at dawn."}, + }, + "doc_id": "d1", + "dataset_id": "ds1", + "tenant_id": "t1", + }); err != nil { + t.Fatalf("Invoke: %v", err) + } + lk.mu.Lock() + key := lk.lastKey + lk.mu.Unlock() + want := "datasetnav:t1:ds1" + if key != want { + t.Fatalf("datasetnav lock key = %q, want %q (lock must be scoped to the dataset, not the document)", key, want) + } +} + +// recordingNavChat echoes each nav-group summary back verbatim (so the child +// summaries are identifiable) and records the root-synthesis user prompt so a +// test can assert which child summaries reached the root overview. +type recordingNavChat struct { + mu sync.Mutex + rootPrompt string +} + +func (r *recordingNavChat) Chat(_ context.Context, req common.ChatRequest) (*common.ChatResponse, error) { + if strings.HasPrefix(req.UserPrompt, "Compose a navigation overview") { + r.mu.Lock() + r.rootPrompt = req.UserPrompt + r.mu.Unlock() + return &common.ChatResponse{Content: "root overview"}, nil + } + // Nav-group summary: echo the group text so the marker survives into the + // root prompt when the root is built from all child summaries. + return &common.ChatResponse{Content: strings.TrimSpace(req.UserPrompt)}, nil +} + +// TestKnowledgeCompiler_Datasetnav_RootIncludesFlushedSummaries is a regression +// for the Medium bug: the root overview used to be built from sink.Products() +// (the not-yet-flushed buffer), so under OnExceed=flush any nav nodes already +// emitted and cleared no longer contributed to the root, making the dataset +// root incomplete precisely in the large-result case where streaming applies. +// Here nav_radius > 1 forces each chunk into its own nav node, a MaxItems=1 +// flush cap drains earlier nodes from the sink before the root is built, and we +// assert every child marker still appears in the recorded root-synthesis prompt. +func TestKnowledgeCompiler_Datasetnav_RootIncludesFlushedSummaries(t *testing.T) { + chat := &recordingNavChat{} + common.SetDepsResolver(func(tenantID, llmID, embeddingModel string) (common.Deps, error) { + return common.Deps{ + Chat: chat, + Embed: mockEmbedder{dim: 8}, + TenantID: tenantID, + }, nil + }) + t.Cleanup(func() { common.SetDepsResolver(nil) }) + + c, err := NewKnowledgeCompilerComponent("KnowledgeCompiler", map[string]any{ + "variant": "datasetnav", "llm_id": "llm1", "embedding_model": "emb1", + // nav_radius > 1 guarantees no two chunks group together (cosine <= 1), + // so each chunk becomes its own nav node. + "extra": map[string]any{"nav_radius": 1.1}, + }) + if err != nil { + t.Fatalf("NewKnowledgeCompilerComponent: %v", err) + } + kc, ok := c.(*KnowledgeCompilerComponent) + if !ok { + t.Fatalf("component is not *KnowledgeCompilerComponent: %T", c) + } + // Tiny cap + flush so earlier nav nodes are emitted and cleared from the + // sink buffer before the root synthesis runs. + kc.Param.Guardrails = common.CapacityGuardrails{MaxItems: 1, OnExceed: common.ExceedFlush} + + markers := []string{"NAVMARKA", "NAVMARKB", "NAVMARKC", "NAVMARKD", "NAVMARKE"} + chunks := make([]any, len(markers)) + for i, m := range markers { + chunks[i] = map[string]any{"id": m, "text": m + " unique section body text"} + } + if _, err := c.Invoke(context.Background(), nil, map[string]any{ + "chunks": chunks, + "doc_id": "d1", + "tenant_id": "t1", + "sink": &sinkRecorder{}, + }); err != nil { + t.Fatalf("Invoke: %v", err) + } + + chat.mu.Lock() + root := chat.rootPrompt + chat.mu.Unlock() + if root == "" { + t.Fatalf("root synthesis prompt was never recorded (root node not built)") + } + for _, m := range markers { + if !strings.Contains(root, m) { + t.Fatalf("root overview is missing flushed child summary %q; the root must be built from ALL child summaries, not just the not-yet-flushed sink buffer.\nroot prompt:\n%s", m, root) + } + } +} + +// fencedChat wraps otherwise-valid extraction JSON in a ```json ... ``` fence, +// the most common way LLMs wrap JSON even when JSONMode is requested. +type fencedChat struct{} + +func (fencedChat) Chat(_ context.Context, req common.ChatRequest) (*common.ChatResponse, error) { + return &common.ChatResponse{Content: "```json\n" + + `{"items":[{"type":"project","name":"Alpha","description":"Alpha project","source_chunk_ids":["c1"]}]}` + + "\n```"}, nil +} + +// proseOnlyChat returns genuine prose with no JSON at all — a model that +// ignored the JSON contract entirely. +type proseOnlyChat struct{} + +func (proseOnlyChat) Chat(_ context.Context, req common.ChatRequest) (*common.ChatResponse, error) { + return &common.ChatResponse{Content: "I could not extract structured facts from this passage."}, nil +} + +type wikiUpdateChat struct{} + +func (wikiUpdateChat) Chat(_ context.Context, req common.ChatRequest) (*common.ChatResponse, error) { + switch { + case req.JSONMode && strings.Contains(req.UserPrompt, "Extract all knowledge"): + return &common.ChatResponse{Content: `{ + "entities": [{"name":"Alpha","type":"person","aliases":["A"],"source_chunk_id":"c1"}], + "concepts": [{"term":"Alpha Protocol","definition_excerpt":"Alpha Protocol is a rollout process","source_chunk_id":"c1"}], + "claims": [{"statement":"Alpha launched a new process in 2026.","subject":"Alpha","confidence":"explicit","source_chunk_id":"c1"}], + "relations": [], + "topics": ["Alpha"] +}`}, nil + case req.JSONMode && strings.Contains(req.UserPrompt, "Return a JSON compilation plan"): + return &common.ChatResponse{Content: `{ + "pages":[{"action":"CREATE","slug":"entity/alpha","title":"Alpha","page_type":"entity","topic":"Alpha","entity_names":["Alpha"],"related_kb_pages":[],"priority":1,"lead":"Alpha overview","sections":[{"heading":"Overview","points":["Alpha overview"]}]}], + "estimated_page_count":1, + "compilation_notes":"ok" +}`}, nil + case strings.Contains(req.SystemPrompt, "wiki page merger"): + return &common.ChatResponse{Content: "# Alpha\n\nExisting fact.\n\nAlpha launched a new process in 2026.\n\n## See also\n\n[Alpha Protocol](artifact/ds1/concept/alpha-protocol)"}, nil + case strings.Contains(req.UserPrompt, "Existing page content"): + return &common.ChatResponse{Content: "# Alpha\n\nAlpha launched a new process in 2026.\n\n## See also\n\n[[concept/alpha-protocol|Alpha Protocol]]"}, nil + default: + return &common.ChatResponse{Content: `{"action":"UPDATE","slug":"entity/alpha","reason":"same entity"}`}, nil + } +} + +type wikiStoreTestStub struct { + page *common.WikiPageCandidate +} + +func (s wikiStoreTestStub) FindSimilarPages(_ context.Context, _, _ string, _ []float32, _ int) ([]common.WikiPageCandidate, error) { + if s.page == nil { + return nil, nil + } + return []common.WikiPageCandidate{*s.page}, nil +} + +func (s wikiStoreTestStub) GetPageBySlug(_ context.Context, _, _, slug string) (*common.WikiPageCandidate, error) { + if s.page != nil && s.page.Slug == slug { + return s.page, nil + } + return nil, nil +} + +// TestKnowledgeCompiler_Structure_FencedJSONNotDropped is a regression for the +// Medium bug: GenJSON used to return {"_raw":...} on any parse failure, which +// parseExtractResult silently treated as "empty extraction", dropping every +// entity/relation. A fenced ```json ... ``` reply is now unwrapped and parsed, +// so the extraction still yields its entities. +func TestKnowledgeCompiler_Structure_FencedJSONNotDropped(t *testing.T) { + common.SetDepsResolver(func(tenantID, llmID, embeddingModel string) (common.Deps, error) { + return common.Deps{Chat: fencedChat{}, Embed: mockEmbedder{dim: 8}, TenantID: tenantID}, nil + }) + t.Cleanup(func() { common.SetDepsResolver(nil) }) + + c, err := NewKnowledgeCompilerComponent("KnowledgeCompiler", map[string]any{ + "variant": "structure", "llm_id": "llm1", "embedding_model": "emb1", + }) + if err != nil { + t.Fatalf("NewKnowledgeCompilerComponent: %v", err) + } + out, err := c.Invoke(context.Background(), nil, map[string]any{ + "chunks": []any{map[string]any{"id": "c1", "text": "Alpha is a Beta"}}, + "doc_id": "d1", + "tenant_id": "t1", + }) + if err != nil { + t.Fatalf("Invoke (fenced JSON should parse): %v", err) + } + var sawAlphaEntity bool + for _, r := range out["chunks"].([]any) { + cm := r.(map[string]any) + if cm["compile_kwd"] == nil { + continue + } + // name_kwd is lowercased (mirrors Python's _struct_to_doc_storage_doc). + if cm["name_kwd"] == "alpha" { + sawAlphaEntity = true + } + } + if !sawAlphaEntity { + t.Fatalf("fenced-JSON reply was dropped (no Alpha entity extracted); GenJSON must unwrap the fence, not treat it as empty extraction") + } +} + +// TestKnowledgeCompiler_Structure_MalformedJSONFailsLoud is the companion +// regression: when the reply is genuinely unparseable (not just fenced), the +// component must fail loudly instead of silently emitting zero knowledge units. +func TestKnowledgeCompiler_Structure_MalformedJSONFailsLoud(t *testing.T) { + common.SetDepsResolver(func(tenantID, llmID, embeddingModel string) (common.Deps, error) { + return common.Deps{Chat: proseOnlyChat{}, Embed: mockEmbedder{dim: 8}, TenantID: tenantID}, nil + }) + t.Cleanup(func() { common.SetDepsResolver(nil) }) + + c, err := NewKnowledgeCompilerComponent("KnowledgeCompiler", map[string]any{ + "variant": "structure", "llm_id": "llm1", "embedding_model": "emb1", + }) + if err != nil { + t.Fatalf("NewKnowledgeCompilerComponent: %v", err) + } + _, err = c.Invoke(context.Background(), nil, map[string]any{ + "chunks": []any{map[string]any{"id": "c1", "text": "Alpha is a Beta"}}, + "doc_id": "d1", + "tenant_id": "t1", + }) + if err == nil { + t.Fatalf("Invoke succeeded on an unparseable LLM reply; it must fail loudly rather than silently drop the extraction") + } +} + +// TestKnowledgeCompiler_PassThroughEnvelope is a regression for the Medium bug: +// the advertised output schema promises name / tenant_id / kb_id are carried +// forward from upstream (pass-through), but mergeChunks only emitted +// output_format + chunks. Headless / manual chaining reads those identity keys +// from the component output (the Tokenizer falls back to globals only in a full +// pipeline), so they must be forwarded when present. +func TestKnowledgeCompiler_PassThroughEnvelope(t *testing.T) { + installMockDeps(t) + + c, err := NewKnowledgeCompilerComponent("KnowledgeCompiler", map[string]any{ + "variant": "structure", "llm_id": "llm1", "embedding_model": "emb1", + }) + if err != nil { + t.Fatalf("NewKnowledgeCompilerComponent: %v", err) + } + out, err := c.Invoke(context.Background(), nil, map[string]any{ + "chunks": []any{map[string]any{"id": "c1", "text": "Alpha is a Beta"}}, + "doc_id": "d1", + "tenant_id": "t1", + "name": "doc.pdf", + "kb_id": "kb9", + }) + if err != nil { + t.Fatalf("Invoke: %v", err) + } + if out["name"] != "doc.pdf" { + t.Fatalf("output missing pass-through name: %v", out["name"]) + } + if out["tenant_id"] != "t1" { + t.Fatalf("output missing pass-through tenant_id: %v", out["tenant_id"]) + } + if out["kb_id"] != "kb9" { + t.Fatalf("output missing pass-through kb_id: %v", out["kb_id"]) + } + if out["output_format"] != "chunks" { + t.Fatalf("output_format = %v, want chunks", out["output_format"]) + } + // Absent keys must NOT be invented. + if _, ok := out["dataset_id"]; ok { + t.Fatalf("output must not invent absent pass-through keys") + } +} + +// groupResolverStub maps every requested group id to a fixed pair of template +// ids, standing in for the production DB-backed group service. +func groupResolverStub(_ context.Context, _ string, groupIDs []string) ([]string, error) { + var out []string + for _, g := range groupIDs { + out = append(out, "tpl-"+g+"-a", "tpl-"+g+"-b") + } + return out, nil +} + +// TestKnowledgeCompiler_GroupIDsResolvedToTemplateIDs is a regression for the +// open question: compilation_template_group_ids were parsed into Param.GroupIDs +// but never resolved or stamped, so group-only configs silently missed +// compilation_template_ids. The component now resolves groups via the installed +// GroupResolver and merges them into the stamped template-id list. +func TestKnowledgeCompiler_GroupIDsResolvedToTemplateIDs(t *testing.T) { + installMockDeps(t) + common.SetGroupResolver(groupResolverStub) + t.Cleanup(func() { common.SetGroupResolver(nil) }) + + c, err := NewKnowledgeCompilerComponent("KnowledgeCompiler", map[string]any{ + "variant": "structure", + "llm_id": "llm1", + "embedding_model": "emb1", + "compilation_template_group_ids": []any{"grp1"}, + }) + if err != nil { + t.Fatalf("NewKnowledgeCompilerComponent: %v", err) + } + out, err := c.Invoke(context.Background(), nil, map[string]any{ + "chunks": []any{map[string]any{"id": "c1", "text": "Alpha is a Beta"}}, + "doc_id": "d1", + "tenant_id": "t1", + }) + if err != nil { + t.Fatalf("Invoke: %v", err) + } + want := map[string]bool{"tpl-grp1-a": true, "tpl-grp1-b": true} + checked := 0 + for _, r := range out["chunks"].([]any) { + cm := r.(map[string]any) + // No parser_config is supplied, so InferType returns "list" and the + // structure variant stamps compile_kwd="list" (not "structure"). + if cm["compile_kwd"] != "list" { + continue + } + checked++ + var got []string + switch v := cm["compilation_template_ids"].(type) { + case []any: + for _, e := range v { + got = append(got, e.(string)) + } + case []string: + got = v + } + if len(got) != 2 || !want[got[0]] || !want[got[1]] { + t.Fatalf("compiled chunk %v: compilation_template_ids = %v, want %v (group ids must resolve to template ids)", cm["id"], got, want) + } + } + if checked == 0 { + t.Fatalf("no compiled chunks inspected (compile_kwd=list expected); assertion was vacuous") + } +} + +// TestKnowledgeCompiler_GroupIDsWithoutResolverFailsLoud is the companion +// regression: a group-only config with no GroupResolver installed must fail +// loudly (surfacing the misconfiguration) rather than silently emitting rows +// that miss compilation_template_ids. +func TestKnowledgeCompiler_GroupIDsWithoutResolverFailsLoud(t *testing.T) { + installMockDeps(t) + common.SetGroupResolver(nil) // ensure no resolver is installed + + c, err := NewKnowledgeCompilerComponent("KnowledgeCompiler", map[string]any{ + "variant": "structure", + "llm_id": "llm1", + "embedding_model": "emb1", + "compilation_template_group_ids": []any{"grp1"}, + }) + if err != nil { + t.Fatalf("NewKnowledgeCompilerComponent: %v", err) + } + _, err = c.Invoke(context.Background(), nil, map[string]any{ + "chunks": []any{map[string]any{"id": "c1", "text": "Alpha is a Beta"}}, + "doc_id": "d1", + "tenant_id": "t1", + }) + if err == nil { + t.Fatalf("Invoke succeeded with group ids but no GroupResolver; it must fail loudly instead of dropping compilation_template_ids") + } + if !strings.Contains(err.Error(), "GroupResolver") { + t.Fatalf("error %q does not mention GroupResolver", err.Error()) + } +} diff --git a/internal/ingestion/component/knowledge_compiler/datasetnav/datasetnav.go b/internal/ingestion/component/knowledge_compiler/datasetnav/datasetnav.go new file mode 100644 index 0000000000..f893ae92ab --- /dev/null +++ b/internal/ingestion/component/knowledge_compiler/datasetnav/datasetnav.go @@ -0,0 +1,788 @@ +// Package datasetnav implements the "datasetnav" variant of KnowledgeCompiler, +// mirroring Python's dataset_nav.py incremental clustering: each input chunk +// (the in-run equivalent of a document leaf) is embedded and placed into the +// nearest nav_cluster via layered KNN descent — merged when the similarity +// clears _MERGE_THRESHOLD, given a fresh sibling cluster above _MIN_SIM, or +// given a new root-level cluster otherwise. Clusters that exceed the fanout +// or doc-count caps split via the same 2-means procedure as Python. +// +// The Go port keeps the whole tree in memory for one Invoke (no ES reads or +// writes): cross-run incremental upsert/removal, which in Python is an +// ES-backed read-modify-write, is the caller's concern. The per-run algorithm +// (thresholds, LLM merge/summary prompts, readable cluster names, split) is +// faithful to the Python original. +package datasetnav + +import ( + "context" + "encoding/json" + "fmt" + "math" + "strings" + + "ragflow/internal/ingestion/component/knowledge_compiler/common" +) + +// Thresholds and caps, mirrored from dataset_nav.py. +const ( + mergeThresholdDefault = 0.80 // merge doc into cluster + recurseThreshold = 0.65 // continue descending into children + minSim = 0.50 // minimum similarity to be considered related + maxFanout = 64 // max child count before triggering rebalance + maxDocsPerCluster = 50 // max docs per leaf cluster before triggering split +) + +// datasetnavLock is the minimal locking surface used to serialize dataset-level +// rebuilds. The production Redis client implements this via SET NX + TTL; tests +// inject an in-memory lock (or nil for no-op). +type datasetnavLock interface { + Acquire(ctx context.Context, key string) (bool, error) + Release(ctx context.Context, key string) error +} + +// navCluster is one internal tree node (Python's nav_cluster row). +type navCluster struct { + Name string + Desc string + Parent string // parent cluster name; "root" for depth-0 clusters + Depth int + DocIDs []string + Vector []float32 +} + +// navDoc is one document leaf (Python's nav_doc row). +type navDoc struct { + ChunkID string + Text string + Parent string // owning cluster name + Depth int + Vector []float32 +} + +// navChild is one child reference under a cluster. +type navChild struct { + name string // cluster name or doc chunk id + isDoc bool + vec []float32 +} + +// navTree holds the in-memory tree built over one Invoke. +type navTree struct { + clusters map[string]*navCluster + order []string // cluster names in creation order + docs []*navDoc + children map[string][]navChild // parent cluster name → children +} + +func newNavTree() *navTree { + return &navTree{clusters: map[string]*navCluster{}, children: map[string][]navChild{}} +} + +func (t *navTree) addCluster(c *navCluster) { + t.clusters[c.Name] = c + t.order = append(t.order, c.Name) +} + +// Run executes the datasetnav variant. +func Run(ctx context.Context, deps common.Deps, param common.Param, inputs common.Inputs) (common.Outputs, error) { + if deps.Embed == nil { + return common.Outputs{}, fmt.Errorf("datasetnav: embedder required") + } + docID := firstNonEmpty(inputs.DocID, deps.DatasetID) + if docID == "" { + docID = "unknown" + } + llmID := firstNonEmpty(param.LLMID, inputs.LLMID) + tenantID := deps.TenantID + + // The rebuild lock is scoped to the dataset, not the document (mirrors + // Python's _nav_lock_key(kb_id)): concurrent runs against the same dataset + // must share one lock. Fall back to docID only when no dataset id is known. + datasetID := firstNonEmpty(deps.DatasetID, docID) + + texts, chunkIDs := chunkTexts(inputs.Chunks) + if len(texts) == 0 { + return common.Outputs{}, nil + } + + if l, ok := deps.Redis.(datasetnavLock); ok && l != nil { + lockKey := "datasetnav:" + tenantID + ":" + datasetID + acquired, err := l.Acquire(ctx, lockKey) + if err != nil { + return common.Outputs{}, err + } + if !acquired { + return common.Outputs{}, fmt.Errorf("datasetnav: lock not acquired for %s", lockKey) + } + defer l.Release(ctx, lockKey) + } + + // Embed every chunk text once (Python embeds each doc summary on upsert). + vectors, err := deps.Embed.Encode(ctx, texts) + if err != nil { + return common.Outputs{}, err + } + if len(vectors) != len(texts) { + return common.Outputs{}, fmt.Errorf("datasetnav: embedding count mismatch (%d vs %d)", len(vectors), len(texts)) + } + + // Sequential placement, mirroring the lock-serialized per-document upserts + // in Python (also what makes the in-run tree deterministic). + tree := newNavTree() + var clusterDescs []string + mergeThreshold := mergeThresholdFor(param) + for i, text := range texts { + if err := ctx.Err(); err != nil { + return common.Outputs{}, err + } + desc, err := tree.place(ctx, deps, llmID, chunkIDs[i], text, vectors[i], mergeThreshold) + if err != nil { + return common.Outputs{}, err + } + clusterDescs = append(clusterDescs, desc...) + } + + // Emit cluster rows in creation order, then nav_doc leaves in placement + // order, then the root overview node (a Go-side convenience: Python's tree + // has no root row, but downstream consumers expect one overview product). + sink := common.NewProductSink(ctx, param.Guardrails, inputs.Sink) + clusterProductID := map[string]string{} + // Pre-compute every cluster id (deterministic from StableRowID) before + // resolving parent edges, so a child emitted before its (later-created) + // parent still gets the correct ParentID instead of "" (reparenting bug). + // Cluster ids are tenant-scoped (as the nav_doc ids already are) so two + // tenants whose datasetID falls back to the same docID cannot collide. + for _, name := range tree.order { + clusterProductID[name] = common.StableRowID("dataset_nav", tenantID, datasetID, "cluster", name) + } + for _, name := range tree.order { + c := tree.clusters[name] + pid := "" + if c.Parent != "" && c.Parent != "root" { + pid = clusterProductID[c.Parent] + } + id := clusterProductID[c.Name] + if err := sink.Add(common.Product{ + ID: id, + DocID: docID, + TenantID: tenantID, + Variant: common.VariantDatasetnav, + Content: payloadJSON(map[string]any{"type": "nav_cluster", "description": c.Desc}), + Vector: c.Vector, + ParentID: pid, + Meta: map[string]any{ + "kind": "nav_cluster", + "type": "nav_cluster", + "name": c.Name, + "parent_name": c.Parent, + "depth": c.Depth, + "doc_ids": append([]string{}, c.DocIDs...), + "size": len(c.DocIDs), + }, + }); err != nil { + return common.Outputs{}, err + } + } + for _, d := range tree.docs { + if err := sink.Add(common.Product{ + // Scope nav_doc ids by tenant and dataset (not just chunk id) so + // synthetic/positional chunk ids from different docs/datasets do not + // collide and overwrite each other across tenants (M5). + ID: common.StableRowID("dataset_nav", tenantID, datasetID, "doc", d.ChunkID), + DocID: docID, + TenantID: tenantID, + Variant: common.VariantDatasetnav, + Content: payloadJSON(map[string]any{"type": "nav_doc", "description": d.Text}), + Vector: d.Vector, + ParentID: clusterProductID[d.Parent], + Meta: map[string]any{ + "kind": "nav_doc", + "type": "nav_doc", + "name": navDocName(d.Parent, d.Text), + "parent_name": d.Parent, + "depth": d.Depth, + "doc_ids": []string{d.ChunkID}, + }, + }); err != nil { + return common.Outputs{}, err + } + } + + // Root overview built from EVERY cluster description (mirrors the retained- + // summaries pattern: already-flushed nodes still contribute to the root). + rootSummary, err := summarize(ctx, deps, llmID, "Compose a navigation overview from these section summaries:\n\n"+formatNavSummaries(clusterDescs)) + if err == nil && rootSummary != "" { + emb, e2 := deps.Embed.Encode(ctx, []string{rootSummary}) + if e2 == nil && len(emb) > 0 { + if err := sink.Add(common.Product{ + ID: common.StableRowID(tenantID, docID, string(common.VariantDatasetnav), "root"), + DocID: docID, + TenantID: tenantID, + Variant: common.VariantDatasetnav, + Content: rootSummary, + Vector: emb[0], + Meta: map[string]any{ + "kind": "root", + "type": "nav_cluster", + "name": "root", + "depth": 0, + "size": len(tree.docs), + }, + }); err != nil { + return common.Outputs{}, err + } + } + } + + out := common.Outputs{ + Products: sink.Products(), + VectorBytes: sink.Bytes(), + Items: sink.TotalItems(), + Flushed: sink.Flushed(), + } + if err := out.EnforceGuardrails(param.Guardrails, inputs.Sink, ctx); err != nil { + return common.Outputs{}, err + } + return out, nil +} + +// place inserts one document (chunk) into the tree, mirroring +// upsert_dataset_nav_doc's placement logic. It returns the descriptions of +// any clusters created for this document (for the root overview). +func (t *navTree) place(ctx context.Context, deps common.Deps, llmID, chunkID, text string, vec []float32, mergeThreshold float64) ([]string, error) { + bestName, bestParent, sim := t.findBestCluster(vec) + + switch { + case bestName != "" && sim >= mergeThreshold: + // ── Merge into the best cluster (mirrors _llm_merge + re-embed) ── + c := t.clusters[bestName] + newDesc, err := llmMerge(ctx, deps, llmID, c.Desc, text) + if err != nil { + return nil, err + } + if newDesc != c.Desc { + c.Desc = newDesc + emb, err := deps.Embed.Encode(ctx, []string{newDesc}) + if err != nil { + return nil, err + } + if len(emb) > 0 { + c.Vector = emb[0] + } + } + if !containsString(c.DocIDs, chunkID) { + c.DocIDs = append(c.DocIDs, chunkID) + } + t.addDoc(c, chunkID, text, vec, c.Depth+1) + if err := t.maybeSplit(ctx, deps, llmID, bestName); err != nil { + return nil, err + } + return nil, nil + + case bestName != "" && sim >= minSim: + // ── Create a sibling/child cluster (mirrors the _MIN_SIM branch) ── + parentForNew := bestParent + if parentForNew == "" { + parentForNew = bestName + } + parentDepth := 1 // Python's default when the parent row is absent + if p, ok := t.clusters[parentForNew]; ok { + parentDepth = p.Depth + } + title, desc, err := llmCreateSummary(ctx, deps, llmID, []string{text}) + if err != nil { + return nil, err + } + name := readableClusterName(title, text) + nc := &navCluster{ + Name: name, + Desc: desc, + Parent: parentForNew, + // Python stamps the new cluster at the PARENT's depth (not +1); + // the nav_doc goes one deeper. Mirrored, quirk included. + Depth: parentDepth, + DocIDs: []string{chunkID}, + Vector: vec, + } + t.addCluster(nc) + t.children[parentForNew] = append(t.children[parentForNew], navChild{name: name, vec: vec}) + t.addDoc(nc, chunkID, text, vec, parentDepth+1) + return []string{desc}, nil + + default: + // ── Create a root-level cluster (mirrors the else branch) ── + title, desc, err := llmCreateSummary(ctx, deps, llmID, []string{text}) + if err != nil { + return nil, err + } + name := readableClusterName(title, text) + nc := &navCluster{ + Name: name, + Desc: desc, + Parent: "root", + Depth: 0, + DocIDs: []string{chunkID}, + Vector: vec, + } + t.addCluster(nc) + t.addDoc(nc, chunkID, text, vec, 1) + return []string{desc}, nil + } +} + +func (t *navTree) addDoc(c *navCluster, chunkID, text string, vec []float32, depth int) { + t.docs = append(t.docs, &navDoc{ChunkID: chunkID, Text: text, Parent: c.Name, Depth: depth, Vector: vec}) + t.children[c.Name] = append(t.children[c.Name], navChild{name: chunkID, isDoc: true, vec: vec}) +} + +// findBestCluster mirrors _find_best_cluster: top-1 root cluster by cosine, +// then descend into the best-matching child while similarity stays above +// _RECURSE_THRESHOLD. Returns (best cluster name, its parent name, sim). +func (t *navTree) findBestCluster(vec []float32) (string, string, float64) { + var best *navCluster + bestSim := 0.0 + for _, c := range t.clusters { + if c.Depth != 0 { + continue + } + s := cosine(vec, c.Vector) + if best == nil || s > bestSim { + best, bestSim = c, s + } + } + if best == nil { + return "", "", 0.0 + } + bestName, bestParent, sim := best.Name, best.Parent, bestSim + visited := map[string]bool{bestName: true} + for sim >= recurseThreshold { + var child *navCluster + childSim := 0.0 + for _, k := range t.children[bestName] { + if k.isDoc || visited[k.name] { + continue + } + c := t.clusters[k.name] + if c == nil { + continue + } + if s := cosine(vec, c.Vector); child == nil || s > childSim { + child, childSim = c, s + } + } + if child == nil || childSim < recurseThreshold { + break + } + bestParent = best.Parent // the (old) best's parent, mirroring Python + bestName = child.Name + sim = childSim + best = child + visited[bestName] = true + } + return bestName, bestParent, sim +} + +// maybeSplit mirrors _maybe_split_cluster: when a cluster exceeds the fanout +// or doc-count caps, split its children into two 2-means groups and reparent +// each group under a fresh sub-cluster. +func (t *navTree) maybeSplit(ctx context.Context, deps common.Deps, llmID, clusterName string) error { + kids := t.children[clusterName] + var clusterKids, docKids []navChild + for _, k := range kids { + if k.isDoc { + docKids = append(docKids, k) + } else { + clusterKids = append(clusterKids, k) + } + } + if len(clusterKids)+len(docKids) <= maxFanout && len(docKids) <= maxDocsPerCluster { + return nil + } + if len(kids) < 4 { + return nil // Python requires >= 4 embeddings to attempt a split + } + + embs := make([][]float32, len(kids)) + for i, k := range kids { + embs[i] = k.vec + } + labels := twoMeans(embs) + + parent := t.clusters[clusterName] + depth := 0 + if parent != nil { + depth = parent.Depth + 1 + } + // navChild collects the new sub-cluster names so the split parent keeps a + // children entry pointing at them; without it findBestCluster can no longer + // descend into the split subtree (M3). + var subChildren []navChild + for gi := 0; gi < 2; gi++ { + var kidIdx []int + for i, lb := range labels { + if lb == gi { + kidIdx = append(kidIdx, i) + } + } + if len(kidIdx) == 0 { + continue + } + var docIDs, descs []string + for _, i := range kidIdx { + k := kids[i] + if k.isDoc { + if !containsString(docIDs, k.name) { + docIDs = append(docIDs, k.name) + } + for _, d := range t.docs { + if d.ChunkID == k.name { + descs = append(descs, d.Text) + } + } + continue + } + if c := t.clusters[k.name]; c != nil { + for _, d := range c.DocIDs { + if !containsString(docIDs, d) { + docIDs = append(docIDs, d) + } + } + descs = append(descs, c.Desc) + } + } + var title, desc string + if len(descs) > 0 { + var err error + title, desc, err = llmCreateSummary(ctx, deps, llmID, descs) + if err != nil { + return err + } + } else { + title, desc = fmt.Sprintf("Group %d", gi+1), fmt.Sprintf("Group %d", gi+1) + } + gname := readableClusterName(title, desc) + emb, err := deps.Embed.Encode(ctx, []string{desc}) + if err != nil { + return err + } + var gvec []float32 + if len(emb) > 0 { + gvec = emb[0] + } + nc := &navCluster{Name: gname, Desc: desc, Parent: clusterName, Depth: depth, DocIDs: docIDs, Vector: gvec} + t.addCluster(nc) + subChildren = append(subChildren, navChild{name: gname}) + + // Reparent the group's children to the new sub-cluster. + for _, i := range kidIdx { + k := kids[i] + if k.isDoc { + for _, d := range t.docs { + if d.ChunkID == k.name { + d.Parent = gname + d.Depth = depth + 1 + } + } + } else if c := t.clusters[k.name]; c != nil { + c.Parent = gname + c.Depth = depth + 1 + } + t.children[gname] = append(t.children[gname], k) + } + } + // The split parent's children become the new sub-clusters so findBestCluster + // can still descend into the split subtree (M3). When no sub-cluster was + // actually created the parent keeps an empty children list. + t.children[clusterName] = subChildren + return nil +} + +// twoMeans mirrors the runtime k-means-like split in _maybe_split_cluster: +// centroids seeded with the first and middle embeddings, 10 refinement +// rounds, squared-euclidean assignment, then a final relabel pass. +func twoMeans(embs [][]float32) []int { + n := len(embs) + if n == 0 { + return nil + } + centroids := [][]float32{append([]float32{}, embs[0]...), append([]float32{}, embs[n/2]...)} + assign := func(e []float32) int { + if sqDist(e, centroids[0]) < sqDist(e, centroids[1]) { + return 0 + } + return 1 + } + for iter := 0; iter < 10; iter++ { + var groups [2][][]float32 + for _, e := range embs { + g := assign(e) + groups[g] = append(groups[g], e) + } + for gi := 0; gi < 2; gi++ { + if len(groups[gi]) == 0 { + continue + } + dim := len(groups[gi][0]) + avg := make([]float32, dim) + for _, e := range groups[gi] { + for d := 0; d < dim; d++ { + avg[d] += e[d] + } + } + for d := range avg { + avg[d] /= float32(len(groups[gi])) + } + centroids[gi] = avg + } + } + labels := make([]int, n) + for i, e := range embs { + labels[i] = assign(e) + } + return labels +} + +func sqDist(a, b []float32) float64 { + var s float64 + for i := 0; i < len(a) && i < len(b); i++ { + d := float64(a[i] - b[i]) + s += d * d + } + return s +} + +func cosine(a, b []float32) float64 { + if len(a) == 0 || len(b) == 0 || len(a) != len(b) { + return 0 + } + var dot, na, nb float64 + for i := range a { + dot += float64(a[i]) * float64(b[i]) + na += float64(a[i]) * float64(a[i]) + nb += float64(b[i]) * float64(b[i]) + } + if na == 0 || nb == 0 { + return 0 + } + return dot / (math.Sqrt(na) * math.Sqrt(nb)) +} + +// ---- LLM helpers (prompts mirrored verbatim from dataset_nav.py) ---- + +var navLLMTemperature = 0.1 + +// llmMerge mirrors _llm_merge: fuse the existing cluster description with the +// new doc summary. The reply may be a JSON object ({"merged"|"result"}) or +// bare text; anything unusable keeps the old description. +func llmMerge(ctx context.Context, deps common.Deps, llmID, clusterDesc, docSummary string) (string, error) { + if deps.Chat == nil { + return clusterDesc, nil + } + prompt := "Merge the following two descriptions of the same topic into " + + "a single concise summary (1-3 sentences):\n\n" + + "Existing: " + clusterDesc + "\n\n" + + "New: " + docSummary + "\n\n" + + "Return ONLY the merged text, no commentary." + resp, err := deps.Chat.Chat(ctx, common.ChatRequest{LLMID: llmID, UserPrompt: prompt, Temperature: &navLLMTemperature}) + if err != nil { + return "", err + } + text := strings.TrimSpace(resp.Content) + if m := parseWholeJSON(text); m != nil { + if s, ok := m["merged"].(string); ok && strings.TrimSpace(s) != "" { + return s, nil + } + if s, ok := m["result"].(string); ok && strings.TrimSpace(s) != "" { + return s, nil + } + return clusterDesc, nil + } + if text != "" { + return text, nil + } + return clusterDesc, nil +} + +// llmCreateSummary mirrors _llm_create_summary: derive a readable (name, +// summary) pair from doc summaries, with the same fallbacks. +func llmCreateSummary(ctx context.Context, deps common.Deps, llmID string, docSummaries []string) (string, string, error) { + fallbackSummary := "" + if len(docSummaries) > 0 { + fallbackSummary = docSummaries[0] + } + if deps.Chat == nil { + return fallbackTitle(fallbackSummary), fallbackSummary, nil + } + prompt := "Given the document excerpts below, produce a short human-readable topic " + + "name and a concise description of their common topic.\n\n" + + strings.Join(docSummaries, "\n---\n") + "\n\n" + + `Return ONLY JSON: {"name": "<2-6 word topic title>", "summary": "<1-3 sentence description>"}` + resp, err := deps.Chat.Chat(ctx, common.ChatRequest{LLMID: llmID, UserPrompt: prompt, Temperature: &navLLMTemperature}) + if err != nil { + return "", "", err + } + text := strings.TrimSpace(resp.Content) + if m := parseWholeJSON(text); m != nil { + summary := fallbackSummary + if s, ok := m["summary"].(string); ok && strings.TrimSpace(s) != "" { + summary = strings.TrimSpace(s) + } else if s, ok := m["result"].(string); ok && strings.TrimSpace(s) != "" { + summary = strings.TrimSpace(s) + } + name := cleanTitle(firstStringOf(m["name"])) + if name == "" { + name = fallbackTitle(summary) + } + return name, summary, nil + } + if text != "" { + return fallbackTitle(text), text, nil + } + return fallbackTitle(fallbackSummary), fallbackSummary, nil +} + +// parseWholeJSON parses the response ONLY when it is a complete JSON object +// (mirrors gen_json's dict path; bare text takes the string fallback path). +func parseWholeJSON(s string) map[string]any { + if !strings.HasPrefix(s, "{") { + return nil + } + var m map[string]any + if err := json.Unmarshal([]byte(s), &m); err != nil { + return nil + } + return m +} + +// cleanTitle mirrors _clean_title: whitespace-normalized, capped at 48 chars. +func cleanTitle(title string) string { + return truncateRunes(strings.Join(strings.Fields(title), " "), 48) +} + +// fallbackTitle mirrors _fallback_title: first 6 words, else "Cluster". +func fallbackTitle(summary string) string { + words := strings.Fields(summary) + if len(words) > 6 { + words = words[:6] + } + if t := strings.Join(words, " "); t != "" { + return t + } + return "Cluster" +} + +// readableClusterName mirrors _readable_cluster_name: " <8-hex>". +func readableClusterName(title, seed string) string { + t := cleanTitle(title) + if t == "" { + t = "Cluster" + } + return t + " " + common.ContentHash(seed)[:8] +} + +// navDocName mirrors _make_nav_doc_row's name field: +// f"{parent_kwd}_{xxh64(summary)[:12]}". +func navDocName(parent, summary string) string { + return parent + "_" + common.ContentHash(summary)[:12] +} + +// summarize is the plain-text LLM helper used for the root overview. +func summarize(ctx context.Context, deps common.Deps, llmID, text string) (string, error) { + if deps.Chat == nil { + return "", nil + } + resp, err := deps.Chat.Chat(ctx, common.ChatRequest{ + LLMID: llmID, + SystemPrompt: navSystemPrompt, + UserPrompt: text, + }) + if err != nil { + return "", err + } + return strings.TrimSpace(resp.Content), nil +} + +// formatNavSummaries renders cluster descriptions as a bullet list for the +// root overview prompt. +func formatNavSummaries(summaries []string) string { + var b strings.Builder + for _, s := range summaries { + b.WriteString("- ") + b.WriteString(s) + b.WriteString("\n") + } + return b.String() +} + +// mergeThresholdFor returns the merge threshold (nav_radius override kept for +// tuning/tests; default _MERGE_THRESHOLD). +func mergeThresholdFor(param common.Param) float64 { + if t, ok := param.Extra["nav_radius"].(float64); ok && t > 0 { + return t + } + return mergeThresholdDefault +} + +// chunkTexts returns the non-empty chunk texts and their ids in parallel +// order (Python skips docs without a summary; we skip empty chunks). +func chunkTexts(chunks []common.Chunk) ([]string, []string) { + var texts, ids []string + for i, c := range chunks { + t := firstNonEmpty(c.Text, c.Content) + if strings.TrimSpace(t) == "" { + continue + } + id := c.ID + if id == "" { + id = fmt.Sprintf("chunk-%d", i) + } + texts = append(texts, t) + ids = append(ids, id) + } + return texts, ids +} + +func payloadJSON(v map[string]any) string { + var b strings.Builder + enc := json.NewEncoder(&b) + enc.SetEscapeHTML(false) + if err := enc.Encode(v); err != nil { + return "{}" + } + return strings.TrimSpace(b.String()) +} + +func truncateRunes(s string, n int) string { + r := []rune(s) + if len(r) <= n { + return s + } + return string(r[:n]) +} + +func firstStringOf(v any) string { + if s, ok := v.(string); ok { + return s + } + return "" +} + +func containsString(haystack []string, needle string) bool { + for _, s := range haystack { + if s == needle { + return true + } + } + return false +} + +func firstNonEmpty(vals ...string) string { + for _, v := range vals { + if v != "" { + return v + } + } + return "" +} + +const navSystemPrompt = `You are a navigation assistant. Summarize the provided text into a concise label and overview that helps a user navigate to the relevant content. Output the summary only.` diff --git a/internal/ingestion/component/knowledge_compiler/datasetnav/datasetnav_test.go b/internal/ingestion/component/knowledge_compiler/datasetnav/datasetnav_test.go new file mode 100644 index 0000000000..848e45f607 --- /dev/null +++ b/internal/ingestion/component/knowledge_compiler/datasetnav/datasetnav_test.go @@ -0,0 +1,289 @@ +package datasetnav + +import ( + "context" + "fmt" + "strings" + "testing" + + "ragflow/internal/ingestion/component/knowledge_compiler/common" +) + +// scriptedEmbedder maps texts to fixed vectors (fallback: halves). +type scriptedEmbedder struct { + vecs map[string][]float32 +} + +func (s scriptedEmbedder) Dimensions() int { return 2 } +func (s scriptedEmbedder) Encode(_ context.Context, texts []string) ([][]float32, error) { + out := make([][]float32, len(texts)) + for i, t := range texts { + if v, ok := s.vecs[t]; ok { + out[i] = v + } else { + out[i] = []float32{0.5, 0.5} + } + } + return out, nil +} + +// navFakeChat answers the three nav LLM calls deterministically: merge keeps +// the existing description, create-summary returns canned JSON, the root +// overview is a fixed string. +type navFakeChat struct { + mergeCalls int + summaryCalls int +} + +func (m *navFakeChat) Chat(_ context.Context, req common.ChatRequest) (*common.ChatResponse, error) { + switch { + case strings.Contains(req.UserPrompt, "Merge the following two descriptions"): + m.mergeCalls++ + existing := "" + if rest, ok := strings.CutPrefix(req.UserPrompt, "Merge the following two descriptions of the same topic into a single concise summary (1-3 sentences):\n\nExisting: "); ok { + existing, _, _ = strings.Cut(rest, "\n\nNew: ") + } + return &common.ChatResponse{Content: fmt.Sprintf(`{"merged": %q}`, existing)}, nil + case strings.Contains(req.UserPrompt, "Given the document excerpts below"): + m.summaryCalls++ + // Distinct summaries per call: readable cluster names hash the + // description, so identical canned summaries would collide (the same + // name-overwrite hazard exists in Python's _nav_cluster_id). + return &common.ChatResponse{Content: fmt.Sprintf(`{"name": "Topic", "summary": "a topic summary %d"}`, m.summaryCalls)}, nil + case strings.HasPrefix(req.UserPrompt, "Compose a navigation overview"): + return &common.ChatResponse{Content: "root overview"}, nil + default: + return &common.ChatResponse{Content: "ok"}, nil + } +} + +func TestTwoMeans(t *testing.T) { + embs := [][]float32{{1, 0}, {0.9, 0.1}, {0, 1}, {0.1, 0.9}} + labels := twoMeans(embs) + if len(labels) != 4 { + t.Fatalf("labels = %v", labels) + } + if labels[0] != labels[1] || labels[2] != labels[3] || labels[0] == labels[2] { + t.Fatalf("expected {x,x,y,y} grouping, got %v", labels) + } +} + +func TestReadableClusterName(t *testing.T) { + got := readableClusterName(" My Topic ", "seed text") + if !strings.HasPrefix(got, "My Topic ") { + t.Fatalf("readableClusterName = %q, want cleaned title + hash suffix", got) + } + if got := readableClusterName("", "seed"); !strings.HasPrefix(got, "Cluster ") { + t.Fatalf("empty title must fall back to Cluster: %q", got) + } +} + +func TestFallbackTitle(t *testing.T) { + if got := fallbackTitle("one two three four five six seven eight"); got != "one two three four five six" { + t.Fatalf("fallbackTitle = %q, want first 6 words", got) + } + if got := fallbackTitle(""); got != "Cluster" { + t.Fatalf("empty fallbackTitle = %q, want Cluster", got) + } +} + +func TestPlace_MergeSiblingRootBranches(t *testing.T) { + vecs := map[string][]float32{ + "A": {1, 0}, + "A2": {0.99, 0.1}, // cosine ≈ 0.995 with A → merge + "B": {0.7, 0.714}, // cosine ≈ 0.7 with A → sibling + "C": {0, 1}, // cosine 0 with A → new root + } + deps := common.Deps{Chat: &navFakeChat{}, Embed: scriptedEmbedder{vecs: vecs}} + tree := newNavTree() + + if _, err := tree.place(context.Background(), deps, "llm", "cA", "A", vecs["A"], mergeThresholdDefault); err != nil { + t.Fatal(err) + } + if _, err := tree.place(context.Background(), deps, "llm", "cA2", "A2", vecs["A2"], mergeThresholdDefault); err != nil { + t.Fatal(err) + } + if _, err := tree.place(context.Background(), deps, "llm", "cB", "B", vecs["B"], mergeThresholdDefault); err != nil { + t.Fatal(err) + } + if _, err := tree.place(context.Background(), deps, "llm", "cC", "C", vecs["C"], mergeThresholdDefault); err != nil { + t.Fatal(err) + } + + if len(tree.order) != 3 { + t.Fatalf("clusters = %d, want 3 (A-root, B-sibling, C-root): %v", len(tree.order), tree.order) + } + // A2 merged into A's cluster. + var clusterA *navCluster + for _, name := range tree.order { + if tree.clusters[name].Depth == 0 && containsString(tree.clusters[name].DocIDs, "cA") { + clusterA = tree.clusters[name] + } + } + if clusterA == nil || len(clusterA.DocIDs) != 2 || clusterA.DocIDs[1] != "cA2" { + t.Fatalf("A2 must merge into cluster A (doc_ids=%v)", clusterA) + } + // B's cluster: sibling under the virtual root (Python depth quirk: parent + // depth default 1 → cluster depth 1, parent "root"). + var clusterB *navCluster + for _, name := range tree.order { + c := tree.clusters[name] + if containsString(c.DocIDs, "cB") { + clusterB = c + } + } + if clusterB == nil || clusterB.Parent != "root" || clusterB.Depth != 1 { + t.Fatalf("B sibling cluster = %+v, want parent root depth 1", clusterB) + } + // C's cluster: a fresh root-level cluster. + var clusterC *navCluster + for _, name := range tree.order { + c := tree.clusters[name] + if containsString(c.DocIDs, "cC") { + clusterC = c + } + } + if clusterC == nil || clusterC.Parent != "root" || clusterC.Depth != 0 { + t.Fatalf("C root cluster = %+v, want parent root depth 0", clusterC) + } + if len(tree.docs) != 4 { + t.Fatalf("nav docs = %d, want 4", len(tree.docs)) + } +} + +func TestMaybeSplit(t *testing.T) { + vecs := map[string][]float32{} + deps := common.Deps{Chat: &navFakeChat{}, Embed: scriptedEmbedder{vecs: vecs}} + tree := newNavTree() + parent := &navCluster{Name: "P", Desc: "parent", Parent: "root", Depth: 0, Vector: []float32{1, 0}} + tree.addCluster(parent) + // 51 docs (> maxDocsPerCluster) alternating between two axes. + for i := 0; i < 51; i++ { + id := fmt.Sprintf("d%d", i) + v := []float32{1, 0} + if i%2 == 1 { + v = []float32{0, 1} + } + parent.DocIDs = append(parent.DocIDs, id) + tree.addDoc(parent, id, "text "+id, v, 1) + } + if err := tree.maybeSplit(context.Background(), deps, "llm", "P"); err != nil { + t.Fatal(err) + } + // Expect two new sub-clusters under P, children reparented, P left with + // no direct children. + if len(tree.order) != 3 { + t.Fatalf("clusters after split = %d, want 3 (P + 2 groups): %v", len(tree.order), tree.order) + } + // The split parent must now point at exactly the two new sub-clusters so + // findBestCluster can descend into the split subtree (M3). + if got := len(tree.children["P"]); got != 2 { + t.Fatalf("split parent children = %d, want 2 (the two new sub-clusters)", got) + } + subDocs := 0 + for _, name := range tree.order[1:] { + c := tree.clusters[name] + if c.Parent != "P" || c.Depth != 1 { + t.Errorf("sub-cluster %s parent/depth = %s/%d, want P/1", name, c.Parent, c.Depth) + } + subDocs += len(c.DocIDs) + } + if subDocs != 51 { + t.Errorf("sub-cluster doc_ids total = %d, want 51", subDocs) + } + for _, d := range tree.docs { + if d.Parent == "P" { + t.Errorf("doc %s not reparented off P", d.ChunkID) + } + if d.Depth != 2 { + t.Errorf("doc %s depth = %d, want 2", d.ChunkID, d.Depth) + } + } +} + +func TestMaybeSplit_BelowCaps(t *testing.T) { + tree := newNavTree() + parent := &navCluster{Name: "P", Desc: "p", Parent: "root", Depth: 0, Vector: []float32{1, 0}} + tree.addCluster(parent) + for i := 0; i < 3; i++ { + tree.addDoc(parent, fmt.Sprintf("d%d", i), "t", []float32{1, 0}, 1) + } + if err := tree.maybeSplit(context.Background(), common.Deps{}, "llm", "P"); err != nil { + t.Fatal(err) + } + if len(tree.order) != 1 { + t.Fatalf("below caps must not split: %v", tree.order) + } +} + +func TestRun_EndToEnd(t *testing.T) { + vecs := map[string][]float32{ + "alpha one": {1, 0}, + "alpha two": {0.99, 0.1}, + "beta one": {0, 1}, + } + deps := common.Deps{Chat: &navFakeChat{}, Embed: scriptedEmbedder{vecs: vecs}, TenantID: "t1", DatasetID: "ds1"} + p := common.Param{}.Defaults() + p.Variant = common.VariantDatasetnav + inputs := common.Inputs{ + DocID: "d1", + Chunks: []common.Chunk{ + {ID: "c1", Text: "alpha one"}, + {ID: "c2", Text: "alpha two"}, + {ID: "c3", Text: "beta one"}, + }, + } + out, err := Run(context.Background(), deps, p, inputs) + if err != nil { + t.Fatalf("Run: %v", err) + } + var clusters, docs, roots int + parentByID := map[string]string{} + for _, p := range out.Products { + switch p.Meta["kind"] { + case "nav_cluster": + clusters++ + parentByID[p.Meta["name"].(string)] = p.Meta["parent_name"].(string) + case "nav_doc": + docs++ + if p.ParentID == "" { + t.Errorf("nav_doc must link to its cluster product") + } + case "root": + roots++ + } + } + if clusters != 2 || docs != 3 || roots != 1 { + t.Fatalf("products = %d clusters + %d docs + %d roots, want 2+3+1", clusters, docs, roots) + } +} + +func TestLLMMergeFallbacks(t *testing.T) { + // Non-JSON text reply is used verbatim (mirrors gen_json's str path). + deps := common.Deps{Chat: chatFunc(func(req common.ChatRequest) (*common.ChatResponse, error) { + return &common.ChatResponse{Content: "merged text"}, nil + })} + got, err := llmMerge(context.Background(), deps, "llm", "old", "new") + if err != nil || got != "merged text" { + t.Fatalf("llmMerge text path = %q, %v", got, err) + } + // Empty reply keeps the old description. + deps2 := common.Deps{Chat: chatFunc(func(req common.ChatRequest) (*common.ChatResponse, error) { + return &common.ChatResponse{Content: ""}, nil + })} + got, _ = llmMerge(context.Background(), deps2, "llm", "old", "new") + if got != "old" { + t.Fatalf("empty reply must keep old desc, got %q", got) + } + // Nil chat keeps the old description without a call. + got, _ = llmMerge(context.Background(), common.Deps{}, "llm", "old", "new") + if got != "old" { + t.Fatalf("nil chat must keep old desc, got %q", got) + } +} + +type chatFunc func(common.ChatRequest) (*common.ChatResponse, error) + +func (f chatFunc) Chat(_ context.Context, req common.ChatRequest) (*common.ChatResponse, error) { + return f(req) +} diff --git a/internal/ingestion/component/knowledge_compiler/golden/corpus.go b/internal/ingestion/component/knowledge_compiler/golden/corpus.go new file mode 100644 index 0000000000..66d18bfbb6 --- /dev/null +++ b/internal/ingestion/component/knowledge_compiler/golden/corpus.go @@ -0,0 +1,55 @@ +// Package golden holds the canonical, committed sample corpus and the +// structural-metric helpers used by the knowledge_compiler golden gates +// (PORT_PLAN.md 缺口 C / 缺口 E). +// +// The corpus is intentionally small and stable so that Go-side product counts +// and clustering structure can be LOCKED as a required CI gate. When a live +// Python baseline is available, the same corpus is fed to +// rag/advanced_rag/knowlege_compile to produce the reference numbers recorded +// beside the fixtures in this package. +package golden + +import ( + "fmt" + + "ragflow/internal/ingestion/component/knowledge_compiler/common" +) + +// FixedCorpus is the canonical sample document. It describes a small, coherent +// ecosystem of named entities (Alpha/Beta/Gamma/Delta/Epsilon) with dense +// cross-references so structure extraction, wiki synthesis, and raptor +// clustering all have meaningful signal. +func FixedCorpus() []common.Chunk { + texts := []string{ + "Alpha is a large language model developed by the Beta research lab.", + "Beta lab also created Gamma, a retrieval engine for knowledge bases.", + "Gamma indexes documents and returns the most relevant passages.", + "Delta is a vector database that stores embeddings produced by Alpha.", + "Epsilon is an evaluation benchmark comparing retrieval systems.", + "Alpha and Delta integrate to power semantic search over documents.", + "Gamma and Epsilon both measure retrieval quality on the same dataset.", + "Beta lab published a survey covering Alpha, Gamma, and Delta.", + "The knowledge graph links Alpha, Beta, Gamma, Delta, and Epsilon.", + "RAG pipelines combine Alpha with Gamma to answer user questions.", + "Embeddings from Alpha are stored in Delta for fast nearest-neighbor search.", + "Epsilon scores show Gamma outperforms naive retrieval on Delta indices.", + } + chunks := make([]common.Chunk, len(texts)) + for i, t := range texts { + chunks[i] = common.Chunk{ + ID: fmt.Sprintf("chunk-%02d", i+1), + Text: t, + } + } + return chunks +} + +// ChunksToAny converts the fixed corpus (or any chunk list) into the runtime +// inputs shape expected by KnowledgeCompilerComponent.Invoke. +func ChunksToAny(chunks []common.Chunk) []any { + out := make([]any, len(chunks)) + for i, c := range chunks { + out[i] = map[string]any{"id": c.ID, "text": c.Text} + } + return out +} diff --git a/internal/ingestion/component/knowledge_compiler/golden/fixtures/raptor_baseline.json b/internal/ingestion/component/knowledge_compiler/golden/fixtures/raptor_baseline.json new file mode 100644 index 0000000000..472c0c3fdd --- /dev/null +++ b/internal/ingestion/component/knowledge_compiler/golden/fixtures/raptor_baseline.json @@ -0,0 +1,12 @@ +{ + "note": "Locked Go-side structural gate for the raptor variant (PORT_PLAN.md 缺口 C). Leaf-cluster counts and tree height are deterministic for the fixed corpus + deterministic signed embedder (see golden_test.go). The exact PSI/AHC leaf-cluster counts below are the regression contract; reconcile NMI/ARI with the Python baseline (rag/advanced_rag/knowlege_compile/raptor.py) when a live run is available.", + "psi_leaf_clusters_min": 4, + "psi_leaf_clusters_max": 4, + "ahc_leaf_clusters_min": 2, + "ahc_leaf_clusters_max": 2, + "gmm_leaf_clusters_min": 6, + "gmm_leaf_clusters_max": 6, + "tree_height_max": 6, + "leaf_coverage_min": 1.0, + "max_depth_max": 6 +} diff --git a/internal/ingestion/component/knowledge_compiler/golden/fixtures/structure_wiki_baseline.json b/internal/ingestion/component/knowledge_compiler/golden/fixtures/structure_wiki_baseline.json new file mode 100644 index 0000000000..eff2c2ac19 --- /dev/null +++ b/internal/ingestion/component/knowledge_compiler/golden/fixtures/structure_wiki_baseline.json @@ -0,0 +1,7 @@ +{ + "note": "Locked Go-side product-count gate for structure/wiki (PORT_PLAN.md 缺口 E). These counts are the golden contract for the fixed corpus with the deterministic mock LLM/embedder. They must be reconciled against the Python baseline (rag/advanced_rag/knowlege_compile/structure.py, wiki.py) when a live run is available; a change here is a deliberate, reviewed contract update. 2026-07-27: structure gate now feeds a graph-kind parser_config (hypergraph path): 12 corpus chunks pack into one batch, the mock extracts 5 entities (Alpha..Epsilon) + 6 relations (fixed Greek pairs) + 1 graph summary = 12 (previously 6 under the legacy tuples contract).", + "schema_version": 1, + "fixed_corpus_chunks": 12, + "structure_products": 12, + "wiki_products": 1 +} diff --git a/internal/ingestion/component/knowledge_compiler/golden/golden_test.go b/internal/ingestion/component/knowledge_compiler/golden/golden_test.go new file mode 100644 index 0000000000..110f3e768b --- /dev/null +++ b/internal/ingestion/component/knowledge_compiler/golden/golden_test.go @@ -0,0 +1,140 @@ +package golden + +import ( + "encoding/json" + "testing" + + "ragflow/internal/ingestion/component/schema" +) + +func TestFixedCorpus_IsStable(t *testing.T) { + first := FixedCorpus() + if len(first) != 12 { + t.Fatalf("FixedCorpus len = %d, want 12 (locked in fixtures)", len(first)) + } + second := FixedCorpus() + if len(first) != len(second) { + t.Fatalf("FixedCorpus not stable: %d vs %d", len(first), len(second)) + } + for i := range first { + if first[i].ID != second[i].ID || first[i].Text != second[i].Text { + t.Errorf("FixedCorpus[%d] differs across calls", i) + break + } + } +} + +func TestChunksToAny(t *testing.T) { + chunks := FixedCorpus() + raw := ChunksToAny(chunks) + if len(raw) != len(chunks) { + t.Fatalf("ChunksToAny len = %d, want %d", len(raw), len(chunks)) + } + for i, item := range raw { + m, ok := item.(map[string]any) + if !ok { + t.Fatalf("ChunksToAny[%d] = %T, want map[string]any", i, item) + } + if m["id"] != chunks[i].ID { + t.Errorf("ChunksToAny[%d].id = %v, want %q", i, m["id"], chunks[i].ID) + } + if m["text"] != chunks[i].Text { + t.Errorf("ChunksToAny[%d].text mismatch", i) + } + } +} + +func TestAnalyzeRaptorProducts_TreeShape(t *testing.T) { + vector := json.RawMessage(`[0.1,0.2,0.3]`) + chunks := []schema.ChunkDoc{ + {Text: "root summary", Extra: mustExtras(t, map[string]any{ + "id": "r1", "doc_id": "d1", "tenant_id": "t1", "compile_kwd": "raptor", + "kc_kind": "root", "kc_level": float64(-1), "q_3_vec": vector, + })}, + {Text: "leaf A", Extra: mustExtras(t, map[string]any{ + "id": "a1", "doc_id": "d1", "tenant_id": "t1", "compile_kwd": "raptor", + "kc_kind": "summary", "kc_level": float64(0), "parent_kwd": "r1", "q_3_vec": vector, + })}, + {Text: "leaf B", Extra: mustExtras(t, map[string]any{ + "id": "b1", "doc_id": "d1", "tenant_id": "t1", "compile_kwd": "raptor", + "kc_kind": "summary", "kc_level": float64(0), "parent_kwd": "r1", "q_3_vec": vector, + })}, + {Text: "mid A", Extra: mustExtras(t, map[string]any{ + "id": "m1", "doc_id": "d1", "tenant_id": "t1", "compile_kwd": "raptor", + "kc_kind": "summary", "kc_level": float64(1), "parent_kwd": "r1", "q_3_vec": vector, + })}, + } + m := AnalyzeRaptorProducts(chunks) + if m.RootCount != 1 { + t.Errorf("RootCount = %d, want 1", m.RootCount) + } + if m.LeafClusters != 2 { + t.Errorf("LeafClusters = %d, want 2 (a1+b1)", m.LeafClusters) + } + if m.MaxDepth != 2 { + t.Errorf("MaxDepth = %d, want 2 (max level 1 + 1)", m.MaxDepth) + } + if !m.AllParented { + t.Error("AllParented = false, want true (every node's parent is in the set)") + } + if !m.VectorOK || !m.SchemaOK { + t.Errorf("VectorOK=%v SchemaOK=%v, want both true", m.VectorOK, m.SchemaOK) + } +} + +func TestAnalyzeRaptorProducts_DetectsDanglingParent(t *testing.T) { + chunks := []schema.ChunkDoc{ + {Text: "root", Extra: mustExtras(t, map[string]any{ + "id": "r1", "doc_id": "d1", "tenant_id": "t1", "compile_kwd": "raptor", + "kc_kind": "root", + })}, + {Text: "orphan", Extra: mustExtras(t, map[string]any{ + "id": "o1", "doc_id": "d1", "tenant_id": "t1", "compile_kwd": "raptor", + "kc_kind": "summary", "kc_level": float64(0), "parent_kwd": "missing-parent", + })}, + } + m := AnalyzeRaptorProducts(chunks) + if m.AllParented { + t.Error("AllParented = true, want false (o1's parent is missing)") + } +} + +func TestCoverageFraction_AllParentedOneRoot(t *testing.T) { + m := TreeMetrics{RootCount: 1, LeafClusters: 3, AllParented: true} + if cov := m.CoverageFraction(12); cov != 1.0 { + t.Errorf("CoverageFraction = %v, want 1.0", cov) + } +} + +func TestCoverageFraction_NoRootIsZero(t *testing.T) { + m := TreeMetrics{LeafClusters: 3, AllParented: true} + if cov := m.CoverageFraction(12); cov != 0.0 { + t.Errorf("CoverageFraction = %v, want 0.0 (no root)", cov) + } +} + +func TestExtraFloat_Roundtrip(t *testing.T) { + doc := schema.ChunkDoc{Extra: mustExtras(t, map[string]any{"kc_level": float64(3)})} + if v, ok := extraFloat(doc, "kc_level"); !ok || v != 3 { + t.Errorf("extraFloat = %v/%v, want 3/true", v, ok) + } + if _, ok := extraFloat(doc, "missing"); ok { + t.Error("extraFloat(missing) = ok, want false") + } +} + +// mustExtras builds a schema.Extra json.RawMessage map from a plain map. The +// test fixtures only need a stable shape; this keeps the golden_test +// independent from the production ChunkDoc constructor details. +func mustExtras(t *testing.T, in map[string]any) map[string]json.RawMessage { + t.Helper() + out := make(map[string]json.RawMessage, len(in)) + for k, v := range in { + b, err := json.Marshal(v) + if err != nil { + t.Fatalf("marshal %s: %v", k, err) + } + out[k] = b + } + return out +} diff --git a/internal/ingestion/component/knowledge_compiler/golden/metrics.go b/internal/ingestion/component/knowledge_compiler/golden/metrics.go new file mode 100644 index 0000000000..4ff03e917c --- /dev/null +++ b/internal/ingestion/component/knowledge_compiler/golden/metrics.go @@ -0,0 +1,113 @@ +package golden + +import ( + "encoding/json" + "strings" + + "ragflow/internal/ingestion/component/schema" +) + +// TreeMetrics summarizes the structural shape of a raptor product tree. +type TreeMetrics struct { + ProductCount int + RootCount int + LeafClusters int // number of level-0 summary nodes (bottom clusters) + MaxDepth int // root(0) .. deepest summary level + AllParented bool + VectorOK bool // every product carries a non-empty vector + SchemaOK bool // every product carries the schema fields +} + +// AnalyzeRaptorProducts validates tree integrity and computes structural +// metrics from a flat chunk list (the compiled raptor output, expressed as +// schema.ChunkDoc values). Used by the 缺口 C golden gate. +func AnalyzeRaptorProducts(chunks []schema.ChunkDoc) TreeMetrics { + ids := make(map[string]bool, len(chunks)) + for _, c := range chunks { + if id, ok := c.GetExtraString("id"); ok { + ids[id] = true + } + } + m := TreeMetrics{ProductCount: len(chunks), AllParented: true, VectorOK: true, SchemaOK: true} + maxLevel := -1 + for _, c := range chunks { + kind, _ := c.GetExtraString("kc_kind") + level := 0 + if lf, ok := extraFloat(c, "kc_level"); ok { + level = int(lf) + } + switch kind { + case "root": + m.RootCount++ + case "summary": + if level == 0 { + m.LeafClusters++ + } + if level > maxLevel { + maxLevel = level + } + } + parent, _ := c.GetExtraString("parent_kwd") + if kind != "root" && parent == "" { + m.AllParented = false + } + if parent != "" && !ids[parent] { + m.AllParented = false + } + if !hasVector(c) { + m.VectorOK = false + } + id, _ := c.GetExtraString("id") + docID, _ := c.GetExtraString("doc_id") + tenant, _ := c.GetExtraString("tenant_id") + ck, _ := c.GetExtraString("compile_kwd") + if id == "" || docID == "" || tenant == "" || c.Text == "" || ck == "" { + m.SchemaOK = false + } + } + m.MaxDepth = maxLevel + 1 + return m +} + +// CoverageFraction reports how completely the input chunks are represented by +// the raptor tree. Every source chunk is assigned to exactly one level-0 +// cluster in buildTree, and each such cluster becomes a leaf summary node, so a +// well-formed tree covers 100% of chunks. nChunks is the input chunk count. +func (m TreeMetrics) CoverageFraction(nChunks int) float64 { + if nChunks <= 0 { + return 0 + } + // Every chunk maps to a level-0 cluster; the number of covered chunks + // equals the total chunk count when at least one leaf cluster exists and + // the tree is well-formed (all parented, single root). + if m.LeafClusters > 0 && m.RootCount == 1 && m.AllParented { + return 1.0 + } + return 0.0 +} + +// extraFloat reads a numeric Extra value by key. +func extraFloat(c schema.ChunkDoc, key string) (float64, bool) { + if c.Extra == nil { + return 0, false + } + raw, ok := c.Extra[key] + if !ok { + return 0, false + } + var f float64 + if err := json.Unmarshal(raw, &f); err != nil { + return 0, false + } + return f, true +} + +// hasVector reports whether the chunk carries any q_<dim>_vec embedding. +func hasVector(c schema.ChunkDoc) bool { + for k := range c.Extra { + if strings.HasPrefix(k, "q_") && strings.HasSuffix(k, "_vec") { + return true + } + } + return false +} diff --git a/internal/ingestion/component/knowledge_compiler/golden_test.go b/internal/ingestion/component/knowledge_compiler/golden_test.go new file mode 100644 index 0000000000..75df1a964e --- /dev/null +++ b/internal/ingestion/component/knowledge_compiler/golden_test.go @@ -0,0 +1,287 @@ +package knowledge_compiler + +import ( + "context" + "encoding/json" + "hash/fnv" + "os" + "path/filepath" + "strings" + "testing" + + "ragflow/internal/ingestion/component/knowledge_compiler/common" + "ragflow/internal/ingestion/component/knowledge_compiler/golden" + "ragflow/internal/ingestion/component/schema" +) + +// chunkHasVector reports whether a compiled chunk carries any q_<dim>_vec embedding. +func chunkHasVector(c schema.ChunkDoc) bool { + for k := range c.Extra { + if strings.HasPrefix(k, "q_") && strings.HasSuffix(k, "_vec") { + return true + } + } + return false +} + +// signedEmbedder is a deterministic embedder whose vectors span both signs. +// The default mockEmbedder emits non-negative vectors in [0,1], which makes +// every pairwise cosine >= 0; under PSI's threshold=0 merge rule that collapses +// the whole corpus into a single cluster and triggers unbounded recursion in +// buildTree. Signed vectors give the clustering real signal so the golden +// corpus separates into multiple small clusters (the production path uses real +// centered embeddings and does not hit this). +type signedEmbedder struct{ dim int } + +func (s signedEmbedder) Dimensions() int { return s.dim } +func (s signedEmbedder) Encode(_ context.Context, texts []string) ([][]float32, error) { + out := make([][]float32, len(texts)) + for i, t := range texts { + out[i] = signedVec(t, s.dim) + } + return out, nil +} +func signedVec(str string, dim int) []float32 { + h := fnv.New32a() + _, _ = h.Write([]byte(str)) + seed := h.Sum32() + v := make([]float32, dim) + for j := 0; j < dim; j++ { + x := int32(seed>>uint(j*3)) % 100 + v[j] = float32(x)/50.0 - 1.0 + } + return v +} + +// installSignedProseDeps wires a prose LLM + signed deterministic embedder. +func installSignedProseDeps(t *testing.T) { + t.Helper() + common.SetDepsResolver(func(tenantID, llmID, embeddingModel string) (common.Deps, error) { + return common.Deps{Chat: proseChat{}, Embed: signedEmbedder{dim: 8}, TenantID: tenantID}, nil + }) + t.Cleanup(func() { common.SetDepsResolver(nil) }) +} + +// runVariantChunks builds the component for variant, runs it over the fixed +// golden corpus, and returns the compiled chunks (schema.ChunkDoc) that were +// merged into the output chunks list. +func runVariantChunks(t *testing.T, variant string, extra map[string]any) []schema.ChunkDoc { + t.Helper() + return runVariantChunksWithInputs(t, variant, extra, nil) +} + +// runVariantChunksWithInputs is runVariantChunks plus extra Invoke-level +// inputs (e.g. parser_config for the structure variant's template shape). +func runVariantChunksWithInputs(t *testing.T, variant string, extra, inputsExtra map[string]any) []schema.ChunkDoc { + t.Helper() + params := map[string]any{"variant": variant, "llm_id": "llm1", "embedding_model": "emb1"} + for k, v := range extra { + params[k] = v + } + c, err := NewKnowledgeCompilerComponent("KnowledgeCompiler", params) + if err != nil { + t.Fatalf("NewKnowledgeCompilerComponent(%s): %v", variant, err) + } + invokeInputs := map[string]any{ + "chunks": golden.ChunksToAny(golden.FixedCorpus()), + "doc_id": "d1", + "tenant_id": "t1", + } + for k, v := range inputsExtra { + invokeInputs[k] = v + } + out, err := c.Invoke(context.Background(), nil, invokeInputs) + if err != nil { + t.Fatalf("Invoke(%s): %v", variant, err) + } + raw, ok := out["chunks"].([]any) + if !ok { + t.Fatalf("Invoke(%s): chunks = %T", variant, out["chunks"]) + } + docs, ok, err := schema.ChunkDocsFromAny(raw) + if !ok || err != nil { + t.Fatalf("Invoke(%s): ChunkDocsFromAny = %v err=%v", variant, ok, err) + } + // The component merges the compiled knowledge units into the upstream input + // chunks; the golden gates only analyze the compiled units (those carrying + // the compile_kwd discriminator), mirroring the old out["products"] surface. + compiled := make([]schema.ChunkDoc, 0, len(docs)) + for _, d := range docs { + if _, has := d.GetExtraString("compile_kwd"); has { + compiled = append(compiled, d) + } + } + return compiled +} + +func loadBaseline(t *testing.T, name string) map[string]any { + t.Helper() + b, err := os.ReadFile(filepath.Join("golden", "fixtures", name)) + if err != nil { + t.Fatalf("read baseline %s: %v", name, err) + } + var m map[string]any + if err := json.Unmarshal(b, &m); err != nil { + t.Fatalf("unmarshal baseline %s: %v", name, err) + } + return m +} + +func num(m map[string]any, key string) int { + switch v := m[key].(type) { + case float64: + return int(v) + case int: + return v + } + return 0 +} + +// TestGolden_Structure_ProductCount is the 缺口 E gate for the structure +// variant: on the fixed corpus with the deterministic mock LLM, the Go product +// count must equal the locked golden baseline. The gate feeds a graph-kind +// parser_config so the hypergraph path (two-stage entity → relation +// extraction + LLM merge dedup) is what gets locked. +func TestGolden_Structure_ProductCount(t *testing.T) { + installMockDeps(t) + prods := runVariantChunksWithInputs(t, "structure", nil, map[string]any{ + "parser_config": graphParserConfig(), + }) + + // Schema invariants: every product carries the schema fields and a vector. + for _, p := range prods { + id, _ := p.GetExtraString("id") + docID, _ := p.GetExtraString("doc_id") + tenant, _ := p.GetExtraString("tenant_id") + if id == "" || docID == "" || tenant == "" || p.Text == "" { + t.Fatalf("structure: product missing schema fields: %+v", p) + } + if !chunkHasVector(p) { + t.Fatalf("structure: product %s missing vector", id) + } + } + + // Exactly one graph product summarizing {nodes, edges}. + graphCount, nodeCount, edgeCount := 0, 0, 0 + for _, p := range prods { + kind, _ := p.GetExtraString("kc_kind") + switch kind { + case "graph": + graphCount++ + case "entity": + nodeCount++ + case "relation": + edgeCount++ + } + } + if graphCount != 1 { + t.Fatalf("structure: graphCount = %d, want 1", graphCount) + } + if nodeCount == 0 || edgeCount == 0 { + t.Fatalf("structure: expected entities+relations, got nodes=%d edges=%d", nodeCount, edgeCount) + } + + baseline := loadBaseline(t, "structure_wiki_baseline.json") + t.Logf("structure: %d products", len(prods)) + want := num(baseline, "structure_products") + if want != 0 && len(prods) != want { + t.Errorf("structure product count = %d, golden baseline = %d (review contract if intentional)", len(prods), want) + } +} + +// TestGolden_Wiki_ProductCount is the 缺口 E gate for the wiki variant. +func TestGolden_Wiki_ProductCount(t *testing.T) { + installProseDeps(t) + prods := runVariantChunks(t, "wiki", nil) + + for _, p := range prods { + id, _ := p.GetExtraString("id") + docID, _ := p.GetExtraString("doc_id") + tenant, _ := p.GetExtraString("tenant_id") + if id == "" || docID == "" || tenant == "" || p.Text == "" { + t.Fatalf("wiki: product missing schema fields: %+v", p) + } + if !chunkHasVector(p) { + t.Fatalf("wiki: product %s missing vector", id) + } + } + + foundPage := false + for _, p := range prods { + if kind, _ := p.GetExtraString("kc_kind"); kind == "page" { + foundPage = true + } + } + if !foundPage { + t.Fatalf("wiki: no 'page' product; got %d products", len(prods)) + } + + baseline := loadBaseline(t, "structure_wiki_baseline.json") + t.Logf("wiki: %d products", len(prods)) + want := num(baseline, "wiki_products") + if want != 0 && len(prods) != want { + t.Errorf("wiki product count = %d, golden baseline = %d (review contract if intentional)", len(prods), want) + } +} + +// TestGolden_Raptor_Structure is the 缺口 C gate for the raptor variant: on the +// fixed corpus, both the PSI and AHC builders must produce a well-formed tree +// (single root, fully parented, vectors + schema intact, full chunk coverage, +// bounded depth/cluster count). The exact NMI/ARI reconciliation with the +// Python baseline is tracked separately; this gate locks the structural +// contract that a regression would violate. +func TestGolden_Raptor_Structure(t *testing.T) { + installSignedProseDeps(t) + baseline := loadBaseline(t, "raptor_baseline.json") + nChunks := len(golden.FixedCorpus()) + + cases := []struct { + method string + minKey string + maxKey string + }{ + {"PSI", "psi_leaf_clusters_min", "psi_leaf_clusters_max"}, + {"AHC", "ahc_leaf_clusters_min", "ahc_leaf_clusters_max"}, + {"GMM", "gmm_leaf_clusters_min", "gmm_leaf_clusters_max"}, + } + for _, tc := range cases { + prods := runVariantChunks(t, "raptor", map[string]any{ + "extra": map[string]any{"clustering_method": tc.method}, + }) + m := golden.AnalyzeRaptorProducts(prods) + t.Logf("raptor(%s): products=%d root=%d leafClusters=%d maxDepth=%d coverage=%.2f", + tc.method, m.ProductCount, m.RootCount, m.LeafClusters, m.MaxDepth, m.CoverageFraction(nChunks)) + + if m.RootCount != 1 { + t.Errorf("raptor(%s): rootCount=%d, want 1", tc.method, m.RootCount) + } + if !m.AllParented { + t.Errorf("raptor(%s): tree has dangling parent_id references", tc.method) + } + if !m.VectorOK { + t.Errorf("raptor(%s): some products missing vectors", tc.method) + } + if !m.SchemaOK { + t.Errorf("raptor(%s): some products missing schema fields", tc.method) + } + if cov := m.CoverageFraction(nChunks); cov < numFloat(baseline, "leaf_coverage_min") { + t.Errorf("raptor(%s): leaf coverage %.2f < baseline %.2f", tc.method, cov, numFloat(baseline, "leaf_coverage_min")) + } + if m.LeafClusters < num(baseline, tc.minKey) || m.LeafClusters > num(baseline, tc.maxKey) { + t.Errorf("raptor(%s): leafClusters=%d outside [%d,%d]", tc.method, m.LeafClusters, num(baseline, tc.minKey), num(baseline, tc.maxKey)) + } + if m.MaxDepth > num(baseline, "max_depth_max") { + t.Errorf("raptor(%s): maxDepth=%d > baseline %d", tc.method, m.MaxDepth, num(baseline, "max_depth_max")) + } + } +} + +func numFloat(m map[string]any, key string) float64 { + switch v := m[key].(type) { + case float64: + return v + case int: + return float64(v) + } + return 0 +} diff --git a/internal/ingestion/component/knowledge_compiler/mindmap/dictify.go b/internal/ingestion/component/knowledge_compiler/mindmap/dictify.go new file mode 100644 index 0000000000..5c55266319 --- /dev/null +++ b/internal/ingestion/component/knowledge_compiler/mindmap/dictify.go @@ -0,0 +1,417 @@ +package mindmap + +import ( + "regexp" + "strings" + + "github.com/yuin/goldmark" + "github.com/yuin/goldmark/ast" + "github.com/yuin/goldmark/text" +) + +// Markdown-outline → mind-map-tree conversion, mirroring the Python pipeline: +// markdown_to_json.dictify (CommonMark nest + Renderer) → _todict/_list_to_kv +// → _merge (multi-batch) → final shaping + _be_children. Node ids keep the +// raw markdown text (asterisks included) until _key strips them, exactly like +// Python. + +// kv is one ordered key/value pair; omap mirrors Python's OrderedDict (key +// order is load-bearing for tree shape and merge order). +type kv struct { + k string + v any +} +type omap []kv + +// Node is one shaped mind-map node (Python's {"id", "children"} shape). +type Node struct { + ID string + Children []*Node +} + +type blockType int + +const ( + headingBlock blockType = iota + listBlock + paraBlock + fencedBlock + otherBlock +) + +// mdBlock is one markdown block in dictify terms: heading (level+text), list +// (flattened items), paragraph/fenced/other (text). +type mdBlock struct { + typ blockType + level int + text string + items []any +} + +// fenceStripRE mirrors re.sub(r"```[^\n]*", "", response): whole fence marker +// lines (``` or ```lang) are removed before parsing. +var fenceStripRE = regexp.MustCompile("```[^\\n]*") + +// StripFences removes markdown code-fence marker lines from an LLM reply. +func StripFences(s string) string { + return fenceStripRE.ReplaceAllString(s, "") +} + +// dictify mirrors markdown_to_json.dictify: nest top-level blocks by the +// MINIMUM top heading level; content before the first heading at each level +// is dropped; when a level has no headings the blocks become a list (the top +// level then renders as {"root": [...]}). +func dictify(md string) omap { + src := []byte(md) + doc := goldmark.DefaultParser().Parse(text.NewReader(src)) + var blocks []mdBlock + for n := doc.FirstChild(); n != nil; n = n.NextSibling() { + blocks = append(blocks, toBlock(n, src)) + } + if len(blocks) == 0 { + return omap{} + } + min := 100000 + for _, b := range blocks { + lvl := 100000 + if b.typ == headingBlock { + lvl = b.level + } + if lvl < min { + min = lvl + } + } + switch t := dictifyBlocks(blocks, min).(type) { + case omap: + out := omap{} + for _, e := range t { + out = append(out, kv{e.k, valuify(e.v)}) + } + return out + case []mdBlock: + // No headings anywhere: {"root": [rendered blocks]} (Renderer's + // non-dict branch). + rendered := make([]any, 0, len(t)) + for _, b := range t { + rendered = append(rendered, renderBlockAny(b)) + } + return omap{{"root", rendered}} + } + return omap{} +} + +// dictifyBlocks mirrors CMarkASTNester._dictify_blocks: split blocks by +// headings at exactly `level`, recursing at level+1. Returns omap or the raw +// []mdBlock (no headings at this level). +func dictifyBlocks(blocks []mdBlock, level int) any { + has := false + for _, b := range blocks { + if b.typ == headingBlock && b.level == level { + has = true + break + } + } + if !has { + return blocks + } + out := omap{} + var cur string + var children []mdBlock + started := false + for _, b := range blocks { + if b.typ == headingBlock && b.level == level { + if started { + out = setOMap(out, cur, children) + } + cur = b.text + children = nil + started = true + continue + } + if started { + children = append(children, b) + } + } + if started { + out = setOMap(out, cur, children) + } + for i := range out { + out[i].v = dictifyBlocks(out[i].v.([]mdBlock), level+1) + } + return out +} + +// valuify mirrors Renderer._valuify: dict → recurse; empty → ""; first child +// a list → its items; otherwise the blocks joined with "\n\n". +func valuify(v any) any { + switch t := v.(type) { + case omap: + out := omap{} + for _, e := range t { + out = append(out, kv{e.k, valuify(e.v)}) + } + return out + case []mdBlock: + if len(t) == 0 { + return "" + } + if t[0].typ == listBlock { + return t[0].items + } + parts := make([]string, 0, len(t)) + for _, b := range t { + parts = append(parts, renderBlockText(b)) + } + return strings.Join(parts, "\n\n") + } + return "" +} + +// renderBlockText renders one block to a string (Renderer._render_*). +// A list appearing mid-children is rendered as its flattened item texts +// (Python str()'s the nested list — a pathological case we render readably). +func renderBlockText(b mdBlock) string { + switch b.typ { + case fencedBlock: + return "```\n" + b.text + "```" + case listBlock: + var parts []string + for _, item := range b.items { + if s, ok := item.(string); ok { + parts = append(parts, s) + } + } + return strings.Join(parts, "\n") + default: + return b.text + } +} + +func renderBlockAny(b mdBlock) any { + if b.typ == listBlock { + return b.items + } + return renderBlockText(b) +} + +// todict mirrors _todict + _list_to_kv: lists under dict keys convert to +// key→definition pairs ([prev scalar] → sublist[0]); a list with no +// definition sub-lists collapses to an empty dict (the bullets are dropped — +// faithful to Python). Recurses into nested dicts. +func todict(m omap) omap { + for i := range m { + switch val := m[i].v.(type) { + case omap: + m[i].v = todict(val) + case []any: + nv := omap{} + for j, item := range val { + lst, isList := item.([]any) + if !isList || j == 0 { + continue + } + key, _ := val[j-1].(string) + var first any + if len(lst) > 0 { + first = lst[0] + } + nv = append(nv, kv{key, first}) + } + m[i].v = nv + } + } + return m +} + +// mergeDicts mirrors _merge: d1 merges INTO d2 (ordered). Both dicts recurse; +// both lists concatenate d2's then d1's items; otherwise d1 wins. Keys absent +// from d2 append in d1's order. +func mergeDicts(d1, d2 omap) omap { + for _, e := range d1 { + idx := indexOMap(d2, e.k) + if idx < 0 { + d2 = append(d2, e) + continue + } + dv := d2[idx].v + d1Dict, d1IsDict := e.v.(omap) + d2Dict, d2IsDict := dv.(omap) + d1List, d1IsList := e.v.([]any) + d2List, d2IsList := dv.([]any) + switch { + case d1IsDict && d2IsDict: + d2[idx].v = mergeDicts(d1Dict, d2Dict) + case d1IsList && d2IsList: + d2[idx].v = append(d2List, d1List...) + default: + d2[idx].v = e.v + } + } + return d2 +} + +// shapeTree mirrors __call__'s final shaping: >1 top-level keys → synthetic +// "root" node whose children are the dict-valued keys (with a keyset seeded +// by those keys); exactly 1 → that key becomes the root id. An empty merged +// dict degrades to a bare root (Python would raise IndexError on +// keys()[0]; the Go port keeps the pipeline alive). +func shapeTree(merged omap) *Node { + if len(merged) > 1 { + keyset := map[string]bool{} + for _, e := range merged { + if _, ok := e.v.(omap); !ok { + continue + } + if k := stripStars(e.k); k != "" { + keyset[k] = true + } + } + var children []*Node + for _, e := range merged { + if _, ok := e.v.(omap); !ok { + continue + } + k := stripStars(e.k) + if k == "" { + continue + } + children = append(children, &Node{ID: k, Children: beChildren(e.v, keyset)}) + } + return &Node{ID: "root", Children: children} + } + if len(merged) == 1 { + k := stripStars(merged[0].k) + return &Node{ID: k, Children: beChildren(merged[0].v, map[string]bool{k: true})} + } + return &Node{ID: "root"} +} + +// beChildren mirrors _be_children: strings/lists become leaf nodes; dict keys +// become nested nodes. keyset is SHARED across the whole subtree walk — a key +// already seen anywhere below is skipped (silent dedup, as in Python). +func beChildren(v any, keyset map[string]bool) []*Node { + switch t := v.(type) { + case string: + keyset[t] = true + if id := stripStars(t); id != "" { + return []*Node{{ID: id}} + } + return nil + case []any: + for _, item := range t { + if s, ok := item.(string); ok { + keyset[s] = true + } + } + var out []*Node + for _, item := range t { + s, ok := item.(string) + if !ok { + continue + } + if id := stripStars(s); id != "" { + out = append(out, &Node{ID: id}) + } + } + return out + case omap: + var out []*Node + for _, e := range t { + k := stripStars(e.k) + if k == "" || keyset[k] { + continue + } + keyset[k] = true + out = append(out, &Node{ID: k, Children: beChildren(e.v, keyset)}) + } + return out + } + return nil +} + +// starsRE mirrors re.compile(r"\*+") used by _key. +var starsRE = regexp.MustCompile(`\*+`) + +// stripStars mirrors _key (re.sub(r"\*+", "", k)). +func stripStars(s string) string { + return starsRE.ReplaceAllString(s, "") +} + +// ---- goldmark AST → mdBlock ---- + +func toBlock(n ast.Node, src []byte) mdBlock { + switch node := n.(type) { + case *ast.Heading: + return mdBlock{typ: headingBlock, level: node.Level, text: headingRawText(node, src)} + case *ast.List: + return mdBlock{typ: listBlock, items: renderList(node, src)} + case *ast.FencedCodeBlock: + return mdBlock{typ: fencedBlock, text: rawLines(n, src, "")} + case *ast.Paragraph: + return mdBlock{typ: paraBlock, text: rawLines(n, src, "\n")} + default: + return mdBlock{typ: otherBlock, text: rawLines(n, src, "\n")} + } +} + +// headingRawText renders the heading's raw source line without the ATX +// markers, keeping inline markdown (e.g. asterisks) intact like CommonMark's +// block.strings (goldmark's inline text would strip emphasis). +func headingRawText(n *ast.Heading, src []byte) string { + raw := strings.TrimSpace(rawLines(n, src, "\n")) + raw = strings.TrimLeft(raw, "#") + return strings.TrimSpace(raw) +} + +// rawLines joins a node's source lines (sep between lines). For fenced code +// the lines already carry their newlines (sep=""). +func rawLines(n ast.Node, src []byte, sep string) string { + lines := n.Lines() + parts := make([]string, 0, lines.Len()) + for i := 0; i < lines.Len(); i++ { + seg := lines.At(i) + v := string(seg.Value(src)) + if sep != "" { + v = strings.TrimRight(v, "\n") + } + parts = append(parts, v) + } + return strings.Join(parts, sep) +} + +// renderList mirrors Renderer._render_List: each ListItem's children render +// (paragraph text / nested list) and the results flatten one level, so an +// item contributes its text and (if any) a nested item list. +func renderList(list *ast.List, src []byte) []any { + var out []any + for item := list.FirstChild(); item != nil; item = item.NextSibling() { + for c := item.FirstChild(); c != nil; c = c.NextSibling() { + switch n := c.(type) { + case *ast.List: + out = append(out, renderList(n, src)) + default: + out = append(out, rawLines(n, src, "\n")) + } + } + } + return out +} + +// setOMap sets key k, replacing the value in place when the key repeats +// (OrderedDict semantics: position kept), else appending. +func setOMap(m omap, k string, v any) omap { + if idx := indexOMap(m, k); idx >= 0 { + m[idx].v = v + return m + } + return append(m, kv{k, v}) +} + +func indexOMap(m omap, k string) int { + for i, e := range m { + if e.k == k { + return i + } + } + return -1 +} diff --git a/internal/ingestion/component/knowledge_compiler/mindmap/mindmap.go b/internal/ingestion/component/knowledge_compiler/mindmap/mindmap.go new file mode 100644 index 0000000000..78c02e4c1d --- /dev/null +++ b/internal/ingestion/component/knowledge_compiler/mindmap/mindmap.go @@ -0,0 +1,250 @@ +// Package mindmap implements the "mindmap" variant of KnowledgeCompiler, +// mirroring Python's MindMapExtractor: the source chunks are packed into +// token-budget batches; each batch gets one LLM call (system = the rendered +// MIND_MAP_EXTRACTION_PROMPT, user = "Output:") whose markdown reply is +// parsed (dictify semantics), list-to-kv converted, merged across batches, +// and shaped into the {"id","children"} mind-map tree. The tree emits as one +// product per node with parent links. +// +// Per PORT_PLAN.md the markdown source is the LLM's reply, NOT the source +// document markdown, so the Parser component is not reused. +package mindmap + +import ( + "context" + "fmt" + "strings" + + "ragflow/internal/ingestion/component/knowledge_compiler/common" + "ragflow/internal/utility" +) + +// Run executes the mindmap variant. +func Run(ctx context.Context, deps common.Deps, param common.Param, inputs common.Inputs) (common.Outputs, error) { + docID := firstNonEmpty(inputs.DocID, deps.DatasetID) + if docID == "" { + docID = "unknown" + } + llmID := firstNonEmpty(param.LLMID, inputs.LLMID) + tenantID := deps.TenantID + + if deps.Chat == nil { + return common.Outputs{}, fmt.Errorf("mindmap: chat model required") + } + + sections := chunkTexts(inputs.Chunks) + if len(sections) == 0 { + return common.Outputs{}, nil + } + + // One LLM task per token-budget batch (mirrors __call__'s task fan-out). + batches := packSections(sections, deps.Tokenizer) + results := make([]omap, len(batches)) + // One LLM task per token-budget batch (mirrors __call__'s task fan-out). + n := param.MaxWorkers + if n <= 0 { + n = 1 + } + wp := utility.NewWorkerPool[func() error, struct{}](n, n, + func(_ context.Context, fn func() error) (struct{}, error) { return struct{}{}, fn() }) + var futs []utility.WorkerPoolFuture[func() error, struct{}] + for i, text := range batches { + i, text := i, text + f, err := wp.Submit(ctx, func() error { + resp, err := deps.Chat.Chat(ctx, common.ChatRequest{ + LLMID: llmID, + SystemPrompt: renderPrompt(text), + UserPrompt: userMessage, + }) + if err != nil { + return err + } + results[i] = todict(dictify(StripFences(resp.Content))) + return nil + }) + if err != nil { + wp.StopWait() + return common.Outputs{}, err + } + futs = append(futs, f) + } + wp.StopWait() + for _, f := range futs { + if res, _ := f.Wait(ctx); res.Err != nil { + return common.Outputs{}, res.Err + } + } + + // Merge batch dicts in batch order (mirrors reduce(self._merge, res)) and + // shape the final tree. Python returns a bare root when nothing parsed. + var merged omap + if len(results) > 0 { + merged = results[0] + for _, r := range results[1:] { + merged = mergeDicts(merged, r) + } + } + root := shapeTree(merged) + + products := treeToProducts(tenantID, docID, root) + + // Batched embedding of each node's content for downstream vector search. + if len(products) > 0 && deps.Embed != nil { + texts := make([]string, len(products)) + for i, p := range products { + texts[i] = p.Content + } + vectors, err := deps.Embed.Encode(ctx, texts) + if err != nil { + return common.Outputs{}, err + } + for i := range products { + if i < len(vectors) { + products[i].Vector = vectors[i] + } + } + } + + // Stream the tree nodes through a ProductSink so the flush policy caps peak + // memory instead of holding the full node set. + sink := common.NewProductSink(ctx, param.Guardrails, inputs.Sink) + for _, p := range products { + if err := sink.Add(p); err != nil { + return common.Outputs{}, err + } + } + out := common.Outputs{ + Products: sink.Products(), + VectorBytes: sink.Bytes(), + Items: sink.TotalItems(), + Flushed: sink.Flushed(), + } + + if err := out.EnforceGuardrails(param.Guardrails, inputs.Sink, ctx); err != nil { + return common.Outputs{}, err + } + return out, nil +} + +// treeToProducts flattens the shaped mind-map tree into Products. The root +// becomes a "root" product whose content is the serialized {"id","children"} +// tree (Python's MindMapResult.output shape); each inner node becomes a +// "node" product linked via parent_id. +func treeToProducts(tenantID, docID string, root *Node) []common.Product { + var out []common.Product + rootID := common.StableRowID(tenantID, docID, string(common.VariantMindmap), "root") + out = append(out, common.Product{ + ID: rootID, + DocID: docID, + TenantID: tenantID, + Variant: common.VariantMindmap, + Content: serializeNode(root), + Meta: map[string]any{ + "kind": "root", + "level": 0, + "name": root.ID, + }, + }) + + type pending struct { + node *Node + parentID string + level int + } + queue := []pending{{root, rootID, 0}} + for len(queue) > 0 { + p := queue[0] + queue = queue[1:] + for _, child := range p.node.Children { + if child.ID == "" { + continue + } + level := p.level + 1 + id := common.StableRowID(tenantID, docID, string(common.VariantMindmap), "node", child.ID) + childTitles := make([]string, 0, len(child.Children)) + for _, gc := range child.Children { + if gc.ID != "" { + childTitles = append(childTitles, gc.ID) + } + } + meta := map[string]any{ + "kind": "node", + "level": level, + "name": child.ID, + } + if len(childTitles) > 0 { + meta["children"] = childTitles + } + out = append(out, common.Product{ + ID: id, + DocID: docID, + TenantID: tenantID, + Variant: common.VariantMindmap, + Content: child.ID, + ParentID: p.parentID, + Meta: meta, + }) + queue = append(queue, pending{child, id, level}) + } + } + return out +} + +// serializeNode renders the tree in Python's MindMapResult.output shape: +// {"id": ..., "children": [...]}. +func serializeNode(node *Node) string { + var b strings.Builder + b.WriteString(`{"id":`) + b.WriteString(quoteJSON(node.ID)) + b.WriteString(`,"children":[`) + for i, c := range node.Children { + if i > 0 { + b.WriteString(",") + } + b.WriteString(serializeNode(c)) + } + b.WriteString("]}") + return b.String() +} + +func quoteJSON(s string) string { + var b strings.Builder + b.WriteByte('"') + for _, r := range s { + switch r { + case '"': + b.WriteString("\\\"") + case '\\': + b.WriteString("\\\\") + case '\n': + b.WriteString("\\n") + case '\t': + b.WriteString("\\t") + default: + b.WriteRune(r) + } + } + b.WriteByte('"') + return b.String() +} + +func chunkTexts(chunks []common.Chunk) []string { + var out []string + for _, c := range chunks { + t := firstNonEmpty(c.Text, c.Content) + if strings.TrimSpace(t) == "" { + continue + } + out = append(out, t) + } + return out +} + +func firstNonEmpty(vals ...string) string { + for _, v := range vals { + if v != "" { + return v + } + } + return "" +} diff --git a/internal/ingestion/component/knowledge_compiler/mindmap/mindmap_test.go b/internal/ingestion/component/knowledge_compiler/mindmap/mindmap_test.go new file mode 100644 index 0000000000..ea78616a64 --- /dev/null +++ b/internal/ingestion/component/knowledge_compiler/mindmap/mindmap_test.go @@ -0,0 +1,239 @@ +package mindmap + +import ( + "strings" + "testing" + + "ragflow/internal/ingestion/component/knowledge_compiler/common" +) + +func TestDictify_HeadingsNest(t *testing.T) { + md := "# Top\n\n## Mid A\n\n## Mid B\n\n### Deep\n" + got := dictify(md) + if len(got) != 1 || got[0].k != "Top" { + t.Fatalf("top-level = %+v, want single key Top", got) + } + mid, ok := got[0].v.(omap) + if !ok || len(mid) != 2 { + t.Fatalf("Top value = %+v, want {Mid A, Mid B}", got[0].v) + } + if mid[0].k != "Mid A" || mid[0].v != "" { + t.Errorf("Mid A = %+v, want empty string value (no children)", mid[0]) + } + deep, ok := mid[1].v.(omap) + if !ok || len(deep) != 1 || deep[0].k != "Deep" { + t.Errorf("Mid B value = %+v, want {Deep}", mid[1].v) + } +} + +func TestDictify_BulletPairsViaTodict(t *testing.T) { + // "- term\n - def" pairs become {term: def}; the unpaired bullet + // ("term2") is dropped — faithful to _list_to_kv. + md := "# T\n\n## S\n- term1\n - def1\n- term2\n" + got := todict(dictify(md)) + if len(got) != 1 || got[0].k != "T" { + t.Fatalf("dictify = %+v", got) + } + s, ok := got[0].v.(omap) + if !ok || len(s) != 1 || s[0].k != "S" { + t.Fatalf("T value = %+v", got[0].v) + } + pairs, ok := s[0].v.(omap) + if !ok || len(pairs) != 1 || pairs[0].k != "term1" || pairs[0].v != "def1" { + t.Fatalf("S value = %+v, want {term1: def1} (term2 dropped)", s[0].v) + } +} + +func TestDictify_PlainListVanishes(t *testing.T) { + // Bullets without definition sub-lists collapse to an empty dict + // (Python _list_to_kv parity). + md := "# T\n\n- a\n- b\n" + got := todict(dictify(md)) + if len(got) != 1 { + t.Fatalf("dictify = %+v", got) + } + if v, ok := got[0].v.(omap); !ok || len(v) != 0 { + t.Fatalf("plain bullets must collapse to empty dict, got %+v", got[0].v) + } +} + +func TestDictify_PlainProse(t *testing.T) { + got := todict(dictify("hello world, no headings here")) + if len(got) != 1 || got[0].k != "root" { + t.Fatalf("prose dictify = %+v, want single root key", got) + } + // The root list has no definition pairs → collapses to empty dict. + if v, ok := got[0].v.(omap); !ok || len(v) != 0 { + t.Fatalf("root value = %+v, want empty dict", got[0].v) + } +} + +func TestDictify_JumpedHeadingBecomesText(t *testing.T) { + // H1 → H3 (no H2): the jumped heading is not treated as a key. + md := "# H1\n\n### H3\n\ntext\n" + got := dictify(md) + if len(got) != 1 || got[0].k != "H1" { + t.Fatalf("dictify = %+v", got) + } + if got[0].v != "H3\n\ntext" { + t.Fatalf("H1 value = %q, want jumped heading rendered as text", got[0].v) + } +} + +func TestDictify_LeadingContentDropped(t *testing.T) { + // Content before the first heading at the minimum level is discarded + // (dictify_list_by parity). + md := "preamble\n\n# H\n\nbody\n" + got := dictify(md) + if len(got) != 1 || got[0].k != "H" || got[0].v != "body" { + t.Fatalf("leading content must be dropped: %+v", got) + } +} + +func TestMergeDicts_Order(t *testing.T) { + d1 := omap{{"A", []any{"x"}}} + d2 := omap{{"A", []any{"y"}}, {"B", "b"}} + got := mergeDicts(d1, d2) + if len(got) != 2 || got[0].k != "A" || got[1].k != "B" { + t.Fatalf("merge = %+v", got) + } + list, _ := got[0].v.([]any) + if len(list) != 2 || list[0] != "y" || list[1] != "x" { + t.Fatalf("list merge = %v, want [y x] (d2 items before d1)", list) + } +} + +func TestShapeTree_MultiTopKeys(t *testing.T) { + merged := omap{ + {"A", omap{{"shared", "x"}}}, + {"B", omap{{"shared", "y"}}}, + } + root := shapeTree(merged) + if root.ID != "root" || len(root.Children) != 2 { + t.Fatalf("shape = %+v, want root with A,B", root) + } + // The second "shared" key must survive (keyset dedups within a subtree + // walk, and A's subtree consumes "shared" first — Python parity: B's + // "shared" is dropped because the keyset is shared across children). + seen := map[string]int{} + var walk func(n *Node) + walk = func(n *Node) { + seen[n.ID]++ + for _, c := range n.Children { + walk(c) + } + } + walk(root) + if seen["shared"] != 1 { + t.Fatalf("shared key count = %d, want 1 (global keyset dedup)", seen["shared"]) + } +} + +func TestShapeTree_SingleKey(t *testing.T) { + merged := omap{{"Only", omap{{"child", "text"}}}} + root := shapeTree(merged) + if root.ID != "Only" { + t.Fatalf("single-key root id = %q, want Only", root.ID) + } + if len(root.Children) != 1 || root.Children[0].ID != "child" { + t.Fatalf("root children = %+v", root.Children) + } +} + +func TestBeChildren_StringLeaf(t *testing.T) { + keyset := map[string]bool{} + out := beChildren("plain string", keyset) + if len(out) != 1 || out[0].ID != "plain string" { + t.Fatalf("string leaf = %+v", out) + } + if !keyset["plain string"] { + t.Fatalf("raw string must enter the keyset") + } +} + +func TestStripStars(t *testing.T) { + if got := stripStars("**Bold** title"); got != "Bold title" { + t.Fatalf("stripStars = %q", got) + } +} + +func TestStripFences(t *testing.T) { + in := "```json\n{\"a\": 1}\n```\n" + if got := StripFences(in); strings.Contains(got, "```") { + t.Fatalf("fences not stripped: %q", got) + } +} + +func TestRenderPrompt(t *testing.T) { + got := renderPrompt("THE TEXT") + if !strings.Contains(got, "-TEXT-\nTHE TEXT") { + t.Fatalf("input_text not substituted: %q", got) + } + if strings.Contains(got, "{input_text}") { + t.Fatalf("placeholder survived: %q", got) + } + // The verbatim prompt body must be present. + if !strings.Contains(got, "Generate a title for user's 'TEXT'。") { + t.Fatalf("prompt body drifted") + } +} + +func TestPackSections_Budget(t *testing.T) { + // budget = max(4096*0.8, 4096-512) = 3584; with the len-based fake + // tokenizer each section of 2000 tokens packs one per batch. + tok := fakeTok{} + sections := []string{strings.Repeat("a", 8000), strings.Repeat("b", 8000), "short"} + got := packSections(sections, tok) + if len(got) != 2 { + t.Fatalf("batches = %d, want 2 (2000+2000 > 3584 splits)", len(got)) + } + // An oversized section is never split: it forms its own batch. + big := packSections([]string{strings.Repeat("x", 20000)}, tok) + if len(big) != 1 { + t.Fatalf("oversized section must stay whole: %d batches", len(big)) + } +} + +type fakeTok struct{} + +func (fakeTok) NumTokens(s string) int { return len(s) / 4 } + +func TestTreeToProducts_ParentLinks(t *testing.T) { + root := &Node{ID: "root", Children: []*Node{ + {ID: "A", Children: []*Node{{ID: "A1"}, {ID: "A2"}}}, + {ID: "B"}, + }} + products := treeToProducts("t1", "d1", root) + if len(products) != 5 { + t.Fatalf("products = %d, want 5 (root+A+A1+A2+B)", len(products)) + } + if products[0].Meta["kind"] != "root" || products[0].ParentID != "" { + t.Errorf("root product malformed: %+v", products[0].Meta) + } + if products[1].ParentID != products[0].ID { + t.Errorf("A parent link = %q, want root id", products[1].ParentID) + } + var aID string + for _, p := range products { + if p.Meta["name"] == "A" { + aID = p.ID + } + } + for _, p := range products { + if n, _ := p.Meta["name"].(string); n == "A1" || n == "A2" { + if p.ParentID != aID { + t.Errorf("%s parent = %q, want A id", n, p.ParentID) + } + } + } +} + +func TestSerializeNode_Shape(t *testing.T) { + root := &Node{ID: "Top \"quoted\"", Children: []*Node{{ID: "child"}}} + js := serializeNode(root) + if !strings.HasPrefix(js, `{"id":"Top \"quoted\""`) { + t.Errorf("serializeNode = %q", js) + } +} + +var _ = common.EstimateTokens // keep common import when unused helpers shift diff --git a/internal/ingestion/component/knowledge_compiler/mindmap/prompt.go b/internal/ingestion/component/knowledge_compiler/mindmap/prompt.go new file mode 100644 index 0000000000..d3ac5e1c5b --- /dev/null +++ b/internal/ingestion/component/knowledge_compiler/mindmap/prompt.go @@ -0,0 +1,88 @@ +package mindmap + +import ( + "strings" + + "ragflow/internal/ingestion/component/knowledge_compiler/common" +) + +// mindMapExtractionPrompt mirrors MIND_MAP_EXTRACTION_PROMPT verbatim +// (rag/graphrag/general/mind_map_prompt.py). It is the SYSTEM message; the +// user turn is the fixed string "Output:" (see userMessage). +const mindMapExtractionPrompt = ` +- Role: You're a talent text processor to summarize a piece of text into a mind map. + +- Step of task: + 1. Generate a title for user's 'TEXT'。 + 2. Classify the 'TEXT' into sections of a mind map. + 3. If the subject matter is really complex, split them into sub-sections and sub-subsections. + 4. Add a shot content summary of the bottom level section. + +- Output requirement: + - Generate at least 4 levels. + - Always try to maximize the number of sub-sections. + - In language of 'Text' + - MUST IN FORMAT OF MARKDOWN + +-TEXT- +{input_text} + +` + +// userMessage is the fixed user turn appended after the rendered system +// prompt (Python: [{"role": "user", "content": "Output:"}]). +const userMessage = "Output:" + +// mindmapMaxLength stands in for chat_mdl.max_length — the ChatInvoker seam +// does not expose the model window, so the batch budget uses the same +// conservative constant the other variants use. Python derives the budget as +// max(max_length*0.8, max_length-512). +const mindmapMaxLength = 4096 + +// batchBudget mirrors max(max_length*0.8, max_length-512). +func batchBudget() int { + a := mindmapMaxLength * 8 / 10 + b := mindmapMaxLength - 512 + if a > b { + return a + } + return b +} + +// renderPrompt mirrors perform_variable_replacements for the single +// {input_text} variable (prompt_variables is empty at the Python call site). +func renderPrompt(text string) string { + return strings.ReplaceAll(mindMapExtractionPrompt, "{input_text}", text) +} + +// packSections mirrors MindMapExtractor.__call__'s batching: sections are +// concatenated WITHOUT a separator into batches whose token count stays +// under the budget; a single section is never split (an oversized section +// forms its own batch). +func packSections(sections []string, tok common.Tokenizer) []string { + budget := batchBudget() + var batches []string + var cur strings.Builder + cnt := 0 + for _, s := range sections { + sc := numTokens(tok, s) + if cnt+sc >= budget && cur.Len() > 0 { + batches = append(batches, cur.String()) + cur.Reset() + cnt = 0 + } + cur.WriteString(s) + cnt += sc + } + if cur.Len() > 0 { + batches = append(batches, cur.String()) + } + return batches +} + +func numTokens(tok common.Tokenizer, s string) int { + if tok != nil { + return tok.NumTokens(s) + } + return common.EstimateTokens(s) +} diff --git a/internal/ingestion/component/knowledge_compiler/raptor/clustering.go b/internal/ingestion/component/knowledge_compiler/raptor/clustering.go new file mode 100644 index 0000000000..993e30d0ef --- /dev/null +++ b/internal/ingestion/component/knowledge_compiler/raptor/clustering.go @@ -0,0 +1,680 @@ +package raptor + +import ( + "math" + "math/rand" + "sort" + + "gonum.org/v1/gonum/mat" +) + +// ClusterSpec configures the clustering backend for the raptor tree. +type ClusterSpec struct { + Method ClusteringMethod + Threshold float64 // AHC: dendrogram gap threshold; GMM: soft-assign threshold + MinClusters int + MaxClusters int +} + +// ClusteringMethod selects the backend. +type ClusteringMethod string + +const ( + // ClusteringAHC is Ward agglomerative + dendrogram gap detection + PCA pre-dim. + ClusteringAHC ClusteringMethod = "AHC" + // ClusteringGMM is a diagonal-covariance Gaussian mixture model selected + // by BIC; implemented in gmmCluster (replaces the deferred ErrNotImplemented). + ClusteringGMM ClusteringMethod = "GMM" + // ClusteringPsi is the Psi builder: cosine matrix + union-find, no clustering. + ClusteringPsi ClusteringMethod = "PSI" +) + +// GapCutHold is the dendrogram-gap weight for Ward AHC. The cluster count is +// chosen where the largest forward gap in the sorted merge heights exceeds +// (gapWeight-1) * previous_gap; values >1 keep the cut conservative, values <1 +// are more aggressive. Mirrors Python's _get_optimal_clusters rule of thumb. +const GapCutHold = 1.0 + +// PCATargetDim is the upper bound for the PCA dimensionality reduction that +// precedes Ward AHC / GMM. It replaces Python's UMAP pre-step (raptor.py runs +// UMAP unconditionally to ≤12 dims). 12 fits "d ≤ 12" stated in PORT_PLAN.md +// §3.3 and is the empirically safe default for 768-dim text embeddings. +const PCATargetDim = 12 + +// GMMSeed pins the kmeans++ initialization so EM trajectories and BIC +// component selection are deterministic across runs. Required for the golden +// gate (缺口 C) and for reproducible production behaviour. +const GMMSeed = 1 + +// RecursionMinClusterSize is the smallest sub-cluster worth recursively +// summarizing: when a cluster has fewer points than this, the recursion skips +// the re-cluster and summarises the cluster's chunks directly at the current +// level. Keeps the tree shallow on small corpora and avoids single-point +// nodes that contribute no new information. +const RecursionMinClusterSize = 4 + +// Cluster dispatch for the spec method; used by both the top-level tree build +// and the recursive sub-clustering so the same method is applied at every +// level (avoids a top-vs-recursion asymmetry that the previous implementation +// had, where recursive levels always used Psi regardless of the user choice). +func clusterByMethod(embeddings [][]float64, spec ClusterSpec) ([]int, [][]float64) { + switch spec.Method { + case ClusteringAHC: + return wardAHC(reducePCA(embeddings, PCATargetDim), spec) + case ClusteringGMM: + return gmmCluster(reducePCA(embeddings, PCATargetDim), spec) + default: + return psiCluster(embeddings, spec.Threshold) + } +} + +// reducePCA projects embeddings (n x d) down to at most targetDim dimensions +// using a thin SVD (gonum SVDThin), keeping the top principal components. It +// stands in for Python's UMAP pre-step (raptor.py runs UMAP unconditionally +// before clustering). PCA preserves most clustering structure for text +// embeddings and is far cheaper; see PORT_PLAN.md §3.3. +func reducePCA(embeddings [][]float64, targetDim int) [][]float64 { + n := len(embeddings) + if n == 0 { + return embeddings + } + d := len(embeddings[0]) + if targetDim <= 0 || targetDim >= d { + return embeddings + } + // Center columns. + data := make([]float64, n*d) + for i, row := range embeddings { + copy(data[i*d:(i+1)*d], row) + } + x := mat.NewDense(n, d, data) + colMean := make([]float64, d) + for j := 0; j < d; j++ { + var s float64 + for i := 0; i < n; i++ { + s += x.At(i, j) + } + colMean[j] = s / float64(n) + } + for i := 0; i < n; i++ { + for j := 0; j < d; j++ { + x.Set(i, j, x.At(i, j)-colMean[j]) + } + } + var svd mat.SVD + ok := svd.Factorize(x, mat.SVDThin) + if !ok { + return embeddings + } + k := targetDim + if k > d { + k = d + } + // With SVDThin, V has only min(n, d) columns, so k must also be bounded by + // the sample count; otherwise vt.Slice below panics on a small corpus + // (e.g. a handful of chunks) feeding high-dimensional embeddings. + if k > n { + k = n + } + if k <= 0 { + return embeddings + } + u, vt := &mat.Dense{}, &mat.Dense{} + svd.UTo(u) + svd.VTo(vt) + // Project: X_centered * V_k (n x k). + proj := mat.NewDense(n, k, nil) + proj.Mul(x, vt.Slice(0, d, 0, k)) + out := make([][]float64, n) + for i := 0; i < n; i++ { + row := make([]float64, k) + for j := 0; j < k; j++ { + row[j] = proj.At(i, j) + } + out[i] = row + } + return out +} + +// wardAHC runs Ward agglomerative clustering on the (already reduced) embeddings +// and returns cluster labels per point plus centroids. The number of clusters is +// chosen by dendrogram gap: merge until the next merge's height exceeds the +// largest gap (scaled by Threshold) or we hit MinClusters. This is a pure-Go +// reimplementation of Python's Ward + _get_optimal_clusters (raptor.py). +func wardAHC(embeddings [][]float64, spec ClusterSpec) ([]int, [][]float64) { + n := len(embeddings) + if n == 0 { + return nil, nil + } + if n == 1 { + return []int{0}, cloneRows(embeddings[:1]) + } + + // Active clusters: each starts as one point. + clusters := make([]*cluster, n) + for i := range embeddings { + clusters[i] = &cluster{ + members: []int{i}, + centroid: cloneRow(embeddings[i]), + size: 1, + id: i, + } + } + + // Precompute pairwise squared Euclidean distances (flat upper triangle). + dist := make([][]float64, n) + for i := range dist { + dist[i] = make([]float64, n) + } + for i := 0; i < n; i++ { + for j := i + 1; j < n; j++ { + d := sqEuclid(embeddings[i], embeddings[j]) + dist[i][j] = d + dist[j][i] = d + } + } + + heights := []float64{} + merges := make([]mergeRec, 0, n-1) + for len(clusters) > 1 { + // Ward linkage: min of (sizeA*sizeB/(sizeA+sizeB)) * ||cA-cB||^2. + bestA, bestB := -1, -1 + bestCost := math.Inf(1) + for a := 0; a < len(clusters); a++ { + for b := a + 1; b < len(clusters); b++ { + da := distBetween(clusters[a].centroid, clusters[b].centroid) + cost := (float64(clusters[a].size) * float64(clusters[b].size) / + float64(clusters[a].size+clusters[b].size)) * da + if cost < bestCost { + bestCost = cost + bestA, bestB = a, b + } + } + } + height := math.Sqrt(bestCost) + heights = append(heights, height) + // Record the merge (before the cluster slice is mutated) and merge + // bestB into bestA. + ca, cb := clusters[bestA], clusters[bestB] + merges = append(merges, mergeRec{ca.id, cb.id, height}) + // ca.centroid / cb.centroid are already per-cluster means, so the + // merged centroid must be the size-weighted mean, not an equal + // average; otherwise Ward merge costs downstream select the wrong + // clusters when sizes differ. + centroid := make([]float64, len(ca.centroid)) + for i := range centroid { + centroid[i] = (float64(ca.size)*ca.centroid[i] + + float64(cb.size)*cb.centroid[i]) / float64(ca.size+cb.size) + } + merged := &cluster{ + members: append(append([]int{}, ca.members...), cb.members...), + centroid: centroid, + size: ca.size + cb.size, + id: ca.id, + } + clusters[bestA] = merged + clusters = append(clusters[:bestB], clusters[bestB+1:]...) + } + + // Dendrogram gap: choose the cluster count from the merge heights, then + // replay the exact recorded merge sequence via union-find (cutTree) instead + // of re-running the O(n^2) greedy scan a second time. + numClusters := chooseClusterCount(heights, spec) + return cutTree(embeddings, merges, numClusters) +} + +type cluster struct { + members []int + centroid []float64 + size int + id int // stable leaf representative, used to record merges for replay +} + +// mergeRec records one Ward merge between two clusters, keyed by their stable +// leaf representatives, so the cut can be replayed without re-scanning pairs. +type mergeRec struct { + a, b int + h float64 +} + +func chooseClusterCount(heights []float64, spec ClusterSpec) int { + if len(heights) == 0 { + return 1 + } + n := len(heights) + 1 + sorted := append([]float64{}, heights...) + sort.Float64s(sorted) + // Largest forward gap (heights sorted ascending; gaps between consecutive). + maxGap := 0.0 + cutIdx := len(sorted) // default: no cut (single cluster) + for i := 1; i < len(sorted); i++ { + g := sorted[i] - sorted[i-1] + if g > maxGap { + maxGap = g + cutIdx = i + } + } + // numClusters = points remaining after performing only the merges below the + // cut point (cutIdx of them); i.e. n - cutIdx. cutIdx == len(sorted) (no gap + // found) => single cluster. This is the gap-driven count; MinClusters/ + // MaxClusters below still clamp it for production tuning. + numClusters := n - cutIdx + if spec.MaxClusters > 0 && numClusters > spec.MaxClusters { + numClusters = spec.MaxClusters + } + if spec.MinClusters > 0 && numClusters < spec.MinClusters { + numClusters = spec.MinClusters + } + if numClusters < 1 { + numClusters = 1 + } + return numClusters +} + +// cutTree replays the recorded Ward merge sequence with a union-find and stops +// once numClusters remain. Replaying the exact recorded order (instead of +// re-running the O(n^2) greedy scan a second time, as the old recluster did) +// yields identical clusters while dropping the duplicate full pass. Centroids +// are recomputed from the final partition over the original embeddings. +func cutTree(embeddings [][]float64, merges []mergeRec, numClusters int) ([]int, [][]float64) { + n := len(embeddings) + if n == 0 { + return nil, nil + } + parent := make([]int, n) + for i := range parent { + parent[i] = i + } + var find func(int) int + find = func(x int) int { + for parent[x] != x { + parent[x] = parent[parent[x]] + x = parent[x] + } + return x + } + stop := n - numClusters + if stop < 0 { + stop = 0 + } + if stop > len(merges) { + stop = len(merges) + } + for i := 0; i < stop; i++ { + ra, rb := find(merges[i].a), find(merges[i].b) + if ra == rb { + continue + } + if ra < rb { + parent[rb] = ra + } else { + parent[ra] = rb + } + } + labels := make([]int, n) + rootID := map[int]int{} + idx := 0 + for i := 0; i < n; i++ { + r := find(i) + if _, ok := rootID[r]; !ok { + rootID[r] = idx + idx++ + } + labels[i] = rootID[r] + } + // Centroids from the final partition (ordered by first appearance). + sums := make([][]float64, idx) + counts := make([]int, idx) + for i := 0; i < n; i++ { + c := labels[i] + if sums[c] == nil { + sums[c] = make([]float64, len(embeddings[i])) + } + for j, v := range embeddings[i] { + sums[c][j] += v + } + counts[c]++ + } + centroidsOut := make([][]float64, 0, idx) + for c := 0; c < idx; c++ { + centroid := make([]float64, len(sums[c])) + for j, s := range sums[c] { + centroid[j] = s / float64(counts[c]) + } + centroidsOut = append(centroidsOut, centroid) + } + return labels, centroidsOut +} + +func sqEuclid(a, b []float64) float64 { + var s float64 + for i := range a { + d := a[i] - b[i] + s += d * d + } + return s +} + +func distBetween(a, b []float64) float64 { + return sqEuclid(a, b) +} + +func cloneRow(r []float64) []float64 { + out := make([]float64, len(r)) + copy(out, r) + return out +} + +func cloneRows(rows [][]float64) [][]float64 { + out := make([][]float64, len(rows)) + for i, r := range rows { + out[i] = cloneRow(r) + } + return out +} + +// gmmRegCovar matches Python sklearn's default reg_covar for a diagonal GMM; it +// keeps per-dimension variances strictly positive and EM numerically stable. +const gmmRegCovar = 1e-4 + +// gmmCluster fits a diagonal-covariance Gaussian mixture model and returns hard +// cluster labels plus centroids (component means). The component count is chosen +// by BIC over [spec.MinClusters, spec.MaxClusters] (clamped to [2, n]); a +// fixed-seed RNG drives kmeans++-style initialization so the result is +// deterministic across runs (required for the golden gate). This is a pure-Go +// reimplementation of Python's covariance_type="diag" GMM (raptor.py), +// replacing the previously deferred ErrNotImplemented backend. +func gmmCluster(embeddings [][]float64, spec ClusterSpec) ([]int, [][]float64) { + n := len(embeddings) + if n == 0 { + return nil, nil + } + if n == 1 { + return []int{0}, cloneRows(embeddings[:1]) + } + d := len(embeddings[0]) + + minK, maxK := spec.MinClusters, spec.MaxClusters + if minK < 2 { + minK = 2 + } + if maxK < minK { + maxK = minK + } + if maxK > n { + maxK = n + } + // Cap the search so GMM does not trivially overfit to one component per + // point (BIC still favours more components on tiny high-dim samples). This + // mirrors raptor.py's capped n_components and yields a meaningful tree + // instead of a flat one-component-per-chunk partition. + if cap := (n + 1) / 2; maxK > cap { + maxK = cap + } + if minK > maxK { + minK = maxK + } + + // Fixed seed -> deterministic model selection and EM trajectory. + rng := rand.New(rand.NewSource(GMMSeed)) + + bestBIC := math.Inf(1) + var bestMeans, bestCovs [][]float64 + var bestWeights []float64 + for k := minK; k <= maxK; k++ { + means, covs, weights, ll := fitDiagonalGMM(embeddings, k, d, rng) + bic := -2*ll + float64(k-1+2*k*d)*math.Log(float64(n)) + if bic < bestBIC { + bestBIC = bic + bestMeans = means + bestCovs = covs + bestWeights = weights + } + } + if bestMeans == nil { + // Degenerate: fall back to a single-cluster partition. + return []int{0}, cloneRows(embeddings[:1]) + } + + // Hard assignment mirrors Python's predict_proba + threshold soft-assign: + // pick the first component whose posterior exceeds the threshold, else the + // argmax. Threshold defaults to 0.1 to match sklearn's GMM behaviour. + thr := spec.Threshold + if thr <= 0 { + thr = 0.1 + } + labels := make([]int, n) + for i := 0; i < n; i++ { + labels[i] = gmmAssign(embeddings[i], bestMeans, bestCovs, bestWeights, thr) + } + return compactLabels(labels), bestMeans +} + +// fitDiagonalGMM runs EM for a diagonal-covariance GMM with kmeans++-style init +// and returns the component means, variances, mixing weights, and the final +// log-likelihood. +func fitDiagonalGMM(embeddings [][]float64, k, d int, rng *rand.Rand) ([][]float64, [][]float64, []float64, float64) { + n := len(embeddings) + means := kmeansPlusPlusInit(embeddings, k, rng) + covs := make([][]float64, k) + globalVar := perDimVariance(embeddings) + for c := 0; c < k; c++ { + covs[c] = cloneRow(globalVar) + } + weights := make([]float64, k) + for c := 0; c < k; c++ { + weights[c] = 1.0 / float64(k) + } + + resp := make([][]float64, n) + for i := range resp { + resp[i] = make([]float64, k) + } + const maxIter = 100 + tol := 1e-4 + prevLL := math.Inf(-1) + for iter := 0; iter < maxIter; iter++ { + // E-step. + ll := 0.0 + for i := 0; i < n; i++ { + logp := make([]float64, k) + for c := 0; c < k; c++ { + logp[c] = diagLogProb(embeddings[i], means[c], covs[c], math.Log(weights[c])) + } + lse := logSumExp(logp) + for c := 0; c < k; c++ { + resp[i][c] = math.Exp(logp[c] - lse) + } + ll += lse + } + if iter > 0 && math.Abs(ll-prevLL) < tol { + prevLL = ll + break + } + prevLL = ll + // M-step. + for c := 0; c < k; c++ { + var nk float64 + mu := make([]float64, d) + for i := 0; i < n; i++ { + w := resp[i][c] + nk += w + for j := 0; j < d; j++ { + mu[j] += w * embeddings[i][j] + } + } + if nk < 1e-9 { + // Empty component: re-seed it to a random point to stay alive. + idx := rng.Intn(n) + means[c] = cloneRow(embeddings[idx]) + covs[c] = cloneRow(globalVar) + weights[c] = 1e-3 + continue + } + for j := 0; j < d; j++ { + mu[j] /= nk + } + varSum := make([]float64, d) + for i := 0; i < n; i++ { + w := resp[i][c] + for j := 0; j < d; j++ { + diff := embeddings[i][j] - mu[j] + varSum[j] += w * diff * diff + } + } + novar := make([]float64, d) + for j := 0; j < d; j++ { + novar[j] = varSum[j]/nk + gmmRegCovar + } + means[c] = mu + covs[c] = novar + weights[c] = nk / float64(n) + } + } + return means, covs, weights, prevLL +} + +// diagLogProb returns log(weight * N(x | mu, diag(cov))) for a diagonal Gaussian. +func diagLogProb(x, mu, cov []float64, logWeight float64) float64 { + s := logWeight + for j := 0; j < len(x); j++ { + v := cov[j] + if v <= 0 { + v = gmmRegCovar + } + diff := x[j] - mu[j] + s += -0.5*math.Log(2*math.Pi*v) - 0.5*diff*diff/v + } + return s +} + +// gmmAssign returns the hard label for x following Python's predict_proba + +// threshold rule: the first component whose posterior exceeds thr, else argmax. +func gmmAssign(x []float64, means, covs [][]float64, weights []float64, thr float64) int { + k := len(means) + logp := make([]float64, k) + for c := 0; c < k; c++ { + logp[c] = diagLogProb(x, means[c], covs[c], math.Log(weights[c])) + } + lse := logSumExp(logp) + resp := make([]float64, k) + for c := 0; c < k; c++ { + resp[c] = math.Exp(logp[c] - lse) + } + for c := 0; c < k; c++ { + if resp[c] > thr { + return c + } + } + return argMax(logp) +} + +// kmeansPlusPlusInit picks k initial means via the k-means++ seeding scheme +// using the provided (deterministic) RNG. +func kmeansPlusPlusInit(embeddings [][]float64, k int, rng *rand.Rand) [][]float64 { + n := len(embeddings) + means := make([][]float64, 0, k) + means = append(means, cloneRow(embeddings[rng.Intn(n)])) + for len(means) < k { + dists := make([]float64, n) + var total float64 + for i := 0; i < n; i++ { + best := math.Inf(1) + for _, m := range means { + if dd := sqEuclid(embeddings[i], m); dd < best { + best = dd + } + } + dists[i] = best + total += best + } + if total <= 0 { + means = append(means, cloneRow(embeddings[rng.Intn(n)])) + continue + } + r := rng.Float64() * total + cum := 0.0 + idx := n - 1 + for i := 0; i < n; i++ { + cum += dists[i] + if cum >= r { + idx = i + break + } + } + means = append(means, cloneRow(embeddings[idx])) + } + return means +} + +// perDimVariance computes the per-dimension population variance of all points +// (used to seed component covariances), with the GMM regularization floor. +func perDimVariance(embeddings [][]float64) []float64 { + n := len(embeddings) + d := len(embeddings[0]) + mean := make([]float64, d) + for i := 0; i < n; i++ { + for j := 0; j < d; j++ { + mean[j] += embeddings[i][j] + } + } + for j := 0; j < d; j++ { + mean[j] /= float64(n) + } + varSum := make([]float64, d) + for i := 0; i < n; i++ { + for j := 0; j < d; j++ { + diff := embeddings[i][j] - mean[j] + varSum[j] += diff * diff + } + } + out := make([]float64, d) + for j := 0; j < d; j++ { + out[j] = varSum[j]/float64(n) + gmmRegCovar + } + return out +} + +// logSumExp computes log(sum(exp(v))) in a numerically stable way. +func logSumExp(v []float64) float64 { + m := v[0] + for _, x := range v[1:] { + if x > m { + m = x + } + } + var s float64 + for _, x := range v { + s += math.Exp(x - m) + } + return m + math.Log(s) +} + +// argMax returns the index of the largest element of v. +func argMax(v []float64) int { + best := 0 + for i := 1; i < len(v); i++ { + if v[i] > v[best] { + best = i + } + } + return best +} + +// compactLabels remaps possibly-sparse label ids to a dense, contiguous range +// [0, m) so downstream clustering consumers see clean cluster identifiers. +func compactLabels(labels []int) []int { + seen := map[int]int{} + out := make([]int, len(labels)) + next := 0 + for i, l := range labels { + id, ok := seen[l] + if !ok { + id = next + seen[l] = next + next++ + } + out[i] = id + } + return out +} diff --git a/internal/ingestion/component/knowledge_compiler/raptor/clustering_test.go b/internal/ingestion/component/knowledge_compiler/raptor/clustering_test.go new file mode 100644 index 0000000000..78d250f272 --- /dev/null +++ b/internal/ingestion/component/knowledge_compiler/raptor/clustering_test.go @@ -0,0 +1,84 @@ +package raptor + +import ( + "math" + "math/rand" + "testing" +) + +// TestGMMCluster_TwoBlobs verifies the diagonal-GMM backend recovers two +// well-separated isotropic Gaussian blobs, produces valid labels/centroids, and +// is deterministic (the golden gate and production behaviour depend on this). +func TestGMMCluster_TwoBlobs(t *testing.T) { + rng := rand.New(rand.NewSource(42)) + blobs := make([][]float64, 0, 20) + for i := 0; i < 10; i++ { + blobs = append(blobs, []float64{1 + rng.NormFloat64()*0.1, 1 + rng.NormFloat64()*0.1}) + } + for i := 0; i < 10; i++ { + blobs = append(blobs, []float64{-1 + rng.NormFloat64()*0.1, -1 + rng.NormFloat64()*0.1}) + } + + spec := ClusterSpec{Method: ClusteringGMM, MinClusters: 2, MaxClusters: 2} + labels, centroids := gmmCluster(blobs, spec) + if len(labels) != len(blobs) { + t.Fatalf("labels len = %d, want %d", len(labels), len(blobs)) + } + if len(centroids) == 0 { + t.Fatal("gmmCluster returned no centroids") + } + for _, c := range centroids { + if len(c) != 2 { + t.Fatalf("centroid dim = %d, want 2", len(c)) + } + } + + // Distinct cluster count must be exactly 2. + counts := map[int]int{} + for _, l := range labels { + counts[l]++ + } + if len(counts) != 2 { + t.Fatalf("recovered %d clusters, want 2 (labels=%v)", len(counts), labels) + } + + // Each recovered centroid must sit near one of the two blob centers. + sum := map[int][]float64{} + cnt := map[int]int{} + for i, l := range labels { + if sum[l] == nil { + sum[l] = make([]float64, 2) + } + sum[l][0] += blobs[i][0] + sum[l][1] += blobs[i][1] + cnt[l]++ + } + for l, s := range sum { + mean := []float64{s[0] / float64(cnt[l]), s[1] / float64(cnt[l])} + d0 := math.Hypot(mean[0]-1, mean[1]-1) + d1 := math.Hypot(mean[0]+1, mean[1]+1) + if math.Min(d0, d1) > 0.3 { + t.Errorf("cluster %d centroid %v is far from both blob centers", l, mean) + } + } + + // Determinism: identical inputs must yield identical labels. + labels2, _ := gmmCluster(blobs, spec) + for i := range labels { + if labels[i] != labels2[i] { + t.Fatalf("non-deterministic GMM labels at %d: %v vs %v", i, labels, labels2) + } + } +} + +// TestGMMCluster_SinglePoint ensures the degenerate n=1 case returns a valid +// single-cluster partition instead of panicking. +func TestGMMCluster_SinglePoint(t *testing.T) { + labels, centroids := gmmCluster([][]float64{{0.5, -0.3}}, ClusterSpec{Method: ClusteringGMM, MinClusters: 2, MaxClusters: 4}) + if len(labels) != 1 || labels[0] != 0 { + t.Fatalf("single-point labels = %v, want [0]", labels) + } + if len(centroids) != 1 { + t.Fatalf("single-point centroids = %d, want 1", len(centroids)) + } +} diff --git a/internal/ingestion/component/knowledge_compiler/raptor/psi.go b/internal/ingestion/component/knowledge_compiler/raptor/psi.go new file mode 100644 index 0000000000..d623729174 --- /dev/null +++ b/internal/ingestion/component/knowledge_compiler/raptor/psi.go @@ -0,0 +1,119 @@ +package raptor + +import ( + "math" +) + +// psiCluster builds the RAPTOR Psi tree: it computes the full pairwise cosine +// similarity matrix and links each node to its most-similar neighbor(s) via +// union-find, forming a forest of connected clusters. No external clustering +// backend is needed (raptor.py Psi path is clustering-free). The returned labels +// group points into connected components; centroids are the mean of each +// component's members. simThreshold controls the edge cutoff (0 = link to the +// single best neighbor only). +func psiCluster(embeddings [][]float64, simThreshold float64) ([]int, [][]float64) { + n := len(embeddings) + if n == 0 { + return nil, nil + } + parent := make([]int, n) + for i := range parent { + parent[i] = i + } + var find func(int) int + find = func(x int) int { + for parent[x] != x { + parent[x] = parent[parent[x]] + x = parent[x] + } + return x + } + union := func(a, b int) { + ra, rb := find(a), find(b) + if ra != rb { + parent[rb] = ra + } + } + + norms := make([]float64, n) + for i, v := range embeddings { + norms[i] = math.Sqrt(dotSelf(v)) + if norms[i] == 0 { + norms[i] = 1 + } + } + + for i := 0; i < n; i++ { + bestJ := -1 + bestScore := simThreshold + for j := 0; j < n; j++ { + if i == j { + continue + } + score := cosine(embeddings[i], embeddings[j], norms[i], norms[j]) + if score > bestScore { + bestScore = score + bestJ = j + } + } + if bestJ >= 0 { + union(i, bestJ) + } + } + + labels := make([]int, n) + rootID := map[int]int{} + idx := 0 + for i := 0; i < n; i++ { + r := find(i) + if _, ok := rootID[r]; !ok { + rootID[r] = idx + idx++ + } + labels[i] = rootID[r] + } + + // Centroids per component. + sums := make([][]float64, len(rootID)) + counts := make([]int, len(rootID)) + d := 0 + if n > 0 { + d = len(embeddings[0]) + } + for i := 0; i < len(sums); i++ { + sums[i] = make([]float64, d) + } + for i := 0; i < n; i++ { + c := labels[i] + counts[c]++ + for j := 0; j < d; j++ { + sums[c][j] += embeddings[i][j] + } + } + centroids := make([][]float64, len(sums)) + for c := range sums { + centroids[c] = make([]float64, d) + if counts[c] > 0 { + for j := 0; j < d; j++ { + centroids[c][j] = sums[c][j] / float64(counts[c]) + } + } + } + return labels, centroids +} + +func cosine(a, b []float64, na, nb float64) float64 { + var s float64 + for i := range a { + s += a[i] * b[i] + } + return s / (na * nb) +} + +func dotSelf(a []float64) float64 { + var s float64 + for _, x := range a { + s += x * x + } + return s +} diff --git a/internal/ingestion/component/knowledge_compiler/raptor/raptor.go b/internal/ingestion/component/knowledge_compiler/raptor/raptor.go new file mode 100644 index 0000000000..e69eb34ff0 --- /dev/null +++ b/internal/ingestion/component/knowledge_compiler/raptor/raptor.go @@ -0,0 +1,341 @@ +// Package raptor implements the "raptor" variant of KnowledgeCompiler: a +// recursive abstractive summarization tree (RAPTOR). It builds a clustering of +// chunk embeddings and summarizes each cluster with the LLM, recursing upward +// until a single root summary remains. Two builders are delivered: +// +// - Psi (ClusteringPsi): cosine-similarity union-find forest, zero clustering +// dependency (lowest risk). +// - classic (ClusteringAHC): Ward agglomerative clustering with dendrogram-gap +// cut, preceded by PCA dimensionality reduction (gonum), standing in for +// Python's unconditional UMAP pre-step. +// - gmm (ClusteringGMM): diagonal-covariance Gaussian mixture model selected +// by BIC (pure-Go reimplementation of Python's covariance_type="diag" GMM). +// +// See PORT_PLAN.md §3.3 and the M6 validation gate (缺口 C): classic must meet +// the NMI/ARI/coverage gate or the caller falls back to Psi-only. +package raptor + +import ( + "context" + "fmt" + "strings" + + "ragflow/internal/ingestion/component/knowledge_compiler/common" +) + +// buildTreeMaxDepth caps recursive summarization depth so a pathological +// clustering input (e.g. all-non-negative embeddings that never separate) can +// never drive unbounded recursion in buildTree (M14). +const buildTreeMaxDepth = 12 + +// Run executes the raptor variant. +func Run(ctx context.Context, deps common.Deps, param common.Param, inputs common.Inputs) (common.Outputs, error) { + if deps.Embed == nil { + return common.Outputs{}, fmt.Errorf("raptor: embedder required") + } + if deps.Chat == nil { + return common.Outputs{}, fmt.Errorf("raptor: chat model required") + } + docID := firstNonEmpty(inputs.DocID, deps.DatasetID) + if docID == "" { + docID = "unknown" + } + llmID := firstNonEmpty(param.LLMID, inputs.LLMID) + tenantID := deps.TenantID + + texts := chunkTexts(inputs.Chunks) + if len(texts) == 0 { + return common.Outputs{}, nil + } + chunkIDs := chunkIDsOf(inputs.Chunks) + + // Embed all source chunks once. + vectors, err := deps.Embed.Encode(ctx, texts) + if err != nil { + return common.Outputs{}, err + } + embeddings := toFloat64Matrix(vectors) + + spec := resolveSpec(param) + + labels, _ := clusterByMethod(embeddings, spec) + + // Build leaf summary products (one per cluster) and recursively summarize up, + // streaming nodes into the sink so the flush policy caps peak memory. + sink := common.NewProductSink(ctx, param.Guardrails, inputs.Sink) + if err := buildTree(ctx, deps, llmID, tenantID, docID, texts, chunkIDs, embeddings, labels, spec, sink); err != nil { + return common.Outputs{}, err + } + + out := common.Outputs{ + Products: sink.Products(), + VectorBytes: sink.Bytes(), + Items: sink.TotalItems(), + Flushed: sink.Flushed(), + } + + if err := out.EnforceGuardrails(param.Guardrails, inputs.Sink, ctx); err != nil { + return common.Outputs{}, err + } + return out, nil +} + +func resolveSpec(param common.Param) ClusterSpec { + spec := ClusterSpec{ + Method: ClusteringPsi, + Threshold: 0.0, + MinClusters: 2, + MaxClusters: 20, + } + if m, ok := param.Extra["clustering_method"].(string); ok && m != "" { + spec.Method = ClusteringMethod(strings.ToUpper(m)) + } + if t, ok := param.Extra["clustering_threshold"].(float64); ok { + spec.Threshold = t + } + if v, ok := param.Extra["min_clusters"].(float64); ok { + spec.MinClusters = int(v) + } + if v, ok := param.Extra["max_clusters"].(float64); ok { + spec.MaxClusters = int(v) + } + if spec.Method == ClusteringAHC && spec.Threshold == 0 { + // AHC uses the dendrogram gap; threshold acts as a *scale* on the gap. + spec.Threshold = 1.0 + } + if spec.Method == ClusteringGMM && spec.Threshold == 0 { + // GMM soft-assign threshold (Python default 0.1). + spec.Threshold = 0.1 + } + return spec +} + +// buildTree summarizes each cluster (level 0) and recurses upward: each parent +// node is the LLM summary of its child cluster's texts. The root is a single +// "raptor" product with a stable id. Nodes are streamed into sink as they are +// produced (so the flush policy can cap peak memory rather than holding the +// whole tree), and only the deepest-level summaries are retained for the root +// synthesis. spec controls both the top-level clustering (already applied by +// the caller via clusterByMethod) and the recursive sub-clustering, so the same +// method is used at every level of the tree. +func buildTree(ctx context.Context, deps common.Deps, llmID, tenantID, docID string, texts, chunkIDs []string, embeddings [][]float64, labels []int, spec ClusterSpec, sink *common.ProductSink) error { + n := len(labels) + // Group point indices by cluster label. + groups := map[int][]int{} + for i := 0; i < n; i++ { + groups[labels[i]] = append(groups[labels[i]], i) + } + + rootID := common.StableRowID(tenantID, docID, string(common.VariantRaptor), "root") + var ( + maxLevel = -1 + topLevelTexts []string + ) + + type nodeTask struct { + pointIdxs []int + parentID string + level int + } + queue := []nodeTask{} + for _, idxs := range groups { + queue = append(queue, nodeTask{pointIdxs: idxs, parentID: rootID, level: 0}) + } + + for len(queue) > 0 { + task := queue[0] + queue = queue[1:] + + // Gather the source text for this cluster's points. + var sb strings.Builder + for _, pi := range task.pointIdxs { + sb.WriteString(texts[pi]) + sb.WriteString("\n\n") + } + summary, err := summarizeTexts(ctx, deps, llmID, sb.String()) + if err != nil { + return err + } + nodeID := common.StableRowID(tenantID, docID, string(common.VariantRaptor), + fmt.Sprintf("L%d", task.level), summary) + embedding, err := deps.Embed.Encode(ctx, []string{summary}) + if err != nil { + return err + } + var vec []float32 + if len(embedding) > 0 { + vec = embedding[0] + } + if err := sink.Add(common.Product{ + ID: nodeID, + DocID: docID, + TenantID: tenantID, + Variant: common.VariantRaptor, + Content: summary, + Vector: vec, + ParentID: task.parentID, + Meta: map[string]any{ + "kind": "summary", + "level": task.level, + "source_chunk_ids": collectIDs(chunkIDs, task.pointIdxs), + }, + }); err != nil { + return err + } + // Track the deepest-level summaries for the root synthesis: only those + // are fed to the root, so the root summarises the top of the tree + // rather than re-summarising every leaf. + if task.level > maxLevel { + maxLevel = task.level + topLevelTexts = topLevelTexts[:0] + } + if task.level == maxLevel { + topLevelTexts = append(topLevelTexts, summary) + } + + // If this cluster is large and we are still below the depth cap, + // recurse: re-cluster its points and enqueue sub-clusters under this + // node. The sub-clustering uses the same method as the top-level pass + // (via clusterByMethod) so a user who picks AHC gets AHC at every level, + // not Psi below AHC. The depth cap guarantees termination even when a + // pathological clustering input (e.g. all-non-negative embeddings that + // never separate) would otherwise recurse without progress (M14). + if task.level < buildTreeMaxDepth && len(task.pointIdxs) > RecursionMinClusterSize { + subEmb := make([][]float64, 0, len(task.pointIdxs)) + for _, pi := range task.pointIdxs { + subEmb = append(subEmb, embeddings[pi]) + } + subLabels, _ := clusterByMethod(subEmb, spec) + subGroups := map[int][]int{} + for i, li := range subLabels { + subGroups[li] = append(subGroups[li], task.pointIdxs[i]) + } + // Only recurse when the sub-cluster was actually split into more than + // one group. If clustering returns a single label covering every point + // (the degenerate case — e.g. PSI threshold=0 on a corpus whose vectors + // all share one sign, or AHC/GMM forced to a single component), the + // subGroups map has exactly one entry equal to task.pointIdxs. Re-enqueuing + // that same work item at level+1 would loop forever with no progress and + // OOM. A single group means the cluster is already atomic, so stop. + if len(subGroups) > 1 { + for _, idxs := range subGroups { + if len(idxs) <= 1 { + continue + } + queue = append(queue, nodeTask{pointIdxs: idxs, parentID: nodeID, level: task.level + 1}) + } + } + } + } + + // Root product ties the tree together. The root summary is the synthesis + // of the highest-level summaries (the leaves of the upper tree), NOT the + // union of all leaf-level summaries — otherwise the root would just + // re-summarize the same text multiple times. + rootSummary, err := summarizeTexts(ctx, deps, llmID, "Synthesize the overall document theme from these section summaries:\n\n"+strings.Join(topLevelTexts, "\n")) + if err != nil { + return fmt.Errorf("raptor: root summarization failed: %w", err) + } + embedding, err := deps.Embed.Encode(ctx, []string{rootSummary}) + if err != nil { + return fmt.Errorf("raptor: root embedding failed: %w", err) + } + if len(embedding) == 0 { + return fmt.Errorf("raptor: root embedding returned no vectors") + } + return sink.Add(common.Product{ + ID: rootID, + DocID: docID, + TenantID: tenantID, + Variant: common.VariantRaptor, + Content: rootSummary, + Vector: embedding[0], + ParentID: "", + Meta: map[string]any{"kind": "root", "level": -1}, + }) +} + +func summarizeTexts(ctx context.Context, deps common.Deps, llmID, text string) (string, error) { + resp, err := deps.Chat.Chat(ctx, common.ChatRequest{ + LLMID: llmID, + SystemPrompt: raptorSystemPrompt, + UserPrompt: text, + }) + if err != nil { + return "", err + } + return strings.TrimSpace(resp.Content), nil +} + +func chunkTexts(chunks []common.Chunk) []string { + var out []string + for _, c := range chunks { + t := firstNonEmpty(c.Text, c.Content) + if strings.TrimSpace(t) == "" { + continue + } + out = append(out, t) + } + return out +} + +// chunkIDsOf returns the id of every non-empty chunk, parallel to chunkTexts. +// When a chunk lacks an id, a synthetic positional id is used so source +// provenance is still recoverable. +func chunkIDsOf(chunks []common.Chunk) []string { + var out []string + for i, c := range chunks { + t := firstNonEmpty(c.Text, c.Content) + if strings.TrimSpace(t) == "" { + continue + } + id := c.ID + if id == "" { + id = fmt.Sprintf("chunk-%d", i) + } + out = append(out, id) + } + return out +} + +// collectIDs returns the chunk ids at the given point indices, de-duplicated +// and in stable order. +func collectIDs(chunkIDs []string, pointIdxs []int) []string { + seen := map[string]bool{} + out := make([]string, 0, len(pointIdxs)) + for _, pi := range pointIdxs { + if pi < 0 || pi >= len(chunkIDs) { + continue + } + id := chunkIDs[pi] + if seen[id] { + continue + } + seen[id] = true + out = append(out, id) + } + return out +} + +func firstNonEmpty(vals ...string) string { + for _, v := range vals { + if v != "" { + return v + } + } + return "" +} + +func toFloat64Matrix(vecs [][]float32) [][]float64 { + out := make([][]float64, len(vecs)) + for i, v := range vecs { + row := make([]float64, len(v)) + for j, x := range v { + row[j] = float64(x) + } + out[i] = row + } + return out +} + +const raptorSystemPrompt = `You are a summarization assistant. Produce a concise, information-dense summary of the provided text that preserves key entities, facts, and relationships. Output the summary only.` diff --git a/internal/ingestion/component/knowledge_compiler/structure/chain.go b/internal/ingestion/component/knowledge_compiler/structure/chain.go new file mode 100644 index 0000000000..c8d19babb7 --- /dev/null +++ b/internal/ingestion/component/knowledge_compiler/structure/chain.go @@ -0,0 +1,473 @@ +package structure + +import ( + "context" + "fmt" + "strings" + + "ragflow/internal/ingestion/component/knowledge_compiler/common" +) + +// Chain-shape validation for ``list`` / ``timeline`` kinds, mirroring the +// Python structure.py block (CHAIN_KINDS, _chain_detect_violations, +// validate_and_correct_chain). Both kinds model a strict linear chain of +// entities (one predecessor, one successor, no cycles). The per-chunk +// extractor is happy to emit branches / cycles when the source text supports +// multiple readings, so the relation set is validated post-extraction and the +// LLM is asked to pick the correct chain out of the offenders. Correction is +// best-effort: any failure returns the input untouched (fail-open). + +// ChainKinds mirrors CHAIN_KINDS. +var ChainKinds = map[Type]bool{TypeList: true, Type("timeline"): true} + +const ( + chainCorrectionMaxChunkChars = 8196 + chainCorrectionMaxChunks = 12 + chainCorrectionMaxRelations = 16 +) + +// chainCorrectionPrompt mirrors CHAIN_CORRECTION_PROMPT verbatim ({} placeholders +// are filled with {kind, bad_relations_json, source_chunks_text}). +const chainCorrectionPrompt = `You are correcting an extracted {kind}-kind structure. + +Constraint: the relations must form a strict linear chain — every entity has +at most one predecessor and at most one successor, and there must be no +cycle. The relations below were flagged by an automated detector as +violating this constraint. Each one carries the issue that was detected. + +Bad relations (review and keep only those supported by the source text): +{bad_relations_json} + +Source chunks the relations were extracted from: +{source_chunks_text} + +Your task: from the bad relations above, pick the subset that should be +kept. Drop the rest. Do not invent new relations. Use only ` + "``from`` and ``to``" + ` slugs that appear verbatim in the bad-relations list. The result +must satisfy the strict-chain constraint. + +Return ONLY a JSON object with this exact shape (no markdown fences, no +commentary): +{ + "keep": [ + {"from": "<slug>", "to": "<slug>"}, + ... + ] +} +` + +// chainCorrectionSystem is the system prompt of the correction call (Python +// passes it inline to gen_json). +const chainCorrectionSystem = "You correct extracted graph relations to satisfy a strict-chain constraint." + +// chainJudgeTemperature mirrors the correction call's gen_conf (temperature 0.0). +var chainJudgeTemperature = 0.0 + +type chainEdge struct{ From, To string } + +// chainExtractEdge mirrors _chain_extract_edge: the relation row's endpoint +// pair, from the authoritative meta columns with a payload fallback. +func chainExtractEdge(row common.Product) (chainEdge, bool) { + if kind, _ := row.Meta["kind"].(string); kind != "relation" { + return chainEdge{}, false + } + from, _ := row.Meta["from"].(string) + to, _ := row.Meta["to"].(string) + from, to = strings.TrimSpace(from), strings.TrimSpace(to) + if from != "" && to != "" { + return chainEdge{from, to}, true + } + payload := parsePayload(row.Content) + if payload == nil { + return chainEdge{}, false + } + from = relationEndpoint(payload, "", "source", "src", "from") + to = relationEndpoint(payload, "", "target", "tgt", "to") + if from == "" || to == "" { + return chainEdge{}, false + } + return chainEdge{from, to}, true +} + +// chainDetectViolations mirrors _chain_detect_violations: it returns +// {edge: [issue strings]} for every edge involved in a self-loop, fan-out, +// fan-in, or a directed cycle (SCC size >= 2). Cycle detection uses iterative +// Kosaraju — no recursion, so the pathologically-deep case Python guards with +// RecursionError cannot occur here. +func chainDetectViolations(edges []chainEdge) map[chainEdge][]string { + issues := map[chainEdge][]string{} + add := func(e chainEdge, reason string) { + issues[e] = append(issues[e], reason) + } + + outGroups := map[string][]chainEdge{} + inGroups := map[string][]chainEdge{} + for _, e := range edges { + if e.From == e.To { + add(e, "self-loop") + } + outGroups[e.From] = append(outGroups[e.From], e) + inGroups[e.To] = append(inGroups[e.To], e) + } + for node, group := range outGroups { + if len(group) > 1 { + siblings := sortedStrings(func() []string { + var out []string + for _, g := range group { + out = append(out, g.To) + } + return out + }()) + reason := fmt.Sprintf("fan-out from '%s' (also points to %s)", node, siblings) + for _, e := range group { + add(e, reason) + } + } + } + for node, group := range inGroups { + if len(group) > 1 { + siblings := sortedStrings(func() []string { + var out []string + for _, g := range group { + out = append(out, g.From) + } + return out + }()) + reason := fmt.Sprintf("fan-in to '%s' (also reached from %s)", node, siblings) + for _, e := range group { + add(e, reason) + } + } + } + + for _, comp := range chainSCCs(edges) { + for _, e := range edges { + if comp[e.From] && comp[e.To] { + add(e, fmt.Sprintf("cycle within %s", sortedKeys(comp))) + } + } + } + return issues +} + +// chainSCCs returns the strongly connected components of size >= 2 (each a +// directed cycle), computed with iterative Kosaraju. +func chainSCCs(edges []chainEdge) []map[string]bool { + adj := map[string][]string{} + radj := map[string][]string{} + nodes := map[string]bool{} + for _, e := range edges { + nodes[e.From] = true + nodes[e.To] = true + adj[e.From] = append(adj[e.From], e.To) + radj[e.To] = append(radj[e.To], e.From) + } + + // Pass 1: finishing order on the forward graph (iterative DFS). + visited := map[string]bool{} + var order []string + for start := range nodes { + if visited[start] { + continue + } + type frame struct { + node string + next int + } + stack := []frame{{node: start}} + visited[start] = true + for len(stack) > 0 { + top := &stack[len(stack)-1] + if top.next < len(adj[top.node]) { + w := adj[top.node][top.next] + top.next++ + if !visited[w] { + visited[w] = true + stack = append(stack, frame{node: w}) + } + continue + } + order = append(order, top.node) + stack = stack[:len(stack)-1] + } + } + + // Pass 2: DFS on the transposed graph in reverse finishing order. + assigned := map[string]bool{} + var sccs []map[string]bool + for i := len(order) - 1; i >= 0; i-- { + root := order[i] + if assigned[root] { + continue + } + comp := map[string]bool{} + stack := []string{root} + assigned[root] = true + for len(stack) > 0 { + v := stack[len(stack)-1] + stack = stack[:len(stack)-1] + comp[v] = true + for _, w := range radj[v] { + if !assigned[w] { + assigned[w] = true + stack = append(stack, w) + } + } + } + if len(comp) >= 2 { + sccs = append(sccs, comp) + } + } + return sccs +} + +// chainGatherChunkText mirrors _chain_gather_chunk_text: deduplicated +// (chunk_id, text) pairs for the correction prompt, capped. +func chainGatherChunkText(badRows []common.Product, chunksByID map[string]string) [][2]string { + seen := map[string]bool{} + var out [][2]string + for _, row := range badRows { + for _, cid := range metaStrings(row.Meta, "source_chunk_ids") { + if cid == "" || seen[cid] { + continue + } + seen[cid] = true + text := strings.TrimSpace(chunksByID[cid]) + if text == "" { + continue + } + if len(text) > chainCorrectionMaxChunkChars { + text = text[:chainCorrectionMaxChunkChars] + } + out = append(out, [2]string{cid, text}) + if len(out) >= chainCorrectionMaxChunks { + return out + } + } + } + return out +} + +// chainCorrectBatch mirrors correct_batch: one LLM correction over a batch of +// bad edges. It fails open — a failed or malformed call retains the batch's +// relations. The returned set holds the edges to keep. +func chainCorrectBatch(ctx context.Context, deps common.Deps, llmID string, kind Type, batchEdges []chainEdge, edgeToRows map[chainEdge][]common.Product, violations map[chainEdge][]string, chunksByID map[string]string) []chainEdge { + batchKeep := append([]chainEdge{}, batchEdges...) + + var badRows []common.Product + relations := make([]map[string]any, 0, len(batchEdges)) + for _, e := range batchEdges { + issue := "cross-batch conflict" + if iss := violations[e]; len(iss) > 0 { + issue = strings.Join(iss, "; ") + } + relations = append(relations, map[string]any{"from": e.From, "to": e.To, "issue": issue}) + badRows = append(badRows, edgeToRows[e]...) + } + chunkPairs := chainGatherChunkText(badRows, chunksByID) + var sourceText strings.Builder + for _, pair := range chunkPairs { + fmt.Fprintf(&sourceText, "[%s]\n%s\n\n", pair[0], pair[1]) + } + if sourceText.Len() == 0 { + sourceText.WriteString("(no source chunks available)") + } + + user := chainCorrectionPrompt + user = strings.ReplaceAll(user, "{kind}", string(kind)) + user = strings.ReplaceAll(user, "{bad_relations_json}", mustJSONList(relations)) + user = strings.ReplaceAll(user, "{source_chunks_text}", strings.TrimSpace(sourceText.String())) + + res, err := common.GenJSON(ctx, deps.Chat, common.ChatRequest{ + LLMID: llmID, SystemPrompt: chainCorrectionSystem, UserPrompt: user, Temperature: &chainJudgeTemperature, + }) + if err != nil { + return batchKeep // fail open + } + keepRaw, ok := res["keep"].([]any) + if !ok { + return batchKeep + } + batchSet := make(map[chainEdge]bool, len(batchEdges)) + for _, e := range batchEdges { + batchSet[e] = true + } + var kept []chainEdge + for _, item := range keepRaw { + m, ok := item.(map[string]any) + if !ok { + continue + } + from, _ := m["from"].(string) + to, _ := m["to"].(string) + e := chainEdge{strings.TrimSpace(from), strings.TrimSpace(to)} + if batchSet[e] { + kept = append(kept, e) + } + } + return kept +} + +// validateAndCorrectChain mirrors validate_and_correct_chain: relations of +// chain-kind rows must form a strict linear chain; offending relations the +// LLM does not keep are dropped from the returned rows. Any failure returns +// the input verbatim (fail-open), so a misbehaving model can never block the +// pipeline. +func validateAndCorrectChain(ctx context.Context, deps common.Deps, llmID string, rows []common.Product, chunksByID map[string]string, kind Type) []common.Product { + if len(rows) == 0 || !ChainKinds[kind] { + return rows + } + + edgeToRows := map[chainEdge][]common.Product{} + var allEdges []chainEdge + for _, row := range rows { + e, ok := chainExtractEdge(row) + if !ok { + continue + } + if _, seen := edgeToRows[e]; !seen { + allEdges = append(allEdges, e) + } + edgeToRows[e] = append(edgeToRows[e], row) + } + violations := chainDetectViolations(allEdges) + if len(violations) == 0 { + return rows + } + badEdges := make([]chainEdge, 0, len(violations)) + for e := range violations { + badEdges = append(badEdges, e) + } + + // Correct in capped batches (sequential: deterministic, and in-run + // violation counts are small — Python parallelises under a semaphore). + keepSet := map[chainEdge]bool{} + for i := 0; i < len(badEdges); i += chainCorrectionMaxRelations { + end := i + chainCorrectionMaxRelations + if end > len(badEdges) { + end = len(badEdges) + } + for _, e := range chainCorrectBatch(ctx, deps, llmID, kind, badEdges[i:end], edgeToRows, violations, chunksByID) { + keepSet[e] = true + } + } + + // Independent corrections can be valid inside each request but conflict + // after their results are combined. Re-check the combined keep set and + // give the model one final decision over the remaining conflicts. + var keepList []chainEdge + for e := range keepSet { + keepList = append(keepList, e) + } + if combined := chainDetectViolations(keepList); len(combined) > 0 { + conflictEdges := make([]chainEdge, 0, len(combined)) + for e := range combined { + conflictEdges = append(conflictEdges, e) + } + finalKeep := chainCorrectBatch(ctx, deps, llmID, kind, conflictEdges, edgeToRows, combined, chunksByID) + for _, e := range conflictEdges { + delete(keepSet, e) + } + for _, e := range finalKeep { + keepSet[e] = true + } + } + + // LLM kept everything → no correction applied. + if len(keepSet) == len(badEdges) { + allKept := true + for _, e := range badEdges { + if !keepSet[e] { + allKept = false + break + } + } + if allKept { + return rows + } + } + + droppedIDs := map[string]bool{} + for _, e := range badEdges { + if keepSet[e] { + continue + } + for _, row := range edgeToRows[e] { + if row.ID != "" { + droppedIDs[row.ID] = true + } + } + } + if len(droppedIDs) == 0 { + return rows + } + corrected := make([]common.Product, 0, len(rows)-len(droppedIDs)) + for _, row := range rows { + if !droppedIDs[row.ID] { + corrected = append(corrected, row) + } + } + return corrected +} + +// dropIsolatedTimelineEntities mirrors cleanup_timeline_isolated_entities: +// for the timeline kind, entity rows not referenced by any relation endpoint +// are dropped. The in-memory port runs it after all merges/rewrites (the +// Python ES version schedules it after every flush for the same reason: a +// later relation can still reference the entity). +func dropIsolatedTimelineEntities(rows []common.Product) []common.Product { + connected := map[string]bool{} + for _, row := range rows { + if e, ok := chainExtractEdge(row); ok { + connected[strings.ToLower(e.From)] = true + connected[strings.ToLower(e.To)] = true + } + } + out := make([]common.Product, 0, len(rows)) + for _, row := range rows { + if kind, _ := row.Meta["kind"].(string); kind == "entity" { + if !connected[strings.ToLower(entityNameValue(row))] { + continue + } + } + out = append(out, row) + } + return out +} + +func sortedStrings(in []string) []string { + seen := map[string]bool{} + var out []string + for _, s := range in { + if !seen[s] { + seen[s] = true + out = append(out, s) + } + } + for i := 1; i < len(out); i++ { + for j := i; j > 0 && out[j] < out[j-1]; j-- { + out[j], out[j-1] = out[j-1], out[j] + } + } + return out +} + +func sortedKeys(m map[string]bool) []string { + var out []string + for k := range m { + out = append(out, k) + } + return sortedStrings(out) +} + +func mustJSONList(items []map[string]any) string { + var b strings.Builder + b.WriteString("[") + for i, item := range items { + if i > 0 { + b.WriteString(", ") + } + b.WriteString(payloadJSON(item)) + } + b.WriteString("]") + return b.String() +} diff --git a/internal/ingestion/component/knowledge_compiler/structure/chain_test.go b/internal/ingestion/component/knowledge_compiler/structure/chain_test.go new file mode 100644 index 0000000000..308660559f --- /dev/null +++ b/internal/ingestion/component/knowledge_compiler/structure/chain_test.go @@ -0,0 +1,210 @@ +package structure + +import ( + "context" + "errors" + "strings" + "testing" + + "ragflow/internal/ingestion/component/knowledge_compiler/common" +) + +func relRow(id, from, to string) common.Product { + return common.Product{ + ID: id, + Content: payloadJSON(map[string]any{"type": "linked", "source": from, "target": to}), + Meta: map[string]any{"kind": "relation", "from": from, "to": to, "source_chunk_ids": []string{"c1"}}, + } +} + +func entRow(id, name string) common.Product { + return common.Product{ + ID: id, + Content: payloadJSON(map[string]any{"type": "letter", "name": name}), + Meta: map[string]any{"kind": "entity", "name": name, "source_chunk_ids": []string{"c1"}}, + } +} + +func TestChainDetectViolations(t *testing.T) { + // Clean chain: no issues. + clean := []chainEdge{{"A", "B"}, {"B", "C"}, {"C", "D"}} + if got := chainDetectViolations(clean); len(got) != 0 { + t.Fatalf("clean chain flagged: %v", got) + } + + bad := []chainEdge{ + {"A", "A"}, // self-loop + {"A", "B"}, // fan-out partner + cycle member + {"A", "C"}, // fan-out partner + {"B", "A"}, // cycle A→B→A + {"X", "C"}, + {"Y", "C"}, // fan-in to C (together with A→C) + } + got := chainDetectViolations(bad) + hasIssue := func(e chainEdge, substr string) bool { + for _, iss := range got[e] { + if strings.Contains(iss, substr) { + return true + } + } + return false + } + if !hasIssue(chainEdge{"A", "A"}, "self-loop") { + t.Errorf("self-loop not detected: %v", got) + } + if !hasIssue(chainEdge{"A", "B"}, "fan-out") || !hasIssue(chainEdge{"A", "C"}, "fan-out") { + t.Errorf("fan-out not detected on both A edges: %v", got) + } + if !hasIssue(chainEdge{"X", "C"}, "fan-in") || !hasIssue(chainEdge{"Y", "C"}, "fan-in") { + t.Errorf("fan-in not detected on both C edges: %v", got) + } + if !hasIssue(chainEdge{"A", "B"}, "cycle") || !hasIssue(chainEdge{"B", "A"}, "cycle") { + t.Errorf("cycle A↔B not detected: %v", got) + } + // Edges uninvolved in any issue must be absent. + if _, ok := got[chainEdge{"A", "C"}]; ok && !hasIssue(chainEdge{"A", "C"}, "fan-out") { + t.Errorf("unexpected issue content for A→C: %v", got) + } +} + +func TestChainSCCs(t *testing.T) { + edges := []chainEdge{{"A", "B"}, {"B", "C"}, {"C", "A"}, {"C", "D"}, {"D", "E"}} + sccs := chainSCCs(edges) + if len(sccs) != 1 { + t.Fatalf("expected exactly one SCC of size>=2, got %d: %v", len(sccs), sccs) + } + comp := sccs[0] + if len(comp) != 3 || !comp["A"] || !comp["B"] || !comp["C"] { + t.Fatalf("SCC = %v, want {A,B,C}", comp) + } +} + +// keepFirstChainChat keeps the Alpha→Beta edge whenever asked to correct. +type keepFirstChainChat struct{ calls int } + +func (m *keepFirstChainChat) Chat(_ context.Context, req common.ChatRequest) (*common.ChatResponse, error) { + m.calls++ + return &common.ChatResponse{Content: `{"keep":[{"from":"Alpha","to":"Beta"}]}`}, nil +} + +type errChainChat struct{} + +func (errChainChat) Chat(_ context.Context, _ common.ChatRequest) (*common.ChatResponse, error) { + return nil, errors.New("llm down") +} + +func TestValidateAndCorrectChainFanOut(t *testing.T) { + rows := []common.Product{ + entRow("e-a", "Alpha"), entRow("e-b", "Beta"), entRow("e-g", "Gamma"), + relRow("r-ab", "Alpha", "Beta"), + relRow("r-ag", "Alpha", "Gamma"), // fan-out partner to be dropped + } + chat := &keepFirstChainChat{} + deps := common.Deps{Chat: chat, Embed: hashEmbedder{dim: 8}} + out := validateAndCorrectChain(context.Background(), deps, "llm1", rows, map[string]string{"c1": "Alpha links Beta and Gamma"}, Type("timeline")) + + var sawAB, sawAG bool + for _, row := range out { + if row.ID == "r-ab" { + sawAB = true + } + if row.ID == "r-ag" { + sawAG = true + } + } + if !sawAB || sawAG { + t.Fatalf("correction must keep Alpha→Beta and drop Alpha→Gamma: got %v", rowIDs(out)) + } + if chat.calls != 1 { + t.Fatalf("expected exactly 1 correction call (single batch, no combined conflicts), got %d", chat.calls) + } +} + +func TestValidateAndCorrectChainFailOpen(t *testing.T) { + rows := []common.Product{ + relRow("r-ab", "Alpha", "Beta"), + relRow("r-ag", "Alpha", "Gamma"), + } + deps := common.Deps{Chat: errChainChat{}, Embed: hashEmbedder{dim: 8}} + out := validateAndCorrectChain(context.Background(), deps, "llm1", rows, nil, TypeList) + if len(out) != len(rows) { + t.Fatalf("fail-open: LLM failure must return the input untouched (%d vs %d)", len(out), len(rows)) + } + + // Non-chain kinds skip validation entirely (no LLM call possible). + deps2 := common.Deps{Chat: errChainChat{}} + out2 := validateAndCorrectChain(context.Background(), deps2, "llm1", rows, nil, TypeHypergraph) + if len(out2) != len(rows) { + t.Fatalf("non-chain kind must pass through untouched") + } +} + +func TestDropIsolatedTimelineEntities(t *testing.T) { + rows := []common.Product{ + entRow("e-a", "Alpha"), entRow("e-b", "Beta"), entRow("e-o", "Orphan"), + relRow("r-ab", "Alpha", "Beta"), + } + out := dropIsolatedTimelineEntities(rows) + if len(out) != 3 { + t.Fatalf("expected 3 survivors (2 entities + 1 relation), got %d", len(out)) + } + for _, row := range out { + if row.ID == "e-o" { + t.Fatalf("orphan timeline entity must be dropped") + } + } +} + +// TestStructureRunTimelineKind exercises the full timeline pipeline end to +// end: extraction → LLM dedup → chain correction (fan-out dropped) → +// orphan-entity cleanup (the dropped relation's target becomes isolated). +func TestStructureRunTimelineKind(t *testing.T) { + deps := common.Deps{Chat: &graphChat{}, Embed: hashEmbedder{dim: 16}, TenantID: "t1", DatasetID: "d1"} + param := testParam() + inputs := common.Inputs{ + DocID: "doc1", + Chunks: []common.Chunk{ + {ID: "c1", Text: "Alpha is a Beta."}, + {ID: "c2", Text: "Alpha also links Gamma."}, + }, + VariantSpecific: map[string]any{"parser_config": map[string]any{ + "kind": "timeline", + "entity": map[string]any{"fields": []any{ + map[string]any{"type": "letter", "description": "a single letter entity"}, + }}, + "relation": map[string]any{"fields": []any{ + map[string]any{"type": "linked", "description": "a link between letters"}, + }}, + }}, + } + out, err := Run(context.Background(), deps, param, inputs) + if err != nil { + t.Fatalf("Run timeline: %v", err) + } + var entities, relations int + for _, p := range out.Products { + switch p.Meta["kind"] { + case "entity": + entities++ + if p.Meta["name"] == "Gamma" { + t.Errorf("Gamma must be dropped: its only relation was corrected away, leaving it isolated") + } + case "relation": + relations++ + if p.Meta["to"] == "Gamma" { + t.Errorf("fan-out relation Alpha→Gamma must be dropped by chain correction") + } + } + } + if entities != 2 || relations != 1 { + t.Fatalf("timeline survivors = %d entities + %d relations, want 2+1 (Alpha, Beta, Alpha→Beta)", entities, relations) + } +} + +func rowIDs(rows []common.Product) []string { + var out []string + for _, r := range rows { + out = append(out, r.ID) + } + return out +} diff --git a/internal/ingestion/component/knowledge_compiler/structure/compile.go b/internal/ingestion/component/knowledge_compiler/structure/compile.go new file mode 100644 index 0000000000..4684c0c6ee --- /dev/null +++ b/internal/ingestion/component/knowledge_compiler/structure/compile.go @@ -0,0 +1,321 @@ +package structure + +import ( + "bytes" + "context" + "encoding/json" + "fmt" + "sort" + "strings" + + "ragflow/internal/ingestion/component/knowledge_compiler/common" +) + +// CompileConfig carries the per-run identity and extraction settings. +type CompileConfig struct { + LLMID string + Type Type // inferred compile kind (list / set / hypergraph) + TenantID string + DocID string + Variant common.Variant + Lang string + ParserConfig map[string]any + // TemplateID is mixed into stable row ids so two templates sharing a + // compile kind don't collide on identical payloads (mirrors Python's + // row_seed_extras in _struct_to_doc_storage_doc). + TemplateID string +} + +// extractionTemperature mirrors Python's gen_conf for structure extraction +// (_struct_extract_hypergraph uses temperature 0.1). +var extractionTemperature = 0.1 + +// extractHypergraph mirrors _struct_extract_hypergraph: stage 1 extracts node +// (entity) items, stage 2 extracts edge (relation) items constrained to the +// stage-1 entities via the {known_nodes} placeholder. The two stages within +// one batch are strictly sequential; batches parallelise at the caller. +func extractHypergraph(ctx context.Context, deps common.Deps, cfg CompileConfig, nodePrompt, edgePromptTmpl, packedText string) (nodes, edges []map[string]any, err error) { + user := UserPrompt(packedText) + nodeRaw, err := common.GenJSON(ctx, deps.Chat, common.ChatRequest{ + LLMID: cfg.LLMID, SystemPrompt: nodePrompt, UserPrompt: user, Temperature: &extractionTemperature, + }) + if err != nil { + return nil, nil, err + } + nodes = unwrapItems(nodeRaw) + + // Known entities: unique values of the config's entity id field, in + // first-seen order (mirrors _struct_extract_hypergraph's known_keys). + idField := EntityIDField(cfg.ParserConfig) + var known []string + for _, n := range nodes { + v := strings.TrimSpace(stringOf(n[idField])) + if v == "" || containsString(known, v) { + continue + } + known = append(known, v) + } + + if strings.TrimSpace(edgePromptTmpl) == "" { + return nodes, nil, nil + } + edgeRaw, err := common.GenJSON(ctx, deps.Chat, common.ChatRequest{ + LLMID: cfg.LLMID, SystemPrompt: fillKnownNodes(edgePromptTmpl, known), UserPrompt: user, Temperature: &extractionTemperature, + }) + if err != nil { + return nil, nil, err + } + return nodes, unwrapItems(edgeRaw), nil +} + +// unwrapItems mirrors _struct_unwrap_items: the contract is +// {"items": [{...}, ...]}; a top-level array is tolerated defensively. +// GenJSON already guarantees a JSON object, so the list form only appears +// under "items". Non-object entries are dropped. +func unwrapItems(raw map[string]any) []map[string]any { + if raw == nil { + return nil + } + arr, ok := raw["items"].([]any) + if !ok { + return nil + } + out := make([]map[string]any, 0, len(arr)) + for _, e := range arr { + if m, ok := e.(map[string]any); ok { + out = append(out, m) + } + } + return out +} + +// payloadChunkIDs mirrors _struct_payload_chunk_ids: keep only model-selected +// chunk IDs that belong to the current batch; fall back to all batch ids when +// the model returned none that qualify. +func payloadChunkIDs(payload map[string]any, batchIDs []string) []string { + var rawIDs []string + switch v := payload["source_chunk_ids"].(type) { + case string: + rawIDs = []string{v} + case []any: + for _, e := range v { + if s, ok := e.(string); ok { + rawIDs = append(rawIDs, s) + } + } + case []string: + rawIDs = v + } + allowed := make(map[string]bool, len(batchIDs)) + for _, id := range batchIDs { + allowed[id] = true + } + var selected []string + seen := map[string]bool{} + for _, id := range rawIDs { + id = strings.TrimSpace(id) + if allowed[id] && !seen[id] { + selected = append(selected, id) + seen[id] = true + } + } + if len(selected) == 0 { + return append([]string{}, batchIDs...) + } + return selected +} + +// payloadDescription mirrors _struct_payload_description: concat the string +// values of every field (lists flattened) with single spaces. Go maps lose +// JSON insertion order, so keys are sorted to keep the embedding input (and +// therefore the row vector) deterministic across runs. +func payloadDescription(payload map[string]any) string { + keys := make([]string, 0, len(payload)) + for k := range payload { + keys = append(keys, k) + } + sort.Strings(keys) + var parts []string + appendValue := func(v any) { + s := strings.TrimSpace(stringOf(v)) + if s != "" { + parts = append(parts, s) + } + } + for _, k := range keys { + switch v := payload[k].(type) { + case []any: + for _, item := range v { + appendValue(item) + } + case []string: + for _, item := range v { + appendValue(item) + } + default: + appendValue(v) + } + } + return strings.Join(parts, " ") +} + +// payloadJSON serialises a payload the way Python's json.dumps(ensure_ascii= +// False) does: Go's json.Marshal escapes <, > and & to < etc., which +// would corrupt entity names containing those characters, so HTML escaping is +// disabled. Keys are alphabetically sorted (map marshal), giving a canonical, +// hash-stable form. +func payloadJSON(payload map[string]any) string { + var b bytes.Buffer + enc := json.NewEncoder(&b) + enc.SetEscapeHTML(false) + if err := enc.Encode(payload); err != nil { + return "{}" + } + return strings.TrimSpace(b.String()) +} + +// parsePayload is the inverse of payloadJSON (nil on malformed content). +func parsePayload(content string) map[string]any { + var m map[string]any + if err := json.Unmarshal([]byte(content), &m); err != nil { + return nil + } + return m +} + +// buildRows converts extracted node/edge payloads into entity/relation +// products, mirroring _struct_process_batch: embedding input is +// payloadDescription(payload); content is the payload JSON (Python's +// content_with_weight); source_chunk_ids are the model's picks filtered to +// the batch; the stable row id hashes (content, doc_id, template_id). +func buildRows(ctx context.Context, deps common.Deps, cfg CompileConfig, nodes, edges []map[string]any, batchIDs []string) ([]common.Product, error) { + srcField, tgtField := RelationMemberFields(cfg.ParserConfig) + + type spec struct { + kind string // "entity" | "relation" + payload map[string]any + } + specs := make([]spec, 0, len(nodes)+len(edges)) + for _, p := range nodes { + specs = append(specs, spec{kind: "entity", payload: p}) + } + for _, p := range edges { + specs = append(specs, spec{kind: "relation", payload: p}) + } + if len(specs) == 0 { + return nil, nil + } + + texts := make([]string, len(specs)) + for i, s := range specs { + texts[i] = payloadDescription(s.payload) + } + vectors, err := deps.Embed.Encode(ctx, texts) + if err != nil { + return nil, err + } + if len(vectors) != len(specs) { + return nil, fmt.Errorf("knowledge_compiler: embedding count mismatch (%d vs %d)", len(vectors), len(specs)) + } + + rows := make([]common.Product, 0, len(specs)) + for i, s := range specs { + content := payloadJSON(s.payload) + idParts := []string{content, cfg.DocID} + if cfg.TemplateID != "" { + idParts = append(idParts, cfg.TemplateID) + } + meta := map[string]any{ + "kind": s.kind, + "source_chunk_ids": payloadChunkIDs(s.payload, batchIDs), + "mention_count": 1, + } + if s.kind == "entity" { + if name := entityName(s.payload); name != "" { + meta["name"] = name + } + if typ := strings.TrimSpace(stringOf(s.payload["type"])); typ != "" { + meta["entity_type"] = typ + } else { + meta["entity_type"] = "other" + } + if desc := strings.TrimSpace(stringOf(s.payload["description"])); desc != "" { + meta["description"] = desc + } + } else { + if from := relationEndpoint(s.payload, srcField, "source", "src", "from"); from != "" { + meta["from"] = from + } + if to := relationEndpoint(s.payload, tgtField, "target", "tgt", "to"); to != "" { + meta["to"] = to + } + if typ := strings.TrimSpace(stringOf(s.payload["type"])); typ != "" { + meta["relation_type"] = typ + } + } + rows = append(rows, common.Product{ + ID: common.StableRowID(idParts...), + DocID: cfg.DocID, + TenantID: cfg.TenantID, + Variant: cfg.Variant, + Content: content, + Vector: vectors[i], + Meta: meta, + }) + } + return rows, nil +} + +// entityName mirrors _struct_graph_entity's name resolution +// (name → text → term → title, "-1" sentinel rejected). +func entityName(payload map[string]any) string { + for _, k := range []string{"name", "text", "term", "title"} { + if s := strings.TrimSpace(stringOf(payload[k])); s != "" && s != "-1" { + return s + } + } + return "" +} + +// relationEndpoint resolves a relation's endpoint: the config-declared member +// field first, then the conventional aliases (mirrors _struct_graph_relation +// combined with _struct_relation_member_fields). +func relationEndpoint(payload map[string]any, declared string, aliases ...string) string { + if declared != "" { + if s := strings.TrimSpace(stringOf(payload[declared])); s != "" && s != "-1" { + return s + } + } + for _, k := range aliases { + if s := strings.TrimSpace(stringOf(payload[k])); s != "" && s != "-1" { + return s + } + } + return "" +} + +// stringOf renders a scalar payload value; maps/slices render as "" (they are +// not scalar text). +func stringOf(v any) string { + switch x := v.(type) { + case nil: + return "" + case string: + return x + case float64: + return strings.TrimSuffix(fmt.Sprintf("%v", x), ".0") + case bool: + return fmt.Sprintf("%v", x) + default: + return "" + } +} + +func containsString(haystack []string, needle string) bool { + for _, s := range haystack { + if s == needle { + return true + } + } + return false +} diff --git a/internal/ingestion/component/knowledge_compiler/structure/graph.go b/internal/ingestion/component/knowledge_compiler/structure/graph.go new file mode 100644 index 0000000000..fdc6b68dc6 --- /dev/null +++ b/internal/ingestion/component/knowledge_compiler/structure/graph.go @@ -0,0 +1,156 @@ +package structure + +import ( + "strings" + + "ragflow/internal/ingestion/component/knowledge_compiler/common" +) + +// Graph projections and rebuild, mirroring Python structure.py's +// _struct_graph_entity / _struct_graph_relation / _struct_merge_graph_entities +// and _struct_rebuild_graph_json. The compact graph is derived from the +// surviving entity/relation payload rows (content_with_weight JSON). + +// graphEntity mirrors _struct_graph_entity: project an entity payload to the +// graph-node shape (mention_count starts at 1; the row's filtered +// source_chunk_ids column is authoritative for provenance). +func graphEntity(payload map[string]any, sourceChunkIDs []string) map[string]any { + name := entityName(payload) + if name == "" { + return nil + } + typ := strings.TrimSpace(stringOf(payload["type"])) + if typ == "" { + typ = "other" + } + var aliases []string + switch v := payload["aliases"].(type) { + case string: + if s := strings.TrimSpace(v); s != "" { + aliases = []string{s} + } + case []any: + for _, e := range v { + if s := strings.TrimSpace(stringOf(e)); s != "" { + aliases = append(aliases, s) + } + } + case []string: + for _, e := range v { + if s := strings.TrimSpace(e); s != "" { + aliases = append(aliases, s) + } + } + } + description := strings.TrimSpace(stringOf(payload["description"])) + if description == "" { + description = strings.TrimSpace(stringOf(payload["definition_excerpt"])) + } + return map[string]any{ + "aliases": aliases, + "mention_count": 1, + "name": name, + "source_chunk_ids": unionOrdered(sourceChunkIDs), + "type": typ, + "description": description, + } +} + +// graphRelation mirrors _struct_graph_relation: project a relation payload to +// the {from, to, type} edge shape. +func graphRelation(payload map[string]any) map[string]any { + src := relationEndpoint(payload, "", "source", "src", "from") + tgt := relationEndpoint(payload, "", "target", "tgt", "to") + if src == "" || tgt == "" { + return nil + } + typ := strings.TrimSpace(stringOf(payload["type"])) + if typ == "" { + typ = "related" + } + return map[string]any{"from": src, "to": tgt, "type": typ} +} + +// mergeGraphEntities mirrors _struct_merge_graph_entities: collapse entities +// sharing (name, type) — mention_count sums, aliases/description/ +// source_chunk_ids union, first-seen order preserved. +func mergeGraphEntities(entities []map[string]any) []map[string]any { + type key struct{ name, typ string } + merged := map[key]int{} + var order []key + var out []map[string]any + for _, entity := range entities { + k := key{name: entity["name"].(string), typ: entity["type"].(string)} + idx, ok := merged[k] + if !ok { + merged[k] = len(out) + order = append(order, k) + out = append(out, entity) + continue + } + target := out[idx] + target["mention_count"] = mentionOf(target) + mentionOf(entity) + target["aliases"] = unionOrdered(toStrings(target["aliases"]), toStrings(entity["aliases"])) + if target["description"] == "" && entity["description"] != "" { + target["description"] = entity["description"] + } + target["source_chunk_ids"] = unionOrdered(toStrings(target["source_chunk_ids"]), toStrings(entity["source_chunk_ids"])) + } + return out +} + +// RebuildStructureGraph mirrors _struct_rebuild_graph_json: rebuild the +// compact {"entities": [...], "relations": [...]} graph from the surviving +// structure products. Entity mention counts sum across rows sharing +// (name, type); relation rows project verbatim (already deduped in-run). +func RebuildStructureGraph(products []common.Product) map[string]any { + var entities []map[string]any + var relations []map[string]any + for _, p := range products { + kind, _ := p.Meta["kind"].(string) + payload := parsePayload(p.Content) + if payload == nil { + continue + } + switch kind { + case "relation": + if rel := graphRelation(payload); rel != nil { + relations = append(relations, rel) + } + case "entity": + if ent := graphEntity(payload, metaStrings(p.Meta, "source_chunk_ids")); ent != nil { + entities = append(entities, ent) + } + } + } + return map[string]any{ + "entities": mergeGraphEntities(entities), + "relations": relations, + } +} + +func mentionOf(entity map[string]any) int { + switch v := entity["mention_count"].(type) { + case int: + return v + case float64: + return int(v) + } + return 1 +} + +func toStrings(v any) []string { + switch x := v.(type) { + case []string: + return x + case []any: + out := make([]string, 0, len(x)) + for _, e := range x { + if s, ok := e.(string); ok { + out = append(out, s) + } + } + return out + } + return nil +} diff --git a/internal/ingestion/component/knowledge_compiler/structure/merge.go b/internal/ingestion/component/knowledge_compiler/structure/merge.go new file mode 100644 index 0000000000..f375875079 --- /dev/null +++ b/internal/ingestion/component/knowledge_compiler/structure/merge.go @@ -0,0 +1,561 @@ +package structure + +import ( + "context" + "fmt" + "strings" + "sync" + + "ragflow/internal/ingestion/component/knowledge_compiler/common" +) + +// Merge prompts, verbatim from Python structure.py. The merge is driven by +// the user-supplied prompts; the decision instruction lets us branch on the +// LLM's verdict via GenJSON. +const mergeSystemPrompt = `You are an intelligent data merging assistant. +You will merge two JSON objects representing the same entity: Item A (existing) and Item B (incoming). + +Merge strategy: +1. Combine information from both items. +2. If fields conflict, use your best judgment to pick the more detailed or recent-looking value. +3. If one item has a null/missing value and the other has data, keep the data. +4. For list fields, combine unique elements from both. +5. Do not invent new information not present in the inputs. +6. Return the result in the exact JSON format of the input items.` + +const mergeUserPrompt = "Item A (existing):\n%s\n\nItem B (incoming):\n%s" + +const mergeDecisionInstruction = `First decide whether Item A and Item B refer to the same logical entity (for entities) or the same logical relation (for relations). Use the merge strategy above only if they are the same. + +Return ONLY a JSON object with this exact structure (no markdown fences, no commentary): +{ + "duplicated": <true | false>, + "merged": <merged JSON object using the same keys as the inputs when duplicated=true; otherwise null> +}` + +// mergeJudgeTemperature mirrors Python's gen_conf for merge judging +// (_struct_merge_pair uses temperature 0.0). +var mergeJudgeTemperature = 0.0 + +// MergeDecision is the variant-level decision a MergeDecider returns for one +// incoming product against its best in-run match. +type MergeDecision int + +const ( + // DecisionKeepBoth keeps the incoming product as a new entry. + DecisionKeepBoth MergeDecision = iota + // DecisionDropIncoming discards the incoming product (duplicate). + DecisionDropIncoming + // DecisionMerge replaces the existing entry with the returned product + // (identity preserved via the existing id). + DecisionMerge +) + +// MergeDecider decides how to treat an incoming product given its best +// in-run cosine match. When it returns DecisionMerge it also returns the +// merged replacement product; for other decisions the product is ignored. +type MergeDecider interface { + Decide(ctx context.Context, existing, incoming common.Product, bestScore float64) (MergeDecision, common.Product, error) +} + +// CosineDecider drops the incoming product when its best cosine match meets +// (>=) Threshold. It performs no LLM call, so it is safe under the store lock. +type CosineDecider struct{ Threshold float64 } + +func (d CosineDecider) Decide(_ context.Context, _, _ common.Product, bestScore float64) (MergeDecision, common.Product, error) { + if bestScore >= d.Threshold { + return DecisionDropIncoming, common.Product{}, nil + } + return DecisionKeepBoth, common.Product{}, nil +} + +// LLMMergeDecider mirrors Python's local-dedup judge (_struct_merge_pair): +// pairs above Threshold are adjudicated by the LLM; on a duplicated verdict +// the merged payload replaces the existing row (re-embedded, provenance +// unioned, id preserved). Entity merges record aliases so relation endpoints +// can be rewritten afterwards (mirrors _struct_local_dedup's entity_aliases). +type LLMMergeDecider struct { + Chat common.ChatInvoker + LLMID string + Embed common.Embedder + Threshold float64 + + mu sync.Mutex + aliases map[string]string +} + +// NewLLMMergeDecider constructs the decider used by the structure variant. +func NewLLMMergeDecider(chat common.ChatInvoker, llmID string, embed common.Embedder, threshold float64) *LLMMergeDecider { + return &LLMMergeDecider{Chat: chat, LLMID: llmID, Embed: embed, Threshold: threshold, aliases: map[string]string{}} +} + +// Decide implements MergeDecider. +func (d *LLMMergeDecider) Decide(ctx context.Context, existing, incoming common.Product, bestScore float64) (MergeDecision, common.Product, error) { + if bestScore < d.Threshold { + return DecisionKeepBoth, common.Product{}, nil + } + merged, err := mergePair(ctx, d.Chat, d.LLMID, existing.Content, incoming.Content) + if err != nil { + return DecisionKeepBoth, common.Product{}, err + } + if merged == nil { + return DecisionKeepBoth, common.Product{}, nil + } + + kind, _ := existing.Meta["kind"].(string) + if kind == "entity" { + oldName := entityNameValue(existing) + incomingName := entityNameValue(incoming) + canonical := entityName(merged) + if canonical == "" { + canonical = oldName + } + for _, alias := range []string{oldName, incomingName} { + if alias != "" && alias != canonical { + d.recordAlias(alias, canonical) + } + } + } + merged = applyMergeInvariants(existing, merged) + + chunkIDs := unionOrdered(metaStrings(existing.Meta, "source_chunk_ids"), metaStrings(incoming.Meta, "source_chunk_ids")) + texts := []string{payloadDescription(merged)} + vecs, err := d.Embed.Encode(ctx, texts) + if err != nil { + return DecisionKeepBoth, common.Product{}, err + } + if len(vecs) == 0 { + return DecisionKeepBoth, common.Product{}, fmt.Errorf("knowledge_compiler: re-embed of merged payload returned no vector") + } + + meta := map[string]any{} + for k, v := range existing.Meta { + meta[k] = v + } + meta["source_chunk_ids"] = chunkIDs + refreshMetaFromPayload(meta, kind, merged) + replacement := common.Product{ + ID: existing.ID, + DocID: existing.DocID, + TenantID: existing.TenantID, + Variant: existing.Variant, + Content: payloadJSON(merged), + Vector: vecs[0], + Meta: meta, + } + return DecisionMerge, replacement, nil +} + +// recordAlias adds one alias→canonical mapping (thread-safe). +func (d *LLMMergeDecider) recordAlias(alias, canonical string) { + d.mu.Lock() + defer d.mu.Unlock() + d.aliases[alias] = canonical +} + +// Aliases returns the accumulated alias→canonical map (a copy). +func (d *LLMMergeDecider) Aliases() map[string]string { + d.mu.Lock() + defer d.mu.Unlock() + out := make(map[string]string, len(d.aliases)) + for k, v := range d.aliases { + out[k] = v + } + return out +} + +// mergePair mirrors _struct_merge_pair: LLM-judged merge of two payload JSON +// documents. It returns the merged payload when the pair is a duplicate, nil +// when not (including unparseable inputs, which Python logs and skips). +func mergePair(ctx context.Context, chat common.ChatInvoker, llmID, existingContent, incomingContent string) (map[string]any, error) { + existingPayload := parsePayload(existingContent) + incomingPayload := parsePayload(incomingContent) + if existingPayload == nil || incomingPayload == nil { + return nil, nil + } + res, err := common.GenJSON(ctx, chat, common.ChatRequest{ + LLMID: llmID, + SystemPrompt: mergeSystemPrompt + "\n\n" + mergeDecisionInstruction, + UserPrompt: fmt.Sprintf(mergeUserPrompt, payloadJSON(existingPayload), payloadJSON(incomingPayload)), + Temperature: &mergeJudgeTemperature, + }) + if err != nil { + return nil, err + } + if dup, _ := res["duplicated"].(bool); !dup { + return nil, nil + } + merged, ok := res["merged"].(map[string]any) + if !ok { + return nil, nil + } + return merged, nil +} + +// applyMergeInvariants mirrors _struct_apply_merge_invariants: for relations, +// force the source/target fields back to the existing payload's values — +// from_entity_kwd / to_entity_kwd must not change across a merge. +func applyMergeInvariants(existing common.Product, mergedPayload map[string]any) map[string]any { + kind, _ := existing.Meta["kind"].(string) + if kind != "relation" { + return mergedPayload + } + existingPayload := parsePayload(existing.Content) + if existingPayload == nil { + return mergedPayload + } + for _, field := range []string{"source", "src", "from"} { + if v, ok := existingPayload[field]; ok { + mergedPayload[field] = v + } + } + for _, field := range []string{"target", "tgt", "to"} { + if v, ok := existingPayload[field]; ok { + mergedPayload[field] = v + } + } + return mergedPayload +} + +// resolveAlias mirrors _struct_resolve_entity_alias: follow the alias chain +// until a non-aliased name, tolerating cycles (stop at the first repeat). +func resolveAlias(name string, aliases map[string]string) string { + current := strings.TrimSpace(name) + seen := map[string]bool{} + for { + next, ok := aliases[current] + if !ok || seen[current] { + return current + } + seen[current] = true + current = next + } +} + +// rewriteRelationPayload mirrors _struct_rewrite_relation_payload: resolve +// the relation's endpoint fields through the alias map in place. Returns +// whether anything changed. +func rewriteRelationPayload(payload map[string]any, aliases map[string]string) bool { + changed := false + for _, fields := range [][]string{{"source", "src", "from"}, {"target", "tgt", "to"}} { + for _, field := range fields { + v, ok := payload[field] + if !ok || v == nil { + continue + } + old := strings.TrimSpace(stringOf(v)) + if old == "" { + continue + } + if newName := resolveAlias(old, aliases); newName != old { + payload[field] = newName + changed = true + } + } + } + return changed +} + +// refreshMetaFromPayload re-derives the name/type/endpoint meta keys of a +// merged row from its new payload (mirrors the column re-derivation in +// _struct_rebuild_doc_storage_doc). +func refreshMetaFromPayload(meta map[string]any, kind string, payload map[string]any) { + if kind == "entity" { + if name := entityName(payload); name != "" { + meta["name"] = name + } else { + delete(meta, "name") + } + if typ := strings.TrimSpace(stringOf(payload["type"])); typ != "" { + meta["entity_type"] = typ + } + if desc := strings.TrimSpace(stringOf(payload["description"])); desc != "" { + meta["description"] = desc + } + return + } + if from := relationEndpoint(payload, "", "source", "src", "from"); from != "" { + meta["from"] = from + } + if to := relationEndpoint(payload, "", "target", "tgt", "to"); to != "" { + meta["to"] = to + } + if typ := strings.TrimSpace(stringOf(payload["type"])); typ != "" { + meta["relation_type"] = typ + } +} + +// entityNameValue reads an entity row's canonical name (payload first, then +// the meta stamp), mirroring _struct_entity_name. +func entityNameValue(p common.Product) string { + if payload := parsePayload(p.Content); payload != nil { + if name := entityName(payload); name != "" { + return name + } + } + name, _ := p.Meta["name"].(string) + return strings.TrimSpace(name) +} + +// metaStrings reads a []string meta value (tolerating []any). +func metaStrings(meta map[string]any, key string) []string { + switch v := meta[key].(type) { + case []string: + return v + case []any: + out := make([]string, 0, len(v)) + for _, e := range v { + if s, ok := e.(string); ok { + out = append(out, s) + } + } + return out + } + return nil +} + +// unionOrdered mirrors _common.union_ordered: concatenated, deduped, +// first-seen order preserved. +func unionOrdered(lists ...[]string) []string { + seen := map[string]bool{} + var out []string + for _, lst := range lists { + for _, v := range lst { + if v == "" || seen[v] { + continue + } + seen[v] = true + out = append(out, v) + } + } + return out +} + +// MergeStats accumulates dedup outcomes across a run. +type MergeStats struct { + Inserted int + Updated int + DuplicatesDropped int +} + +// MergeIntoStore folds the incoming rows into store, using DedupeAdd so the +// check-then-act is atomic and the decider (which may call the LLM) runs +// outside the store lock. It returns counts of inserted/merged/dropped rows. +// +// When a row is dropped as a duplicate, its source_chunk_ids are still merged +// into the surviving entry's source_chunk_ids (order-preserving union), so an +// entity that appears across multiple chunks accumulates provenance from all +// of them — mirroring Python's _struct_merge_graph_entities. +func MergeIntoStore(ctx context.Context, store *common.MemStore, decider MergeDecider, rows []common.Product) (MergeStats, error) { + var stats MergeStats + for _, row := range rows { + var droppedExisting common.Product + var dropped bool + action, err := store.DedupeAdd(row, 0, func(existing common.Product, score float64) (common.KeepAction, common.Product, error) { + d, replacement, e := decider.Decide(ctx, existing, row, score) + if e != nil { + return common.KeepAdd, common.Product{}, e + } + switch d { + case DecisionDropIncoming: + // Capture the existing entry so we can fold the incoming + // source_chunk_ids into it after the drop (provenance union, + // mirroring Python's _struct_merge_graph_entities). + droppedExisting = existing + dropped = true + return common.KeepDrop, common.Product{}, nil + case DecisionMerge: + return common.KeepMerge, replacement, nil + default: + return common.KeepAdd, common.Product{}, nil + } + }) + if err != nil { + return stats, err + } + switch action { + case common.KeepDrop: + stats.DuplicatesDropped++ + // Fold the incoming source_chunk_ids into the surviving entry. + if dropped && droppedExisting.ID != "" { + if ids := metaStrings(row.Meta, "source_chunk_ids"); len(ids) > 0 { + store.MergeSourceChunkIDs(droppedExisting.ID, ids) + } + } + case common.KeepMerge: + stats.Updated++ + stats.DuplicatesDropped++ + default: + stats.Inserted++ + } + } + return stats, nil +} + +// GroupedDeduper mirrors _struct_local_dedup's group-by-filter-key behaviour: +// rows are bucketed by their relation endpoints (entities share one bucket, +// mirroring _struct_filter_key where entity rows carry no from/to), so an +// entity never merges with a relation and relations only merge with +// same-endpoint relations. One MemStore per bucket keeps cosine candidates +// group-scoped. doc/compile/template are constant per run and therefore not +// part of the key (they are in Python's key for the cross-document ES case). +type GroupedDeduper struct { + decider MergeDecider + + mu sync.Mutex + order []string + stores map[string]*common.MemStore + stats MergeStats +} + +// NewGroupedDeduper constructs a deduper around the given decider. +func NewGroupedDeduper(decider MergeDecider) *GroupedDeduper { + return &GroupedDeduper{decider: decider, stores: map[string]*common.MemStore{}} +} + +// groupKey mirrors _struct_filter_key minus the per-run constants. +func groupKey(row common.Product) string { + from, _ := row.Meta["from"].(string) + to, _ := row.Meta["to"].(string) + return from + "\x00" + to +} + +// Add folds one row into its endpoint bucket. Safe for concurrent use, but +// callers that need deterministic merge outcomes (mirroring Python's +// sequential _struct_local_dedup) should call it sequentially in batch order. +func (g *GroupedDeduper) Add(ctx context.Context, row common.Product) error { + key := groupKey(row) + g.mu.Lock() + store, ok := g.stores[key] + if !ok { + store = common.NewMemStore() + g.stores[key] = store + g.order = append(g.order, key) + } + g.mu.Unlock() + + s, err := MergeIntoStore(ctx, store, g.decider, []common.Product{row}) + if err != nil { + return err + } + g.mu.Lock() + g.stats.Inserted += s.Inserted + g.stats.Updated += s.Updated + g.stats.DuplicatesDropped += s.DuplicatesDropped + g.mu.Unlock() + return nil +} + +// Stats returns the accumulated dedup counters. +func (g *GroupedDeduper) Stats() MergeStats { + g.mu.Lock() + defer g.mu.Unlock() + return g.stats +} + +// Rows concatenates every bucket's surviving rows in first-seen bucket order. +func (g *GroupedDeduper) Rows() []common.Product { + g.mu.Lock() + defer g.mu.Unlock() + var out []common.Product + for _, key := range g.order { + out = append(out, g.stores[key].Snapshot()...) + } + return out +} + +// RewriteRelations mirrors the alias-rewrite tail of _struct_local_dedup: +// when entity merges recorded aliases, every relation row's endpoints are +// resolved through them (re-embedding the changed payloads, preserving row +// ids), then ALL relation rows are re-folded through the deduper so relations +// that became identical collapse (Python's second pass with +// rewrite_relations=False). It is a no-op without aliases. +func (g *GroupedDeduper) RewriteRelations(ctx context.Context, aliases map[string]string, embed common.Embedder) error { + if len(aliases) == 0 { + return nil + } + + // Detach every relation row from its bucket (entity rows stay put). + g.mu.Lock() + var relations []common.Product + keptOrder := make([]string, 0, len(g.order)) + for _, key := range g.order { + store := g.stores[key] + var remaining []common.Product + for _, row := range store.Snapshot() { + if kind, _ := row.Meta["kind"].(string); kind == "relation" { + relations = append(relations, row) + } else { + remaining = append(remaining, row) + } + } + if len(remaining) == len(store.Snapshot()) { + keptOrder = append(keptOrder, key) + continue + } + delete(g.stores, key) + if len(remaining) > 0 { + fresh := common.NewMemStore() + for _, row := range remaining { + fresh.Add(row) + } + g.stores[key] = fresh + keptOrder = append(keptOrder, key) + } + } + g.order = keptOrder + g.mu.Unlock() + + if len(relations) == 0 { + return nil + } + + // Rewrite endpoints; re-embed only the rows whose payload changed + // (mirrors _struct_rewrite_relation_doc). Row ids are preserved. + var changedIdx []int + for i, row := range relations { + payload := parsePayload(row.Content) + if payload == nil { + continue + } + if !rewriteRelationPayload(payload, aliases) { + continue + } + row.Content = payloadJSON(payload) + if row.Meta == nil { + row.Meta = map[string]any{} + } + if from := relationEndpoint(payload, "", "source", "src", "from"); from != "" { + row.Meta["from"] = from + } + if to := relationEndpoint(payload, "", "target", "tgt", "to"); to != "" { + row.Meta["to"] = to + } + relations[i] = row + changedIdx = append(changedIdx, i) + } + if len(changedIdx) > 0 { + texts := make([]string, len(changedIdx)) + for i, idx := range changedIdx { + texts[i] = payloadDescription(parsePayload(relations[idx].Content)) + } + vecs, err := embed.Encode(ctx, texts) + if err != nil { + return err + } + if len(vecs) != len(changedIdx) { + return fmt.Errorf("knowledge_compiler: re-embed after relation rewrite returned %d vectors for %d rows", len(vecs), len(changedIdx)) + } + for i, idx := range changedIdx { + relations[idx].Vector = vecs[i] + } + } + + // Second dedup pass over the rewritten relation set. + for _, row := range relations { + if err := g.Add(ctx, row); err != nil { + return err + } + } + return nil +} diff --git a/internal/ingestion/component/knowledge_compiler/structure/prompt.go b/internal/ingestion/component/knowledge_compiler/structure/prompt.go new file mode 100644 index 0000000000..006b956793 --- /dev/null +++ b/internal/ingestion/component/knowledge_compiler/structure/prompt.go @@ -0,0 +1,384 @@ +package structure + +import ( + "fmt" + "strings" + "time" + + "ragflow/internal/ingestion/component/knowledge_compiler/common" +) + +// Prompt construction for the structure (graph) variant, mirroring Python +// structure.py: `_struct_infer_type`, `_struct_render_fields`, +// `_struct_render_type_fields` and `_struct_hypergraph_prompts`. All three +// compile kinds (list / set / hypergraph) share the same two-stage +// node → edge prompts and the {"items": [...]} response contract; they differ +// only in the Auto-type label and (for set) the uniqueness rule. + +// Type identifies the structure compile kind. +type Type string + +const ( + TypeList Type = "list" + TypeSet Type = "set" + TypeHypergraph Type = "hypergraph" +) + +// typeAliases mirrors Python's _STRUCT_TYPE_ALIASES: "graph" and +// "knowledge_graph" both mean the hypergraph compile. +var typeAliases = map[string]Type{ + "graph": TypeHypergraph, + "knowledge_graph": TypeHypergraph, +} + +// normalizeKind mirrors _struct_normalize_kind. +func normalizeKind(raw any) string { + s, ok := raw.(string) + if !ok { + return "" + } + return strings.ReplaceAll(strings.ToLower(strings.TrimSpace(s)), "-", "_") +} + +// InferType mirrors _struct_infer_type: explicit compile_type restricted to +// the known compile kinds; then kind returned VERBATIM (aliases resolved) — +// Python stamps arbitrary kinds ("timeline", "page_index") as the autotype; +// then the legacy output.entities+output.relations shape; default "list". +func InferType(parserConfig map[string]any) Type { + if t, ok := resolveTypeAlias(normalizeKind(common.Get(parserConfig, "compile_type"))); ok { + return t + } + kind := normalizeKind(common.Get(parserConfig, "kind")) + if t, ok := typeAliases[kind]; ok { + return t + } + if kind != "" { + return Type(kind) + } + output := common.GetMap(parserConfig, "output") + if common.Get(output, "entities") != nil && common.Get(output, "relations") != nil { + return TypeHypergraph + } + return TypeList +} + +func resolveTypeAlias(normalized string) (Type, bool) { + if t, ok := typeAliases[normalized]; ok { + return t, true + } + switch Type(normalized) { + case TypeList, TypeSet, TypeHypergraph: + return Type(normalized), true + } + return "", false +} + +// renderFields mirrors _struct_render_fields (legacy output.entities / +// output.relations field shape). Returns (bulleted descriptions, JSON skeleton). +func renderFields(fields []any, lang string) (string, string) { + var lines []string + var skel []string + for _, raw := range fields { + f, ok := raw.(map[string]any) + if !ok { + continue + } + name, _ := f["name"].(string) + ftype, _ := f["type"].(string) + if ftype == "" { + ftype = "str" + } + desc := common.Localize(f["description"], lang) + reqLabel := "required" + if req, ok := f["required"].(bool); ok && !req { + reqLabel = "optional" + } + lines = append(lines, fmt.Sprintf("- %s (%s, %s): %s", name, ftype, reqLabel, desc)) + var placeholder string + switch ftype { + case "list": + placeholder = "[<string>, ...]" + case "int": + placeholder = "<int>" + case "float": + placeholder = "<float>" + case "bool": + placeholder = "<true|false>" + default: + placeholder = "<string>" + } + skel = append(skel, fmt.Sprintf("%q: %s", name, placeholder)) + } + return strings.Join(lines, "\n"), "{ " + strings.Join(skel, ", ") + " }" +} + +// renderTypeFields mirrors _struct_render_type_fields (new compilation-template +// shape: allowed item `type` values with descriptions/rules). +func renderTypeFields(fields []any, lang, kind string) (string, string) { + var lines []string + var typeValues []string + for _, raw := range fields { + f, ok := raw.(map[string]any) + if !ok { + continue + } + typ := strings.TrimSpace(firstStringOf(f["type"])) + if typ == "" { + continue + } + typeValues = append(typeValues, typ) + lines = append(lines, "- type: "+typ) + if desc := common.Localize(f["description"], lang); desc != "" { + lines = append(lines, " description: "+desc) + } + if rule := common.Localize(f["rule"], lang); rule != "" { + lines = append(lines, " rule: "+rule) + } + } + if len(typeValues) == 0 { + typeValues = append(typeValues, "other") + lines = append(lines, "- type: other") + } + oneOf := strings.Join(typeValues, "|") + var skeleton string + if kind == "relation" { + skeleton = `{ "type": "<one of: ` + oneOf + `>", "source": "<known entity name>", "target": "<known entity name>", "description": "<evidence or relation description>", "source_chunk_ids": ["<source chunk id>", ...] }` + } else { + skeleton = `{ "type": "<one of: ` + oneOf + `>", "name": "<exact extracted item text>", "description": "<evidence, definition, or detail from the source>", "source_chunk_ids": ["<source chunk id>", ...] }` + } + return strings.Join(lines, "\n"), skeleton +} + +// HypergraphPrompts mirrors _struct_hypergraph_prompts: it renders the node +// (entity) and edge (relation) system prompts from the parser_config template. +// The returned edge prompt still contains the "{known_nodes}" placeholder, +// filled per batch from the stage-1 entities (see fillKnownNodes). An empty +// edge prompt means the config declares no relations and stage 2 is skipped. +func HypergraphPrompts(parserConfig map[string]any, lang string) (nodePrompt, edgePrompt string) { + autotype := InferType(parserConfig) + guideline := common.GetMap(parserConfig, "guideline") + output := common.GetMap(parserConfig, "output") + options := common.GetMap(parserConfig, "options") + + target := common.Localize(common.Get(guideline, "target"), lang) + rulesE := common.Localize(common.Get(guideline, "rules_for_entities"), lang) + rulesR := common.Localize(common.Get(guideline, "rules_for_relations"), lang) + rulesT := common.Localize(common.Get(guideline, "rules_for_time"), lang) + globalRules := common.Localize(common.Get(parserConfig, "global_rules"), lang) + + observationTime, _ := common.Get(options, "observation_time").(string) + if strings.TrimSpace(observationTime) == "" { + observationTime = time.Now().Format("2006-01-02") + } + if strings.Contains(rulesT, "{observation_time}") { + rulesT = strings.ReplaceAll(rulesT, "{observation_time}", observationTime) + } + + usesTemplateShape := common.Get(parserConfig, "entity") != nil || common.Get(parserConfig, "relation") != nil + var entitiesCfg, relationsCfg map[string]any + if usesTemplateShape { + entitiesCfg = common.GetMap(parserConfig, "entity") + relationsCfg = common.GetMap(parserConfig, "relation") + } else { + entitiesCfg = common.GetMap(output, "entities") + relationsCfg = common.GetMap(output, "relations") + } + entDesc := common.Localize(common.Get(entitiesCfg, "description"), lang) + relDesc := common.Localize(common.Get(relationsCfg, "description"), lang) + entFields := configFields(entitiesCfg) + relFields := configFields(relationsCfg) + + var entFieldsText, entSkel, relFieldsText, relSkel string + if usesTemplateShape { + entFieldsText, entSkel = renderTypeFields(entFields, lang, "entity") + relFieldsText, relSkel = renderTypeFields(relFields, lang, "relation") + } else { + entFieldsText, entSkel = renderFields(entFields, lang) + relFieldsText, relSkel = renderFields(relFields, lang) + } + + var nodeParts []string + if target != "" { + nodeParts = append(nodeParts, "# Role and Task:\n"+target) + } + if globalRules != "" { + nodeParts = append(nodeParts, "## Global Rules:\n"+globalRules) + } + if rulesE != "" { + nodeParts = append(nodeParts, "## Entity Extraction Rules:\n"+rulesE) + } + if entDesc != "" { + nodeParts = append(nodeParts, "## Entity Description:\n"+entDesc) + } + nodeParts = append(nodeParts, "## Entity Fields:\n"+entFieldsText) + unique := "" + if autotype == TypeSet { + unique = "Items must be unique. " + } + nodeParts = append(nodeParts, + "## Response Format:\n"+ + "Reply with a single JSON object of the form: "+ + fmt.Sprintf(`{"items": [%s, ...]}.`, entSkel)+"\n"+ + fmt.Sprintf(`Auto-type: %q. `, string(autotype))+unique+"Return JSON only, no commentary.") + nodePrompt = strings.Join(nodeParts, "\n\n") + + if len(relationsCfg) == 0 { + return nodePrompt, "" + } + + var edgeParts []string + if target != "" { + edgeParts = append(edgeParts, "# Role and Task:\n"+target) + } + if globalRules != "" { + edgeParts = append(edgeParts, "## Global Rules:\n"+globalRules) + } + if rulesR != "" { + edgeParts = append(edgeParts, "## Relation Extraction Rules:\n"+rulesR) + } + if rulesT != "" { + edgeParts = append(edgeParts, "## Time Rules:\n"+rulesT) + } + if relDesc != "" { + edgeParts = append(edgeParts, "## Relation Description:\n"+relDesc) + } + edgeParts = append(edgeParts, "## Relation Fields:\n"+relFieldsText) + edgeParts = append(edgeParts, "## Known Entities:\n{known_nodes}") + edgeParts = append(edgeParts, + "## Response Format:\n"+ + "Reply with a single JSON object of the form: "+ + fmt.Sprintf(`{"items": [%s, ...]}.`, relSkel)+"\n"+ + "Only create relations between entities listed in 'Known Entities'. "+ + "Return JSON only, no commentary.") + edgePrompt = strings.Join(edgeParts, "\n\n") + return nodePrompt, edgePrompt +} + +// fillKnownNodes renders the stage-2 edge prompt by substituting the +// {known_nodes} placeholder with the bulleted known-entity list ("(none)" +// when stage 1 found nothing), mirroring _struct_extract_hypergraph. +func fillKnownNodes(edgePromptTemplate string, knownKeys []string) string { + knownStr := "(none)" + if len(knownKeys) > 0 { + knownStr = "- " + strings.Join(knownKeys, "\n- ") + } + return strings.ReplaceAll(edgePromptTemplate, "{known_nodes}", knownStr) +} + +// UserPrompt wraps one batch's packed chunk text exactly as +// _struct_extract_hypergraph does; the same user prompt serves both stages. +func UserPrompt(packedText string) string { + return "## Source Text:\n" + + "Each source chunk is enclosed by [CHUNK_ID: ...] and [END_CHUNK]. " + + "For every entity and relation, return source_chunk_ids containing only " + + "the IDs of chunks that support that item.\n" + + packedText + "\n\n## Output (JSON only):" +} + +// PackBatch renders one batch's chunks as [CHUNK_ID: id]...[END_CHUNK] +// segments (mirrors _struct_process_batch) and returns the batch chunk ids in +// order. Chunks without text are skipped, matching Python's isinstance check. +func PackBatch(batch []common.Chunk) (string, []string) { + segments := make([]string, 0, len(batch)) + ids := make([]string, 0, len(batch)) + for i, c := range batch { + id := strings.TrimSpace(c.ID) + if id == "" { + id = fmt.Sprintf("chunk-%d", i+1) + } + text := common.FirstNonEmpty(c.Text, c.Content) + if strings.TrimSpace(text) == "" { + continue + } + ids = append(ids, id) + segments = append(segments, "[CHUNK_ID: "+id+"]\n"+text+"\n[END_CHUNK]") + } + return strings.Join(segments, "\n\n"), ids +} + +// EntityIDField mirrors _struct_entity_id_field. +func EntityIDField(parserConfig map[string]any) string { + if common.Get(parserConfig, "entity") != nil { + return "name" + } + identifiers := common.GetMap(parserConfig, "identifiers") + if entityID, ok := common.Get(identifiers, "entity_id").(string); ok { + if s := strings.TrimSpace(entityID); s != "" && !strings.Contains(s, "{") { + return s + } + } + output := common.GetMap(parserConfig, "output") + entitiesCfg := common.GetMap(output, "entities") + for _, raw := range configFields(entitiesCfg) { + if f, ok := raw.(map[string]any); ok { + if req, isBool := f["required"].(bool); !isBool || req { + if name, _ := f["name"].(string); name != "" { + return name + } + } + } + } + return "name" +} + +// RelationMemberFields mirrors _struct_relation_member_fields: it returns the +// payload keys carrying a relation's endpoints, or ("","") when the schema +// declares none. +func RelationMemberFields(parserConfig map[string]any) (string, string) { + identifiers := common.GetMap(parserConfig, "identifiers") + if members := common.GetMap(identifiers, "relation_members"); members != nil { + src, _ := members["source"].(string) + if src == "" { + src, _ = members["src"].(string) + } + tgt, _ := members["target"].(string) + if tgt == "" { + tgt, _ = members["tgt"].(string) + } + if src != "" || tgt != "" { + return src, tgt + } + } + if common.Get(parserConfig, "relation") != nil { + return "source", "target" + } + output := common.GetMap(parserConfig, "output") + relationsCfg := common.GetMap(output, "relations") + names := map[string]bool{} + for _, raw := range configFields(relationsCfg) { + if f, ok := raw.(map[string]any); ok { + if name, _ := f["name"].(string); name != "" { + names[name] = true + } + } + } + if names["source"] && names["target"] { + return "source", "target" + } + return "", "" +} + +// configFields extracts cfg["fields"] as a []any (nil when absent). +func configFields(cfg map[string]any) []any { + if cfg == nil { + return nil + } + switch v := cfg["fields"].(type) { + case []any: + return v + case []map[string]any: + out := make([]any, 0, len(v)) + for _, f := range v { + out = append(out, f) + } + return out + } + return nil +} + +func firstStringOf(v any) string { + if s, ok := v.(string); ok { + return s + } + return "" +} diff --git a/internal/ingestion/component/knowledge_compiler/structure/structure.go b/internal/ingestion/component/knowledge_compiler/structure/structure.go new file mode 100644 index 0000000000..a36d64391e --- /dev/null +++ b/internal/ingestion/component/knowledge_compiler/structure/structure.go @@ -0,0 +1,205 @@ +// Package structure implements the "structure" variant of KnowledgeCompiler: +// document-level structure compilation (list / set / hypergraph — the graph +// kind) as a two-stage entity → relation LLM extraction with template-driven +// prompts, followed by LLM-judged in-run merge dedup. Stage semantics and +// prompts mirror Python's rag/advanced_rag/knowlege_compile/structure.py; the +// Go port keeps all intermediate state in memory (no ES reads/writes). +package structure + +import ( + "context" + "fmt" + + "ragflow/internal/ingestion/component/knowledge_compiler/common" + "ragflow/internal/utility" +) + +// structureBatchTokenBudget caps one extraction batch's packed chunk tokens. +// Python derives the budget from chat_mdl.max_length minus the prompt +// overhead; the Go ChatInvoker seam does not expose the model window, so we +// use the same conservative constant the wiki variant uses. +const structureBatchTokenBudget = 4096 + +// Run executes the structure variant: +// 1. MAP — per-batch two-stage (node → edge) extraction, parallel across +// batches, results kept in batch order (mirrors _run_chunked_pipeline). +// 2. DEDUP — sequential LLM-judged merge in batch order, grouped by +// relation endpoints, then a relation-rewrite pass for entity aliases +// (mirrors _struct_local_dedup). +// 3. KIND POST-PROCESSING — chain validation for list/timeline (LLM +// correction, fail-open) and the timeline orphan-entity filter (mirrors +// validate_and_correct_chain + cleanup_timeline_isolated_entities). +// 4. GRAPH — one compact {"entities","relations"} summary row (mirrors +// _struct_rebuild_graph_json). +// +// It never writes ES; the downstream writer persists the returned products. +func Run(ctx context.Context, deps common.Deps, param common.Param, inputs common.Inputs) (common.Outputs, error) { + parserConfig, _ := inputs.VariantSpecific["parser_config"].(map[string]any) + compileType := InferType(parserConfig) + docID := common.FirstNonEmpty(inputs.DocID, deps.DatasetID, "unknown") + llmID := common.FirstNonEmpty(param.LLMID, inputs.LLMID) + cfg := CompileConfig{ + LLMID: llmID, + Type: compileType, + TenantID: deps.TenantID, + DocID: docID, + Variant: common.VariantStructure, + Lang: param.Language, + ParserConfig: parserConfig, + TemplateID: common.FirstNonEmpty(param.TemplateIDs...), + } + + nodePrompt, edgePromptTmpl := HypergraphPrompts(parserConfig, param.Language) + + // ---- MAP ---- + batches := common.PackBatches(inputs.Chunks, structureBatchTokenBudget, deps.Tokenizer) + perBatch := make([][]common.Product, len(batches)) + n := param.MaxWorkers + if n <= 0 { + n = 1 + } + wp := utility.NewWorkerPool[func() error, struct{}](n, n, + func(_ context.Context, fn func() error) (struct{}, error) { return struct{}{}, fn() }) + var futs []utility.WorkerPoolFuture[func() error, struct{}] + for i, batch := range batches { + i, batch := i, batch + f, err := wp.Submit(ctx, func() error { + packed, batchIDs := PackBatch(batch) + if len(batchIDs) == 0 { + return nil + } + nodes, edges, err := extractHypergraph(ctx, deps, cfg, nodePrompt, edgePromptTmpl, packed) + if err != nil { + return err + } + rows, err := buildRows(ctx, deps, cfg, nodes, edges, batchIDs) + if err != nil { + return err + } + perBatch[i] = rows + return nil + }) + if err != nil { + wp.StopWait() + return common.Outputs{}, err + } + futs = append(futs, f) + } + wp.StopWait() + for _, f := range futs { + if res, _ := f.Wait(ctx); res.Err != nil { + return common.Outputs{}, res.Err + } + } + + // ---- DEDUP ---- + // Sequential in batch order so merge outcomes are deterministic and match + // Python's _struct_local_dedup (which folds docs in list order). + decider := NewLLMMergeDecider(deps.Chat, llmID, deps.Embed, param.SimilarityThreshold) + deduper := NewGroupedDeduper(decider) + for _, rows := range perBatch { + for _, row := range rows { + if err := deduper.Add(ctx, row); err != nil { + return common.Outputs{}, err + } + } + } + if err := deduper.RewriteRelations(ctx, decider.Aliases(), deps.Embed); err != nil { + return common.Outputs{}, err + } + stats := deduper.Stats() + prods := deduper.Rows() + + // ---- KIND POST-PROCESSING ---- + // Chain kinds (list/timeline): relations must form a strict linear chain; + // offending relations the LLM does not keep are dropped (fail-open). + // Timeline additionally drops entity rows no surviving relation references. + // (Mirrors Python's validate_and_correct_chain — which runs right after + // local dedup — and cleanup_timeline_isolated_entities.) + if ChainKinds[compileType] { + chunksByID := make(map[string]string, len(inputs.Chunks)) + for _, ch := range inputs.Chunks { + if id := ch.ID; id != "" { + chunksByID[id] = common.FirstNonEmpty(ch.Text, ch.Content) + } + } + prods = validateAndCorrectChain(ctx, deps, llmID, prods, chunksByID, compileType) + } + if compileType == Type("timeline") { + prods = dropIsolatedTimelineEntities(prods) + } + + // Python stamps the inferred compile kind (list/set/hypergraph) as each + // row's compile_kwd; the chunk converter picks it up from Meta. + for i := range prods { + prods[i].Meta["compile_kwd"] = string(compileType) + } + + // ---- GRAPH ---- + graphProduct, err := buildGraphProduct(ctx, deps, cfg, prods) + if err != nil { + return common.Outputs{}, err + } + + // Stream the merged products (plus the graph) through a ProductSink so the + // flush policy caps peak memory instead of holding the full result set. + sink := common.NewProductSink(ctx, param.Guardrails, inputs.Sink) + for _, p := range prods { + if err := sink.Add(p); err != nil { + return common.Outputs{}, err + } + } + if err := sink.Add(graphProduct); err != nil { + return common.Outputs{}, err + } + + out := common.Outputs{ + Products: sink.Products(), + VectorBytes: sink.Bytes(), + Items: sink.TotalItems(), + Flushed: sink.Flushed(), + DuplicatesDropped: stats.DuplicatesDropped, + } + // Capacity guardrails (centralised in common.EnforceGuardrails so every + // variant honours the same policy — error/flush/none — without drift). + if err := out.EnforceGuardrails(param.Guardrails, inputs.Sink, ctx); err != nil { + return common.Outputs{}, err + } + return out, nil +} + +// buildGraphProduct rebuilds the compact graph JSON from the surviving +// entity/relation rows and wraps it as a single "graph" product so the +// downstream writer has a ready structure to persist. The row id mirrors +// Python's _struct_graph_row_id (doc : structure_graph : compile : template). +func buildGraphProduct(ctx context.Context, deps common.Deps, cfg CompileConfig, prods []common.Product) (common.Product, error) { + if deps.Embed == nil { + return common.Product{}, fmt.Errorf("knowledge_compiler: embedding model is required to build the graph product") + } + graph := RebuildStructureGraph(prods) + graphContent := payloadJSON(graph) + vecs, err := deps.Embed.Encode(ctx, []string{graphContent}) + if err != nil { + return common.Product{}, err + } + if len(vecs) == 0 { + return common.Product{}, fmt.Errorf("knowledge_compiler: embedding the graph summary returned no vector") + } + idParts := []string{cfg.DocID, "structure_graph", string(cfg.Type)} + if cfg.TemplateID != "" { + idParts = append(idParts, cfg.TemplateID) + } + return common.Product{ + ID: common.StableRowID(idParts...), + DocID: cfg.DocID, + TenantID: cfg.TenantID, + Variant: cfg.Variant, + Content: graphContent, + Vector: vecs[0], + Meta: map[string]any{ + "kind": "graph", + "compile_kwd": string(cfg.Type), + "source_doc_ids": []string{cfg.DocID}, + }, + }, nil +} diff --git a/internal/ingestion/component/knowledge_compiler/structure/structure_test.go b/internal/ingestion/component/knowledge_compiler/structure/structure_test.go new file mode 100644 index 0000000000..2815f98ff0 --- /dev/null +++ b/internal/ingestion/component/knowledge_compiler/structure/structure_test.go @@ -0,0 +1,613 @@ +package structure + +import ( + "context" + "encoding/json" + "fmt" + "math" + "strings" + "testing" + + "github.com/cespare/xxhash/v2" + + "ragflow/internal/ingestion/component/knowledge_compiler/common" +) + +// ---- mocks ---- + +// graphChat answers the three LLM call shapes of the structure variant: +// node extraction, edge extraction, and merge judging. It emits items for +// every fixture phrase present in the user prompt so assertions hold whether +// the chunks pack into one batch or several. +type graphChat struct { + mergeCalls int + nodeCalls int + edgeCalls int +} + +var graphFixtures = []struct { + phrase string + entities []map[string]any + relations []map[string]any +}{ + { + phrase: "Alpha is a Beta", + entities: []map[string]any{ + {"type": "letter", "name": "Alpha", "description": "the letter Alpha"}, + {"type": "letter", "name": "Beta", "description": "the letter Beta"}, + }, + relations: []map[string]any{ + {"type": "linked", "source": "Alpha", "target": "Beta", "description": "Alpha is a Beta"}, + }, + }, + { + phrase: "Beta related to Gamma", + entities: []map[string]any{ + {"type": "letter", "name": "Beta", "description": "the letter Beta"}, + {"type": "letter", "name": "Gamma", "description": "the letter Gamma"}, + }, + relations: []map[string]any{ + {"type": "linked", "source": "Beta", "target": "Gamma", "description": "Beta related to Gamma"}, + }, + }, + { + phrase: "Al partners with Gamma", + entities: []map[string]any{ + {"type": "letter", "name": "Al", "description": "short for Alpha"}, + {"type": "letter", "name": "Gamma", "description": "the letter Gamma"}, + }, + relations: []map[string]any{ + {"type": "partners_with", "source": "Al", "target": "Gamma", "description": "Al partners with Gamma"}, + }, + }, + { + phrase: "Alpha also links Gamma", + entities: []map[string]any{ + {"type": "letter", "name": "Alpha", "description": "the letter Alpha"}, + {"type": "letter", "name": "Gamma", "description": "the letter Gamma"}, + }, + relations: []map[string]any{ + {"type": "linked", "source": "Alpha", "target": "Gamma", "description": "Alpha also links Gamma"}, + }, + }, +} + +func itemsJSON(items []map[string]any, chunkIDs []string) string { + for _, item := range items { + item["source_chunk_ids"] = chunkIDs + } + b, _ := json.Marshal(map[string]any{"items": items}) + return string(b) +} + +func chunkIDsFromPrompt(prompt string) []string { + var ids []string + for _, id := range []string{"c1", "c2", "c3"} { + if strings.Contains(prompt, "[CHUNK_ID: "+id+"]") { + ids = append(ids, id) + } + } + if len(ids) == 0 { + return []string{"c1"} + } + return ids +} + +func (m *graphChat) Chat(_ context.Context, req common.ChatRequest) (*common.ChatResponse, error) { + ids := chunkIDsFromPrompt(req.UserPrompt) + switch { + case strings.HasPrefix(req.UserPrompt, "Item A (existing):"): + m.mergeCalls++ + a, b := parseMergeItems(req.UserPrompt) + if sameLogicalItem(a, b) { + return &common.ChatResponse{Content: fmt.Sprintf(`{"duplicated":true,"merged":%s}`, payloadJSON(a))}, nil + } + return &common.ChatResponse{Content: `{"duplicated":false,"merged":null}`}, nil + case strings.Contains(req.SystemPrompt, "strict-chain constraint"): + // Chain correction: keep Alpha→Beta, drop the fan-out partner. + return &common.ChatResponse{Content: `{"keep":[{"from":"Alpha","to":"Beta"}]}`}, nil + case strings.Contains(req.SystemPrompt, "## Known Entities:"): + m.edgeCalls++ + var rels []map[string]any + for _, f := range graphFixtures { + if strings.Contains(req.UserPrompt, f.phrase) { + for _, r := range f.relations { + rels = append(rels, copyPayload(r)) + } + } + } + return &common.ChatResponse{Content: itemsJSON(rels, ids)}, nil + default: + m.nodeCalls++ + var ents []map[string]any + for _, f := range graphFixtures { + if strings.Contains(req.UserPrompt, f.phrase) { + for _, e := range f.entities { + ents = append(ents, copyPayload(e)) + } + } + } + return &common.ChatResponse{Content: itemsJSON(ents, ids)}, nil + } +} + +func copyPayload(p map[string]any) map[string]any { + out := map[string]any{} + for k, v := range p { + out[k] = v + } + return out +} + +// parseMergeItems extracts the Item A / Item B payload JSONs from a merge +// judge user prompt. +func parseMergeItems(prompt string) (map[string]any, map[string]any) { + body := strings.TrimPrefix(prompt, "Item A (existing):\n") + parts := strings.SplitN(body, "\n\nItem B (incoming):\n", 2) + if len(parts) != 2 { + return nil, nil + } + return parsePayload(parts[0]), parsePayload(parts[1]) +} + +// sameLogicalItem mirrors what a reasonable judge would decide on the +// fixtures: entities are duplicates when their names match (or the known +// Al≡Alpha alias pair); relations when source+target+type all match. +func sameLogicalItem(a, b map[string]any) bool { + if a == nil || b == nil { + return false + } + aName, bName := entityName(a), entityName(b) + if aName != "" && bName != "" { + if aName == bName { + return true + } + pair := map[string]bool{aName: true, bName: true} + return pair["Al"] && pair["Alpha"] + } + as, bs := relationEndpoint(a, "", "source"), relationEndpoint(b, "", "source") + at, bt := relationEndpoint(a, "", "target"), relationEndpoint(b, "", "target") + return as != "" && as == bs && at == bt && stringOf(a["type"]) == stringOf(b["type"]) +} + +// hashEmbedder is content-deterministic: identical text => identical (unit) +// vector => cosine 1.0. Near-identical texts do NOT collide, so only exact +// duplicate payloads reach the LLM judge. +type hashEmbedder struct{ dim int } + +func (m hashEmbedder) Dimensions() int { return m.dim } + +func (m hashEmbedder) Encode(_ context.Context, texts []string) ([][]float32, error) { + out := make([][]float32, len(texts)) + for i, t := range texts { + v := make([]float32, m.dim) + for j := 0; j < m.dim; j++ { + h := xxhash.Sum64String(fmt.Sprintf("%d:%s", j, t)) + v[j] = float32(int64(h%(1<<20))-(1<<19)) / float32(1<<19) + } + var s float64 + for _, x := range v { + s += float64(x) * float64(x) + } + if s = math.Sqrt(s); s > 0 { + for k := range v { + v[k] = float32(float64(v[k]) / s) + } + } + out[i] = v + } + return out, nil +} + +// constEmbedder maps every text to the same vector, so every pair in a dedup +// bucket reaches the LLM judge (used to exercise alias merges deterministically). +type constEmbedder struct{ dim int } + +func (m constEmbedder) Dimensions() int { return m.dim } +func (m constEmbedder) Encode(_ context.Context, texts []string) ([][]float32, error) { + v := make([]float32, m.dim) + for i := range v { + v[i] = 1 + } + out := make([][]float32, len(texts)) + for i := range out { + out[i] = v + } + return out, nil +} + +func graphParserConfig() map[string]any { + return map[string]any{ + "kind": "graph", + "entity": map[string]any{"fields": []any{ + map[string]any{"type": "letter", "description": "a single letter entity"}, + }}, + "relation": map[string]any{"fields": []any{ + map[string]any{"type": "linked", "description": "a link between letters"}, + map[string]any{"type": "partners_with", "description": "a partnership"}, + }}, + } +} + +func testParam() common.Param { + p := common.Param{}.Defaults() + p.Variant = common.VariantStructure + p.SimilarityThreshold = 0.99 + p.MaxWorkers = 1 + return p +} + +// ---- prompt alignment tests ---- + +func TestHypergraphPromptsTemplateShape(t *testing.T) { + cfg := map[string]any{ + "kind": "graph", + "guideline": map[string]any{"target": "Extract the letter graph.", "rules_for_entities": "Be precise.", "rules_for_relations": "Only stated links.", "rules_for_time": "As of {observation_time}."}, + "global_rules": "No inventions.", + "entity": map[string]any{ + "description": "Letter entities.", + "fields": []any{map[string]any{"type": "letter", "description": "a letter", "rule": "single glyph"}}, + }, + "relation": map[string]any{ + "description": "Letter links.", + "fields": []any{map[string]any{"type": "linked", "description": "a link"}}, + }, + } + node, edge := HypergraphPrompts(cfg, "en") + + for _, want := range []string{ + "# Role and Task:\nExtract the letter graph.", + "## Global Rules:\nNo inventions.", + "## Entity Extraction Rules:\nBe precise.", + "## Entity Description:\nLetter entities.", + "## Entity Fields:\n- type: letter\n description: a letter\n rule: single glyph", + `Auto-type: "hypergraph". `, + "Return JSON only, no commentary.", + `"name": "<exact extracted item text>"`, + `"source_chunk_ids": ["<source chunk id>", ...]`, + } { + if !strings.Contains(node, want) { + t.Errorf("node prompt missing %q\n---\n%s", want, node) + } + } + if strings.Contains(node, "Items must be unique.") { + t.Errorf("hypergraph (not set) must not demand uniqueness") + } + + for _, want := range []string{ + "## Relation Extraction Rules:\nOnly stated links.", + "## Time Rules:\nAs of ", + "## Relation Description:\nLetter links.", + "## Relation Fields:\n- type: linked\n description: a link", + "## Known Entities:\n{known_nodes}", + "Only create relations between entities listed in 'Known Entities'.", + `"source": "<known entity name>"`, + } { + if !strings.Contains(edge, want) { + t.Errorf("edge prompt missing %q\n---\n%s", want, edge) + } + } + if strings.Contains(edge, "{observation_time}") { + t.Errorf("observation_time placeholder was not substituted") + } +} + +func TestHypergraphPromptsSetUniquenessAndNoRelations(t *testing.T) { + cfg := map[string]any{ + "compile_type": "set", + "entity": map[string]any{"fields": []any{map[string]any{"type": "letter"}}}, + } + node, edge := HypergraphPrompts(cfg, "en") + if !strings.Contains(node, "Items must be unique. ") { + t.Errorf("set kind must demand uniqueness:\n%s", node) + } + if !strings.Contains(node, `Auto-type: "set". `) { + t.Errorf("set Auto-type missing:\n%s", node) + } + if edge != "" { + t.Errorf("config without relations must skip the edge stage, got:\n%s", edge) + } + if got := InferType(cfg); got != TypeSet { + t.Errorf("InferType = %q, want set", got) + } + // "graph" / "knowledge_graph" aliases resolve to hypergraph. + for _, alias := range []string{"graph", "knowledge_graph"} { + if got := InferType(map[string]any{"kind": alias}); got != TypeHypergraph { + t.Errorf("InferType(kind=%q) = %q, want hypergraph", alias, got) + } + } + // Arbitrary kinds (timeline, page_index) are returned verbatim — Python + // stamps them as the autotype/compile_kwd. + if got := InferType(map[string]any{"kind": "timeline"}); got != Type("timeline") { + t.Errorf("InferType(kind=timeline) = %q, want verbatim timeline", got) + } + // An unknown compile_type is NOT a compile kind and falls through to kind. + if got := InferType(map[string]any{"compile_type": "timeline"}); got != TypeList { + t.Errorf("InferType(compile_type=timeline) = %q, want list (falls through)", got) + } + if got := InferType(map[string]any{}); got != TypeList { + t.Errorf("InferType(empty) = %q, want list", got) + } +} + +func TestFillKnownNodes(t *testing.T) { + tmpl := "## Known Entities:\n{known_nodes}" + if got := fillKnownNodes(tmpl, nil); !strings.Contains(got, "(none)") { + t.Fatalf("empty known list must render (none): %q", got) + } + got := fillKnownNodes(tmpl, []string{"Alpha", "Beta"}) + if !strings.Contains(got, "- Alpha\n- Beta") { + t.Fatalf("known list not rendered: %q", got) + } +} + +// ---- payload helpers ---- + +func TestPayloadChunkIDs(t *testing.T) { + payload := map[string]any{"source_chunk_ids": []any{"c1", "nope", "c1"}} + got := payloadChunkIDs(payload, []string{"c1", "c2"}) + if len(got) != 1 || got[0] != "c1" { + t.Fatalf("filter+dedupe failed: %v", got) + } + // Fallback: none of the model's ids belong to the batch -> all batch ids. + payload = map[string]any{"source_chunk_ids": []any{"nope"}} + got = payloadChunkIDs(payload, []string{"c1", "c2"}) + if len(got) != 2 { + t.Fatalf("fallback to batch ids failed: %v", got) + } + // A bare string is accepted (mirrors Python). + payload = map[string]any{"source_chunk_ids": "c2"} + got = payloadChunkIDs(payload, []string{"c1", "c2"}) + if len(got) != 1 || got[0] != "c2" { + t.Fatalf("string form failed: %v", got) + } +} + +func TestPayloadDescriptionSortedAndFlattened(t *testing.T) { + payload := map[string]any{ + "name": "Beta", + "type": "letter", + "source_chunk_ids": []any{"c1", "c2"}, + } + got := payloadDescription(payload) + // Keys are sorted for determinism: name < source_chunk_ids < type. + if got != "Beta c1 c2 letter" { + t.Fatalf("payloadDescription = %q, want %q", got, "Beta c1 c2 letter") + } +} + +// ---- Run: extraction + LLM dedup + graph ---- + +func TestStructureRunGraphKind(t *testing.T) { + deps := common.Deps{Chat: &graphChat{}, Embed: hashEmbedder{dim: 16}, TenantID: "t1", DatasetID: "d1"} + param := testParam() + inputs := common.Inputs{ + DocID: "doc1", + Chunks: []common.Chunk{ + {ID: "c1", Text: "Alpha is a Beta."}, + {ID: "c2", Text: "Beta related to Gamma."}, + }, + VariantSpecific: map[string]any{"parser_config": graphParserConfig()}, + } + + out, err := Run(context.Background(), deps, param, inputs) + if err != nil { + t.Fatalf("Run: %v", err) + } + + var entities, relations, graphs int + var betaProduct *common.Product + for i, p := range out.Products { + switch p.Meta["kind"] { + case "entity": + entities++ + if p.Meta["name"] == "Beta" { + betaProduct = &out.Products[i] + } + if p.Meta["entity_type"] != "letter" { + t.Errorf("entity %v missing entity_type stamp", p.Meta["name"]) + } + if p.Meta["compile_kwd"] != "hypergraph" { + t.Errorf("entity row compile_kwd = %v, want hypergraph", p.Meta["compile_kwd"]) + } + case "relation": + relations++ + case "graph": + graphs++ + } + } + if entities != 3 || relations != 2 || graphs != 1 { + t.Fatalf("products = %d entities + %d relations + %d graph, want 3+2+1", entities, relations, graphs) + } + if out.DuplicatesDropped != 1 { + t.Fatalf("DuplicatesDropped = %d, want 1 (the cross-chunk Beta)", out.DuplicatesDropped) + } + if betaProduct == nil { + t.Fatal("no Beta entity product") + } + ids := metaStrings(betaProduct.Meta, "source_chunk_ids") + if len(ids) != 2 { + t.Fatalf("Beta provenance = %v, want {c1,c2} unioned through the merge", ids) + } + // The merged entity keeps the LLM-merged payload content (parseable JSON). + if parsePayload(betaProduct.Content) == nil { + t.Fatalf("entity content is not payload JSON: %q", betaProduct.Content) + } + + // Graph summary mirrors Python's {entities, relations} shape. + g := parsePayload(graphContentOf(out)) + if g == nil { + t.Fatal("graph product missing/unparseable") + } + gEnts, _ := g["entities"].([]any) + gRels, _ := g["relations"].([]any) + if len(gEnts) != 3 || len(gRels) != 2 { + t.Fatalf("graph = %d entities + %d relations, want 3+2", len(gEnts), len(gRels)) + } +} + +func TestStructureListKindSkipsRelations(t *testing.T) { + deps := common.Deps{Chat: &graphChat{}, Embed: hashEmbedder{dim: 16}, TenantID: "t1", DatasetID: "d1"} + param := testParam() + inputs := common.Inputs{ + DocID: "doc1", + Chunks: []common.Chunk{{ID: "c1", Text: "Alpha is a Beta."}}, + // No parser_config: InferType yields "list", so no edge stage runs. + } + out, err := Run(context.Background(), deps, param, inputs) + if err != nil { + t.Fatalf("Run: %v", err) + } + var relations int + for _, p := range out.Products { + if p.Meta["kind"] == "relation" { + relations++ + } + if p.Meta["compile_kwd"] != "list" { + t.Errorf("compile_kwd = %v, want list", p.Meta["compile_kwd"]) + } + } + if relations != 0 { + t.Fatalf("list kind must not extract relations, got %d", relations) + } +} + +func TestStructureAliasRewrite(t *testing.T) { + // constEmbedder makes every pair a judge candidate; the mock judge merges + // Al into Alpha (the known alias pair), so the relation Al→Gamma must be + // rewritten to Alpha→Gamma by the alias pass. + deps := common.Deps{Chat: &graphChat{}, Embed: constEmbedder{dim: 4}, TenantID: "t1", DatasetID: "d1"} + param := testParam() + inputs := common.Inputs{ + DocID: "doc1", + Chunks: []common.Chunk{ + {ID: "c1", Text: "Alpha is a Beta."}, + {ID: "c2", Text: "Al partners with Gamma."}, + }, + VariantSpecific: map[string]any{"parser_config": graphParserConfig()}, + } + out, err := Run(context.Background(), deps, param, inputs) + if err != nil { + t.Fatalf("Run: %v", err) + } + + var sawAlEntity, sawRewritten bool + for _, p := range out.Products { + if p.Meta["kind"] == "entity" && p.Meta["name"] == "Al" { + sawAlEntity = true + } + if p.Meta["kind"] == "relation" { + payload := parsePayload(p.Content) + if payload == nil { + t.Fatalf("relation content not payload JSON: %q", p.Content) + } + if relationEndpoint(payload, "", "source") == "Alpha" && relationEndpoint(payload, "", "target") == "Gamma" { + sawRewritten = true + if p.Meta["from"] != "Alpha" || p.Meta["to"] != "Gamma" { + t.Errorf("rewritten relation meta from/to = %v/%v", p.Meta["from"], p.Meta["to"]) + } + } + } + } + if sawAlEntity { + t.Errorf("Al entity should have been merged into Alpha") + } + if !sawRewritten { + t.Errorf("relation Al→Gamma was not rewritten to Alpha→Gamma") + } +} + +// ---- merge unit tests ---- + +func TestLLMMergeDeciderContracts(t *testing.T) { + chat := &graphChat{} + d := NewLLMMergeDecider(chat, "llm1", hashEmbedder{dim: 8}, 0.99) + + // Below threshold: no LLM call, keep both. + before := chat.mergeCalls + dec, _, err := d.Decide(context.Background(), common.Product{}, common.Product{}, 0.5) + if err != nil || dec != DecisionKeepBoth || chat.mergeCalls != before { + t.Fatalf("below-threshold pair must be kept without an LLM call: dec=%v err=%v", dec, err) + } + + existing := common.Product{ + ID: "row-alpha", + Content: payloadJSON(map[string]any{"type": "letter", "name": "Alpha", "description": "the letter Alpha"}), + Meta: map[string]any{"kind": "entity", "name": "Alpha", "source_chunk_ids": []string{"c1"}}, + } + incoming := common.Product{ + ID: "row-alpha-2", + Content: payloadJSON(map[string]any{"type": "letter", "name": "Al", "description": "short for Alpha"}), + Meta: map[string]any{"kind": "entity", "name": "Al", "source_chunk_ids": []string{"c2"}}, + } + dec, merged, err := d.Decide(context.Background(), existing, incoming, 1.0) + if err != nil { + t.Fatalf("Decide: %v", err) + } + if dec != DecisionMerge { + t.Fatalf("alias pair must merge: dec=%v", dec) + } + if merged.ID != "row-alpha" { + t.Errorf("merged row must preserve the existing id, got %q", merged.ID) + } + if got := entityName(parsePayload(merged.Content)); got != "Alpha" { + t.Errorf("merged canonical name = %q, want Alpha", got) + } + if ids := metaStrings(merged.Meta, "source_chunk_ids"); len(ids) != 2 { + t.Errorf("merged provenance = %v, want {c1,c2}", ids) + } + if aliases := d.Aliases(); aliases["Al"] != "Alpha" { + t.Errorf("alias map = %v, want Al→Alpha", aliases) + } +} + +func TestApplyMergeInvariants(t *testing.T) { + existing := common.Product{ + Content: payloadJSON(map[string]any{"type": "linked", "source": "Alpha", "target": "Beta"}), + Meta: map[string]any{"kind": "relation"}, + } + merged := map[string]any{"type": "linked", "source": "WRONG", "target": "ALSO_WRONG", "description": "richer"} + out := applyMergeInvariants(existing, merged) + if out["source"] != "Alpha" || out["target"] != "Beta" { + t.Fatalf("relation endpoints must be pinned to the existing payload: %v", out) + } + if out["description"] != "richer" { + t.Fatalf("non-endpoint fields must survive the merge: %v", out) + } +} + +func TestResolveAliasToleratesCycles(t *testing.T) { + aliases := map[string]string{"A": "B", "B": "A"} + if got := resolveAlias("A", aliases); got != "A" && got != "B" { + t.Fatalf("cycle must terminate: %q", got) + } + if got := resolveAlias("x", map[string]string{"x": "y", "y": "z"}); got != "z" { + t.Fatalf("chain resolution failed: %q", got) + } +} + +func TestCosineDecider(t *testing.T) { + d := CosineDecider{Threshold: 0.99} + got, _, err := d.Decide(context.Background(), common.Product{}, common.Product{}, 1.0) + if err != nil { + t.Fatal(err) + } + if got != DecisionDropIncoming { + t.Fatalf("expected drop at 1.0, got %v", got) + } + got, _, _ = d.Decide(context.Background(), common.Product{}, common.Product{}, 0.5) + if got != DecisionKeepBoth { + t.Fatalf("expected keep at 0.5, got %v", got) + } +} + +// ---- helpers ---- + +func graphContentOf(out common.Outputs) string { + for _, p := range out.Products { + if p.Meta["kind"] == "graph" { + return p.Content + } + } + return "" +} diff --git a/internal/ingestion/component/knowledge_compiler/wiki/page.go b/internal/ingestion/component/knowledge_compiler/wiki/page.go new file mode 100644 index 0000000000..e8cd10ea5a --- /dev/null +++ b/internal/ingestion/component/knowledge_compiler/wiki/page.go @@ -0,0 +1,312 @@ +package wiki + +import ( + "context" + "fmt" + "hash/fnv" + "strings" + "unicode" + + "ragflow/internal/ingestion/component/knowledge_compiler/common" +) + +type wikiPageResult struct { + Slug string + Title string + PageType string + Topic string + Action string + EntityNames []string + RelatedKBPages []string + ContentRaw string + Content string + Summary string + Outlinks []string + SourceChunkIDs []string + SourceDocIDs []string +} + +// mapFromChunks runs the MAP stage: extract raw facts/notes from the source. +func mapFromChunks(ctx context.Context, deps common.Deps, llmID, sections string) (string, error) { + resp, err := deps.Chat.Chat(ctx, common.ChatRequest{ + LLMID: llmID, + SystemPrompt: wikiMapSystem, + UserPrompt: sections, + }) + if err != nil { + return "", err + } + return resp.Content, nil +} + +// reduceFromExtracts consolidates the MAP output into topic clusters. +func reduceFromExtracts(ctx context.Context, deps common.Deps, llmID, raw string) (string, error) { + resp, err := deps.Chat.Chat(ctx, common.ChatRequest{ + LLMID: llmID, + SystemPrompt: wikiReduceSystem, + UserPrompt: raw, + }) + if err != nil { + return "", err + } + return resp.Content, nil +} + +// planFromReduction arranges topics into a page outline. +func planFromReduction(ctx context.Context, deps common.Deps, llmID, reduced string) (string, error) { + resp, err := deps.Chat.Chat(ctx, common.ChatRequest{ + LLMID: llmID, + SystemPrompt: wikiPlanSystem, + UserPrompt: reduced, + }) + if err != nil { + return "", err + } + return resp.Content, nil +} + +// refineFromPlan writes the final wiki page from the outline. +func refineFromPlan(ctx context.Context, deps common.Deps, llmID, plan string) (string, error) { + resp, err := deps.Chat.Chat(ctx, common.ChatRequest{ + LLMID: llmID, + SystemPrompt: wikiRefineSystem, + UserPrompt: plan, + }) + if err != nil { + return "", err + } + return resp.Content, nil +} + +// Section is one heading block parsed out of a wiki page's markdown body. The +// page itself is a separate product; sections live as children linked via +// parent_id so the retrieval-side can scope queries to a section (e.g. user +// asks "what does the page say about X?" and the SIDE picks the page + the X +// section). +type Section struct { + Title string + Level int // 1..6, ATX heading depth + Content string // section body until the next heading of equal-or-lower depth +} + +// ParseOutline splits a wiki page's markdown body into a flat section list. +// Headings are ATX style ("#", "##", ...). The page is treated as a single +// implicit "preamble" section when it starts with non-heading text. The page +// itself is *not* in the result; buildPageProducts emits it separately so +// callers can attach the section list under it. +func ParseOutline(md string) []Section { + var sections []Section + var current *Section + flush := func() { + if current != nil { + current.Content = strings.TrimSpace(current.Content) + if current.Title != "" || current.Content != "" { + sections = append(sections, *current) + } + current = nil + } + } + for _, line := range strings.Split(md, "\n") { + trimmed := strings.TrimRight(line, "\r") + if lvl := headingLevel(trimmed); lvl > 0 { + flush() + title := strings.TrimSpace(trimmed[lvl:]) + current = &Section{Title: strings.TrimSpace(stripMarkdownEmphasis(title)), Level: lvl} + continue + } + if current == nil { + // Preamble: treat leading prose as a synthetic level-0 section so + // the page body is not silently discarded. + current = &Section{Title: "", Level: 0} + } + current.Content += trimmed + "\n" + } + flush() + return sections +} + +// headingLevel returns the ATX heading level of s (1..6) or 0 if not a heading. +func headingLevel(s string) int { + if !strings.HasPrefix(s, "#") { + return 0 + } + n := 0 + for n < len(s) && s[n] == '#' { + n++ + } + if n > 6 || n >= len(s) { + return 0 + } + if s[n] != ' ' { + return 0 + } + return n +} + +func stripMarkdownEmphasis(s string) string { + return strings.Trim(s, "`*_") +} + +// buildPageProducts turns the final wiki page into one "page" product plus one +// "section" product per heading in the page body. Sections carry their title, +// level, and content; vector embedding is applied by the caller in a single +// batched Encode. Sections are linked to the page via parent_id. +func buildPageProducts(tenantID, docID, page string, sourceChunkIDs []string) []common.Product { + return buildWikiPageProducts(tenantID, docID, []wikiPageResult{{ + Slug: slugify(firstHeading(page)), + Title: firstHeading(page), + PageType: "concept", + Topic: firstParagraph(page), + Action: "CREATE", + ContentRaw: page, + Content: page, + Summary: firstParagraph(page), + SourceChunkIDs: sourceChunkIDs, + SourceDocIDs: []string{docID}, + }}) +} + +// buildWikiPageProducts turns multiple wiki page drafts into page + section +// products. Each page result becomes one page product and one section product +// per heading in its markdown body. +func buildWikiPageProducts(tenantID, docID string, pages []wikiPageResult) []common.Product { + if len(pages) == 0 { + return nil + } + out := make([]common.Product, 0, len(pages)*2) + for _, page := range pages { + slug := strings.TrimSpace(page.Slug) + if slug == "" { + slug = slugify(page.Title) + } + if slug == "" { + continue + } + title := strings.TrimSpace(page.Title) + if title == "" { + title = slug + } + pageType := strings.TrimSpace(page.PageType) + if pageType == "" { + pageType = "concept" + } + content := page.Content + if strings.TrimSpace(content) == "" { + content = page.ContentRaw + } + summary := strings.TrimSpace(page.Summary) + if summary == "" { + summary = firstParagraph(content) + } + sourceChunkIDs := uniqueStrings(page.SourceChunkIDs) + sourceDocIDs := uniqueStrings(page.SourceDocIDs) + if len(sourceDocIDs) == 0 && docID != "" { + sourceDocIDs = []string{docID} + } + pageID := common.StableRowID(tenantID, docID, string(common.VariantWiki), "page", slug) + out = append(out, common.Product{ + ID: pageID, + DocID: docID, + TenantID: tenantID, + Variant: common.VariantWiki, + Content: content, + ParentID: "", + Meta: map[string]any{ + "kind": "page", + "slug": slug, + "title": title, + "page_type": pageType, + "topic": strings.TrimSpace(page.Topic), + "summary": summary, + "entity_names": uniqueStrings(page.EntityNames), + "related_kb_pages": uniqueStrings(page.RelatedKBPages), + "outlinks": uniqueStrings(page.Outlinks), + "source_chunk_ids": sourceChunkIDs, + "source_doc_ids": sourceDocIDs, + "content_md_raw": page.ContentRaw, + }, + }) + for i, sec := range ParseOutline(content) { + if strings.TrimSpace(sec.Title) == "" { + continue + } + sectionContent := strings.TrimSpace(sec.Content) + if sectionContent == "" { + continue + } + out = append(out, common.Product{ + ID: common.StableRowID(tenantID, docID, string(common.VariantWiki), "section", slug, sec.Title, pageID), + DocID: docID, + TenantID: tenantID, + Variant: common.VariantWiki, + Content: sectionContent, + ParentID: pageID, + Meta: map[string]any{ + "kind": "section", + "slug": slugify(sec.Title), + "title": sec.Title, + "page_slug": slug, + "page_title": title, + "page_type": pageType, + "section_level": sec.Level, + "section_index": i, + "source_chunk_ids": sourceChunkIDs, + "source_doc_ids": sourceDocIDs, + }, + }) + } + } + return out +} + +// firstHeading returns the first markdown "# "-prefixed line, trimmed. +func firstHeading(md string) string { + for _, line := range strings.Split(md, "\n") { + s := strings.TrimSpace(line) + if strings.HasPrefix(s, "# ") { + return strings.TrimSpace(strings.TrimPrefix(s, "# ")) + } + } + return "" +} + +// firstParagraph returns the first non-empty, non-heading line, capped. +func firstParagraph(md string) string { + for _, line := range strings.Split(md, "\n") { + s := strings.TrimSpace(line) + if s == "" || strings.HasPrefix(s, "#") { + continue + } + if len(s) > 300 { + return s[:300] + } + return s + } + return "" +} + +// slugify produces a filesystem-safe slug from a heading. ASCII letters/digits +// and hyphens are kept verbatim, and unicode letters/digits (e.g. CJK titles) +// are preserved too so distinct non-Latin titles don't all collapse to "page" +// (M17). When the result is empty (punctuation-only or whitespace-only title) +// a stable hash of the input is used so different such titles stay distinct. +func slugify(s string) string { + trimmed := strings.TrimSpace(s) + if trimmed == "" { + return "page" + } + s = strings.ToLower(trimmed) + s = strings.ReplaceAll(s, " ", "-") + var b strings.Builder + for _, r := range s { + if (r >= 'a' && r <= 'z') || (r >= '0' && r <= '9') || r == '-' || unicode.IsLetter(r) || unicode.IsDigit(r) { + b.WriteRune(r) + } + } + if b.Len() == 0 { + h := fnv.New32a() + _, _ = h.Write([]byte(trimmed)) + return "page-" + fmt.Sprintf("%08x", h.Sum32()) + } + return b.String() +} diff --git a/internal/ingestion/component/knowledge_compiler/wiki/page_test.go b/internal/ingestion/component/knowledge_compiler/wiki/page_test.go new file mode 100644 index 0000000000..5b8c563c7d --- /dev/null +++ b/internal/ingestion/component/knowledge_compiler/wiki/page_test.go @@ -0,0 +1,232 @@ +package wiki + +import ( + "strings" + "testing" +) + +func TestParseOutline_Basic(t *testing.T) { + md := `# Main Title + +Preamble paragraph that introduces the page. + +## Background + +This is the background section. + +## Methods + +This is the methods section. + +### Subsection + +Deeper detail. + +## Conclusion + +Final thoughts.` + sections := ParseOutline(md) + if len(sections) != 5 { + t.Fatalf("ParseOutline produced %d sections, want 5", len(sections)) + } + titles := []string{"Main Title", "Background", "Methods", "Subsection", "Conclusion"} + levels := []int{1, 2, 2, 3, 2} + for i, sec := range sections { + if sec.Title != titles[i] { + t.Errorf("section %d title = %q, want %q", i, sec.Title, titles[i]) + } + if sec.Level != levels[i] { + t.Errorf("section %d level = %d, want %d", i, sec.Level, levels[i]) + } + if strings.TrimSpace(sec.Content) == "" { + t.Errorf("section %d content is empty", i) + } + } +} + +func TestParseOutline_Empty(t *testing.T) { + if got := ParseOutline(""); len(got) != 0 { + t.Fatalf("ParseOutline(\"\") = %d sections, want 0", len(got)) + } + if got := ParseOutline(" \n\n \n"); len(got) != 0 { + t.Fatalf("ParseOutline(whitespace) = %d sections, want 0", len(got)) + } +} + +func TestParseOutline_PreambleOnly(t *testing.T) { + // Markdown with no headings at all: the prose still becomes a synthetic + // level-0 section so the page body is not silently dropped. + sections := ParseOutline("Just some prose, no headings here.") + if len(sections) != 1 { + t.Fatalf("preamble-only produced %d sections, want 1", len(sections)) + } + if sections[0].Level != 0 || sections[0].Title != "" { + t.Errorf("synthetic preamble = level=%d title=%q, want level=0 title=\"\"", sections[0].Level, sections[0].Title) + } + if !strings.Contains(sections[0].Content, "Just some prose") { + t.Errorf("synthetic preamble content missing prose: %q", sections[0].Content) + } +} + +func TestParseOutline_StripsEmphasis(t *testing.T) { + md := "## **Bold heading**" + sections := ParseOutline(md) + if len(sections) != 1 || sections[0].Title != "Bold heading" { + t.Fatalf("emphasis strip: got %#v, want title \"Bold heading\"", sections) + } +} + +func TestBuildPageProducts_PageAndSections(t *testing.T) { + md := `# The Page Title + +Lead paragraph. + +## Section A + +A body. + +## Section B + +B body. + +### Subsection B1 + +Deeper.` + products := buildPageProducts("t1", "d1", md, []string{"c1", "c2"}) + if len(products) < 2 { + t.Fatalf("buildPageProducts = %d products, want page + 3 sections (>= 4)", len(products)) + } + // First must be the page itself. + page := products[0] + if kind, _ := page.Meta["kind"].(string); kind != "page" { + t.Fatalf("products[0].kind = %q, want \"page\"", page.Meta["kind"]) + } + if title, _ := page.Meta["title"].(string); title != "The Page Title" { + t.Errorf("page title = %q, want \"The Page Title\"", title) + } + if slug, _ := page.Meta["slug"].(string); slug != "the-page-title" { + t.Errorf("page slug = %q, want \"the-page-title\"", slug) + } + if ids, _ := page.Meta["source_chunk_ids"].([]string); len(ids) != 2 { + t.Errorf("page source_chunk_ids = %#v, want 2 ids", page.Meta["source_chunk_ids"]) + } + // Sections are linked via parent_id = page.ID. + for i, p := range products[1:] { + if p.ParentID != page.ID { + t.Errorf("section[%d].parent_id = %q, want %q", i, p.ParentID, page.ID) + } + if kind, _ := p.Meta["kind"].(string); kind != "section" { + t.Errorf("section[%d].kind = %q, want \"section\"", i, p.Meta["kind"]) + } + if _, has := p.Meta["section_level"]; !has { + t.Errorf("section[%d] missing section_level meta", i) + } + if ids, _ := p.Meta["source_chunk_ids"].([]string); len(ids) != 2 { + t.Errorf("section[%d] source_chunk_ids = %#v, want 2 ids", i, p.Meta["source_chunk_ids"]) + } + if strings.TrimSpace(p.Content) == "" { + t.Errorf("section[%d] content is empty", i) + } + } +} + +func TestBuildPageProducts_NoSections(t *testing.T) { + // A page that contains only preamble: still emits a page product, with + // zero section children (the preamble is rolled into the page summary). + products := buildPageProducts("t1", "d1", "Just prose, no headings.", []string{"c1"}) + if len(products) != 1 { + t.Fatalf("preamble-only page = %d products, want 1", len(products)) + } + if kind, _ := products[0].Meta["kind"].(string); kind != "page" { + t.Errorf("kind = %q, want page", products[0].Meta["kind"]) + } +} + +func TestBuildPageProducts_StableIDs(t *testing.T) { + md := "# Title\n\n## A\n\nbody\n" + first := buildPageProducts("t1", "d1", md, []string{"c1"}) + second := buildPageProducts("t1", "d1", md, []string{"c1"}) + if len(first) != len(second) { + t.Fatalf("length mismatch: %d vs %d", len(first), len(second)) + } + for i := range first { + if first[i].ID != second[i].ID { + t.Errorf("product[%d] id not stable: %q vs %q", i, first[i].ID, second[i].ID) + } + } +} + +func TestBuildWikiPageProducts_MultiPage(t *testing.T) { + pages := []wikiPageResult{ + { + Slug: "entity/alpha", + Title: "Alpha", + PageType: "entity", + ContentRaw: "# Alpha\n\nSee [[concept/beta]].\n\n## Details\n\nAlpha body.", + Content: "# Alpha\n\nSee [Beta](artifact/kb/concept/beta).\n\n## Details\n\nAlpha body.", + SourceChunkIDs: []string{"c1"}, + SourceDocIDs: []string{"d1"}, + Outlinks: []string{"concept/beta"}, + }, + { + Slug: "concept/beta", + Title: "Beta", + PageType: "concept", + ContentRaw: "# Beta\n\nBody.", + Content: "# Beta\n\nBody.", + SourceChunkIDs: []string{"c2"}, + SourceDocIDs: []string{"d1"}, + }, + } + products := buildWikiPageProducts("t1", "d1", pages) + if len(products) != 5 { + t.Fatalf("buildWikiPageProducts = %d products, want 5", len(products)) + } + page := products[0] + if kind, _ := page.Meta["kind"].(string); kind != "page" { + t.Fatalf("page kind = %q, want page", page.Meta["kind"]) + } + if outlinks, _ := page.Meta["outlinks"].([]string); len(outlinks) != 1 || outlinks[0] != "concept/beta" { + t.Fatalf("page outlinks = %#v, want concept/beta", page.Meta["outlinks"]) + } + if ids, _ := page.Meta["source_chunk_ids"].([]string); len(ids) != 1 || ids[0] != "c1" { + t.Fatalf("page source_chunk_ids = %#v, want c1", page.Meta["source_chunk_ids"]) + } + if ids, _ := page.Meta["source_doc_ids"].([]string); len(ids) != 1 || ids[0] != "d1" { + t.Fatalf("page source_doc_ids = %#v, want d1", page.Meta["source_doc_ids"]) + } +} + +func TestTransformWikiLinks(t *testing.T) { + rendered, outlinks := transformWikiLinks( + "See [[concept/beta|Beta]] and [[entity/alpha]]. Also [Beta](artifact/kb/concept/beta).", + "kb", + map[string]string{"entity/alpha": "Alpha", "concept/beta": "Beta"}, + ) + if !strings.Contains(rendered, "(artifact/kb/concept/beta)") || !strings.Contains(rendered, "(artifact/kb/entity/alpha)") { + t.Fatalf("rendered links not rewritten: %q", rendered) + } + if len(outlinks) != 2 { + t.Fatalf("outlinks = %#v, want 2 unique slugs", outlinks) + } +} + +func TestSlugify(t *testing.T) { + cases := []struct{ in, want string }{ + {"Hello World", "hello-world"}, + {" Trim Me ", "trim-me"}, + {"Punctuation!@#Stripped", "punctuationstripped"}, + {"", "page"}, + // Non-ASCII (CJK) letters are preserved (M17), so the slug keeps the + // Chinese characters and the trailing ASCII word. + {"中文 title", "中文-title"}, + {"a-b-c", "a-b-c"}, + // Punctuation-only title yields a hash fallback, not the bare "page". + {"!@#", "page-8b024aa5"}, + } + for _, c := range cases { + if got := slugify(c.in); got != c.want { + t.Errorf("slugify(%q) = %q, want %q", c.in, got, c.want) + } + } +} diff --git a/internal/ingestion/component/knowledge_compiler/wiki/prompt.go b/internal/ingestion/component/knowledge_compiler/wiki/prompt.go new file mode 100644 index 0000000000..6da9738c75 --- /dev/null +++ b/internal/ingestion/component/knowledge_compiler/wiki/prompt.go @@ -0,0 +1,459 @@ +package wiki + +import ( + "fmt" + "strings" + + "ragflow/internal/ingestion/component/knowledge_compiler/common" +) + +const wikiMapSystem = `You are a knowledge extraction engine. Extract structured knowledge from the provided document sections. Return ONLY valid JSON matching the schema exactly. Never include text outside the JSON object. Keep the source language of the document.` + +const wikiReduceSystem = `You are a knowledge synthesis engine. Normalize and deduplicate the structured extracts while preserving source_chunk_ids provenance. Claims, relations, and topics are pass-through facts; do not invent new ones. Return ONLY valid JSON matching the same schema.` + +const wikiPlanSystem = `You are a knowledge compilation planner. Given structured knowledge, produce a wiki page plan. Return ONLY valid JSON.` + +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 wikiMapUserTemplate = `## Document context +Document id: {doc_id} +Batch contains {chunk_count} packed chunk(s). Each chunk is introduced by a +"[CHUNK_ID <id>]" line. The chunk_id values to choose from are: +{chunk_id_list} + +## Packed chunks +{packed_chunks} + +--- + +Extract all knowledge from every chunk and return a single JSON object with this +exact schema: + +{ + "entities": [ + { + "name": "string", + "type": "string", + "aliases": ["string"], + "source_chunk_id": "string" + } + ], + "concepts": [ + { + "term": "string", + "definition_excerpt": "string", + "source_chunk_id": "string" + } + ], + "claims": [ + { + "statement": "string", + "subject": "string", + "confidence": "explicit", + "source_chunk_id": "string" + } + ], + "relations": [ + { + "from": "string", + "to": "string", + "type": "string", + "source_chunk_id": "string" + } + ], + "topics": ["string"] +} + +## Customization +The following rules and schema customizations override the defaults above and +must be honoured when extracting from this knowledge base: + +- Entity types (override default): {entity_type_rules} +- Relation types (override default): {relation_type_rules} +- Concept term guidance: {concept_term} +- Concept definition guidance: {concept_definition_excerpt} +- Claim statement guidance: {claim_statement} +- Claim subject guidance: {claim_subject} +{custom_rules} + +Rules: +- source_chunk_id MUST be one of the chunk_id values listed above. +- Do not invent opaque identifiers, hashes, or prompt scaffolding as names. +- Extract entities and concepts from human-readable text only. +- Claims should be liberal: every factual sentence about an entity or concept is a claim. +- Relations should be explicit links only. +- Topics should capture coherent subtopics that could become wiki sections. +- Return empty arrays [] for categories with no findings. +- Return ONLY the JSON object, no markdown fences, no commentary.` + +const wikiReduceUserTemplate = `## Reduced input +Document id: {doc_id} +{ + "entities": {entities}, + "concepts": {concepts}, + "claims": {claims}, + "relations": {relations}, + "topics": {topics} +} + +Normalize duplicates, merge aliases and provenance, and return the same schema.` + +const wikiPlanBatchUserTemplate = `## Knowledge base context +Document id: {doc_id} +Batch: {batch_index}/{batch_total} + +## Reduced input +{ + "entities": {entities}, + "concepts": {concepts}, + "claims": {claims}, + "relations": {relations}, + "topics": {topics} +} + +Return a JSON compilation plan with one or more page entries: +{ + "pages": [ + { + "action": "CREATE", + "slug": "string", + "title": "string", + "page_type": "entity | concept | topic", + "topic": "string", + "entity_names": ["string"], + "related_kb_pages": ["string"], + "priority": 1, + "lead": "string", + "sections": [ + { "heading": "string", "points": ["string"] } + ] + } + ], + "estimated_page_count": 1, + "compilation_notes": "string" +} + +Rules: +- Prefer one page per high-signal entity or concept when the batch supports it. +- 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 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 +{planned_page} + +## Candidate existing pages +{candidates} + +Return JSON: +{ + "action": "UPDATE | CREATE", + "slug": "string", + "reason": "string" +} + +Rules: +- Choose UPDATE only when the planned page and candidate refer to the same underlying entity, concept, or topic. +- Prefer CREATE when the overlap is weak, ambiguous, or just generally related. +- If action is UPDATE, slug must be exactly one candidate slug. +- Return only JSON.` + +const wikiRefineWriterExample = `Each page must be a proper encyclopedic article, not a flat bullet list: +1. Opening paragraph that defines the page subject. +2. H2 sections with prose before any bullets. +3. Bold key terms on first use. +4. Link related pages with [[slug]] or [[slug|display text]]. +5. End with a short "See also" section when relevant.` + +const wikiRefineWriterSystemTemplate = `You are an enterprise knowledge compilation writer. Your job is to write a single, high-quality wiki page by reading the SOURCE TEXT and using the evidence checklist as guidance for what to cover. + +Write in the same language as the source text. Do not translate content. + +# Page structure +{template_example} + +# Wikilinks +- Use [[slug]] or [[slug|display text]] to cross-link. +- Only link to slugs from the available-pages list. + +Return only markdown.` + +const wikiRefineWriterUserTemplate = `## Task +{action} the following wiki page. + +## Page specification +- Slug: {slug} +- Title: {title} +- Type: {page_type} + +## Available pages (ONLY use these slugs for [[wikilinks]]) +{all_plan_slugs} + +{existing_section} + +## Source document text +{source_context} + +## Evidence checklist ({evidence_count} items) +{evidence_blocks} + +## Instructions +Write the complete wiki page in markdown based on the source text above. +Cross-link to other pages using [[slug]] or [[slug|display text]] and only use slugs from the available pages list. +Return only the markdown content.` + +const wikiRefineMergeSystem = `You are a wiki page merger. You receive two versions of the same wiki page: +- EXISTING: the current version in the knowledge base. +- INCOMING: a new version generated from a different source document. + +Produce one unified page that preserves all factual content from both versions. +- Keep all facts, numbers, procedures, and named entities from both versions. +- Remove exact duplicates. +- Preserve and normalize [[wikilinks]]. +- Keep the same language as the existing page. +- Do not summarize or condense away factual detail. +- Return only markdown.` + +func formatWikiChunkBatch(docID string, batch []common.Chunk) (string, []string) { + var b strings.Builder + ids := make([]string, 0, len(batch)) + for i, c := range batch { + id := strings.TrimSpace(c.ID) + if id == "" { + id = fmt.Sprintf("chunk-%d", i+1) + } + ids = append(ids, id) + text := firstNonEmpty(c.Text, c.Content) + if strings.TrimSpace(text) == "" { + continue + } + b.WriteString("[CHUNK_ID ") + b.WriteString(id) + b.WriteString("]\n") + b.WriteString(text) + b.WriteString("\n\n") + } + _ = docID + return b.String(), ids +} + +func buildWikiMapPrompt(docID string, batch []common.Chunk, parserConfig any, language string) (string, []string) { + body, ids := formatWikiChunkBatch(docID, batch) + entityTypeRules := "person|org|product|regulation|location|system|equipment|other" + relationTypeRules := "include|ordered|owns|part_of|caused_by|regulates|uses|located_in|other" + conceptTerm := "named term or topic" + conceptDef := "short definition excerpt from the source text" + claimStatement := "factual statement" + claimSubject := "entity or concept that the claim is about" + customRules := "" + + if cfg, ok := parserConfig.(map[string]any); ok { + if rules := wikiTemplateCustomRules(cfg, language); rules != "" { + customRules = rules + } + if ent := wikiTemplateFields(cfg, "entity"); len(ent) > 0 { + if s := wikiRenderSchemaBody(ent, language, defaultEntitySchemaBody); s != "" { + entityTypeRules = s + } + } + if rel := wikiTemplateFields(cfg, "relation"); len(rel) > 0 { + if s := wikiRenderSchemaBody(rel, language, defaultRelationSchemaBody); s != "" { + relationTypeRules = s + } + } + if conceptFields := wikiTemplateFields(cfg, "concept"); len(conceptFields) > 0 { + if s := wikiPipeJoin(conceptFields, "term"); s != "" { + conceptTerm = s + } + if s := wikiColonJoin(conceptFields, "term", "definition_excerpt"); s != "" { + conceptDef = s + } + } + if claimFields := wikiTemplateFields(cfg, "claim"); len(claimFields) > 0 { + if s := wikiNamedFieldDescription(claimFields, "statement"); s != "" { + claimStatement = s + } + if s := wikiNamedFieldDescription(claimFields, "subject"); s != "" { + claimSubject = s + } + } + } + + user := renderWikiTemplate(wikiMapUserTemplate, map[string]string{ + "doc_id": docID, + "chunk_count": fmt.Sprintf("%d", len(batch)), + "chunk_id_list": strings.Join(prefixWithDash(ids), "\n"), + "packed_chunks": body, + "entity_type_rules": entityTypeRules, + "relation_type_rules": relationTypeRules, + "concept_term": conceptTerm, + "concept_definition_excerpt": conceptDef, + "claim_statement": claimStatement, + "claim_subject": claimSubject, + "custom_rules": customRules, + }) + return user, ids +} + +func buildWikiRefineWriterSystem(example string) string { + body := strings.TrimSpace(example) + if body == "" { + body = wikiRefineWriterExample + } + return renderWikiTemplate(wikiRefineWriterSystemTemplate, map[string]string{ + "template_example": body, + }) +} + +const defaultEntitySchemaBody = `"name": "string — entity canonical name as it appears in text", + "type": "string — one of: person|org|product|regulation|location|system|equipment|other", + "aliases": ["string"], + "source_chunk_id": "string — exact value from the chunk_id list above"` + +const defaultRelationSchemaBody = `"from": "string — source entity/concept name", + "to": "string — target entity/concept name", + "type": "string — e.g. owns|part_of|caused_by|regulates|uses|located_in|other", + "source_chunk_id": "string — exact value from the chunk_id list above"` + +func wikiTemplateCustomRules(parserConfig map[string]any, language string) string { + guideline := common.GetMap(parserConfig, "guideline") + if len(guideline) == 0 { + return "" + } + var sections []string + if rules := common.Localize(common.Get(guideline, "rules_for_entities"), language); strings.TrimSpace(rules) != "" { + sections = append(sections, "## Entity extraction rules (from knowledge base config):\n"+rules) + } + if rules := common.Localize(common.Get(guideline, "rules_for_relations"), language); strings.TrimSpace(rules) != "" { + sections = append(sections, "## Relation extraction rules (from knowledge base config):\n"+rules) + } + if len(sections) == 0 { + return "" + } + return "\n" + strings.Join(sections, "\n\n") + "\n" +} + +func wikiTemplateFields(parserConfig map[string]any, section string) []any { + cfg := common.GetMap(parserConfig, section) + if len(cfg) == 0 { + return nil + } + if fields, ok := cfg["fields"].([]any); ok { + return fields + } + if fields, ok := cfg["fields"].([]map[string]any); ok { + out := make([]any, 0, len(fields)) + for _, f := range fields { + out = append(out, f) + } + return out + } + return nil +} + +func wikiRenderSchemaBody(fields []any, language, defaultBody string) string { + if len(fields) == 0 { + return defaultBody + } + lines := make([]string, 0, len(fields)+1) + seen := map[string]bool{} + for _, raw := range fields { + field, ok := raw.(map[string]any) + if !ok { + continue + } + name := strings.TrimSpace(firstString(field["name"])) + if name == "" || name == "source_chunk_id" || seen[name] { + continue + } + seen[name] = true + ftype := strings.TrimSpace(firstString(field["type"])) + if ftype == "" { + ftype = "str" + } + desc := common.Localize(field["description"], language) + var placeholder string + switch ftype { + case "list": + placeholder = `["string"]` + case "int": + placeholder = "0" + case "float": + placeholder = "0.0" + case "bool": + placeholder = "false" + default: + if strings.TrimSpace(desc) != "" { + desc = strings.NewReplacer("\n", " ", "{", "(", "}", ")").Replace(desc) + placeholder = fmt.Sprintf(`"string — %s"`, strings.TrimSpace(desc)) + } else { + placeholder = `"string"` + } + } + lines = append(lines, fmt.Sprintf(` "%s": %s`, name, placeholder)) + } + if len(lines) == 0 { + return defaultBody + } + lines = append(lines, ` "source_chunk_id": "string — exact value from the chunk_id list above"`) + return strings.Join(lines, ",\n") +} + +func wikiPipeJoin(fields []any, key string) string { + var vals []string + for _, raw := range fields { + field, ok := raw.(map[string]any) + if !ok { + continue + } + if s := strings.TrimSpace(firstString(field[key])); s != "" { + vals = append(vals, s) + } + } + return strings.Join(vals, "|") +} + +func wikiColonJoin(fields []any, leftKey, rightKey string) string { + var vals []string + for _, raw := range fields { + field, ok := raw.(map[string]any) + if !ok { + continue + } + left := strings.TrimSpace(firstString(field[leftKey])) + right := strings.TrimSpace(firstString(field[rightKey])) + if left != "" || right != "" { + vals = append(vals, left+":"+right) + } + } + return strings.Join(vals, "\n") +} + +func wikiNamedFieldDescription(fields []any, name string) string { + for _, raw := range fields { + field, ok := raw.(map[string]any) + if !ok { + continue + } + if strings.EqualFold(strings.TrimSpace(firstString(field["name"])), name) { + if s := strings.TrimSpace(common.Localize(field["description"], "en")); s != "" { + return s + } + } + } + return "" +} diff --git a/internal/ingestion/component/knowledge_compiler/wiki/wiki.go b/internal/ingestion/component/knowledge_compiler/wiki/wiki.go new file mode 100644 index 0000000000..8ffe6c3e33 --- /dev/null +++ b/internal/ingestion/component/knowledge_compiler/wiki/wiki.go @@ -0,0 +1,1972 @@ +// Package wiki implements the "wiki" variant of KnowledgeCompiler: a +// document artifact pipeline MAP -> REDUCE -> PLAN -> REFINE that produces a +// wiki-style page (and supporting section products) with in-memory state only. +// The stage semantics are aligned with the Python wiki.py design, but the Go +// port keeps all intermediate artifacts in memory instead of persisting them to +// ES between stages. +package wiki + +import ( + "bytes" + "context" + "encoding/json" + "fmt" + "math" + "net/url" + "regexp" + "sort" + "strings" + + "ragflow/internal/ingestion/component/knowledge_compiler/common" + "ragflow/internal/ingestion/component/knowledge_compiler/structure" +) + +type wikiPipeline struct { + ctx context.Context + deps common.Deps + param common.Param + inputs common.Inputs + docID string + tenantID string + datasetID string + llmID string + + mapExtracts []wikiExtract + reduced wikiExtract + plan wikiPlan + pages []wikiPageResult +} + +type wikiExtract struct { + Entities []wikiEntity `json:"entities"` + Concepts []wikiConcept `json:"concepts"` + Claims []wikiClaim `json:"claims"` + Relations []wikiRelation `json:"relations"` + Topics []string `json:"topics"` +} + +type wikiEntity struct { + Name string `json:"name"` + Type string `json:"type"` + Aliases []string `json:"aliases,omitempty"` + SourceChunkIDs []string `json:"source_chunk_ids,omitempty"` +} + +type wikiConcept struct { + Term string `json:"term"` + Definition string `json:"definition_excerpt"` + SourceChunkIDs []string `json:"source_chunk_ids,omitempty"` +} + +type wikiClaim struct { + Statement string `json:"statement"` + Subject string `json:"subject"` + Confidence string `json:"confidence"` + SourceChunkIDs []string `json:"source_chunk_ids,omitempty"` +} + +type wikiRelation struct { + From string `json:"from"` + To string `json:"to"` + Type string `json:"type"` + SourceChunkIDs []string `json:"source_chunk_ids,omitempty"` +} + +type wikiPlan struct { + Title string `json:"title"` + Slug string `json:"slug"` + Lead string `json:"lead"` + Sections []wikiPlanSection `json:"sections"` + PageType string `json:"page_type,omitempty"` + Topic string `json:"topic,omitempty"` + Entities []string `json:"entity_names,omitempty"` + Related []string `json:"related_kb_pages,omitempty"` + Pages []wikiPlanPage `json:"pages,omitempty"` +} + +type wikiPlanSection struct { + Heading string `json:"heading"` + Points []string `json:"points"` +} + +type wikiPlanPage struct { + Action string `json:"action"` + Slug string `json:"slug"` + Title string `json:"title"` + PageType string `json:"page_type"` + Topic string `json:"topic"` + EntityNames []string `json:"entity_names"` + RelatedKB []string `json:"related_kb_pages"` + Priority int `json:"priority"` + Lead string `json:"lead"` + Sections []wikiPlanSection `json:"sections"` +} + +const ( + wikiPlanUpdateThreshold = 0.92 + wikiPlanMaybeThreshold = 0.78 + wikiPlanSearchTopK = 3 + wikiMergeShrinkRatio = 0.7 +) + +// Run executes the wiki variant. +func Run(ctx context.Context, deps common.Deps, param common.Param, inputs common.Inputs) (common.Outputs, error) { + docID := inputs.DocID + llmID := firstNonEmpty(inputs.LLMID, param.LLMID) + if docID == "" { + docID = "unknown" + } + p := &wikiPipeline{ + ctx: ctx, + deps: deps, + param: param, + inputs: inputs, + docID: docID, + tenantID: deps.TenantID, + datasetID: deps.DatasetID, + llmID: llmID, + } + if err := p.run(); err != nil { + return common.Outputs{}, err + } + + products := buildWikiPageProducts(p.tenantID, p.docID, p.pages) + for i := range products { + if products[i].Meta == nil { + products[i].Meta = map[string]any{} + } + } + if len(products) > 0 { + if deps.Chat == nil { + return common.Outputs{}, fmt.Errorf("wiki: chat model required for page generation") + } + if deps.Embed == nil { + return common.Outputs{}, fmt.Errorf("wiki: embedder required (page + sections must carry vectors)") + } + texts := make([]string, len(products)) + for i, p := range products { + texts[i] = p.Content + } + vectors, err := deps.Embed.Encode(ctx, texts) + if err != nil { + return common.Outputs{}, err + } + for i := range products { + if i < len(vectors) { + products[i].Vector = vectors[i] + } + } + } + + store := common.NewMemStore() + decider := structure.CosineDecider{Threshold: param.SimilarityThreshold} + stats := structure.MergeStats{} + for _, p := range products { + s, err := structure.MergeIntoStore(ctx, store, decider, []common.Product{p}) + if err != nil { + return common.Outputs{}, err + } + stats.Inserted += s.Inserted + stats.Updated += s.Updated + stats.DuplicatesDropped += s.DuplicatesDropped + } + + prod := store.Snapshot() + if param.EnableHistoricalDedup || len(inputs.HistoricalCandidates) > 0 { + survivors, dropped, err := dedupHistorical(ctx, deps, param, p.tenantID, p.datasetID, inputs.HistoricalCandidates, prod) + if err != nil { + return common.Outputs{}, err + } + prod = survivors + stats.DuplicatesDropped += dropped + } + + sink := common.NewProductSink(ctx, param.Guardrails, inputs.Sink) + for _, p := range prod { + if err := sink.Add(p); err != nil { + return common.Outputs{}, err + } + } + out := common.Outputs{ + Products: sink.Products(), + VectorBytes: sink.Bytes(), + Items: sink.TotalItems(), + Flushed: sink.Flushed(), + DuplicatesDropped: stats.DuplicatesDropped, + } + if err := out.EnforceGuardrails(param.Guardrails, inputs.Sink, ctx); err != nil { + return common.Outputs{}, err + } + return out, nil +} + +func (p *wikiPipeline) run() error { + if len(p.inputs.Chunks) == 0 { + return nil + } + if err := p.runMap(); err != nil { + return err + } + p.reduced = reduceExtracts(p.mapExtracts) + plan, err := p.runPlan() + if err != nil { + return err + } + p.plan = plan + pages, err := p.runRefine() + if err != nil { + return err + } + p.pages = pages + return nil +} + +func (p *wikiPipeline) runMap() error { + batches := common.PackBatches(p.inputs.Chunks, 4096, p.deps.Tokenizer) + for _, batch := range batches { + if err := p.ctx.Err(); err != nil { + return err + } + extract, err := p.mapBatch(batch) + if err != nil { + return err + } + p.mapExtracts = append(p.mapExtracts, extract) + } + return 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) + raw, err := common.GenJSON(p.ctx, p.deps.Chat, common.ChatRequest{ + LLMID: p.llmID, + SystemPrompt: wikiMapSystem, + UserPrompt: user, + }) + if err != nil { + return wikiExtract{}, err + } + return parseWikiExtract(raw), nil +} + +func (p *wikiPipeline) runPlan() (wikiPlan, error) { + batches := packWikiPlanBatches(p.reduced, wikiPlanTokenBudget) + if len(batches) == 0 { + batches = []wikiExtract{p.reduced} + } + plans := make([]wikiPlan, 0, len(batches)) + for i, batch := range batches { + plan, err := p.runPlanBatch(batch, i+1, len(batches)) + if err != nil { + return wikiPlan{}, err + } + plans = append(plans, plan) + } + if len(plans) == 1 { + plan := normalizeWikiPlan(plans[0], p.docID, p.reduced) + return p.reconcilePlan(plan) + } + plan, err := p.mergePlanCandidates(plans) + if err != nil { + return wikiPlan{}, err + } + return p.reconcilePlan(plan) +} + +func (p *wikiPipeline) runRefine() ([]wikiPageResult, error) { + pages := normalizeWikiPlanPages(p.plan.Pages, p.reduced) + if len(pages) == 0 { + return nil, nil + } + pageTitles := map[string]string{} + allPlanSlugs := make([]string, 0, len(pages)) + for _, page := range pages { + if page.Slug == "" { + continue + } + allPlanSlugs = append(allPlanSlugs, page.Slug) + pageTitles[page.Slug] = page.Title + } + entityLookup := buildWikiEntityLookup(p.reduced.Entities) + conceptLookup := buildWikiConceptLookup(p.reduced.Concepts) + results := make([]wikiPageResult, 0, len(pages)) + for _, planItem := range pages { + if p.ctx.Err() != nil { + return nil, p.ctx.Err() + } + evidence := assembleWikiPageEvidence(planItem, p.reduced.Claims, entityLookup, conceptLookup) + sourceChunkIDs := collectWikiEvidenceChunkIDs(evidence) + sourceContext := buildSourceContext(p.inputs.Chunks, sourceChunkIDs) + if strings.TrimSpace(sourceContext) == "" { + sourceContext = buildSourceContext(p.inputs.Chunks, p.reduced.sourceChunkIDs()) + } + available := make([]string, 0, len(allPlanSlugs)) + for _, slug := range allPlanSlugs { + if slug != planItem.Slug { + available = append(available, "- [["+slug+"]]") + } + } + 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 nil, err + } + } + existingSection := "" + existingRaw := "" + if existing != nil { + existingRaw = firstNonEmpty(existing.ContentMDRaw, existing.ContentMD) + if strings.TrimSpace(existingRaw) != "" { + existingSection = "## Existing page content (UPDATE — integrate new evidence into this)\n\n" + existingRaw + "\n" + } + } + user := renderWikiTemplate(wikiRefineWriterUserTemplate, map[string]string{ + "action": firstNonEmpty(planItem.Action, "CREATE"), + "slug": planItem.Slug, + "title": firstNonEmpty(planItem.Title, planItem.Slug), + "page_type": firstNonEmpty(planItem.PageType, "concept"), + "all_plan_slugs": strings.Join(available, "\n"), + "existing_section": existingSection, + "source_context": sourceContext, + "evidence_count": fmt.Sprintf("%d", len(evidence)), + "evidence_blocks": formatWikiEvidenceBlocks(evidence), + }) + resp, err := p.deps.Chat.Chat(p.ctx, common.ChatRequest{ + LLMID: p.llmID, + SystemPrompt: buildWikiRefineWriterSystem(""), + UserPrompt: user, + }) + if err != nil { + return nil, err + } + if resp == nil { + return nil, fmt.Errorf("knowledge_compiler: wiki refine returned no response") + } + contentRaw := strings.TrimSpace(firstNonEmpty(resp.Content)) + if contentRaw == "" { + contentRaw = "# " + firstNonEmpty(planItem.Title, planItem.Slug) + "\n\n(Page generation produced no content.)" + } + if strings.TrimSpace(existingRaw) != "" { + contentRaw, err = p.mergeWikiPageContent(existingRaw, contentRaw, planItem.Slug) + if err != nil { + return nil, err + } + } + contentRendered, outlinks := transformWikiLinks(contentRaw, firstNonEmpty(p.datasetID, p.docID), pageTitles) + sourceDocIDs := collectWikiSourceDocIDs(p.inputs.Chunks, sourceChunkIDs, p.docID) + summary := firstParagraph(contentRendered) + if summary == "" { + summary = firstNonEmpty(planItem.Title, planItem.Slug) + } + topic := firstNonEmpty(planItem.Topic, planItem.Title, planItem.Slug) + results = append(results, wikiPageResult{ + Slug: planItem.Slug, + Title: firstNonEmpty(planItem.Title, planItem.Slug), + PageType: firstNonEmpty(planItem.PageType, "concept"), + Topic: topic, + Action: firstNonEmpty(planItem.Action, "CREATE"), + EntityNames: uniqueStrings(planItem.EntityNames), + RelatedKBPages: uniqueStrings(planItem.RelatedKB), + ContentRaw: contentRaw, + Content: contentRendered, + Summary: summary, + Outlinks: outlinks, + SourceChunkIDs: sourceChunkIDs, + SourceDocIDs: sourceDocIDs, + }) + } + return results, nil +} + +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), + "entities": mustJSON(batch.Entities), + "concepts": mustJSON(batch.Concepts), + "claims": mustJSON(batch.Claims), + "relations": mustJSON(batch.Relations), + "topics": mustJSON(batch.Topics), + }) + raw, err := common.GenJSON(p.ctx, p.deps.Chat, common.ChatRequest{ + LLMID: p.llmID, + SystemPrompt: wikiPlanSystem, + UserPrompt: user, + }) + if err != nil { + return wikiPlan{}, err + } + return parseWikiPlan(raw, p.docID, batch), nil +} + +func (p *wikiPipeline) mergePlanCandidates(plans []wikiPlan) (wikiPlan, error) { + user := renderWikiTemplate(wikiPlanMergeUserTemplate, map[string]string{ + "doc_id": p.docID, + "candidates": mustPrettyJSON(plans), + }) + raw, err := common.GenJSON(p.ctx, p.deps.Chat, common.ChatRequest{ + LLMID: p.llmID, + SystemPrompt: wikiPlanSystem, + UserPrompt: user, + }) + if err != nil { + return wikiPlan{}, err + } + return normalizeWikiPlan(parseWikiPlan(raw, p.docID, p.reduced), p.docID, p.reduced), nil +} + +func (p *wikiPipeline) reconcilePlan(plan wikiPlan) (wikiPlan, error) { + if p.deps.WikiPages == nil || p.deps.Embed == nil || strings.TrimSpace(p.datasetID) == "" || len(plan.Pages) == 0 { + return plan, nil + } + queryTexts := make([]string, 0, len(plan.Pages)) + pageIdx := make([]int, 0, len(plan.Pages)) + for i, page := range plan.Pages { + page = normalizeWikiPlanPage(page) + plan.Pages[i] = page + query := buildWikiPageQueryText(page) + if strings.TrimSpace(query) == "" { + continue + } + queryTexts = append(queryTexts, query) + pageIdx = append(pageIdx, i) + } + if len(queryTexts) == 0 { + return plan, nil + } + vectors, err := p.deps.Embed.Encode(p.ctx, queryTexts) + if err != nil { + return wikiPlan{}, err + } + for i, idx := range pageIdx { + if idx >= len(plan.Pages) || i >= len(vectors) { + continue + } + page := plan.Pages[idx] + existing, err := p.reconcilePlanPage(page, vectors[i]) + if err != nil { + return wikiPlan{}, err + } + if existing == nil { + continue + } + page.Action = "UPDATE" + page.Slug = firstNonEmpty(existing.Slug, page.Slug) + page.Title = firstNonEmpty(existing.Title, page.Title) + page.PageType = firstNonEmpty(existing.PageType, page.PageType) + page.Topic = firstNonEmpty(existing.Topic, page.Topic) + page.RelatedKB = mergeStrings(page.RelatedKB, existing.RelatedKBPages) + page.EntityNames = mergeStrings(page.EntityNames, existing.EntityNames) + plan.Pages[idx] = page + } + plan.Pages = normalizeWikiPlanPages(plan.Pages, p.reduced) + return plan, nil +} + +func (p *wikiPipeline) reconcilePlanPage(page wikiPlanPage, queryVec []float32) (*common.WikiPageCandidate, error) { + if p.deps.WikiPages == nil { + return nil, nil + } + if hit, err := p.deps.WikiPages.GetPageBySlug(p.ctx, p.tenantID, p.datasetID, page.Slug); err != nil { + return nil, err + } else if hit != nil { + return hit, nil + } + cands, err := p.deps.WikiPages.FindSimilarPages(p.ctx, p.tenantID, p.datasetID, queryVec, wikiPlanSearchTopK) + if err != nil || len(cands) == 0 { + return nil, err + } + cands = rerankWikiPlanCandidates(page, cands) + targetNames := normalizedStringSet(page.EntityNames) + pageTitle := normKey(page.Title) + pageTopic := normKey(page.Topic) + maybe := make([]common.WikiPageCandidate, 0, len(cands)) + for i := range cands { + cand := cands[i] + if cand.Score >= wikiPlanUpdateThreshold { + return &cand, nil + } + if cand.Score < wikiPlanMaybeThreshold { + continue + } + if normKey(cand.Title) == pageTitle || normKey(cand.Topic) == pageTopic { + return &cand, nil + } + if intersectsNormalized(targetNames, cand.EntityNames) { + return &cand, nil + } + maybe = append(maybe, cand) + } + if len(maybe) > 0 { + return p.resolveMaybePlanPage(page, maybe) + } + return nil, nil +} + +func rerankWikiPlanCandidates(page wikiPlanPage, candidates []common.WikiPageCandidate) []common.WikiPageCandidate { + if len(candidates) <= 1 { + return candidates + } + type scored struct { + candidate common.WikiPageCandidate + boost float64 + } + targetNames := normalizedStringSet(page.EntityNames) + pageType := normKey(page.PageType) + pageTitle := normKey(page.Title) + pageTopic := normKey(page.Topic) + scoredCands := make([]scored, 0, len(candidates)) + for _, cand := range candidates { + boost := cand.Score + if pageType != "" && normKey(cand.PageType) == pageType { + boost += 0.08 + } + if pageTitle != "" && normKey(cand.Title) == pageTitle { + boost += 0.12 + } + if pageTopic != "" && normKey(cand.Topic) == pageTopic { + boost += 0.08 + } + if overlap := normalizedOverlapCount(targetNames, cand.EntityNames); overlap > 0 { + boost += 0.05 * float64(overlap) + } + scoredCands = append(scoredCands, scored{candidate: cand, boost: boost}) + } + sort.SliceStable(scoredCands, func(i, j int) bool { + if scoredCands[i].boost == scoredCands[j].boost { + return scoredCands[i].candidate.Slug < scoredCands[j].candidate.Slug + } + return scoredCands[i].boost > scoredCands[j].boost + }) + out := make([]common.WikiPageCandidate, 0, len(scoredCands)) + for _, item := range scoredCands { + out = append(out, item.candidate) + } + return out +} + +func (p *wikiPipeline) resolveMaybePlanPage(page wikiPlanPage, candidates []common.WikiPageCandidate) (*common.WikiPageCandidate, error) { + if len(candidates) == 0 { + return nil, nil + } + raw, err := common.GenJSON(p.ctx, p.deps.Chat, common.ChatRequest{ + LLMID: p.llmID, + SystemPrompt: wikiPlanReconcileSystem, + UserPrompt: renderWikiTemplate(wikiPlanReconcileUserTemplate, map[string]string{ + "planned_page": mustPrettyJSON(page), + "candidates": mustPrettyJSON(candidates), + }), + }) + if err != nil { + return nil, err + } + action := strings.ToUpper(strings.TrimSpace(firstString(raw["action"]))) + if action != "UPDATE" { + return nil, nil + } + slug := strings.TrimSpace(firstString(raw["slug"])) + if slug == "" { + return nil, nil + } + for i := range candidates { + if candidates[i].Slug == slug { + return &candidates[i], nil + } + } + return nil, nil +} + +func buildWikiPageQueryText(page wikiPlanPage) string { + parts := []string{page.Title, page.Topic, strings.Join(page.EntityNames, " ")} + var out []string + for _, part := range parts { + if s := strings.TrimSpace(part); s != "" { + out = append(out, s) + } + } + return strings.Join(out, "\n") +} + +func (p *wikiPipeline) mergeWikiPageContent(existingMD, newMD, slug string) (string, error) { + if strings.TrimSpace(existingMD) == "" { + return newMD, nil + } + if strings.TrimSpace(newMD) == "" { + return existingMD, nil + } + if strings.TrimSpace(existingMD) == strings.TrimSpace(newMD) { + return newMD, nil + } + fallback := conservativeMergeWikiMarkdown(existingMD, newMD) + user := fmt.Sprintf("Merge these two versions of wiki page `%s`:\n\n## EXISTING VERSION\n\n%s\n\n---\n\n## INCOMING VERSION\n\n%s\n\n---\n\nProduce the merged page now. Return only the markdown content.", slug, existingMD, newMD) + resp, err := p.deps.Chat.Chat(p.ctx, common.ChatRequest{ + LLMID: p.llmID, + SystemPrompt: wikiRefineMergeSystem, + UserPrompt: user, + }) + if err != nil { + return "", err + } + if resp == nil { + return fallback, nil + } + merged := strings.TrimSpace(firstNonEmpty(resp.Content)) + if merged == "" { + return fallback, nil + } + minAcceptable := int(float64(maxInt(len(existingMD), len(newMD))) * wikiMergeShrinkRatio) + if len(merged) < minAcceptable { + return fallback, nil + } + if wikiMarkdownDropsContent(merged, existingMD) || wikiMarkdownDropsContent(merged, newMD) { + return fallback, nil + } + return merged, nil +} + +func conservativeMergeWikiMarkdown(existingMD, newMD string) string { + if strings.Contains(existingMD, newMD) { + return strings.TrimSpace(existingMD) + } + if strings.Contains(newMD, existingMD) { + return strings.TrimSpace(newMD) + } + title := firstNonEmpty(wikiMarkdownTitle(existingMD), wikiMarkdownTitle(newMD)) + blocks := make([]string, 0, 8) + seen := map[string]bool{} + appendBlocks := func(body string) { + for _, block := range splitWikiMarkdownBlocks(body) { + key := normKey(block) + if key == "" || seen[key] { + continue + } + seen[key] = true + blocks = append(blocks, block) + } + } + appendBlocks(wikiMarkdownBody(existingMD)) + appendBlocks(wikiMarkdownBody(newMD)) + if title != "" { + if len(blocks) == 0 { + return "# " + title + } + return "# " + title + "\n\n" + strings.Join(blocks, "\n\n") + } + return strings.Join(blocks, "\n\n") +} + +func wikiMarkdownDropsContent(mergedMD, sourceMD string) bool { + mergedNorm := normKey(mergedMD) + for _, block := range splitWikiMarkdownBlocks(wikiMarkdownBody(sourceMD)) { + if !wikiMarkdownContentBlock(block) { + continue + } + if !strings.Contains(mergedNorm, normKey(block)) { + return true + } + } + return false +} + +func wikiMarkdownTitle(md string) string { + for _, line := range strings.Split(md, "\n") { + line = strings.TrimSpace(line) + if strings.HasPrefix(line, "# ") { + return strings.TrimSpace(strings.TrimPrefix(line, "# ")) + } + if line != "" { + break + } + } + return "" +} + +func wikiMarkdownBody(md string) string { + lines := strings.Split(md, "\n") + for i, line := range lines { + line = strings.TrimSpace(line) + if strings.HasPrefix(line, "# ") { + return strings.TrimSpace(strings.Join(lines[i+1:], "\n")) + } + if line != "" { + break + } + } + return strings.TrimSpace(md) +} + +func splitWikiMarkdownBlocks(md string) []string { + parts := strings.Split(strings.TrimSpace(md), "\n\n") + out := make([]string, 0, len(parts)) + for _, part := range parts { + part = strings.TrimSpace(part) + if part != "" { + out = append(out, part) + } + } + return out +} + +func wikiMarkdownContentBlock(block string) bool { + block = strings.TrimSpace(block) + if block == "" { + return false + } + if strings.HasPrefix(block, "#") { + return false + } + return true +} + +func reduceExtracts(extracts []wikiExtract) wikiExtract { + out := wikiExtract{} + type entityKey struct { + Name string + Type string + } + type conceptKey struct { + Term string + } + entities := map[entityKey]*wikiEntity{} + concepts := map[conceptKey]*wikiConcept{} + seenRelations := map[string]bool{} + seenTopics := map[string]bool{} + + for _, ex := range extracts { + for _, e := range ex.Entities { + key := entityKey{Name: normKey(e.Name), Type: normKey(e.Type)} + if cur, ok := entities[key]; ok { + cur.Aliases = mergeStrings(cur.Aliases, e.Aliases) + cur.SourceChunkIDs = mergeStrings(cur.SourceChunkIDs, e.SourceChunkIDs) + if cur.Type == "" { + cur.Type = e.Type + } + continue + } + item := e + item.Name = strings.TrimSpace(item.Name) + item.Type = strings.TrimSpace(item.Type) + item.Aliases = uniqueStrings(item.Aliases) + item.SourceChunkIDs = uniqueStrings(item.SourceChunkIDs) + entities[key] = &item + } + for _, c := range ex.Concepts { + key := conceptKey{Term: normKey(c.Term)} + if cur, ok := concepts[key]; ok { + if cur.Definition == "" { + cur.Definition = c.Definition + } + cur.SourceChunkIDs = mergeStrings(cur.SourceChunkIDs, c.SourceChunkIDs) + continue + } + item := c + item.Term = strings.TrimSpace(item.Term) + item.Definition = strings.TrimSpace(item.Definition) + item.SourceChunkIDs = uniqueStrings(item.SourceChunkIDs) + concepts[key] = &item + } + for _, c := range ex.Claims { + c.SourceChunkIDs = uniqueStrings(c.SourceChunkIDs) + out.Claims = append(out.Claims, c) + } + for _, r := range ex.Relations { + key := normKey(r.From) + "\x00" + normKey(r.Type) + "\x00" + normKey(r.To) + if seenRelations[key] { + continue + } + r.SourceChunkIDs = uniqueStrings(r.SourceChunkIDs) + seenRelations[key] = true + out.Relations = append(out.Relations, r) + } + for _, t := range ex.Topics { + t = strings.TrimSpace(t) + if t == "" { + continue + } + key := normKey(t) + if seenTopics[key] { + continue + } + seenTopics[key] = true + out.Topics = append(out.Topics, t) + } + } + + out.Entities = make([]wikiEntity, 0, len(entities)) + keys := make([]entityKey, 0, len(entities)) + for k := range entities { + keys = append(keys, k) + } + sort.Slice(keys, func(i, j int) bool { + if keys[i].Name == keys[j].Name { + return keys[i].Type < keys[j].Type + } + return keys[i].Name < keys[j].Name + }) + for _, k := range keys { + out.Entities = append(out.Entities, *entities[k]) + } + out.Concepts = make([]wikiConcept, 0, len(concepts)) + ckeys := make([]conceptKey, 0, len(concepts)) + for k := range concepts { + ckeys = append(ckeys, k) + } + sort.Slice(ckeys, func(i, j int) bool { return ckeys[i].Term < ckeys[j].Term }) + for _, k := range ckeys { + out.Concepts = append(out.Concepts, *concepts[k]) + } + return out +} + +func parseWikiExtract(raw map[string]any) wikiExtract { + out := wikiExtract{} + out.Entities = parseWikiEntities(raw["entities"]) + out.Concepts = parseWikiConcepts(raw["concepts"]) + out.Claims = parseWikiClaims(raw["claims"]) + out.Relations = parseWikiRelations(raw["relations"]) + out.Topics = parseWikiStrings(raw["topics"]) + return out +} + +func parseWikiPlan(raw map[string]any, docID string, reduced wikiExtract) wikiPlan { + plan := wikiPlan{ + Title: firstString(raw["title"]), + Slug: firstString(raw["slug"]), + Lead: firstString(raw["lead"]), + Sections: parseWikiPlanSections(raw["sections"]), + PageType: firstString(raw["page_type"]), + Topic: firstString(raw["topic"]), + Entities: parseWikiStrings(raw["entity_names"]), + Related: parseWikiStrings(raw["related_kb_pages"]), + Pages: parseWikiPlanPages(raw["pages"]), + } + if len(plan.Pages) > 0 { + first := plan.Pages[0] + if plan.Title == "" { + plan.Title = first.Title + } + if plan.Slug == "" { + plan.Slug = first.Slug + } + if plan.Lead == "" { + plan.Lead = first.Lead + } + if plan.PageType == "" { + plan.PageType = first.PageType + } + if plan.Topic == "" { + plan.Topic = first.Topic + } + if len(plan.Entities) == 0 { + plan.Entities = first.EntityNames + } + if len(plan.Related) == 0 { + plan.Related = first.RelatedKB + } + if len(plan.Sections) == 0 { + plan.Sections = first.Sections + } + } + if plan.Title == "" { + if len(reduced.Topics) > 0 { + plan.Title = reduced.Topics[0] + } else if len(reduced.Entities) > 0 { + plan.Title = reduced.Entities[0].Name + } else if len(reduced.Concepts) > 0 { + plan.Title = reduced.Concepts[0].Term + } else { + plan.Title = docID + } + } + if plan.Slug == "" { + plan.Slug = slugify(plan.Title) + } + if plan.Lead == "" { + plan.Lead = plan.Title + } + if len(plan.Sections) == 0 { + plan.Sections = []wikiPlanSection{{Heading: "Overview", Points: []string{plan.Lead}}} + } + return plan +} + +func normalizeWikiPlan(plan wikiPlan, docID string, reduced wikiExtract) wikiPlan { + plan.Pages = normalizeWikiPlanPages(plan.Pages, reduced) + if plan.Title == "" { + plan.Title = firstPlanTitle(plan.Pages, reduced, docID) + } + if plan.Slug == "" { + plan.Slug = slugify(plan.Title) + } + if plan.Lead == "" { + plan.Lead = plan.Title + } + if len(plan.Sections) == 0 { + plan.Sections = []wikiPlanSection{{Heading: "Overview", Points: []string{plan.Lead}}} + } + return plan +} + +func normalizeWikiPlanPages(pages []wikiPlanPage, reduced wikiExtract) []wikiPlanPage { + existing := map[string]wikiPlanPage{} + out := make([]wikiPlanPage, 0, len(pages)) + for _, page := range pages { + page = normalizeWikiPlanPage(page) + if page.Slug == "" { + continue + } + if _, ok := existing[page.Slug]; ok { + continue + } + existing[page.Slug] = page + out = append(out, page) + } + fallback := buildWikiFallbackPages(reduced) + if len(out) == 0 { + for _, page := range fallback { + if page.Slug == "" { + continue + } + existing[page.Slug] = page + out = append(out, page) + } + } + sort.SliceStable(out, func(i, j int) bool { + if out[i].Priority == out[j].Priority { + return out[i].Slug < out[j].Slug + } + return out[i].Priority < out[j].Priority + }) + for i := range out { + if out[i].Priority <= 0 { + out[i].Priority = i + 1 + } + } + return out +} + +func normalizeWikiPlanPage(page wikiPlanPage) wikiPlanPage { + page.Action = strings.ToUpper(strings.TrimSpace(page.Action)) + if page.Action == "" { + page.Action = "CREATE" + } + page.Slug = strings.TrimSpace(page.Slug) + if page.Slug == "" { + page.Slug = slugify(page.Title) + } + page.Title = strings.TrimSpace(page.Title) + if page.Title == "" { + page.Title = page.Slug + } + page.PageType = strings.TrimSpace(page.PageType) + if page.PageType == "" { + page.PageType = "concept" + } + page.Topic = strings.TrimSpace(page.Topic) + if page.Topic == "" { + page.Topic = page.Title + } + page.EntityNames = uniqueStrings(page.EntityNames) + page.RelatedKB = uniqueStrings(page.RelatedKB) + page.Sections = normalizeWikiPlanSections(page.Sections) + if page.Priority <= 0 { + page.Priority = 99 + } + return page +} + +func normalizeWikiPlanSections(sections []wikiPlanSection) []wikiPlanSection { + if len(sections) == 0 { + return []wikiPlanSection{{Heading: "Overview", Points: nil}} + } + out := make([]wikiPlanSection, 0, len(sections)) + for _, sec := range sections { + sec.Heading = strings.TrimSpace(sec.Heading) + sec.Points = uniqueStrings(sec.Points) + if sec.Heading == "" { + continue + } + out = append(out, sec) + } + if len(out) == 0 { + return []wikiPlanSection{{Heading: "Overview", Points: nil}} + } + return out +} + +func buildWikiFallbackPages(reduced wikiExtract) []wikiPlanPage { + var out []wikiPlanPage + seen := map[string]bool{} + appendPage := func(page wikiPlanPage) { + page = normalizeWikiPlanPage(page) + if page.Slug == "" || seen[page.Slug] { + return + } + seen[page.Slug] = true + out = append(out, page) + } + for _, e := range reduced.Entities { + appendPage(wikiPlanPage{ + Action: "CREATE", + Slug: "entity/" + slugify(e.Name), + Title: e.Name, + PageType: "entity", + Topic: e.Name, + EntityNames: append([]string{e.Name}, e.Aliases...), + Priority: len(out) + 1, + Lead: e.Name, + Sections: []wikiPlanSection{{Heading: "Overview", Points: []string{e.Name}}}, + }) + } + for _, c := range reduced.Concepts { + appendPage(wikiPlanPage{ + Action: "CREATE", + Slug: "concept/" + slugify(c.Term), + Title: c.Term, + PageType: "concept", + Topic: c.Term, + EntityNames: []string{c.Term}, + Priority: len(out) + 1, + Lead: c.Definition, + Sections: []wikiPlanSection{{Heading: "Overview", Points: []string{c.Term}}}, + }) + } + if len(out) == 0 { + for _, t := range reduced.Topics { + appendPage(wikiPlanPage{ + Action: "CREATE", + Slug: "topic/" + slugify(t), + Title: t, + PageType: "topic", + Topic: t, + Priority: len(out) + 1, + Lead: t, + Sections: []wikiPlanSection{{Heading: "Overview", Points: []string{t}}}, + }) + } + } + if len(out) == 0 { + appendPage(wikiPlanPage{ + Action: "CREATE", + Slug: slugify(firstPlanTitle(nil, reduced, "")), + Title: firstPlanTitle(nil, reduced, ""), + PageType: "concept", + Topic: firstPlanTitle(nil, reduced, ""), + Priority: 1, + Lead: firstPlanTitle(nil, reduced, ""), + Sections: []wikiPlanSection{{Heading: "Overview", Points: []string{firstPlanTitle(nil, reduced, "")}}}, + }) + } + return out +} + +func firstPlanTitle(pages []wikiPlanPage, reduced wikiExtract, docID string) string { + for _, page := range pages { + if s := strings.TrimSpace(page.Title); s != "" { + return s + } + } + if len(reduced.Topics) > 0 { + return reduced.Topics[0] + } + if len(reduced.Entities) > 0 { + return reduced.Entities[0].Name + } + if len(reduced.Concepts) > 0 { + return reduced.Concepts[0].Term + } + if docID != "" { + return docID + } + return "wiki" +} + +func (p wikiPlan) firstPageCandidate() wikiPlanPage { + if len(p.Pages) > 0 { + return p.Pages[0] + } + return wikiPlanPage{ + Action: "CREATE", + Slug: p.Slug, + Title: p.Title, + PageType: p.PageType, + Topic: p.Topic, + EntityNames: uniqueStrings(p.Entities), + RelatedKB: uniqueStrings(p.Related), + Priority: 1, + Lead: p.Lead, + Sections: p.Sections, + } +} + +func packWikiPlanBatches(reduced wikiExtract, budget int) []wikiExtract { + if budget <= 0 { + budget = wikiPlanTokenBudget + } + var batches []wikiExtract + cur := wikiExtract{} + curTokens := 0 + flush := func() { + if len(cur.Entities)+len(cur.Concepts)+len(cur.Claims)+len(cur.Relations)+len(cur.Topics) > 0 { + batches = append(batches, cur) + cur = wikiExtract{} + curTokens = 0 + } + } + add := func(itemTokens int, addFn func()) { + if curTokens > 0 && curTokens+itemTokens > budget { + flush() + } + addFn() + curTokens += itemTokens + } + for _, e := range reduced.Entities { + add(common.EstimateTokens(mustJSON(e)), func() { cur.Entities = append(cur.Entities, e) }) + } + for _, c := range reduced.Concepts { + add(common.EstimateTokens(mustJSON(c)), func() { cur.Concepts = append(cur.Concepts, c) }) + } + for _, c := range reduced.Claims { + add(common.EstimateTokens(mustJSON(c)), func() { cur.Claims = append(cur.Claims, c) }) + } + for _, r := range reduced.Relations { + add(common.EstimateTokens(mustJSON(r)), func() { cur.Relations = append(cur.Relations, r) }) + } + for _, t := range reduced.Topics { + add(common.EstimateTokens(mustJSON(t)), func() { cur.Topics = append(cur.Topics, t) }) + } + flush() + return batches +} + +func parseWikiPlanPages(raw any) []wikiPlanPage { + arr, ok := raw.([]any) + if !ok { + return nil + } + out := make([]wikiPlanPage, 0, len(arr)) + for _, item := range arr { + m, ok := item.(map[string]any) + if !ok { + continue + } + page := wikiPlanPage{ + Action: strings.TrimSpace(firstString(m["action"])), + Slug: strings.TrimSpace(firstString(m["slug"])), + Title: strings.TrimSpace(firstString(m["title"])), + PageType: strings.TrimSpace(firstString(m["page_type"])), + Topic: strings.TrimSpace(firstString(m["topic"])), + EntityNames: parseWikiStrings(m["entity_names"]), + RelatedKB: parseWikiStrings(m["related_kb_pages"]), + Priority: int(firstNumber(m["priority"])), + Lead: strings.TrimSpace(firstString(m["lead"])), + Sections: parseWikiPlanSections(m["sections"]), + } + if page.Slug != "" || page.Title != "" { + out = append(out, page) + } + } + return out +} + +func parseWikiEntities(raw any) []wikiEntity { + arr, ok := raw.([]any) + if !ok { + return nil + } + out := make([]wikiEntity, 0, len(arr)) + for _, item := range arr { + m, ok := item.(map[string]any) + if !ok { + continue + } + e := wikiEntity{ + Name: strings.TrimSpace(firstString(m["name"])), + Type: strings.TrimSpace(firstString(m["type"])), + Aliases: parseWikiStrings(m["aliases"]), + SourceChunkIDs: parseWikiStrings(m["source_chunk_ids"]), + } + if len(e.SourceChunkIDs) == 0 { + if s := strings.TrimSpace(firstString(m["source_chunk_id"])); s != "" { + e.SourceChunkIDs = []string{s} + } + } + if e.Name != "" { + out = append(out, e) + } + } + return out +} + +func parseWikiConcepts(raw any) []wikiConcept { + arr, ok := raw.([]any) + if !ok { + return nil + } + out := make([]wikiConcept, 0, len(arr)) + for _, item := range arr { + m, ok := item.(map[string]any) + if !ok { + continue + } + c := wikiConcept{ + Term: strings.TrimSpace(firstString(m["term"])), + Definition: strings.TrimSpace(firstString(m["definition_excerpt"])), + SourceChunkIDs: parseWikiStrings(m["source_chunk_ids"]), + } + if len(c.SourceChunkIDs) == 0 { + if s := strings.TrimSpace(firstString(m["source_chunk_id"])); s != "" { + c.SourceChunkIDs = []string{s} + } + } + if c.Term != "" { + out = append(out, c) + } + } + return out +} + +func parseWikiClaims(raw any) []wikiClaim { + arr, ok := raw.([]any) + if !ok { + return nil + } + out := make([]wikiClaim, 0, len(arr)) + for _, item := range arr { + m, ok := item.(map[string]any) + if !ok { + continue + } + c := wikiClaim{ + Statement: strings.TrimSpace(firstString(m["statement"])), + Subject: strings.TrimSpace(firstString(m["subject"])), + Confidence: strings.TrimSpace(firstString(m["confidence"])), + SourceChunkIDs: parseWikiStrings(m["source_chunk_ids"]), + } + if len(c.SourceChunkIDs) == 0 { + if s := strings.TrimSpace(firstString(m["source_chunk_id"])); s != "" { + c.SourceChunkIDs = []string{s} + } + } + if c.Statement != "" { + out = append(out, c) + } + } + return out +} + +func parseWikiRelations(raw any) []wikiRelation { + arr, ok := raw.([]any) + if !ok { + return nil + } + out := make([]wikiRelation, 0, len(arr)) + for _, item := range arr { + m, ok := item.(map[string]any) + if !ok { + continue + } + r := wikiRelation{ + From: strings.TrimSpace(firstString(m["from"])), + To: strings.TrimSpace(firstString(m["to"])), + Type: strings.TrimSpace(firstString(m["type"])), + SourceChunkIDs: parseWikiStrings(m["source_chunk_ids"]), + } + if len(r.SourceChunkIDs) == 0 { + if s := strings.TrimSpace(firstString(m["source_chunk_id"])); s != "" { + r.SourceChunkIDs = []string{s} + } + } + if r.From != "" && r.To != "" { + out = append(out, r) + } + } + return out +} + +func parseWikiPlanSections(raw any) []wikiPlanSection { + arr, ok := raw.([]any) + if !ok { + return nil + } + out := make([]wikiPlanSection, 0, len(arr)) + for _, item := range arr { + m, ok := item.(map[string]any) + if !ok { + continue + } + sec := wikiPlanSection{ + Heading: strings.TrimSpace(firstString(m["heading"])), + Points: parseWikiStrings(m["points"]), + } + if sec.Heading != "" { + out = append(out, sec) + } + } + return out +} + +func firstNumber(v any) float64 { + switch x := v.(type) { + case int: + return float64(x) + case int64: + return float64(x) + case float32: + return float64(x) + case float64: + return x + default: + return 0 + } +} + +func parseWikiStrings(raw any) []string { + switch v := raw.(type) { + case []string: + return uniqueStrings(v) + case []any: + out := make([]string, 0, len(v)) + for _, item := range v { + if s := strings.TrimSpace(firstString(item)); s != "" { + out = append(out, s) + } + } + return uniqueStrings(out) + case string: + if s := strings.TrimSpace(v); s != "" { + return []string{s} + } + } + return nil +} + +func buildSourceContext(chunks []common.Chunk, sourceChunkIDs []string) string { + if len(sourceChunkIDs) == 0 { + var b strings.Builder + for _, ch := range chunks { + text := firstNonEmpty(ch.Text, ch.Content) + if strings.TrimSpace(text) == "" { + continue + } + b.WriteString(text) + b.WriteString("\n\n") + } + return b.String() + } + lookup := map[string]string{} + for i, ch := range chunks { + id := strings.TrimSpace(ch.ID) + if id == "" { + id = fmt.Sprintf("chunk-%d", i+1) + } + lookup[id] = firstNonEmpty(ch.Text, ch.Content) + } + var b strings.Builder + for _, id := range sourceChunkIDs { + if text := strings.TrimSpace(lookup[id]); text != "" { + b.WriteString("[CHUNK_ID ") + b.WriteString(id) + b.WriteString("]\n") + b.WriteString(text) + b.WriteString("\n\n") + } + } + return b.String() +} + +func buildEvidenceChecklist(reduced wikiExtract) string { + var lines []string + for _, e := range reduced.Entities { + lines = append(lines, fmt.Sprintf("- entity: %s (%s)", e.Name, e.Type)) + } + for _, c := range reduced.Concepts { + lines = append(lines, fmt.Sprintf("- concept: %s", c.Term)) + } + for _, c := range reduced.Claims { + lines = append(lines, fmt.Sprintf("- claim: %s", c.Statement)) + } + for _, r := range reduced.Relations { + lines = append(lines, fmt.Sprintf("- relation: %s -> %s (%s)", r.From, r.To, r.Type)) + } + for _, t := range reduced.Topics { + lines = append(lines, fmt.Sprintf("- topic: %s", t)) + } + if len(lines) == 0 { + return "- (none)" + } + return strings.Join(lines, "\n") +} + +type wikiEvidenceItem struct { + Statement string + Subject string + Confidence string + ChunkIDs []string + Synthetic bool +} + +func assembleWikiPageEvidence(planItem wikiPlanPage, claims []wikiClaim, entityByName, conceptByTerm map[string]wikiExtractItem) []wikiEvidenceItem { + rawNames := uniqueStrings(append([]string{}, planItem.EntityNames...)) + if len(rawNames) == 0 { + if t := strings.TrimSpace(planItem.Title); t != "" { + rawNames = []string{t} + } else if s := strings.TrimSpace(planItem.Slug); s != "" { + rawNames = []string{s} + } + } + if len(rawNames) == 0 { + return nil + } + + namesLower := make([]string, 0, len(rawNames)) + patterns := make([]*regexp.Regexp, 0, len(rawNames)) + for _, n := range rawNames { + namesLower = append(namesLower, strings.ToLower(n)) + patterns = append(patterns, regexp.MustCompile(`(?i)\b`+regexp.QuoteMeta(n)+`\b`)) + } + + var evidence []wikiEvidenceItem + for _, claim := range claims { + subjRaw := strings.TrimSpace(claim.Subject) + if subjRaw == "" { + continue + } + subjLower := strings.ToLower(subjRaw) + matched := false + for _, n := range namesLower { + if subjLower == n { + matched = true + break + } + } + if !matched { + for _, re := range patterns { + if re.MatchString(subjRaw) { + matched = true + break + } + } + } + if !matched { + continue + } + evidence = append(evidence, wikiEvidenceItem{ + Statement: claim.Statement, + Subject: claim.Subject, + Confidence: firstNonEmpty(claim.Confidence, "explicit"), + ChunkIDs: uniqueStrings(claim.SourceChunkIDs), + }) + } + if len(evidence) > 0 { + return evidence + } + + var fallbackChunkIDs []string + matchedNames := make([]string, 0, len(rawNames)) + for _, name := range rawNames { + key := strings.ToLower(strings.TrimSpace(name)) + var hit wikiExtractItem + var ok bool + if entityByName != nil { + hit, ok = entityByName[key] + } + if !ok && conceptByTerm != nil { + hit, ok = conceptByTerm[key] + } + if !ok { + continue + } + fallbackChunkIDs = mergeStrings(fallbackChunkIDs, hit.SourceChunkIDs) + matchedNames = append(matchedNames, name) + } + if len(fallbackChunkIDs) == 0 { + return nil + } + subject := rawNames[0] + if len(matchedNames) > 0 { + subject = matchedNames[0] + } + return []wikiEvidenceItem{{ + Statement: "", + Subject: subject, + Confidence: "inferred", + ChunkIDs: fallbackChunkIDs, + Synthetic: true, + }} +} + +type wikiExtractItem struct { + SourceChunkIDs []string +} + +func buildWikiEntityLookup(items []wikiEntity) map[string]wikiExtractItem { + out := map[string]wikiExtractItem{} + for _, item := range items { + hit := wikiExtractItem{SourceChunkIDs: uniqueStrings(item.SourceChunkIDs)} + name := strings.TrimSpace(item.Name) + if name != "" { + out[strings.ToLower(name)] = hit + } + for _, alias := range item.Aliases { + if alias = strings.TrimSpace(alias); alias != "" { + out[strings.ToLower(alias)] = hit + } + } + } + return out +} + +func buildWikiConceptLookup(items []wikiConcept) map[string]wikiExtractItem { + out := map[string]wikiExtractItem{} + for _, item := range items { + hit := wikiExtractItem{SourceChunkIDs: uniqueStrings(item.SourceChunkIDs)} + term := strings.TrimSpace(item.Term) + if term != "" { + out[strings.ToLower(term)] = hit + } + } + return out +} + +func formatWikiEvidenceBlocks(evidence []wikiEvidenceItem) string { + var lines []string + for _, ev := range evidence { + if ev.Synthetic { + continue + } + confidence := strings.ToUpper(firstNonEmpty(ev.Confidence, "explicit")) + lines = append(lines, fmt.Sprintf("%d. [%s] %s\n %s", len(lines)+1, confidence, ev.Subject, ev.Statement)) + } + if len(lines) == 0 { + return "(no pre-extracted evidence — extract facts directly from the source document text above)" + } + return strings.Join(lines, "\n\n") +} + +func collectWikiEvidenceChunkIDs(evidence []wikiEvidenceItem) []string { + var out []string + for _, ev := range evidence { + out = mergeStrings(out, ev.ChunkIDs) + } + return out +} + +func collectWikiSourceDocIDs(chunks []common.Chunk, chunkIDs []string, fallback string) []string { + if len(chunkIDs) == 0 { + if fallback != "" { + return []string{fallback} + } + return nil + } + lookup := map[string][]string{} + for _, ch := range chunks { + id := strings.TrimSpace(ch.ID) + if id == "" { + continue + } + docID := strings.TrimSpace(firstString(ch.Meta["doc_id"])) + if docID == "" { + docID = strings.TrimSpace(firstString(ch.Meta["doc_id_kwd"])) + } + if docID != "" { + lookup[id] = []string{docID} + } + } + out := make([]string, 0, len(chunkIDs)) + for _, cid := range chunkIDs { + if ids := lookup[strings.TrimSpace(cid)]; len(ids) > 0 { + out = mergeStrings(out, ids) + } + } + if len(out) == 0 && fallback != "" { + return []string{fallback} + } + return out +} + +var ( + wikiWikilinkPipeRe = regexp.MustCompile(`\[\[([^\[\]\|]+?)\|([^\[\]]+?)\]\]`) + wikiWikilinkSimpleRe = regexp.MustCompile(`\[\[([^\[\]\|]+?)\]\]`) + wikiArtifactMarkdownLink = regexp.MustCompile(`\[([^\]]+)\]\(([^)]+)\)`) +) + +func transformWikiLinks(content, kbID string, pageTitles map[string]string) (string, []string) { + kbID = strings.TrimSpace(kbID) + seen := map[string]bool{} + var outlinks []string + track := func(slug string) { + slug = strings.TrimSpace(slug) + if slug == "" || seen[slug] { + return + } + seen[slug] = true + outlinks = append(outlinks, slug) + } + displayText := func(label, slug string) string { + label = strings.TrimSpace(label) + if label != slug && label != slugify(lastPathSlug(slug)) { + return label + } + if title := strings.TrimSpace(pageTitles[slug]); title != "" { + return title + } + readable := strings.ReplaceAll(lastPathSlug(slug), "-", " ") + readable = strings.ReplaceAll(readable, "_", " ") + readable = strings.TrimSpace(readable) + if readable == "" { + return label + } + return strings.Title(readable) + } + artifactSlug := func(href string) string { + parsed, err := url.Parse(href) + if err != nil { + return "" + } + path := parsed.Path + if path == "" && parsed.Opaque != "" { + path = parsed.Opaque + } + parts := strings.Split(strings.Trim(path, "/"), "/") + if len(parts) >= 2 && parts[0] == "artifact" && parts[1] == kbID { + return strings.Join(parts[2:], "/") + } + if len(parts) >= 2 && parts[0] == kbID { + return strings.Join(parts[1:], "/") + } + return "" + } + rewriteMD := func(label, href string) string { + slug := artifactSlug(href) + if slug == "" { + return "[" + label + "](" + href + ")" + } + track(slug) + return "[" + displayText(label, slug) + "](artifact/" + kbID + "/" + slug + ")" + } + out := wikiArtifactMarkdownLink.ReplaceAllStringFunc(content, func(match string) string { + sub := wikiArtifactMarkdownLink.FindStringSubmatch(match) + if len(sub) != 3 { + return match + } + return rewriteMD(sub[1], sub[2]) + }) + out = wikiWikilinkPipeRe.ReplaceAllStringFunc(out, func(match string) string { + sub := wikiWikilinkPipeRe.FindStringSubmatch(match) + if len(sub) != 3 { + return match + } + slug := strings.TrimSpace(sub[1]) + track(slug) + return "[" + sub[2] + "](artifact/" + kbID + "/" + slug + ")" + }) + out = wikiWikilinkSimpleRe.ReplaceAllStringFunc(out, func(match string) string { + sub := wikiWikilinkSimpleRe.FindStringSubmatch(match) + if len(sub) != 2 { + return match + } + slug := strings.TrimSpace(sub[1]) + track(slug) + return "[" + displayText(slug, slug) + "](artifact/" + kbID + "/" + slug + ")" + }) + return out, outlinks +} + +func lastPathSlug(slug string) string { + slug = strings.TrimSpace(slug) + if idx := strings.LastIndex(slug, "/"); idx >= 0 && idx < len(slug)-1 { + return slug[idx+1:] + } + return slug +} + +func mustJSON(v any) string { + b, err := json.Marshal(v) + if err != nil { + return "[]" + } + return string(b) +} + +func mustPrettyJSON(v any) string { + b, err := json.MarshalIndent(v, "", " ") + if err != nil { + return "{}" + } + return string(b) +} + +func renderWikiTemplate(tmpl string, values map[string]string) string { + out := tmpl + keys := make([]string, 0, len(values)) + for k := range values { + keys = append(keys, k) + } + sort.Strings(keys) + for _, k := range keys { + out = strings.ReplaceAll(out, "{"+k+"}", values[k]) + } + return out +} + +func prefixWithDash(items []string) []string { + out := make([]string, 0, len(items)) + for _, item := range items { + out = append(out, "- "+item) + } + return out +} + +func firstNonEmpty(vals ...string) string { + for _, v := range vals { + if s := strings.TrimSpace(v); s != "" { + return s + } + } + return "" +} + +func uniqueStrings(in []string) []string { + if len(in) == 0 { + return nil + } + seen := map[string]bool{} + out := make([]string, 0, len(in)) + for _, s := range in { + s = strings.TrimSpace(s) + if s == "" || seen[s] { + continue + } + seen[s] = true + out = append(out, s) + } + return out +} + +func mergeStrings(a, b []string) []string { + return uniqueStrings(append(append([]string{}, a...), b...)) +} + +func normalizedStringSet(items []string) map[string]struct{} { + out := make(map[string]struct{}, len(items)) + for _, item := range items { + if key := normKey(item); key != "" { + out[key] = struct{}{} + } + } + return out +} + +func intersectsNormalized(needle map[string]struct{}, haystack []string) bool { + if len(needle) == 0 { + return false + } + for _, item := range haystack { + if _, ok := needle[normKey(item)]; ok { + return true + } + } + return false +} + +func normalizedOverlapCount(needle map[string]struct{}, haystack []string) int { + if len(needle) == 0 { + return 0 + } + count := 0 + seen := map[string]struct{}{} + for _, item := range haystack { + key := normKey(item) + if key == "" { + continue + } + if _, dup := seen[key]; dup { + continue + } + seen[key] = struct{}{} + if _, ok := needle[key]; ok { + count++ + } + } + return count +} + +func normKey(s string) string { + return strings.ToLower(strings.Join(strings.Fields(strings.TrimSpace(s)), " ")) +} + +func firstString(v any) string { + switch x := v.(type) { + case string: + return x + case fmt.Stringer: + return x.String() + default: + return "" + } +} + +func (e wikiExtract) sourceChunkIDs() []string { + seen := map[string]bool{} + var out []string + add := func(s string) { + s = strings.TrimSpace(s) + if s == "" || seen[s] { + return + } + seen[s] = true + out = append(out, s) + } + for _, item := range e.Entities { + for _, id := range item.SourceChunkIDs { + add(id) + } + } + for _, item := range e.Concepts { + for _, id := range item.SourceChunkIDs { + add(id) + } + } + for _, item := range e.Claims { + for _, id := range item.SourceChunkIDs { + add(id) + } + } + for _, item := range e.Relations { + for _, id := range item.SourceChunkIDs { + add(id) + } + } + return out +} + +func (e wikiExtract) asPlanPayload() map[string]any { + return map[string]any{ + "entities": e.Entities, + "concepts": e.Concepts, + "claims": e.Claims, + "relations": e.Relations, + "topics": e.Topics, + } +} + +func compactJSON(v any) string { + b, err := json.Marshal(v) + if err != nil { + return "{}" + } + return string(bytes.TrimSpace(b)) +} + +func (p *wikiPipeline) maybeSourceIDs() []string { + if ids := p.reduced.sourceChunkIDs(); len(ids) > 0 { + return ids + } + seen := map[string]bool{} + var out []string + for _, ch := range p.inputs.Chunks { + id := strings.TrimSpace(ch.ID) + if id == "" || seen[id] { + continue + } + seen[id] = true + out = append(out, id) + } + 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. +func dedupHistorical(ctx context.Context, deps common.Deps, param common.Param, tenantID, datasetID string, override []common.Candidate, products []common.Product) (survivors []common.Product, dropped int, err error) { + if len(products) == 0 { + return products, 0, nil + } + threshold := param.SimilarityThreshold + + if len(override) > 0 { + for _, p := range products { + if nearDuplicateOf(p.Vector, override, threshold) { + dropped++ + continue + } + survivors = append(survivors, p) + } + return survivors, dropped, nil + } + + if deps.HistoricalKNN == nil { + return products, 0, nil + } + for _, p := range products { + hits, e := deps.HistoricalKNN.TopKHistory(ctx, tenantID, datasetID, string(common.VariantWiki), p.Vector, wikiHistoricalK, threshold) + if e != nil { + return survivors, dropped, e + } + if len(hits) > 0 { + dropped++ + continue + } + survivors = append(survivors, p) + } + return survivors, dropped, nil +} + +func nearDuplicateOf(vec []float32, cands []common.Candidate, threshold float64) bool { + qn := l2Norm32(vec) + if qn == 0 { + qn = 1 + } + for _, c := range cands { + cn := l2Norm32(c.Vector) + if cn == 0 { + cn = 1 + } + var dot float64 + for i := 0; i < len(vec) && i < len(c.Vector); i++ { + dot += float64(vec[i]) * float64(c.Vector[i]) + } + if dot/(qn*cn) >= threshold { + return true + } + } + return false +} + +func l2Norm32(v []float32) float64 { + var s float64 + for _, x := range v { + s += float64(x) * float64(x) + } + return math.Sqrt(s) +} + +func maxInt(a, b int) int { + if a > b { + return a + } + return b +} + +// wikiHistoricalK is the K used for the historical-dedup KNN lookup. +const wikiHistoricalK = 5 + +// wikiPlanTokenBudget caps one planning round's approximate token load. +const wikiPlanTokenBudget = 3500 diff --git a/internal/ingestion/component/knowledge_compiler/wiki/wiki_test.go b/internal/ingestion/component/knowledge_compiler/wiki/wiki_test.go new file mode 100644 index 0000000000..5f233ba648 --- /dev/null +++ b/internal/ingestion/component/knowledge_compiler/wiki/wiki_test.go @@ -0,0 +1,185 @@ +package wiki + +import ( + "context" + "strings" + "testing" + + "ragflow/internal/ingestion/component/knowledge_compiler/common" +) + +func TestReduceExtracts_MergesProvenance(t *testing.T) { + reduced := reduceExtracts([]wikiExtract{ + { + Entities: []wikiEntity{{Name: "Alpha", Type: "thing", SourceChunkIDs: []string{"c1"}}}, + Claims: []wikiClaim{{Statement: "Alpha exists", Subject: "Alpha", SourceChunkIDs: []string{"c1"}}}, + }, + { + Entities: []wikiEntity{{Name: "Alpha", Type: "thing", SourceChunkIDs: []string{"c2"}}}, + Claims: []wikiClaim{{Statement: "Alpha exists", Subject: "Alpha", SourceChunkIDs: []string{"c2"}}}, + }, + }) + if len(reduced.Entities) != 1 { + t.Fatalf("entities=%d, want 1", len(reduced.Entities)) + } + if ids := reduced.Entities[0].SourceChunkIDs; len(ids) != 2 { + t.Fatalf("entity provenance = %#v, want 2 chunk ids", ids) + } + if len(reduced.Claims) != 2 { + t.Fatalf("claims=%d, want 2", len(reduced.Claims)) + } +} + +func TestPackWikiPlanBatches_SplitsLargeInput(t *testing.T) { + reduced := wikiExtract{ + Entities: []wikiEntity{ + {Name: strings.Repeat("a", 1000)}, + {Name: strings.Repeat("b", 1000)}, + {Name: strings.Repeat("c", 1000)}, + }, + } + batches := packWikiPlanBatches(reduced, 1) + if len(batches) < 2 { + t.Fatalf("expected multiple batches, got %d", len(batches)) + } +} + +func TestBuildSourceContext_SelectsKnownChunks(t *testing.T) { + ctx := buildSourceContext([]common.Chunk{ + {ID: "c1", Text: "alpha text"}, + {ID: "c2", Text: "beta text"}, + {ID: "c3", Text: "gamma text"}, + }, []string{"c2", "c3"}) + if strings.Contains(ctx, "alpha text") { + t.Fatalf("source context leaked unselected chunk: %q", ctx) + } + if !strings.Contains(ctx, "beta text") || !strings.Contains(ctx, "gamma text") { + t.Fatalf("source context missing selected chunks: %q", ctx) + } +} + +func TestNormalizeWikiPlanPages_FallbacksToEntitiesAndConcepts(t *testing.T) { + plan := normalizeWikiPlan(wikiPlan{}, "doc-1", wikiExtract{ + Entities: []wikiEntity{{Name: "Alpha", Aliases: []string{"A"}}}, + Concepts: []wikiConcept{{Term: "Beta"}}, + }) + if len(plan.Pages) < 2 { + t.Fatalf("normalizeWikiPlan generated %d pages, want at least 2", len(plan.Pages)) + } + if plan.Pages[0].Slug == "" || plan.Pages[1].Slug == "" { + t.Fatalf("fallback pages missing slugs: %#v", plan.Pages) + } +} + +type reconcileChatStub struct { + resp string +} + +func (s reconcileChatStub) Chat(_ context.Context, req common.ChatRequest) (*common.ChatResponse, error) { + return &common.ChatResponse{Content: s.resp}, nil +} + +type reconcileEmbedStub struct{} + +func (reconcileEmbedStub) Encode(_ context.Context, texts []string) ([][]float32, error) { + out := make([][]float32, len(texts)) + for i := range texts { + out[i] = []float32{0.1, 0.2, 0.3} + } + return out, nil +} + +func (reconcileEmbedStub) Dimensions() int { return 3 } + +type wikiStoreStub struct { + slugHit *common.WikiPageCandidate + similar []common.WikiPageCandidate +} + +func (s wikiStoreStub) FindSimilarPages(_ context.Context, _, _ string, _ []float32, _ int) ([]common.WikiPageCandidate, error) { + return s.similar, nil +} + +func (s wikiStoreStub) GetPageBySlug(_ context.Context, _, _, slug string) (*common.WikiPageCandidate, error) { + if s.slugHit != nil && s.slugHit.Slug == slug { + return s.slugHit, nil + } + return nil, nil +} + +func TestReconcilePlanPage_MaybeUsesLLMDecision(t *testing.T) { + p := &wikiPipeline{ + ctx: context.Background(), + tenantID: "t1", + datasetID: "kb1", + llmID: "llm1", + deps: common.Deps{ + Chat: reconcileChatStub{resp: `{"action":"UPDATE","slug":"entity/existing","reason":"same entity"}`}, + Embed: reconcileEmbedStub{}, + WikiPages: wikiStoreStub{similar: []common.WikiPageCandidate{{Slug: "entity/existing", Title: "Existing", Score: 0.81}}}, + }, + } + got, err := p.reconcilePlanPage(wikiPlanPage{ + Slug: "entity/new-alpha", + Title: "Alpha", + PageType: "entity", + Topic: "Alpha", + EntityNames: []string{"Alpha Prime"}, + }, []float32{0.1, 0.2, 0.3}) + if err != nil { + t.Fatalf("reconcilePlanPage err = %v", err) + } + if got == nil || got.Slug != "entity/existing" { + t.Fatalf("reconcilePlanPage = %#v, want entity/existing", got) + } +} + +func TestReconcilePlanPage_LowScoreSkipsLLM(t *testing.T) { + p := &wikiPipeline{ + ctx: context.Background(), + tenantID: "t1", + datasetID: "kb1", + llmID: "llm1", + deps: common.Deps{ + Chat: reconcileChatStub{resp: `{"action":"UPDATE","slug":"entity/existing","reason":"same entity"}`}, + Embed: reconcileEmbedStub{}, + WikiPages: wikiStoreStub{similar: []common.WikiPageCandidate{{Slug: "entity/existing", Title: "Existing", Score: 0.6}}}, + }, + } + got, err := p.reconcilePlanPage(wikiPlanPage{ + Slug: "entity/new-alpha", + Title: "Alpha", + PageType: "entity", + Topic: "Alpha", + EntityNames: []string{"Alpha Prime"}, + }, []float32{0.1, 0.2, 0.3}) + if err != nil { + t.Fatalf("reconcilePlanPage err = %v", err) + } + if got != nil { + t.Fatalf("reconcilePlanPage = %#v, want nil", got) + } +} + +func TestMergeWikiPageContent_PreservesShortExistingPage(t *testing.T) { + p := &wikiPipeline{ + ctx: context.Background(), + deps: common.Deps{ + Chat: reconcileChatStub{resp: "# Alpha\n\nAlpha launched a new process in 2026.\n"}, + }, + } + merged, err := p.mergeWikiPageContent( + "# Alpha\n\nExisting fact.\n", + "# Alpha\n\nAlpha launched a new process in 2026.\n", + "entity/alpha", + ) + if err != nil { + t.Fatalf("mergeWikiPageContent err = %v", err) + } + if !strings.Contains(merged, "Existing fact.") { + t.Fatalf("merged page dropped existing content: %q", merged) + } + if !strings.Contains(merged, "Alpha launched a new process in 2026.") { + t.Fatalf("merged page dropped incoming content: %q", merged) + } +} diff --git a/internal/ingestion/component/schema/extractor.go b/internal/ingestion/component/schema/extractor.go index c5f7733914..783defc253 100644 --- a/internal/ingestion/component/schema/extractor.go +++ b/internal/ingestion/component/schema/extractor.go @@ -16,6 +16,8 @@ package schema +import "ragflow/internal/common" + // TagLabel is a single labeled record from the tag definition file: // a piece of content and the tags associated with it. type TagLabel struct { @@ -114,19 +116,34 @@ type ExtractorParam struct { // storage. Used only when AutoTags > 0 and no inline tag // source text is wired in. TagFileID string `json:"tag_file_id"` + + // EnableMetadata enables automatic structured metadata extraction + // with a fixed prompt. When > 0, runEnableMetadata asks the LLM to + // fill the fields listed in Metadata as a JSON object and + // merges the result into chunk["metadata"] (mirrors Python's + // gen_metadata_task path; see doc_metadata_go_port_research.md). + EnableMetadata int `json:"enable_metadata,omitempty"` + + // Metadata lists the target field definitions for + // EnableMetadata (mirrors parser_config.metadata / built_in_metadata: + // {key, type, description, enum}). When empty, EnableMetadata is a + // no-op (nothing to extract). + Metadata []common.MetadataFieldDef `json:"metadata,omitempty"` } // Defaults returns the default ExtractorParam. func (ExtractorParam) Defaults() ExtractorParam { return ExtractorParam{ - FieldName: "", - LLMID: "", - SystemPrompt: "", - Prompt: "", - AutoKeywords: 0, - AutoQuestions: 0, - AutoTags: 0, - TagFileID: "", + FieldName: "", + LLMID: "", + SystemPrompt: "", + Prompt: "", + AutoKeywords: 0, + AutoQuestions: 0, + AutoTags: 0, + TagFileID: "", + EnableMetadata: 0, + Metadata: nil, } } diff --git a/internal/ingestion/knowledge_compile/consumer.go b/internal/ingestion/knowledge_compile/consumer.go new file mode 100644 index 0000000000..e7b2a1b9aa --- /dev/null +++ b/internal/ingestion/knowledge_compile/consumer.go @@ -0,0 +1,285 @@ +// +// 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 knowledge_compile + +import ( + "context" + "sync" + "time" +) + +// Consumer is the dataset-level post-processing worker (§11.5). Multiple +// instances compete on the MySQL scheduling rows; each KB is processed by at +// most one instance at a time via the per-KB claim (so the same KB is handled +// by a single worker that serializes the batch). The MySQL row — not the +// broker — is the scheduling system of record and the source of same-KB +// serialization. +type Consumer struct { + scheduler Scheduler + reader Reader + writer Writer + factory DeduperFactory + + batchSize int + ttl time.Duration + heartbeat time.Duration + pollInterval time.Duration + sweepInterval time.Duration + + mu sync.Mutex + seqs map[string]map[string]uint64 // dataset -> docID -> last applied seq (per-doc out-of-order guard) + tombs map[string]map[string]uint64 // dataset -> docID -> delete event seq (tombstone) +} + +// NewConsumer constructs a Consumer driven by the given Scheduler. Tests pass a +// FakeScheduler and override the Reader/Writer/Deduper via options. +func NewConsumer(scheduler Scheduler, opts ...Option) *Consumer { + c := &Consumer{ + scheduler: scheduler, + reader: infinityReader{}, + writer: infinityWriter{}, + factory: defaultDeduperFactory, + batchSize: 32, + ttl: 2 * time.Minute, + heartbeat: 20 * time.Second, + pollInterval: 2 * time.Second, + sweepInterval: 30 * time.Second, + seqs: map[string]map[string]uint64{}, + tombs: map[string]map[string]uint64{}, + } + for _, o := range opts { + o(c) + } + return c +} + +// Run is one owned worker loop (Option E §11.5/§11.7): it wakes on NATS notify +// and otherwise polls the scheduling table for a claimable KB, then claims the +// closed batch, processes it, and acks. It returns when ctx is cancelled. +func (c *Consumer) Run(ctx context.Context) { + if c.scheduler == nil { + return + } + notifyCh, _ := c.scheduler.SubscribeNotify(ctx) + poll := time.NewTicker(c.pollInterval) + sweep := time.NewTicker(c.sweepInterval) + defer poll.Stop() + defer sweep.Stop() + for { + select { + case <-ctx.Done(): + return + case ds := <-notifyCh: + if ds != "" { + c.processDataset(ctx, ds) + } + case <-poll.C: + c.claimOne(ctx) + case <-sweep.C: + // Recover inflight left by crashed workers (crash recovery, §11.5). + if _, err := c.scheduler.ReclaimExpired(ctx, time.Now()); err != nil { + // best-effort; next tick retries + _ = err + } + } + } +} + +// claimOne finds a claimable KB and processes it. FindClaimable returns at most +// one dataset so the worker handles it before looking for more. +func (c *Consumer) claimOne(ctx context.Context) { + ids, err := c.scheduler.FindClaimable(ctx, 1) + if err != nil || len(ids) == 0 { + return + } + for _, ds := range ids { + c.processDataset(ctx, ds) + } +} + +// processDataset claims the closed batch for datasetID, processes it, and acks. +// It is the Option E replacement for the old processOnce (lease + drain + merge): +// the claim transaction returns a frozen batch boundary, so there is no moving +// target and no Nak-churn routing. +func (c *Consumer) processDataset(ctx context.Context, datasetID string) { + cr, ok, err := c.scheduler.Claim(ctx, datasetID, c.batchSize) + if err != nil { + return + } + if !ok || len(cr.Entries) == 0 { + return // race lost or nothing to claim + } + + // Heartbeat refreshes the claim TTL while we process; a failed touch means + // the lease was taken over (or reclaimed) and we must abort without acking. + stopHb := make(chan struct{}) + hbFailed := make(chan struct{}, 1) + go func() { + t := time.NewTicker(c.heartbeat) + defer t.Stop() + for { + select { + case <-stopHb: + return + case <-ctx.Done(): + return + case <-t.C: + alive, e := c.scheduler.TouchClaim(ctx, datasetID, cr.Token, c.ttl) + if e != nil || !alive { + select { + case hbFailed <- struct{}{}: + default: + } + return + } + } + } + }() + + done := make(chan struct{}) + var batchErr error + go func() { + defer close(done) + batchErr = c.processBatch(ctx, cr.TenantID, datasetID, cr.Entries) + }() + + select { + case <-ctx.Done(): + close(stopHb) + <-done + return // graceful shutdown: leave inflight for reclamation, do not ack + case <-hbFailed: + close(stopHb) + <-done + return // lease lost: do not ack; sweeper/redelivery reprocesses (idempotent) + case <-done: + close(stopHb) + } + + // Ack only on success. A batch error (reader/dedup/writer failure) means we + // 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 { + return + } + if _, err := c.scheduler.Ack(ctx, datasetID, cr.Token, cr.Entries); err != nil { + _ = err + } +} + +// processBatch applies out-of-order / tombstone handling, then recomputes and +// writes the dataset-level merged products for the claimed closed batch. It +// 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 { + c.mu.Lock() + if c.tombs == nil { + c.tombs = map[string]map[string]uint64{} + } + if c.seqs == nil { + c.seqs = map[string]map[string]uint64{} + } + tomb := c.tombs[kb] + docSeqs := c.seqs[kb] + if docSeqs == nil { + docSeqs = map[string]uint64{} + c.seqs[kb] = docSeqs + } + var completed []BacklogEntry + var deleted []string + for _, e := range entries { + switch EventType(e.EventType) { + case EventTypeDeleted: + if tomb == nil { + tomb = map[string]uint64{} + c.tombs[kb] = tomb + } + tomb[e.DocID] = e.Seq + deleted = append(deleted, e.DocID) + case EventTypeCompleted: + // A tombstone means the doc was deleted; a completion with a seq + // <= the delete seq is the stale original completion (skip it). + // A completion with a higher seq is a genuine re-ingest after + // deletion: clear the tombstone so the doc is processed again + // (otherwise the tombstone would skip it forever and grow + // unbounded across the consumer's lifetime). + if delSeq, ok := tomb[e.DocID]; ok { + if e.Seq <= delSeq { + continue // deleted (or a stale completion) before it completed + } + delete(tomb, e.DocID) + } + // Seq is per-document, so the stale/duplicate check must be scoped + // to the document, not the whole dataset (C4). + if prev, ok := docSeqs[e.DocID]; ok && e.Seq <= prev { + continue // stale / duplicate completion for this doc + } + docSeqs[e.DocID] = e.Seq + completed = append(completed, e) + } + } + c.mu.Unlock() + + if len(deleted) == 0 && len(completed) == 0 { + return nil + } + + deduper, err := c.factory(tenant) + if err != nil || deduper == nil { + deduper = NewNoopDeduper() + } + + products, err := c.reader.LoadCompiledProducts(ctx, tenant, kb) + if err != nil { + return err + } + if len(products) == 0 { + // Nothing to merge. Still drop fully-orphaned merged products for + // deleted docs, then treat as success (nothing to do). + for _, d := range deleted { + _ = c.writer.DeleteMergedForDoc(ctx, tenant, kb, d) + } + return nil + } + + // With deletions present, recompute the whole KB; otherwise scope to the + // contributing documents of this batch (efficiency). + if len(deleted) == 0 { + docSet := make(map[string]bool, len(completed)) + for _, e := range completed { + docSet[e.DocID] = true + } + scoped := products[:0] + for _, p := range products { + if docSet[p.DocID] { + scoped = append(scoped, p) + } + } + products = scoped + } + + merged, err := deduper.Dedup(ctx, products) + if err != nil { + return err + } + if err := c.writer.WriteMerged(ctx, tenant, kb, merged); err != nil { + return err + } + for _, d := range deleted { + _ = c.writer.DeleteMergedForDoc(ctx, tenant, kb, d) + } + return nil +} diff --git a/internal/ingestion/knowledge_compile/consumer_test.go b/internal/ingestion/knowledge_compile/consumer_test.go new file mode 100644 index 0000000000..fb2c1ee988 --- /dev/null +++ b/internal/ingestion/knowledge_compile/consumer_test.go @@ -0,0 +1,191 @@ +// +// 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 knowledge_compile + +import ( + "context" + "sync" + "testing" + "time" + + kccommon "ragflow/internal/ingestion/component/knowledge_compiler/common" +) + +// fakeReader returns a fixed product set. +type fakeReader struct { + mu sync.Mutex + products []kccommon.Product + calls int +} + +func (r *fakeReader) LoadCompiledProducts(_ context.Context, _, _ string) ([]kccommon.Product, error) { + r.mu.Lock() + defer r.mu.Unlock() + r.calls++ + out := make([]kccommon.Product, len(r.products)) + copy(out, r.products) + return out, nil +} + +// fakeWriter captures written merged products. +type fakeWriter struct { + mu sync.Mutex + written [][]kccommon.Product + deleted []string +} + +func (w *fakeWriter) WriteMerged(_ context.Context, _, _ string, products []kccommon.Product) error { + w.mu.Lock() + defer w.mu.Unlock() + cp := make([]kccommon.Product, len(products)) + copy(cp, products) + w.written = append(w.written, cp) + return nil +} + +func (w *fakeWriter) DeleteMergedForDoc(_ context.Context, _, _, docID string) error { + w.mu.Lock() + defer w.mu.Unlock() + w.deleted = append(w.deleted, docID) + return nil +} + +func sampleProducts() []kccommon.Product { + return []kccommon.Product{ + {ID: "p1", DocID: "d1", TenantID: "t1", Variant: kccommon.Variant("structure"), + Content: `{"name":"X"}`, Meta: map[string]any{"name": "X", "kind": "entity"}}, + {ID: "p2", DocID: "d1", TenantID: "t1", Variant: kccommon.Variant("structure"), + Content: `{"name":"Y"}`, Meta: map[string]any{"name": "Y", "kind": "entity"}}, + } +} + +func newTestConsumer(sch *FakeScheduler, r *fakeReader, w *fakeWriter, factory DeduperFactory) *Consumer { + return NewConsumer(sch, + WithReader(r), + WithWriter(w), + WithDeduperFactory(factory), + WithBatchSize(32), + ) +} + +func TestConsumerCompletedWritesMerged(t *testing.T) { + sch := NewFakeScheduler() + r := &fakeReader{products: sampleProducts()} + w := &fakeWriter{} + c := newTestConsumer(sch, r, w, func(string) (Deduper, error) { return NewNoopDeduper(), nil }) + + if err := sch.AppendBacklog(context.Background(), "t1", "kb1", "d1", string(EventTypeCompleted), 1); err != nil { + t.Fatalf("append: %v", err) + } + c.processDataset(context.Background(), "kb1") + + w.mu.Lock() + defer w.mu.Unlock() + if len(w.written) != 1 { + t.Fatalf("expected 1 WriteMerged call, got %d", len(w.written)) + } + if len(w.written[0]) != 2 { + t.Fatalf("expected 2 merged products, got %d", len(w.written[0])) + } + // After ack, the claim row must be cleared (no live lease left behind). + if got, _ := sch.FindClaimable(context.Background(), 1); len(got) != 0 { + t.Fatalf("expected no claimable row after ack, got %v", got) + } +} + +func TestConsumerTombstoneSkipsCompletedBeforeDeleted(t *testing.T) { + sch := NewFakeScheduler() + r := &fakeReader{products: sampleProducts()} + w := &fakeWriter{} + c := newTestConsumer(sch, r, w, func(string) (Deduper, error) { return NewNoopDeduper(), nil }) + + // deleted before completed -> completed must be skipped. + if err := sch.AppendBacklog(context.Background(), "t1", "kb1", "d1", string(EventTypeDeleted), 1); err != nil { + t.Fatalf("append deleted: %v", err) + } + if err := sch.AppendBacklog(context.Background(), "t1", "kb1", "d1", string(EventTypeCompleted), 2); err != nil { + t.Fatalf("append completed: %v", err) + } + c.processDataset(context.Background(), "kb1") + + w.mu.Lock() + defer w.mu.Unlock() + if len(w.deleted) != 1 { + t.Fatalf("expected 1 orphan-delete call for d1, got %d", len(w.deleted)) + } +} + +func TestSchedulerClaimClosedBatch(t *testing.T) { + sch := NewFakeScheduler() + for i := 0; i < 40; i++ { + docID := "d" + string(rune('a'+i%26)) + string(rune('0'+i/26)) + if err := sch.AppendBacklog(context.Background(), "t1", "kb1", docID, string(EventTypeCompleted), uint64(i)); err != nil { + t.Fatalf("append: %v", err) + } + } + // First claim returns the bounded prefix (batchSize=32), not all 40. + cr1, ok, err := sch.Claim(context.Background(), "kb1", 32) + if err != nil || !ok { + t.Fatalf("claim1: ok=%v err=%v", ok, err) + } + if len(cr1.Entries) != 32 { + t.Fatalf("expected 32 entries in first claim, got %d", len(cr1.Entries)) + } + // A second claim by the same holder (still live lease) must not re-claim + // the same dataset until the first batch is acked. + _, ok2, _ := sch.Claim(context.Background(), "kb1", 32) + if ok2 { + t.Fatalf("second claim should have lost the race (live lease)") + } + // Ack the first batch, then the remaining 8 become claimable. + if _, err := sch.Ack(context.Background(), "kb1", cr1.Token, cr1.Entries); err != nil { + t.Fatalf("ack: %v", err) + } + cr3, ok3, err := sch.Claim(context.Background(), "kb1", 32) + if err != nil || !ok3 { + t.Fatalf("claim3: ok=%v err=%v", ok3, err) + } + if len(cr3.Entries) != 8 { + t.Fatalf("expected 8 remaining entries, got %d", len(cr3.Entries)) + } +} + +func TestSchedulerReclaimExpired(t *testing.T) { + sch := NewFakeScheduler() + if err := sch.AppendBacklog(context.Background(), "t1", "kb1", "d1", string(EventTypeCompleted), 1); err != nil { + t.Fatalf("append: %v", err) + } + _, ok, err := sch.Claim(context.Background(), "kb1", 32) + if err != nil || !ok { + t.Fatalf("claim: ok=%v err=%v", ok, err) + } + // Simulate a crash: the inflight is never acked. Reclaim moves it back. + n, err := sch.ReclaimExpired(context.Background(), time.Now().Add(10*time.Minute)) + if err != nil { + t.Fatalf("reclaim: %v", err) + } + if n != 1 { + t.Fatalf("expected 1 reclaimed entry, got %d", n) + } + // The entry is claimable again. + cr2, ok2, err := sch.Claim(context.Background(), "kb1", 32) + if err != nil || !ok2 { + t.Fatalf("reclaim claim: ok=%v err=%v", ok2, err) + } + if len(cr2.Entries) != 1 { + t.Fatalf("expected 1 entry after reclaim, got %d", len(cr2.Entries)) + } +} diff --git a/internal/ingestion/knowledge_compile/dedup.go b/internal/ingestion/knowledge_compile/dedup.go new file mode 100644 index 0000000000..fabc7412db --- /dev/null +++ b/internal/ingestion/knowledge_compile/dedup.go @@ -0,0 +1,74 @@ +// +// 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 knowledge_compile + +import ( + "context" + + kccommon "ragflow/internal/ingestion/component/knowledge_compiler/common" + "ragflow/internal/ingestion/component/knowledge_compiler/structure" +) + +// Deduper folds a set of per-document compiled Products into the dataset-level +// merged set. The LLM-backed implementation reuses the component's +// GroupedDeduper + LLMMergeDecider (§11.6). +type Deduper interface { + Dedup(ctx context.Context, rows []kccommon.Product) ([]kccommon.Product, error) +} + +// DeduperFactory builds a per-tenant Deduper. It is invoked once per batch so +// the LLM deps can be resolved for the owning tenant. +type DeduperFactory func(tenant string) (Deduper, error) + +// llmDeduper wraps the component's GroupedDeduper (which internally uses +// LLMMergeDecider for duplicate-judging), scoped to the whole KB batch. +type llmDeduper struct { + group *structure.GroupedDeduper + decider *structure.LLMMergeDecider + embed kccommon.Embedder +} + +// NewLLMDeduper builds a KB-scoped deduper from the runtime chat/embed deps. +func NewLLMDeduper(chat kccommon.ChatInvoker, embed kccommon.Embedder, llmID string, threshold float64) Deduper { + decider := structure.NewLLMMergeDecider(chat, llmID, embed, threshold) + return &llmDeduper{group: structure.NewGroupedDeduper(decider), decider: decider, embed: embed} +} + +func (x *llmDeduper) Dedup(ctx context.Context, rows []kccommon.Product) ([]kccommon.Product, error) { + for _, r := range rows { + if err := x.group.Add(ctx, r); err != nil { + return nil, err + } + } + // Apply the aliases recorded by the LLM merge decider to relation endpoints + // so merged entities collapse consistently with the per-document dedup path. + if err := x.group.RewriteRelations(ctx, x.decider.Aliases(), x.embed); err != nil { + return nil, err + } + return x.group.Rows(), nil +} + +// noopDeduper performs no LLM merge; it returns the input rows unchanged so +// the writer still emits dataset-level products (without cross-document merging). +// Used as a safe fallback when LLM deps are unavailable. +type noopDeduper struct{} + +func (noopDeduper) Dedup(_ context.Context, rows []kccommon.Product) ([]kccommon.Product, error) { + return rows, nil +} + +// NewNoopDeduper builds the fallback deduper. +func NewNoopDeduper() Deduper { return noopDeduper{} } diff --git a/internal/ingestion/knowledge_compile/event.go b/internal/ingestion/knowledge_compile/event.go new file mode 100644 index 0000000000..25e1f3248c --- /dev/null +++ b/internal/ingestion/knowledge_compile/event.go @@ -0,0 +1,115 @@ +// +// Copyright 2026 The InfiniFlow Authors. All Rights Reserved. +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +// Package knowledge_compile implements the dataset-level post-processing consumer described in +// docs/develop/knowledge_compile_design.md §11 (Option E). +// +// Pipeline (KnowledgeCompiler, per document) writes compiled chunks with +// available_int=0. Its completion/deletion is recorded by appending a +// BacklogEntry to the KB's durable MySQL scheduling row +// (knowledge_compile_docs), then waking idle workers over NATS. A cluster of +// competing workers claims a closed batch (backlog -> inflight) per KB, runs +// dataset-level dedup on that batch, and writes the merged dataset-level +// products with available_int=1. The MySQL row — not the broker — is the +// scheduling system of record and the source of same-KB serialization. +package knowledge_compile + +import ( + "context" + "encoding/json" +) + +// Subjects and event types for the knowledge-compile stream. Both sit under the +// knowledge.compile.events.> prefix declared on the NATS stream/consumer. +const ( + SubjectCompleted = "knowledge.compile.events.completed" + SubjectDeleted = "knowledge.compile.events.deleted" +) + +// EventType enumerates the KC event kinds. +type EventType string + +const ( + EventTypeCompleted EventType = "doc_completed" + EventTypeDeleted EventType = "doc_deleted" +) + +// KCCompileEvent is the payload published when a document's pipeline finishes +// (doc_completed) or is deleted (doc_deleted). It is intentionally decoupled +// from common.TaskMessage because the consumer reads the raw JSON body. +type KCCompileEvent struct { + TenantID string `json:"tenant_id"` + DatasetID string `json:"dataset_id"` // the KB scope + DocID string `json:"doc_id"` // the contributing document + EventType string `json:"event_type"` // EventType value + Seq uint64 `json:"seq"` // per-doc monotonic sequence, for out-of-order correction + Timestamp int64 `json:"ts"` +} + +// Subject returns the NATS subject for this event. +func (e KCCompileEvent) Subject() string { + if EventType(e.EventType) == EventTypeDeleted { + return SubjectDeleted + } + return SubjectCompleted +} + +// Marshal serializes the event to JSON. +func (e KCCompileEvent) Marshal() ([]byte, error) { return json.Marshal(e) } + +// ParseEvent deserializes a KCCompileEvent from raw message bytes. +func ParseEvent(data []byte) (KCCompileEvent, error) { + var e KCCompileEvent + if err := json.Unmarshal(data, &e); err != nil { + return KCCompileEvent{}, err + } + return e, nil +} + +// defaultScheduler is the package-level Scheduler used by the publishing path +// (PublishCompleted / PublishDeleted). It is installed by Provision (called +// once by the owning Ingestor at startup). Until then, publishing is a no-op. +var defaultScheduler Scheduler + +// SetScheduler installs the package-level Scheduler used for publishing. +func SetScheduler(s Scheduler) { defaultScheduler = s } + +// DefaultScheduler returns the package-level Scheduler (nil until Provision). +func DefaultScheduler() Scheduler { return defaultScheduler } + +// PublishCompleted records a doc_completed event: it appends the doc to the +// KB's durable MySQL backlog and wakes idle workers over NATS. It is a no-op +// when no Scheduler has been installed (e.g. DB unavailable). A failure is +// returned so callers can log but never fail the pipeline on it. +func PublishCompleted(ctx context.Context, tenantID, datasetID, docID string, seq uint64) error { + if defaultScheduler == nil { + return nil + } + if err := defaultScheduler.AppendBacklog(ctx, tenantID, datasetID, docID, string(EventTypeCompleted), seq); err != nil { + return err + } + return defaultScheduler.Notify(ctx, datasetID) +} + +// PublishDeleted records a doc_deleted event the same way (append + notify). +func PublishDeleted(ctx context.Context, tenantID, datasetID, docID string, seq uint64) error { + if defaultScheduler == nil { + return nil + } + if err := defaultScheduler.AppendBacklog(ctx, tenantID, datasetID, docID, string(EventTypeDeleted), seq); err != nil { + return err + } + return defaultScheduler.Notify(ctx, datasetID) +} diff --git a/internal/ingestion/knowledge_compile/reader.go b/internal/ingestion/knowledge_compile/reader.go new file mode 100644 index 0000000000..1b6d9d43bb --- /dev/null +++ b/internal/ingestion/knowledge_compile/reader.go @@ -0,0 +1,168 @@ +// +// 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 knowledge_compile + +import ( + "context" + "fmt" + "strings" + + "ragflow/internal/engine" + "ragflow/internal/engine/types" + kccommon "ragflow/internal/ingestion/component/knowledge_compiler/common" +) + +// Reader loads the per-document compiled products (available_int=0) of a KB so +// the consumer can recompute the dataset-level merged set (§11.6 step 1, §11.7 +// incremental re-dedup). +type Reader interface { + LoadCompiledProducts(ctx context.Context, tenant, kb string) ([]kccommon.Product, error) +} + +type infinityReader struct{} + +func (infinityReader) LoadCompiledProducts(ctx context.Context, tenant, kb string) ([]kccommon.Product, error) { + eng := engine.Get() + if eng == nil { + return nil, nil + } + baseName := fmt.Sprintf("ragflow_%s", tenant) + const batchSize = 5000 + var out []kccommon.Product + offset := 0 + for { + res, err := eng.Search(ctx, &types.SearchRequest{ + IndexNames: []string{baseName}, + KbIDs: []string{kb}, + Filter: map[string]interface{}{"available_int": 0}, + SelectFields: []string{ + "id", "doc_id", "tenant_id", "compile_kwd", + "content_with_weight", "kc_payload", + "source_chunk_ids", "source_doc_ids", + "name_kwd", "entity_type_kwd", "from_entity_kwd", "to_entity_kwd", + "slug_kwd", "type", + }, + Limit: batchSize, + Offset: offset, + }) + if err != nil { + return nil, err + } + for _, c := range res.Chunks { + // Only compiled products carry compile_kwd; skip ordinary source chunks. + if _, ok := c["compile_kwd"]; !ok { + continue + } + if p, ok := productFromChunkMap(c, tenant); ok { + out = append(out, p) + } + } + // The KB-wide scan is paginated: keep fetching until a page returns + // fewer than batchSize rows, so a KB larger than the cap merges against + // the full compiled set instead of a truncated slice. + if len(res.Chunks) < batchSize { + break + } + offset += batchSize + } + return out, nil +} + +// productFromChunkMap reconstructs a kccommon.Product from a stored compiled +// chunk document. It reads the payload from kc_payload (falling back to +// content_with_weight) and the embedding from the q_<dim>_vec column. +func productFromChunkMap(c map[string]interface{}, tenant string) (kccommon.Product, bool) { + content, _ := c["kc_payload"].(string) + if content == "" { + content, _ = c["content_with_weight"].(string) + } + if content == "" { + return kccommon.Product{}, false + } + id, _ := c["id"].(string) + docID, _ := c["doc_id"].(string) + variant, _ := c["compile_kwd"].(string) + + meta := map[string]any{} + if v, ok := c["name_kwd"].(string); ok && v != "" { + meta["name"] = v + } + if v, ok := c["entity_type_kwd"].(string); ok && v != "" { + meta["entity_type"] = v + } + if v, ok := c["from_entity_kwd"].(string); ok && v != "" { + meta["from"] = v + meta["kind"] = "relation" + } + if v, ok := c["to_entity_kwd"].(string); ok && v != "" { + meta["to"] = v + meta["kind"] = "relation" + } + if v, ok := c["slug_kwd"].(string); ok && v != "" { + meta["slug"] = v + } + if v, ok := c["type"].(string); ok && v != "" { + meta["type"] = v + } + if _, ok := meta["kind"]; !ok { + if _, hasName := meta["name"]; hasName { + meta["kind"] = "entity" + } + } + if v := metaStringSlice(c, "source_chunk_ids"); len(v) > 0 { + meta["source_chunk_ids"] = v + } + if v := metaStringSlice(c, "source_doc_ids"); len(v) > 0 { + meta["source_doc_ids"] = v + } + + return kccommon.Product{ + ID: id, + DocID: docID, + TenantID: tenant, + Variant: kccommon.Variant(variant), + Content: content, + Vector: vectorFromChunkMap(c), + Meta: meta, + }, true +} + +// vectorFromChunkMap extracts the embedding from the q_<dim>_vec column, whose +// exact name depends on the embedding dimension. +func vectorFromChunkMap(c map[string]interface{}) []float32 { + for k, v := range c { + if strings.HasPrefix(k, "q_") && strings.HasSuffix(k, "_vec") { + return toFloat32Slice(v) + } + } + return nil +} + +func toFloat32Slice(v interface{}) []float32 { + switch arr := v.(type) { + case []float32: + return arr + case []interface{}: + out := make([]float32, 0, len(arr)) + for _, e := range arr { + if f, ok := e.(float64); ok { + out = append(out, float32(f)) + } + } + return out + } + return nil +} diff --git a/internal/ingestion/knowledge_compile/scheduler.go b/internal/ingestion/knowledge_compile/scheduler.go new file mode 100644 index 0000000000..43abe90edc --- /dev/null +++ b/internal/ingestion/knowledge_compile/scheduler.go @@ -0,0 +1,534 @@ +// +// 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 knowledge_compile + +import ( + "context" + "encoding/json" + "errors" + "fmt" + "sync" + "time" + + "ragflow/internal/engine" + "ragflow/internal/entity" + + "gorm.io/gorm" + "gorm.io/gorm/clause" +) + +// ErrClaimTokenMismatch is returned by Ack/TouchClaim when the supplied token +// no longer identifies the live claim (another worker took over, or the +// sweeper reclaimed it). The caller must abort without writing. +var ErrClaimTokenMismatch = errors.New("knowledge_compile: claim token mismatch") + +// notifySubject is the Option E wake-up channel (§11.4): publishers push a +// {dataset_id} payload here after appending to a KB's backlog; workers +// subscribe to wake promptly. It carries no routing payload — MySQL is the +// scheduling truth. +const notifySubject = "notify.kc.workers" + +// 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 +// broker-based design without re-reading the queue. +type BacklogEntry struct { + DocID string `json:"doc_id"` + EventType string `json:"event_type"` + Seq uint64 `json:"seq"` +} + +// ClaimResult is returned by Scheduler.Claim: the KB's tenant plus the closed +// batch of entries moved into inflight, and the token identifying this claim. +type ClaimResult struct { + TenantID string + Entries []BacklogEntry + Token string +} + +// Scheduler owns the claim/ack lifecycle against the durable scheduling store +// (Option E §11.4/§11.5). MySQL is the system of record; NATS notify is only a +// wake-up. The claim is a move (backlog -> inflight), never a copy, so a doc id +// is visible to at most one worker at a time; same-KB serialization follows +// from the inflight set, not from the broker. +type Scheduler interface { + // Provision ensures the backing store (table + notify subject) exists. + Provision(ctx context.Context) error + // AppendBacklog adds one doc event to the KB's backlog. The append is + // transactional and preserves any existing backlog (concurrent publishers + // do not clobber each other). + AppendBacklog(ctx context.Context, tenantID, datasetID, docID, eventType string, seq uint64) error + // Claim moves a bounded prefix of backlog -> inflight for datasetID and + // returns the closed batch. acquired=false when the row is held by another + // live lease (the race was lost) or the backlog is empty. + Claim(ctx context.Context, datasetID string, batchSize int) (ClaimResult, bool, error) + // Ack removes the claimed batch from inflight; clears the claim metadata + // only when backlog is also empty. Returns the remaining backlog size. + Ack(ctx context.Context, datasetID, token string, batch []BacklogEntry) (backlogRemaining int, err error) + // TouchClaim extends the lease while processing (heartbeat). Returns false + // when the lease is gone (taken over / reclaimed) so the worker must abort. + TouchClaim(ctx context.Context, datasetID, token string, ttl time.Duration) (bool, error) + // FindClaimable returns up to limit dataset ids with non-empty backlog and + // no live lease (a free token to claim). + FindClaimable(ctx context.Context, limit int) ([]string, error) + // Notify wakes idle workers about a dataset that just got backlog. + Notify(ctx context.Context, datasetID string) error + // SubscribeNotify returns a channel of dataset ids pushed by Notify, or nil + // when the implementation has no push wake-up (callers fall back to polling). + SubscribeNotify(ctx context.Context) (<-chan string, error) + // ReclaimExpired moves expired inflight entries back to backlog (sweeper). + ReclaimExpired(ctx context.Context, now time.Time) (int, error) +} + +// newScheduler is the production constructor: a MySQL-backed scheduler with an +// optional NATS wake-up. holder identifies this ingestor instance; ttl is the +// claim lease duration. +func newScheduler(db *gorm.DB, mq engine.MessageQueue, holder string, ttl time.Duration) Scheduler { + if ttl <= 0 { + ttl = 2 * time.Minute + } + return &mysqlScheduler{db: db, mq: mq, holder: holder, leaseTTL: ttl} +} + +// ---- JSON helpers (the *_doc_ids columns are TEXT holding []BacklogEntry) ---- + +func marshalEntries(es []BacklogEntry) string { + if len(es) == 0 { + return "[]" + } + b, _ := json.Marshal(es) + return string(b) +} + +func parseEntries(s string) []BacklogEntry { + if s == "" || s == "[]" { + return nil + } + var es []BacklogEntry + if err := json.Unmarshal([]byte(s), &es); err != nil || len(es) == 0 { + return nil + } + return es +} + +func minInt(a, b int) int { + if a < b { + return a + } + return b +} + +// ---- MySQL-backed scheduler ---- + +type mysqlScheduler struct { + db *gorm.DB + mq engine.MessageQueue + holder string + leaseTTL time.Duration +} + +func (s *mysqlScheduler) Provision(ctx context.Context) error { + if s.db == nil { + return nil + } + return s.db.WithContext(ctx).AutoMigrate(&entity.KnowledgeCompileDoc{}) +} + +func (s *mysqlScheduler) AppendBacklog(ctx context.Context, tenantID, datasetID, docID, eventType string, seq uint64) error { + if s.db == nil { + return nil + } + entry := BacklogEntry{DocID: docID, EventType: eventType, Seq: seq} + // KnowledgeCompileDoc.dataset_id is the PRIMARY KEY, so the SELECT ... FOR + // UPDATE below also takes an InnoDB gap lock for a not-yet-existing dataset; + // concurrent publishers for the same dataset serialize on that lock and the + // loser observes the row already present. FirstOrCreate then reliably finds + // or inserts exactly one row, and the append below runs under the same row + // lock so concurrent appends from different publishers cannot interleave. + err := s.db.WithContext(ctx).Transaction(func(tx *gorm.DB) error { + // Set the key/tenant fields on the struct itself (not only via a string + // Where) so the inserted row is fully populated and later queries by + // dataset_id find it (M19: a raw-string Where alone leaves DatasetID + // blank on the created row). + row := entity.KnowledgeCompileDoc{ + DatasetID: datasetID, + TenantID: tenantID, + BacklogDocIDs: "[]", + InflightDocIDs: "[]", + } + if err := tx.Clauses(clause.Locking{Strength: "UPDATE"}). + Where("dataset_id = ?", datasetID). + FirstOrCreate(&row).Error; err != nil { + return err + } + if row.TenantID == "" { + row.TenantID = tenantID + } + backlog := parseEntries(row.BacklogDocIDs) + backlog = append(backlog, entry) + row.BacklogDocIDs = marshalEntries(backlog) + return tx.Save(&row).Error + }) + if err != nil { + return fmt.Errorf("knowledge_compile: append backlog %s: %w", datasetID, err) + } + return nil +} + +func (s *mysqlScheduler) Claim(ctx context.Context, datasetID string, batchSize int) (ClaimResult, bool, error) { + if s.db == nil { + return ClaimResult{}, false, nil + } + var res ClaimResult + var acquired bool + err := s.db.WithContext(ctx).Transaction(func(tx *gorm.DB) error { + var row entity.KnowledgeCompileDoc + if err := tx.Clauses(clause.Locking{Strength: "UPDATE"}). + Where("dataset_id = ?", datasetID).First(&row).Error; err != nil { + if errors.Is(err, gorm.ErrRecordNotFound) { + return nil // nothing to claim + } + return err + } + now := time.Now() + liveLease := row.ClaimOwner != "" && row.ClaimExpiresAt != nil && row.ClaimExpiresAt.After(now) + if liveLease { + return nil // a live lease means some worker is already processing this KB + } + backlog := parseEntries(row.BacklogDocIDs) + if len(backlog) == 0 { + return nil + } + n := minInt(batchSize, len(backlog)) + batch := backlog[:n] + inflight := parseEntries(row.InflightDocIDs) + inflight = append(inflight, batch...) + row.BacklogDocIDs = marshalEntries(backlog[n:]) + row.InflightDocIDs = marshalEntries(inflight) + row.ClaimOwner = s.holder + row.ClaimToken = generateHolder() + exp := now.Add(s.leaseTTL) + row.ClaimExpiresAt = &exp + if err := tx.Save(&row).Error; err != nil { + return err + } + res = ClaimResult{TenantID: row.TenantID, Entries: batch, Token: row.ClaimToken} + acquired = true + return nil + }) + if err != nil { + return ClaimResult{}, false, err + } + return res, acquired, nil +} + +func (s *mysqlScheduler) Ack(ctx context.Context, datasetID, token string, batch []BacklogEntry) (int, error) { + if s.db == nil { + return 0, nil + } + var remaining int + err := s.db.WithContext(ctx).Transaction(func(tx *gorm.DB) error { + var row entity.KnowledgeCompileDoc + if err := tx.Clauses(clause.Locking{Strength: "UPDATE"}). + Where("dataset_id = ?", datasetID).First(&row).Error; err != nil { + return err + } + if row.ClaimToken != token { + return ErrClaimTokenMismatch + } + inflight := parseEntries(row.InflightDocIDs) + inflight = removeEntries(inflight, batch) + row.InflightDocIDs = marshalEntries(inflight) + if len(inflight) == 0 { + row.ClaimOwner = "" + row.ClaimToken = "" + row.ClaimExpiresAt = nil + } + if err := tx.Save(&row).Error; err != nil { + return err + } + remaining = len(parseEntries(row.BacklogDocIDs)) + return nil + }) + if err != nil { + return 0, err + } + return remaining, nil +} + +func (s *mysqlScheduler) TouchClaim(ctx context.Context, datasetID, token string, ttl time.Duration) (bool, error) { + if s.db == nil { + return false, nil + } + res := s.db.WithContext(ctx).Model(&entity.KnowledgeCompileDoc{}). + Where("dataset_id = ? AND claim_token = ?", datasetID, token). + Update("claim_expires_at", time.Now().Add(ttl)) + if res.Error != nil { + return false, res.Error + } + return res.RowsAffected > 0, nil +} + +func (s *mysqlScheduler) FindClaimable(ctx context.Context, limit int) ([]string, error) { + if s.db == nil || limit <= 0 { + return nil, nil + } + var ids []string + err := s.db.WithContext(ctx).Model(&entity.KnowledgeCompileDoc{}). + Where("(claim_expires_at IS NULL OR claim_expires_at <= ?) AND backlog_doc_ids <> '[]' AND backlog_doc_ids <> '' AND backlog_doc_ids IS NOT NULL", time.Now()). + Order("priority DESC, updated_at ASC"). + Limit(limit). + Pluck("dataset_id", &ids).Error + if err != nil { + return nil, err + } + return ids, nil +} + +func (s *mysqlScheduler) Notify(ctx context.Context, datasetID string) error { + if s.mq == nil { + return nil + } + payload, _ := json.Marshal(map[string]string{"dataset_id": datasetID}) + return s.mq.PublishKnowledgeCompile(notifySubject, payload) +} + +func (s *mysqlScheduler) SubscribeNotify(ctx context.Context) (<-chan string, error) { + if s.mq == nil { + return nil, nil + } + return s.mq.SubscribeNotify(ctx) +} + +func (s *mysqlScheduler) ReclaimExpired(ctx context.Context, now time.Time) (int, error) { + if s.db == nil { + return 0, nil + } + var rows []entity.KnowledgeCompileDoc + if err := s.db.WithContext(ctx).Model(&entity.KnowledgeCompileDoc{}). + Where("claim_expires_at IS NOT NULL AND claim_expires_at <= ? AND inflight_doc_ids <> '[]' AND inflight_doc_ids <> ''", now). + Find(&rows).Error; err != nil { + return 0, err + } + total := 0 + for _, row := range rows { + inflight := parseEntries(row.InflightDocIDs) + if len(inflight) == 0 { + continue + } + total += len(inflight) + err := s.db.WithContext(ctx).Transaction(func(tx *gorm.DB) error { + var cur entity.KnowledgeCompileDoc + if err := tx.Clauses(clause.Locking{Strength: "UPDATE"}). + Where("dataset_id = ?", row.DatasetID).First(&cur).Error; err != nil { + return err + } + // Re-check liveness under the lock so we don't clobber a lease that + // was refreshed between the scan and now. + if cur.ClaimExpiresAt != nil && cur.ClaimExpiresAt.After(now) { + return nil + } + backlog := parseEntries(cur.BacklogDocIDs) + backlog = append(backlog, parseEntries(cur.InflightDocIDs)...) + cur.BacklogDocIDs = marshalEntries(backlog) + cur.InflightDocIDs = "[]" + cur.ClaimOwner = "" + cur.ClaimToken = "" + cur.ClaimExpiresAt = nil + return tx.Save(&cur).Error + }) + if err != nil { + return total, err + } + } + return total, nil +} + +// removeEntries drops every entry in batch from the inflight slice (match by +// doc_id + event_type + seq). The batch is the exact set the worker claimed, so +// this is a precise removal, not a "clear all" — any inflight added by a +// concurrent claim is preserved. +func removeEntries(inflight, batch []BacklogEntry) []BacklogEntry { + if len(batch) == 0 { + return inflight + } + drop := make(map[BacklogEntry]bool, len(batch)) + for _, e := range batch { + drop[e] = true + } + out := make([]BacklogEntry, 0, len(inflight)) + for _, e := range inflight { + if drop[e] { + continue + } + out = append(out, e) + } + return out +} + +// ---- in-memory scheduler for tests ---- + +type fakeRow struct { + tenant string + backlog []BacklogEntry + inflight []BacklogEntry + owner string + token string + expires *time.Time +} + +// FakeScheduler is an in-memory Scheduler used by tests. It mirrors the +// MySQL semantics (move-not-copy claim, token-checked ack, lease takeover). +type FakeScheduler struct { + mu sync.Mutex + rows map[string]*fakeRow + notifyCh chan string + holder string + leaseTTL time.Duration +} + +// NewFakeScheduler constructs an in-memory scheduler for tests. +func NewFakeScheduler() *FakeScheduler { + return &FakeScheduler{ + rows: map[string]*fakeRow{}, + notifyCh: make(chan string, 64), + holder: generateHolder(), + leaseTTL: 2 * time.Minute, + } +} + +func (f *FakeScheduler) Provision(_ context.Context) error { return nil } + +func (f *FakeScheduler) AppendBacklog(_ context.Context, tenantID, datasetID, docID, eventType string, seq uint64) error { + f.mu.Lock() + defer f.mu.Unlock() + r, ok := f.rows[datasetID] + if !ok { + r = &fakeRow{tenant: tenantID} + f.rows[datasetID] = r + } + if r.tenant == "" { + r.tenant = tenantID + } + r.backlog = append(r.backlog, BacklogEntry{DocID: docID, EventType: eventType, Seq: seq}) + return nil +} + +func (f *FakeScheduler) Claim(_ context.Context, datasetID string, batchSize int) (ClaimResult, bool, error) { + f.mu.Lock() + defer f.mu.Unlock() + r, ok := f.rows[datasetID] + if !ok { + return ClaimResult{}, false, nil + } + now := time.Now() + live := r.owner != "" && r.expires != nil && r.expires.After(now) + if live { + return ClaimResult{}, false, nil + } + if len(r.backlog) == 0 { + return ClaimResult{}, false, nil + } + n := minInt(batchSize, len(r.backlog)) + batch := append([]BacklogEntry{}, r.backlog[:n]...) + r.inflight = append(r.inflight, batch...) + r.backlog = append([]BacklogEntry{}, r.backlog[n:]...) + r.owner = f.holder + r.token = generateHolder() + exp := now.Add(f.leaseTTL) + r.expires = &exp + return ClaimResult{TenantID: r.tenant, Entries: batch, Token: r.token}, true, nil +} + +func (f *FakeScheduler) Ack(_ context.Context, datasetID, token string, batch []BacklogEntry) (int, error) { + f.mu.Lock() + defer f.mu.Unlock() + r, ok := f.rows[datasetID] + if !ok { + return 0, ErrClaimTokenMismatch + } + if r.token != token { + return 0, ErrClaimTokenMismatch + } + r.inflight = removeEntries(r.inflight, batch) + if len(r.inflight) == 0 { + r.owner, r.token, r.expires = "", "", nil + } + return len(r.backlog), nil +} + +func (f *FakeScheduler) TouchClaim(_ context.Context, datasetID, token string, _ time.Duration) (bool, error) { + f.mu.Lock() + defer f.mu.Unlock() + r, ok := f.rows[datasetID] + if !ok || r.token != token { + return false, nil + } + exp := time.Now().Add(f.leaseTTL) + r.expires = &exp + return true, nil +} + +func (f *FakeScheduler) FindClaimable(_ context.Context, limit int) ([]string, error) { + f.mu.Lock() + defer f.mu.Unlock() + var ids []string + now := time.Now() + for id, r := range f.rows { + if len(r.backlog) == 0 { + continue + } + if r.owner != "" && r.expires != nil && r.expires.After(now) { + continue + } + ids = append(ids, id) + if limit > 0 && len(ids) >= limit { + break + } + } + return ids, nil +} + +func (f *FakeScheduler) Notify(_ context.Context, datasetID string) error { + select { + case f.notifyCh <- datasetID: + default: + } + return nil +} + +func (f *FakeScheduler) SubscribeNotify(_ context.Context) (<-chan string, error) { + return f.notifyCh, nil +} + +func (f *FakeScheduler) ReclaimExpired(_ context.Context, now time.Time) (int, error) { + f.mu.Lock() + defer f.mu.Unlock() + total := 0 + for _, r := range f.rows { + if r.owner != "" && r.expires != nil && r.expires.After(now) { + continue + } + if len(r.inflight) == 0 { + continue + } + total += len(r.inflight) + r.backlog = append(r.backlog, r.inflight...) + r.inflight = nil + r.owner, r.token, r.expires = "", "", nil + } + return total, nil +} diff --git a/internal/ingestion/knowledge_compile/service.go b/internal/ingestion/knowledge_compile/service.go new file mode 100644 index 0000000000..e7401047fb --- /dev/null +++ b/internal/ingestion/knowledge_compile/service.go @@ -0,0 +1,146 @@ +// +// 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 knowledge_compile + +import ( + "context" + "crypto/rand" + "encoding/hex" + "fmt" + "time" + + "gorm.io/gorm" + "ragflow/internal/engine" + kccommon "ragflow/internal/ingestion/component/knowledge_compiler/common" +) + +// Option configures a Consumer. +type Option func(*Consumer) + +// WithBatchSize sets the per-KB batch size trigger (closed-batch boundary). +func WithBatchSize(n int) Option { + return func(c *Consumer) { + if n > 0 { + c.batchSize = n + } + } +} + +// WithTTL sets the per-KB claim lease TTL. +func WithTTL(d time.Duration) Option { + return func(c *Consumer) { + if d > 0 { + c.ttl = d + } + } +} + +// WithHeartbeat sets the claim heartbeat period (must be < TTL). +func WithHeartbeat(d time.Duration) Option { + return func(c *Consumer) { + if d > 0 { + c.heartbeat = d + } + } +} + +// WithPollInterval sets how often the worker poll loop looks for a claimable KB +// when no NATS notify arrives. +func WithPollInterval(d time.Duration) Option { + return func(c *Consumer) { + if d > 0 { + c.pollInterval = d + } + } +} + +// WithSweepInterval sets how often the worker runs ReclaimExpired to recover +// inflight left by crashed workers. +func WithSweepInterval(d time.Duration) Option { + return func(c *Consumer) { + if d > 0 { + c.sweepInterval = d + } + } +} + +// WithReader overrides the chunk Reader (used in tests). +func WithReader(r Reader) Option { return func(c *Consumer) { c.reader = r } } + +// WithWriter overrides the chunk Writer (used in tests). +func WithWriter(w Writer) Option { return func(c *Consumer) { c.writer = w } } + +// WithDeduperFactory overrides the per-tenant Deduper factory. +func WithDeduperFactory(f DeduperFactory) Option { return func(c *Consumer) { c.factory = f } } + +// defaultLLMID / defaultEmbedding are the model ids used when resolving LLM +// deps for the dataset-level deduper. Set via SetModelConfig (typically from the +// server bootstrap). Empty strings make the factory fall back to the noop +// deduper until production wiring supplies real values. +var ( + defaultLLMID string + defaultEmbedding string +) + +// SetModelConfig records the model ids used to build the LLM deduper. +func SetModelConfig(llmID, embedding string) { + defaultLLMID = llmID + defaultEmbedding = embedding +} + +// defaultDeduperFactory resolves the per-tenant LLM deps and builds the +// KB-scoped deduper. On any failure it returns an error so the caller falls +// back to the noop deduper (merged products are still written, just without +// cross-document LLM merging). +func defaultDeduperFactory(tenant string) (Deduper, error) { + deps, err := kccommon.ResolveDeps(tenant, defaultLLMID, defaultEmbedding) + if err != nil { + return nil, err + } + return NewLLMDeduper(deps.Chat, deps.Embed, defaultLLMID, 0.99), nil +} + +func generateHolder() string { + b := make([]byte, 8) + if _, err := rand.Read(b); err != nil { + return fmt.Sprintf("knowledgecompile-%d", time.Now().UnixNano()) + } + return "knowledgecompile-" + hex.EncodeToString(b) +} + +// Provision initializes the dataset-level compile scheduling store and installs +// the package-level Scheduler used by the publishing path. It is called once by +// the owning Ingestor (which then drives the consumer from its own worker pool), +// so the dataset-level consumer shares the ingestor's lifecycle instead of +// being a standalone service. Best-effort: provisioning errors are returned so +// the caller can log and continue (the pipeline still writes available_int=0 +// compiled chunks; they just won't be merged until a scheduler is available). +func Provision(ctx context.Context, mq engine.MessageQueue, db *gorm.DB) error { + if db == nil { + return nil + } + s := newScheduler(db, mq, generateHolder(), 2*time.Minute) + // Bound the startup AutoMigrate so a slow/unreachable DB cannot block + // startup indefinitely. The caller's ctx is also honoured (cancelled on + // shutdown). + ctx, cancel := context.WithTimeout(ctx, 30*time.Second) + defer cancel() + if err := s.Provision(ctx); err != nil { + return err + } + SetScheduler(s) + return nil +} diff --git a/internal/ingestion/knowledge_compile/writer.go b/internal/ingestion/knowledge_compile/writer.go new file mode 100644 index 0000000000..ec4631b146 --- /dev/null +++ b/internal/ingestion/knowledge_compile/writer.go @@ -0,0 +1,177 @@ +// +// 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 knowledge_compile + +import ( + "context" + "crypto/sha256" + "encoding/hex" + "fmt" + + "ragflow/internal/engine" + "ragflow/internal/engine/types" + kccommon "ragflow/internal/ingestion/component/knowledge_compiler/common" +) + +// Writer persists dataset-level merged products and removes them on document +// deletion (§11.7). +type Writer interface { + // WriteMerged upserts the dataset-level merged products (available_int=1). + WriteMerged(ctx context.Context, tenant, kb string, products []kccommon.Product) error + // DeleteMergedForDoc removes merged products that became fully orphaned by + // the deletion of docID. Multi-doc merged products are recomputed by the + // caller (via the Reader + Deduper), not here. + DeleteMergedForDoc(ctx context.Context, tenant, kb, docID string) error +} + +type infinityWriter struct{} + +func (infinityWriter) WriteMerged(ctx context.Context, tenant, kb string, products []kccommon.Product) error { + if len(products) == 0 { + return nil + } + eng := engine.Get() + if eng == nil { + return nil + } + baseName := fmt.Sprintf("ragflow_%s", tenant) + chunks := make([]map[string]interface{}, 0, len(products)) + for _, p := range products { + chunks = append(chunks, mergedChunkMap(tenant, kb, p)) + } + _, err := eng.InsertChunks(ctx, chunks, baseName, kb) + return err +} + +// mergedChunkMap builds the chunk-index document for a dataset-level merged +// product. It uses the dataset-level idempotency key (§11.6) as `id`, never the +// per-doc key, and is always available_int=1 (searchable). +func mergedChunkMap(tenant, kb string, p kccommon.Product) map[string]interface{} { + srcDocIDs := metaStringSlice(p.Meta, "source_doc_ids") + srcChunkIDs := metaStringSlice(p.Meta, "source_chunk_ids") + m := map[string]interface{}{ + "id": datasetLevelID(tenant, kb, p), + "doc_id": kb, + "tenant_id": tenant, + "kb_id": kb, + "available_int": 1, + "kc_merged": 1, + "compile_kwd": string(p.Variant), + "content_with_weight": p.Content, + "kc_payload": p.Content, // raw payload, for Reader reconstruction + "source_doc_ids": srcDocIDs, + "source_chunk_ids": srcChunkIDs, + } + // Persist the merged product's embedding under the dimension-suffixed column + // used elsewhere in the index, so dataset-level rows remain vector-searchable + // and Reader.LoadCompiledProducts can reconstruct them (otherwise the vector + // is silently dropped and search returns nothing for merged rows). + if dim := len(p.Vector); dim > 0 { + m[fmt.Sprintf("q_%d_vec", dim)] = p.Vector + } + return m +} + +func (infinityWriter) DeleteMergedForDoc(ctx context.Context, tenant, kb, docID string) error { + eng := engine.Get() + if eng == nil { + return nil + } + baseName := fmt.Sprintf("ragflow_%s", tenant) + const batchSize = 2000 + // Re-query from the start each iteration (deletions shift the result set), + // deleting every fully-orphaned merged row we find, until a page returns + // fewer than batchSize candidates. Without pagination, orphaned merged rows + // beyond the 2000 cap would never be cleaned up. + for { + res, err := eng.Search(ctx, &types.SearchRequest{ + IndexNames: []string{baseName}, + KbIDs: []string{kb}, + Filter: map[string]interface{}{"available_int": 1, "kc_merged": 1}, + SelectFields: []string{"id", "source_doc_ids"}, + Limit: batchSize, + Offset: 0, + }) + if err != nil { + return err + } + if len(res.Chunks) == 0 { + break + } + for _, c := range res.Chunks { + ids := metaStringSlice(c, "source_doc_ids") + // Fully orphaned: the only contributor is the deleted document. + if len(ids) == 1 && ids[0] == docID { + if _, err := eng.DeleteChunks(ctx, map[string]interface{}{ + "id": c["id"], + "kb_id": kb, + "available_int": 1, + }, baseName, kb); err != nil { + return err + } + } + } + if len(res.Chunks) < batchSize { + break + } + } + return nil +} + +// canonicalKey derives a stable cluster key for a merged product. +func canonicalKey(p kccommon.Product) string { + if slug, ok := p.Meta["slug"].(string); ok && slug != "" { + return slug + } + if p.Meta["name"] != nil { + name, _ := p.Meta["name"].(string) + typ, _ := p.Meta["entity_type"].(string) + if typ == "" { + typ, _ = p.Meta["type"].(string) + } + if name != "" { + return hashStr(name + "\x00" + typ) + } + } + return hashStr(p.Content) +} + +// datasetLevelID is the dataset-level idempotency key (§11.6): a stable hash of +// (tenant, kb, variant, canonical cluster key). +func datasetLevelID(tenant, kb string, p kccommon.Product) string { + return hashStr(tenant + "\x00" + kb + "\x00" + string(p.Variant) + "\x00" + canonicalKey(p)) +} + +func hashStr(s string) string { + sum := sha256.Sum256([]byte(s)) + return hex.EncodeToString(sum[:]) +} + +func metaStringSlice(m map[string]any, key string) []string { + switch v := m[key].(type) { + case []string: + return v + case []any: + out := make([]string, 0, len(v)) + for _, e := range v { + if s, ok := e.(string); ok { + out = append(out, s) + } + } + return out + } + return nil +} diff --git a/internal/ingestion/pipeline/pipeline_knowledge_compiler_test.go b/internal/ingestion/pipeline/pipeline_knowledge_compiler_test.go new file mode 100644 index 0000000000..57483a1c5d --- /dev/null +++ b/internal/ingestion/pipeline/pipeline_knowledge_compiler_test.go @@ -0,0 +1,70 @@ +// +// 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 pipeline + +import ( + "os" + "path/filepath" + "testing" + + // Blank-import the KnowledgeCompiler component so its runtime.MustRegister + // init hook fires and the component name referenced by the template is + // resolvable in the default runtime factory. + "ragflow/internal/agent/runtime" + _ "ragflow/internal/ingestion/component/knowledge_compiler" +) + +// TestKnowledgeCompilerTemplate_RegisteredAndDecodable verifies the +// knowledge_compiler ingestion template is loaded by the built-in registry, +// carries the expected title, references runtime-registered components, and +// decodes into a pipeline via NewPipelineFromDSL. A full Run is intentionally +// not exercised here: it needs real LLM/embedder/ES wiring (covered by the +// knowledge_compiler component E2E tests) and is skipped in the headless +// chunks-smoke test. +func TestKnowledgeCompilerTemplate_RegisteredAndDecodable(t *testing.T) { + r, err := DefaultRegistry() + if err != nil { + t.Fatalf("DefaultRegistry: %v", err) + } + tpl, ok := r.Get("knowledge_compiler") + if !ok { + t.Fatal("knowledge_compiler template not registered in built-in registry") + } + if tpl.Title != "Knowledge Compiler" { + t.Fatalf("title = %q, want 'Knowledge Compiler'", tpl.Title) + } + + // Ensure the default runtime factory is installed, then confirm every + // component_name the template references is actually registered. + runtime.InstallDefaultRegistryFactory() + if runtime.DefaultFactory() == nil { + t.Fatal("default runtime factory not installed") + } + 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) + } + } + + path := filepath.Join(repoRootFromPipelineTest(t), "internal", "ingestion", "pipeline", "template", "ingestion_pipeline_knowledge_compiler.json") + raw, err := os.ReadFile(path) + if err != nil { + t.Fatalf("read template: %v", err) + } + if _, err := NewPipelineFromDSL(raw, "kc-validation"); err != nil { + t.Fatalf("NewPipelineFromDSL: %v", err) + } +} diff --git a/internal/ingestion/pipeline/template/ingestion_pipeline_audio.json b/internal/ingestion/pipeline/template/ingestion_pipeline_audio.json index 3a38646f4d..1a20de075c 100644 --- a/internal/ingestion/pipeline/template/ingestion_pipeline_audio.json +++ b/internal/ingestion/pipeline/template/ingestion_pipeline_audio.json @@ -135,7 +135,9 @@ "auto_keywords": 0, "auto_questions": 0, "llm_id": "", - "auto_tags": 0 + "auto_tags": 0, + "enable_metadata": 0, + "metadata": [] } } } diff --git a/internal/ingestion/pipeline/template/ingestion_pipeline_book.json b/internal/ingestion/pipeline/template/ingestion_pipeline_book.json index 537a41acf5..1b6336f78c 100644 --- a/internal/ingestion/pipeline/template/ingestion_pipeline_book.json +++ b/internal/ingestion/pipeline/template/ingestion_pipeline_book.json @@ -196,7 +196,9 @@ "auto_keywords": 0, "auto_questions": 0, "llm_id": "", - "auto_tags": 0 + "auto_tags": 0, + "enable_metadata": 0, + "metadata": [] } } } diff --git a/internal/ingestion/pipeline/template/ingestion_pipeline_email.json b/internal/ingestion/pipeline/template/ingestion_pipeline_email.json index a04f3fadf5..25947234e1 100644 --- a/internal/ingestion/pipeline/template/ingestion_pipeline_email.json +++ b/internal/ingestion/pipeline/template/ingestion_pipeline_email.json @@ -140,7 +140,9 @@ "auto_keywords": 0, "auto_questions": 0, "llm_id": "", - "auto_tags": 0 + "auto_tags": 0, + "enable_metadata": 0, + "metadata": [] } } } diff --git a/internal/ingestion/pipeline/template/ingestion_pipeline_general.json b/internal/ingestion/pipeline/template/ingestion_pipeline_general.json index 46b6655beb..b540ff1b57 100644 --- a/internal/ingestion/pipeline/template/ingestion_pipeline_general.json +++ b/internal/ingestion/pipeline/template/ingestion_pipeline_general.json @@ -224,7 +224,9 @@ "auto_keywords": 0, "auto_questions": 0, "llm_id": "", - "auto_tags": 0 + "auto_tags": 0, + "enable_metadata": 0, + "metadata": [] } } } diff --git a/internal/ingestion/pipeline/template/ingestion_pipeline_knowledge_compiler.json b/internal/ingestion/pipeline/template/ingestion_pipeline_knowledge_compiler.json new file mode 100644 index 0000000000..6b9bf3f5da --- /dev/null +++ b/internal/ingestion/pipeline/template/ingestion_pipeline_knowledge_compiler.json @@ -0,0 +1,310 @@ +{ + "id": 47, + "title": { + "en": "Knowledge Compiler", + "de": "Wissens-Kompilierer", + "zh": "知识编译" + }, + "description": { + "en": "Compiles parsed chunks into structured knowledge units (graph/wiki/raptor/mindmap/datasetnav) via the KnowledgeCompiler component, emitting them as chunks merged into the upstream chunk stream. Ideal for building a retrievable knowledge layer on top of chunked documents.", + "de": "Kompiliert geparste Chunks über die KnowledgeCompiler-Komponente in strukturierte Wissenseinheiten und gibt sie als Chunks im upstream-Chunk-Strom zurück.", + "zh": "通过 KnowledgeCompiler 组件将解析后的分块编译为结构化知识单元（图谱/百科/RAPTOR/思维导图/数据集导航），以 chunks 形式合并进上游分块流输出，适合在分块文档之上构建可检索的知识层。" + }, + "canvas_type": "Ingestion Pipeline", + "canvas_category": "dataflow_canvas", + "dsl": { + "components": { + "File": { + "downstream": [ + "Parser:HipSignsRhyme" + ], + "obj": { + "component_name": "File", + "params": {} + }, + "upstream": [] + }, + "Parser:HipSignsRhyme": { + "downstream": [ + "TokenChunker:SixApplesFall" + ], + "obj": { + "component_name": "Parser", + "params": { + "outputs": { + "html": { + "type": "string", + "value": "" + }, + "json": { + "type": "Array<object>", + "value": [] + }, + "markdown": { + "type": "string", + "value": "" + }, + "text": { + "type": "string", + "value": "" + } + }, + "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" + ] + } + } + } + }, + "upstream": [ + "File" + ] + }, + "TokenChunker:SixApplesFall": { + "downstream": [ + "KnowledgeCompiler:KnownSwiftLions" + ], + "obj": { + "component_name": "TokenChunker", + "params": { + "children_delimiters": [], + "chunk_token_size": 512, + "delimiter_mode": "token_size", + "delimiters": [ + "\n", + "!", + "?", + "。", + "；", + "！", + "？" + ], + "image_context_size": 0, + "outputs": { + "chunks": { + "type": "Array<Object>", + "value": [] + } + }, + "overlapped_percent": 0, + "table_context_size": 0 + } + }, + "upstream": [ + "Parser:HipSignsRhyme" + ] + }, + "KnowledgeCompiler:KnownSwiftLions": { + "downstream": [], + "obj": { + "component_name": "KnowledgeCompiler", + "params": { + "variant": "structure", + "language": "English" + } + }, + "upstream": [ + "TokenChunker:SixApplesFall" + ] + } + }, + "globals": { + "sys.history": [] + }, + "graph": { + "edges": [ + { + "id": "xy-edge__Filestart-Parser:HipSignsRhymeend", + "source": "File", + "sourceHandle": "start", + "target": "Parser:HipSignsRhyme", + "targetHandle": "end" + }, + { + "id": "xy-edge__Parser:HipSignsRhymestart-TokenChunker:SixApplesFallend", + "source": "Parser:HipSignsRhyme", + "sourceHandle": "start", + "target": "TokenChunker:SixApplesFall", + "targetHandle": "end" + }, + { + "id": "xy-edge__TokenChunker:SixApplesFallstart-KnowledgeCompiler:KnownSwiftLionsend", + "source": "TokenChunker:SixApplesFall", + "sourceHandle": "start", + "target": "KnowledgeCompiler:KnownSwiftLions", + "targetHandle": "end" + } + ], + "nodes": [ + { + "data": { + "label": "File", + "name": "File" + }, + "id": "File", + "measured": { + "height": 50, + "width": 200 + }, + "position": { + "x": 50, + "y": 200 + }, + "sourcePosition": "left", + "targetPosition": "right", + "type": "beginNode" + }, + { + "data": { + "label": "Parser", + "name": "Parser_0" + }, + "id": "Parser:HipSignsRhyme", + "measured": { + "height": 57, + "width": 200 + }, + "position": { + "x": 316.99524094206413, + "y": 195.39629819663406 + }, + "sourcePosition": "right", + "targetPosition": "left", + "type": "parserNode" + }, + { + "data": { + "label": "TokenChunker", + "name": "Token Chunker_0" + }, + "id": "TokenChunker:SixApplesFall", + "measured": { + "height": 74, + "width": 200 + }, + "position": { + "x": 616.9952409420641, + "y": 195.39629819663406 + }, + "sourcePosition": "right", + "targetPosition": "left", + "type": "chunkerNode" + }, + { + "data": { + "label": "KnowledgeCompiler", + "name": "Knowledge Compiler_0" + }, + "id": "KnowledgeCompiler:KnownSwiftLions", + "measured": { + "height": 74, + "width": 200 + }, + "position": { + "x": 916.9952409420641, + "y": 195.39629819663406 + }, + "sourcePosition": "right", + "targetPosition": "left", + "type": "componentNode" + } + ] + }, + "history": [], + "messages": [], + "path": [], + "retrieval": [], + "variables": [] + }, + "parser_ids": [ + "knowledge_compiler" + ] +} diff --git a/internal/ingestion/pipeline/template/ingestion_pipeline_laws.json b/internal/ingestion/pipeline/template/ingestion_pipeline_laws.json index 7ab36a2fac..7b6e5ca7b5 100644 --- a/internal/ingestion/pipeline/template/ingestion_pipeline_laws.json +++ b/internal/ingestion/pipeline/template/ingestion_pipeline_laws.json @@ -209,7 +209,9 @@ "auto_keywords": 0, "auto_questions": 0, "llm_id": "", - "auto_tags": 0 + "auto_tags": 0, + "enable_metadata": 0, + "metadata": [] } } } diff --git a/internal/ingestion/pipeline/template/ingestion_pipeline_manual.json b/internal/ingestion/pipeline/template/ingestion_pipeline_manual.json index 8a82d00a8d..c227ccd188 100644 --- a/internal/ingestion/pipeline/template/ingestion_pipeline_manual.json +++ b/internal/ingestion/pipeline/template/ingestion_pipeline_manual.json @@ -176,7 +176,9 @@ "auto_keywords": 0, "auto_questions": 0, "llm_id": "", - "auto_tags": 0 + "auto_tags": 0, + "enable_metadata": 0, + "metadata": [] } } } diff --git a/internal/ingestion/pipeline/template/ingestion_pipeline_one.json b/internal/ingestion/pipeline/template/ingestion_pipeline_one.json index 87b6e003a0..3a0005c821 100644 --- a/internal/ingestion/pipeline/template/ingestion_pipeline_one.json +++ b/internal/ingestion/pipeline/template/ingestion_pipeline_one.json @@ -189,7 +189,9 @@ "auto_keywords": 0, "auto_questions": 0, "llm_id": "", - "auto_tags": 0 + "auto_tags": 0, + "enable_metadata": 0, + "metadata": [] } } } diff --git a/internal/ingestion/pipeline/template/ingestion_pipeline_paper.json b/internal/ingestion/pipeline/template/ingestion_pipeline_paper.json index 490e1935d0..b90d09b8b3 100644 --- a/internal/ingestion/pipeline/template/ingestion_pipeline_paper.json +++ b/internal/ingestion/pipeline/template/ingestion_pipeline_paper.json @@ -157,7 +157,9 @@ "auto_keywords": 0, "auto_questions": 0, "llm_id": "", - "auto_tags": 0 + "auto_tags": 0, + "enable_metadata": 0, + "metadata": [] } } } diff --git a/internal/ingestion/pipeline/template/ingestion_pipeline_picture.json b/internal/ingestion/pipeline/template/ingestion_pipeline_picture.json index f04f47a273..cf1e9dbea4 100644 --- a/internal/ingestion/pipeline/template/ingestion_pipeline_picture.json +++ b/internal/ingestion/pipeline/template/ingestion_pipeline_picture.json @@ -149,7 +149,9 @@ "auto_keywords": 0, "auto_questions": 0, "llm_id": "", - "auto_tags": 0 + "auto_tags": 0, + "enable_metadata": 0, + "metadata": [] } } } diff --git a/internal/ingestion/pipeline/template/ingestion_pipeline_presentation.json b/internal/ingestion/pipeline/template/ingestion_pipeline_presentation.json index 341ee68e45..14e9db6912 100644 --- a/internal/ingestion/pipeline/template/ingestion_pipeline_presentation.json +++ b/internal/ingestion/pipeline/template/ingestion_pipeline_presentation.json @@ -135,7 +135,9 @@ "auto_keywords": 0, "auto_questions": 0, "llm_id": "", - "auto_tags": 0 + "auto_tags": 0, + "enable_metadata": 0, + "metadata": [] } } } diff --git a/internal/ingestion/pipeline/template/ingestion_pipeline_resume.json b/internal/ingestion/pipeline/template/ingestion_pipeline_resume.json index b0a4be7c5b..7f449c4d34 100644 --- a/internal/ingestion/pipeline/template/ingestion_pipeline_resume.json +++ b/internal/ingestion/pipeline/template/ingestion_pipeline_resume.json @@ -49,7 +49,9 @@ "top_p": 0.3, "auto_keywords": 0, "auto_questions": 0, - "auto_tags": 0 + "auto_tags": 0, + "enable_metadata": 0, + "metadata": [] } }, "upstream": [ diff --git a/internal/ingestion/pipeline/template_integration_test.go b/internal/ingestion/pipeline/template_integration_test.go index 59ead62fe8..dee9f7ab22 100644 --- a/internal/ingestion/pipeline/template_integration_test.go +++ b/internal/ingestion/pipeline/template_integration_test.go @@ -811,6 +811,9 @@ func TestPipelineRun_AllIngestionTemplates_RealComponentsSmoke(t *testing.T) { if templateUsesComponent(t, templateBytes, "TagChunker") { t.Skip("template uses TagChunker which requires tag-structured content and parser setups not available for generic .md input; covered separately") } + if templateUsesComponent(t, templateBytes, "KnowledgeCompiler") { + t.Skip("template uses KnowledgeCompiler which requires LLM/embedder/ES wiring not available in the headless smoke run; covered by the knowledge_compiler component E2E tests") + } terminalIDs := terminalComponentIDsFromTemplate(t, templateBytes) if len(terminalIDs) != 1 { t.Fatalf("terminal ids = %v, want exactly 1 terminal", terminalIDs) diff --git a/internal/ingestion/service/doc_state.go b/internal/ingestion/service/doc_state.go index 64530f60a9..92b83fa3aa 100644 --- a/internal/ingestion/service/doc_state.go +++ b/internal/ingestion/service/doc_state.go @@ -23,6 +23,7 @@ import ( "ragflow/internal/common" taskpkg "ragflow/internal/ingestion/task" documentpkg "ragflow/internal/service/document" + "ragflow/internal/utility" ) // docStateSvc is the subset of *service.DocumentService needed to finalize a @@ -65,10 +66,14 @@ func (u *docStateUpdater) apply(ctx context.Context, r *taskpkg.PipelineResult) } } -// mergeDocMetadata reads existing metadata, fills in keys not already present -// (existing keys are preserved, not overwritten), and writes the merged map back. -// A read failure aborts the merge: SetDocumentMetadata is a full overwrite, so -// writing with an empty baseline would destroy existing keys. +// mergeDocMetadata reads existing metadata, unions it with the freshly +// aggregated doc metadata (list values merged + de-duplicated, scalars from the +// stored map winning — matching Python task_executor.py:572 +// update_metadata_to(metadata, existing_meta)), then splits combined values +// before writing the merged map back (Python doc_metadata_service.py:468 +// _split_combined_values). A read failure aborts the merge: SetDocumentMetadata +// is a full overwrite, so writing with an empty baseline would destroy existing +// keys. func mergeDocMetadata(ctx context.Context, svc docStateSvc, docID string, metadata map[string]any) error { existing, err := svc.GetDocumentMetadataByID(ctx, docID) if err != nil { @@ -77,10 +82,7 @@ func mergeDocMetadata(ctx context.Context, svc docStateSvc, docID string, metada if existing == nil { existing = map[string]any{} } - for k, v := range metadata { - if _, exists := existing[k]; !exists { - existing[k] = v - } - } - return svc.SetDocumentMetadata(ctx, docID, existing) + merged := utility.UpdateMetadataTo(metadata, existing) + merged = common.SplitCombinedMetadataValues(merged) + return svc.SetDocumentMetadata(ctx, docID, merged) } diff --git a/internal/ingestion/service/doc_state_test.go b/internal/ingestion/service/doc_state_test.go index f403d3a090..5ea98f4cdc 100644 --- a/internal/ingestion/service/doc_state_test.go +++ b/internal/ingestion/service/doc_state_test.go @@ -111,6 +111,56 @@ func TestDocStateUpdater_PreservesExistingKey(t *testing.T) { } } +func TestDocStateUpdater_UnionsListValues(t *testing.T) { + // Stored doc metadata already has a list value; a new run extracts more + // (possibly overlapping) values. The merge must union + de-dupe, matching + // Python update_metadata_to(metadata, existing_meta). + svc := &stubDocStateSvc{metaData: map[string]any{"people": []string{"关羽", "张辽"}}} + u := &docStateUpdater{docSvc: svc} + ctx := t.Context() + u.apply(ctx, &taskpkg.PipelineResult{ + DocID: "doc-1", + Metadata: map[string]any{"people": []string{"张辽", "刘备"}}, + ChunkCount: 1, TokenConsumption: 10, + }) + got, ok := svc.metaData["people"].([]string) + if !ok { + t.Fatalf("people should be []string, got %T", svc.metaData["people"]) + } + want := []string{"张辽", "刘备", "关羽"} + if len(got) != len(want) { + t.Fatalf("people=%v, want union %v", got, want) + } + seen := map[string]bool{} + for _, p := range got { + if seen[p] { + t.Fatalf("duplicate %q in %v", p, got) + } + seen[p] = true + } + for _, w := range want { + if !seen[w] { + t.Fatalf("missing %q in %v", w, got) + } + } +} + +func TestDocStateUpdater_PreservesExistingScalar(t *testing.T) { + // Python update_metadata_to keeps the stored scalar when both sides carry + // a scalar for the same key (stored wins). + svc := &stubDocStateSvc{metaData: map[string]any{"author": "Alice"}} + u := &docStateUpdater{docSvc: svc} + ctx := t.Context() + u.apply(ctx, &taskpkg.PipelineResult{ + DocID: "doc-1", + Metadata: map[string]any{"author": "Bob"}, + ChunkCount: 1, TokenConsumption: 10, + }) + if svc.metaData["author"] != "Alice" { + t.Fatalf("stored scalar must win: got %q", svc.metaData["author"]) + } +} + func TestDocStateUpdater_IncrementArgs(t *testing.T) { svc := &stubDocStateSvc{} u := &docStateUpdater{docSvc: svc} diff --git a/internal/ingestion/service/ingestion_service.go b/internal/ingestion/service/ingestion_service.go index fa8c78ed2f..d4b2b4bdeb 100644 --- a/internal/ingestion/service/ingestion_service.go +++ b/internal/ingestion/service/ingestion_service.go @@ -28,9 +28,11 @@ import ( "time" "ragflow/internal/common" + "ragflow/internal/dao" "ragflow/internal/engine" redis2 "ragflow/internal/engine/redis" "ragflow/internal/entity" + "ragflow/internal/ingestion/knowledge_compile" pipelinepkg "ragflow/internal/ingestion/pipeline" taskpkg "ragflow/internal/ingestion/task" servicepkg "ragflow/internal/service" @@ -70,6 +72,23 @@ type Ingestor struct { ingestionTaskSvc *servicepkg.IngestionTaskService docState *docStateUpdater + // knowledgeCompile is the dataset-level post-processing consumer (§11, + // Option E) owned by this ingestor. It is driven by kcConcurrency owned + // worker goroutines, each running the Consumer's Run loop (poll MySQL + // scheduling rows + NATS notify → claim closed batch → merge → ack). The + // MySQL knowledge_compile_docs table — not the broker — is the scheduling + // system of record and the source of same-KB serialization, so different + // datasets compile in parallel while the same dataset is serialized by its + // claim row. The workers are started/joined within the ingestor (via + // compileWg), so they share its lifecycle and goroutine set instead of + // running as a separate service goroutine. + knowledgeCompile *knowledge_compile.Consumer + kcLLMID string + kcEmbedding string + kcConcurrency int32 // number of parallel dataset-level compile workers + + compileWg sync.WaitGroup + // runDocumentTask dispatches to the migrated task handler path. // Tests may override this to verify branch routing without invoking // the full downstream stack. @@ -106,6 +125,7 @@ func NewIngestor(name string, maxConcurrency int32, supportedTypes []string) *In } ingestor.runDocumentTask = ingestor.defaultRunDocumentTask ingestor.cancelCheck = ingestor.defaultCancelCheck + ingestor.kcConcurrency = maxConcurrency // parallel dataset-level compile workers default to the task width return ingestor } @@ -120,21 +140,52 @@ const consumeErrorBackoff = 1 * time.Second func (e *Ingestor) Start() error { common.Info(fmt.Sprintf("Ingestor %s initialized", e.id)) + var startErr error + e.startOnce.Do(func() { + startErr = e.start() + }) + return startErr +} + +// start runs the full startup sequence. It is invoked at most once (guarded by +// startOnce in Start) so repeated Start calls cannot launch duplicate worker +// pools, compile consumers, or consume loops, and the first initialization +// error is retained and returned to every later caller. +func (e *Ingestor) start() error { msgQueueEngine := engine.GetMessageQueueEngine() - err := msgQueueEngine.InitConsumer("tasks.RAGFLOW") - if err != nil { + if err := msgQueueEngine.InitConsumer("tasks.RAGFLOW"); err != nil { return err } - // Ensure worker pool is started on first task - go e.startWorkerPool() + // Start the task worker pool and the dataset-level compile consumer as + // owned goroutines joined by Stop via workerWg/compileWg. Start follows + // the standard lifecycle contract: it returns immediately after kicking + // these off rather than blocking on the consume loop itself. + e.startWorkerPool() + e.startDatasetKnowledgeCompile() + // Run the main tasks.RAGFLOW consume loop off the caller's goroutine so + // Start returns promptly; it is joined by Stop via workerWg. + e.workerWg.Add(1) + go e.consumeLoop() + return nil +} + +// consumeLoop is the main tasks.RAGFLOW consume loop. It runs until e.ctx is +// cancelled (graceful shutdown); per-message failures are settled by +// processMessage and never terminate the consumer. Dataset-level compile work +// is owned by the kcConcurrency compileLoop workers (each running the +// Consumer's Run loop), so this loop stays focused on tasks.RAGFLOW and is +// never held up by compile work. +func (e *Ingestor) consumeLoop() { + defer e.workerWg.Done() + msgQueueEngine := engine.GetMessageQueueEngine() for { // Graceful shutdown is the only condition under which the consume // loop exits. Per-message processing failures never terminate the // consumer: processMessage settles (ack/nack) each message itself. if err := e.ctx.Err(); err != nil { - return nil + return } taskHandles, err := msgQueueEngine.GetMessages(4) if err != nil { @@ -142,13 +193,84 @@ func (e *Ingestor) Start() error { select { case <-time.After(consumeErrorBackoff): case <-e.ctx.Done(): - return nil + return } continue } - for _, taskHandle := range taskHandles { - e.processMessage(taskHandle) + if len(taskHandles) > 0 { + for _, taskHandle := range taskHandles { + e.processMessage(taskHandle) + } + continue } + // tasks.RAGFLOW idle: dataset-level compile work is owned by the + // kcConcurrency compileLoop workers (started in + // startDatasetKnowledgeCompile), so the task loop stays focused on + // tasks.RAGFLOW and is never held up by compile work. + } +} + +// SetKnowledgeCompileModelConfig supplies the default LLM/embedding model ids +// used by the dataset-level compile consumer's deduper. Call it before Start. +func (e *Ingestor) SetKnowledgeCompileModelConfig(llmID, embedding string) { + e.kcLLMID = llmID + e.kcEmbedding = embedding +} + +// SetKnowledgeCompileConcurrency sets the number of parallel dataset-level +// compile workers. Multiple datasets compile concurrently (each worker runs its +// own Run loop that claims a KB's closed batch and merges it); the per-KB claim +// row serializes the same dataset, preserving the §11 O(unique pairs) +// merged-dedup invariant. Call before Start. A value <= 0 falls back to +// runtime.NumCPU() at start time. +func (e *Ingestor) SetKnowledgeCompileConcurrency(n int32) { + e.kcConcurrency = n +} + +// startDatasetKnowledgeCompile provisions the dataset-level compile scheduling +// store (MySQL knowledge_compile_docs + NATS notify subject), builds the +// Consumer, and starts the owned compile-worker pool. The consumer is driven by +// the ingestor's own goroutine set — kcConcurrency compileLoop workers, each +// running the Consumer's Run loop (see §11.7 of knowledge_compile_design.md) — +// so its lifecycle shares the ingestor's. Best-effort: any provisioning failure +// is logged and the pipeline continues to write available_int=0 compiled +// chunks; they just won't be merged until a scheduler is available. +func (e *Ingestor) startDatasetKnowledgeCompile() { + mq := engine.GetMessageQueueEngine() + if mq == nil { + common.Warn("message queue not initialized; dataset-level compile consumer disabled") + return + } + knowledge_compile.SetModelConfig(e.kcLLMID, e.kcEmbedding) + if err := knowledge_compile.Provision(e.ctx, mq, dao.DB); err != nil { + common.Warn(fmt.Sprintf("dataset-level compile consumer unavailable; compiled chunks will not be merged: %v", err)) + return + } + e.knowledgeCompile = knowledge_compile.NewConsumer(knowledge_compile.DefaultScheduler()) + n := e.kcConcurrency + if n <= 0 { + n = int32(runtime.NumCPU()) + } + // kcConcurrency owned workers, each running its own Run loop. Concurrency is + // bounded by the goroutine count itself (no extra semaphore). Different + // datasets compile in parallel while the same dataset is serialized by its + // MySQL claim row — preserving the §11 invariant that merged dedup stays + // O(unique pairs) instead of O(N). + for i := int32(0); i < n; i++ { + e.compileWg.Add(1) + go e.compileLoop(i) + } +} + +// compileLoop is one of the owned dataset-level compile workers. It runs the +// Consumer's Run loop (poll scheduling rows + NATS notify → claim closed batch +// → merge → ack). kcConcurrency such goroutines run in parallel, so different +// datasets compile concurrently while the same dataset is serialized by its +// claim row. All are joined in Stop via compileWg. +func (e *Ingestor) compileLoop(id int32) { + defer e.compileWg.Done() + if e.knowledgeCompile != nil { + e.knowledgeCompile.Run(e.ctx) } } @@ -754,6 +876,7 @@ func (e *Ingestor) Stop(ctx context.Context) { waitDone := make(chan struct{}) go func() { e.workerWg.Wait() + e.compileWg.Wait() close(waitDone) }() diff --git a/internal/ingestion/task/knowledge_compiler_wiring.go b/internal/ingestion/task/knowledge_compiler_wiring.go new file mode 100644 index 0000000000..244bef3e3e --- /dev/null +++ b/internal/ingestion/task/knowledge_compiler_wiring.go @@ -0,0 +1,291 @@ +// +// 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 ( + "context" + "fmt" + "strings" + "sync/atomic" + + "ragflow/internal/dao" + "ragflow/internal/engine" + enginetypes "ragflow/internal/engine/types" + "ragflow/internal/entity/models" + kc "ragflow/internal/ingestion/component/knowledge_compiler/common" + "ragflow/internal/service" +) + +// This file is the composition-root wiring for the KnowledgeCompiler ingestion +// component. The component package (internal/ingestion/component/knowledge_compiler) +// is deliberately DB-independent: it owns the compile schema but not the model +// resolution or the storage engine. The DepsResolver seam is injected here, at +// the task-package level, so the component never imports internal/service +// directly (which would invert the dependency direction — see PORT_PLAN.md §4). +// +// The component returns its compiled knowledge units as chunk-aligned docs merged +// into the upstream chunk stream, so it needs no separate writer: the caller +// (pipeline / downstream tokenizer) handles any persistence, exactly as it does +// for ordinary chunks. + +func init() { + kc.SetDepsResolver(newKnowledgeCompilerDepsResolver()) + kc.SetGroupResolver(newKnowledgeCompilerGroupResolver()) +} + +// newKnowledgeCompilerGroupResolver builds the production GroupResolver backed by +// the compilation_template DAO. Without it, any config carrying +// compilation_template_group_ids would fail loud at runtime (the component +// refuses to silently drop the compilation_template_ids stamp). It resolves each +// group id to its child template ids so group-based configs stamp the full set +// on every compiled unit. +func newKnowledgeCompilerGroupResolver() kc.GroupResolver { + tmplDAO := dao.NewCompilationTemplateDAO() + return func(ctx context.Context, tenantID string, groupIDs []string) ([]string, error) { + return tmplDAO.ResolveGroupTemplateIDs(ctx, tenantID, groupIDs) + } +} + +// newKnowledgeCompilerDepsResolver builds the production DepsResolver. Each call +// yields a fresh Deps whose ChatInvoker / Embedder are bound to the resolved +// tenant + model ids (captured in the closure), mirroring how the Tokenizer +// component resolves its embedder. +func newKnowledgeCompilerDepsResolver() kc.DepsResolver { + svc := service.NewModelProviderService() + return func(tenantID, llmID, embeddingModel string) (kc.Deps, error) { + if strings.TrimSpace(llmID) == "" { + return kc.Deps{}, fmt.Errorf("knowledge_compiler: llm_id is required for production deps resolution") + } + return kc.Deps{ + Chat: &kcChatInvoker{svc: svc, tenantID: tenantID, llmID: llmID}, + Embed: &kcEmbedder{svc: svc, tenantID: tenantID, embdID: embeddingModel}, + WikiPages: &kcWikiPageStore{docEngine: engine.Get()}, + // HistoricalKNN / Redis are optional (wiki historical dedup, + // datasetnav lock). They are wired separately when the + // surrounding pipeline supplies the backing services. + }, nil + } +} + +// kcChatInvoker adapts service.ModelProviderService.Chat to the +// knowledge_compiler ChatInvoker seam. +type kcChatInvoker struct { + svc *service.ModelProviderService + tenantID string + llmID string +} + +func (c *kcChatInvoker) Chat(ctx context.Context, req kc.ChatRequest) (*kc.ChatResponse, error) { + llmID := c.llmID + if req.LLMID != "" { + llmID = req.LLMID + } + msgs := []models.Message{ + {Role: "system", Content: req.SystemPrompt}, + {Role: "user", Content: req.UserPrompt}, + } + // Python's knowledge compilation pins per-call-site temperatures + // (extraction 0.1, merge judging 0.0); nil leaves the driver default. + var config *models.ChatConfig + if req.Temperature != nil { + config = &models.ChatConfig{Temperature: req.Temperature} + } + resp, err := c.svc.Chat(ctx, c.tenantID, llmID, msgs, config) + if err != nil { + return nil, err + } + content := "" + if resp != nil && resp.Answer != nil { + content = *resp.Answer + } + return &kc.ChatResponse{Content: content}, nil +} + +// kcEmbedder adapts service.ModelProviderService.GetEmbeddingModel to the +// knowledge_compiler Embedder seam. Vectors are returned as []float32 to match +// the component's product schema. +type kcEmbedder struct { + svc *service.ModelProviderService + tenantID string + embdID string + dim atomic.Int64 +} + +func (e *kcEmbedder) Encode(ctx context.Context, texts []string) ([][]float32, error) { + if len(texts) == 0 { + return nil, nil + } + embdID := strings.TrimSpace(e.embdID) + if embdID == "" { + return nil, fmt.Errorf("knowledge_compiler: embedding_model is required for production embedding") + } + mdl, err := e.svc.GetEmbeddingModel(ctx, e.tenantID, embdID) + 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) + } + config := &models.EmbeddingConfig{} + // Embed expects *string for the model name; nil ModelUsage (not tracked here). + embeds, err := mdl.ModelDriver.Embed(ctx, mdl.ModelName, texts, mdl.APIConfig, config, nil) + if err != nil { + return nil, fmt.Errorf("knowledge_compiler: embed: %w", err) + } + out := make([][]float32, len(embeds)) + for i, v := range embeds { + out[i] = float64sToFloat32(v.Embedding) + } + if len(out) > 0 { + e.dim.CompareAndSwap(0, int64(len(out[0]))) + } + return out, nil +} + +func (e *kcEmbedder) Dimensions() int { return int(e.dim.Load()) } + +// float64sToFloat32 converts an embedding vector to the product schema's +// []float32 representation. +func float64sToFloat32(in []float64) []float32 { + out := make([]float32, len(in)) + for i, x := range in { + out[i] = float32(x) + } + return out +} + +type kcWikiPageStore struct { + docEngine engine.DocEngine +} + +func (s *kcWikiPageStore) FindSimilarPages(ctx context.Context, tenantID, datasetID string, queryVec []float32, k int) ([]kc.WikiPageCandidate, error) { + if s == nil || s.docEngine == nil || len(queryVec) == 0 || k <= 0 || strings.TrimSpace(datasetID) == "" { + return nil, nil + } + vec := make([]float64, len(queryVec)) + for i, v := range queryVec { + vec[i] = float64(v) + } + req := &enginetypes.SearchRequest{ + IndexNames: []string{fmt.Sprintf("ragflow_%s", tenantID)}, + 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"}, + Filter: map[string]interface{}{ + "compile_kwd": "artifact_page", + "kc_kind": "page", + }, + MatchExprs: []interface{}{&enginetypes.MatchDenseExpr{ + VectorColumnName: fmt.Sprintf("q_%d_vec", len(vec)), + EmbeddingData: vec, + EmbeddingDataType: "float", + DistanceType: "cosine", + TopN: k, + ExtraOptions: map[string]interface{}{"similarity": 0.0}, + }}, + } + res, err := s.docEngine.Search(ctx, req) + if err != nil || res == nil { + return nil, err + } + out := make([]kc.WikiPageCandidate, 0, len(res.Chunks)) + for _, row := range res.Chunks { + out = append(out, wikiPageCandidateFromRow(row)) + } + return out, nil +} + +func (s *kcWikiPageStore) GetPageBySlug(ctx context.Context, tenantID, datasetID, slug string) (*kc.WikiPageCandidate, error) { + if s == nil || s.docEngine == nil || strings.TrimSpace(datasetID) == "" || strings.TrimSpace(slug) == "" { + return nil, nil + } + req := &enginetypes.SearchRequest{ + IndexNames: []string{fmt.Sprintf("ragflow_%s", tenantID)}, + KbIDs: []string{datasetID}, + Limit: 1, + SelectFields: []string{"id", "slug_kwd", "title_kwd", "page_type_kwd", "topic_kwd", "summary_with_weight", "content_with_weight", "entity_names_kwd", "related_kb_pages_kwd", "outlinks_kwd", "kc_content_md_raw", "_score"}, + Filter: map[string]interface{}{ + "compile_kwd": "artifact_page", + "slug_kwd": slug, + "kc_kind": "page", + }, + } + res, err := s.docEngine.Search(ctx, req) + if err != nil || res == nil || len(res.Chunks) == 0 { + return nil, err + } + page := wikiPageCandidateFromRow(res.Chunks[0]) + return &page, nil +} + +func wikiPageCandidateFromRow(row map[string]interface{}) kc.WikiPageCandidate { + return kc.WikiPageCandidate{ + ID: strings.TrimSpace(anyString(row["id"])), + Slug: strings.TrimSpace(anyString(row["slug_kwd"])), + Title: strings.TrimSpace(anyString(row["title_kwd"])), + PageType: strings.TrimSpace(anyString(row["page_type_kwd"])), + Topic: strings.TrimSpace(anyString(row["topic_kwd"])), + Summary: strings.TrimSpace(anyString(row["summary_with_weight"])), + ContentMD: strings.TrimSpace(anyString(row["content_with_weight"])), + ContentMDRaw: strings.TrimSpace(anyString(row["kc_content_md_raw"])), + EntityNames: anyStrings(row["entity_names_kwd"]), + RelatedKBPages: anyStrings(row["related_kb_pages_kwd"]), + Outlinks: anyStrings(row["outlinks_kwd"]), + Score: anyFloat(row["_score"]), + } +} + +func anyString(v interface{}) string { + switch x := v.(type) { + case string: + return x + default: + return "" + } +} + +func anyStrings(v interface{}) []string { + switch x := v.(type) { + case []string: + return x + case []interface{}: + out := make([]string, 0, len(x)) + for _, item := range x { + if s, ok := item.(string); ok && strings.TrimSpace(s) != "" { + out = append(out, strings.TrimSpace(s)) + } + } + return out + default: + return nil + } +} + +func anyFloat(v interface{}) float64 { + switch x := v.(type) { + case float64: + return x + case float32: + return float64(x) + case int: + return float64(x) + case int64: + return float64(x) + default: + return 0 + } +} diff --git a/internal/ingestion/task/pipeline_executor.go b/internal/ingestion/task/pipeline_executor.go index 07d11d2eb8..dfd187a550 100644 --- a/internal/ingestion/task/pipeline_executor.go +++ b/internal/ingestion/task/pipeline_executor.go @@ -28,6 +28,7 @@ import ( "ragflow/internal/dao" "ragflow/internal/engine" "ragflow/internal/entity" + "ragflow/internal/ingestion/knowledge_compile" pipelinepkg "ragflow/internal/ingestion/pipeline" "gorm.io/gorm" @@ -182,6 +183,7 @@ func (s *PipelineExecutor) Execute(ctx context.Context) (*PipelineResult, error) if pipelineDSL != "" { s.recordPipelineLog(ctx, dao.DB, s.taskCtx.Doc.ID, pipelineDSL, "done") } + return result, nil } @@ -222,10 +224,27 @@ func (s *PipelineExecutor) processOutput(ctx context.Context, pipelineOutput map } } - if err = s.indexWriter.Write(ctx, chunks); err != nil { + // Per-document compiled knowledge products (those emitted by the + // KnowledgeCompiler component and stamped with `compile_kwd`) are persisted + // as available_int=0: invisible to the normal retriever until the dataset-level + // post-processing consumer (§11) merges them into available_int=1 products. + // Ordinary source chunks stay available_int=1 (the index default). + markCompiledProductsHidden(chunks) + + if err := s.indexWriter.Write(ctx, chunks); err != nil { return nil, err } + // All chunks are now persisted. Notify the dataset-level post-processing consumer + // (§11) that this document is complete: its compiled products were written + // available_int=0 and the consumer later merges them into dataset-level products + // (available_int=1). The notification is sent only after a successful persist + // and is best-effort / non-fatal — a delivery failure is logged but does not + // fail the pipeline task. + if err := knowledge_compile.PublishCompleted(ctx, s.taskCtx.Tenant.ID, s.taskCtx.Doc.KbID, s.taskCtx.Doc.ID, 0); err != nil { + common.Logger.Warn(fmt.Sprintf("knowledge_compile: publish doc_completed for %s failed: %v", s.taskCtx.Doc.ID, err)) + } + chunkCount := countDistinctChunkIDs(chunks) return &PipelineResult{ @@ -256,6 +275,23 @@ func countDistinctChunkIDs(chunks []map[string]any) int { return len(seen) } +// markCompiledProductsHidden sets available_int=0 on the per-document compiled +// knowledge products so they are hidden from the retriever until the dataset-level +// post-processing consumer merges them into available_int=1 products (§11). A +// chunk is a compiled product iff it carries the compile_kwd discriminator the +// KnowledgeCompiler component stamps; ordinary source chunks (no compile_kwd) +// keep the index default available_int=1 and remain immediately searchable. +// Merged dataset-level products are written by the consumer, never here, so they are +// never double-marked. +func markCompiledProductsHidden(chunks []map[string]any) { + for _, ck := range chunks { + if _, ok := ck["compile_kwd"]; !ok { + continue + } + ck["available_int"] = 0 + } +} + func (s *PipelineExecutor) recordPipelineLog(ctx context.Context, db *gorm.DB, docID, dsl, status string) { var dslMap entity.JSONMap if err := json.Unmarshal([]byte(dsl), &dslMap); err != nil { @@ -340,6 +376,13 @@ func (s *PipelineExecutor) runPipelineWithDSL(ctx context.Context, dsl string) ( parserConfig := map[string]interface{}(s.taskCtx.Doc.ParserConfig) common.InjectExtractorLLMID(parserConfig, s.taskCtx.Tenant.LLMID) + // When the dataset enables auto-metadata, ensure the Extractor node(s) + // carry the enable_metadata mode + field schema so the LLM extraction fires + // (mirrors Python task_executor.py:519 enabling gen_metadata_task). The + // dataset flag is authoritative: a node that already has enable_metadata + // turned on keeps its own config, but a shipped DSL defaulting it to 0 is + // still overridden so auto-metadata can activate. + common.InjectExtractorEnableMetadata(parserConfig) // Surface component params whose cpnID is absent from the DSL. The // runtime merge (override_params) silently drops such entries; diff --git a/internal/ingestion/task/pipeline_executor_defaults_test.go b/internal/ingestion/task/pipeline_executor_defaults_test.go index ceaa44fdc5..acc2113a8e 100644 --- a/internal/ingestion/task/pipeline_executor_defaults_test.go +++ b/internal/ingestion/task/pipeline_executor_defaults_test.go @@ -38,19 +38,20 @@ import ( // Comparison is done per-component (see assertComponentsMatch), so the JSON // serialization order inside each entry is irrelevant. var builtinComponentParamsGolden = map[string]string{ - "audio": "{\"File\": {}, \"Parser:SongsFillAir\": {\"audio\": {\"output_format\": \"text\", \"preprocess\": [\"main_content\"], \"suffix\": [\"aac\", \"aiff\", \"ape\", \"au\", \"da\", \"flac\", \"midi\", \"mp3\", \"ogg\", \"oggvorbis\", \"realaudio\", \"vqf\", \"wav\", \"wave\", \"wma\"]}}, \"TokenChunker:BlueSkiesLaugh\": {}, \"Tokenizer:KindEyesWatch\": {\"fields\": \"text\", \"filename_embd_weight\": 0.1, \"search_method\": [\"embedding\", \"full_text\"]}, \"Extractor:AutoExtractDefault\": {\"field_name\": \"\", \"auto_keywords\": 0, \"auto_questions\": 0, \"llm_id\": \"\", \"auto_tags\": 0}}", - "book": "{\"File\": {}, \"Parser:HipSignsRhyme\": {\"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\"]}, \"pdf\": {\"flatten_media_to_text\": false, \"output_format\": \"json\", \"parse_method\": \"DeepDOC\", \"preprocess\": [\"main_content\"], \"remove_toc\": true, \"suffix\": [\"pdf\"], \"vlm\": {}}, \"text&code\": {\"output_format\": \"json\", \"preprocess\": [\"main_content\"], \"suffix\": [\"txt\"]}}, \"TitleChunker:GrumpyGarlicsBake\": {\"hierarchy\": 5, \"include_heading_content\": true, \"levels\": [[\"^#[^#]\", \"^##[^#]\", \"^###[^#]\", \"^####[^#]\"], [\"第[零一二三四五六七八九十百0-9]+(分?编|部分)\", \"第[零一二三四五六七八九十百0-9]+章\", \"第[零一二三四五六七八九十百0-9]+节\", \"第[零一二三四五六七八九十百0-9]+条\", \"[\\\\(（][零一二三四五六七八九十百]+[\\\\)）]\"], [\"第[0-9]+章\", \"第[0-9]+节\", \"[0-9]{1,2}[\\\\. 、]\", \"[0-9]{1,2}\\\\.[0-9]{1,2}($|[^a-zA-Z/%~.-])\", \"[0-9]{1,2}\\\\.[0-9]{1,2}\\\\.[0-9]{1,2}\"], [\"第[零一二三四五六七八九十百0-9]+章\", \"第[零一二三四五六七八九十百0-9]+节\", \"[零一二三四五六七八九十百]+[ 、]\", \"[\\\\(（][零一二三四五六七八九十百]+[\\\\)）]\", \"[\\\\(（][0-9]{,2}[\\\\)）]\"], [\"PART (ONE|TWO|THREE|FOUR|FIVE|SIX|SEVEN|EIGHT|NINE|TEN)\", \"Chapter (I+V?|VI*|XI|IX|X)\", \"Section [0-9]+\", \"Article [0-9]+\"]], \"method\": \"hierarchy\"}, \"Tokenizer:HotDonutsRing\": {\"fields\": \"text\", \"filename_embd_weight\": 0.1, \"search_method\": [\"embedding\", \"full_text\"]}, \"Extractor:AutoExtractDefault\": {\"field_name\": \"\", \"auto_keywords\": 0, \"auto_questions\": 0, \"llm_id\": \"\", \"auto_tags\": 0}}", - "email": "{\"File\": {}, \"Parser:BirdsFlutterHigh\": {\"email\": {\"fields\": [\"from\", \"to\", \"cc\", \"bcc\", \"date\", \"subject\", \"body\", \"attachments\"], \"output_format\": \"text\", \"preprocess\": [\"main_content\"], \"suffix\": [\"eml\"]}}, \"TokenChunker:WarmBreadSmells\": {\"children_delimiters\": [], \"chunk_token_size\": 512, \"delimiter_mode\": \"token_size\", \"delimiters\": [\"\\n\", \"!\", \"?\", \"。\", \"；\", \"！\", \"？\"], \"image_context_size\": 0, \"overlapped_percent\": 0, \"table_context_size\": 0}, \"Tokenizer:NiceWordsSpoken\": {\"fields\": \"text\", \"filename_embd_weight\": 0.1, \"search_method\": [\"embedding\", \"full_text\"]}, \"Extractor:AutoExtractDefault\": {\"field_name\": \"\", \"auto_keywords\": 0, \"auto_questions\": 0, \"llm_id\": \"\", \"auto_tags\": 0}}", - "general": "{\"File\": {}, \"Parser:HipSignsRhyme\": {\"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\", \"pages\": [[1, 100000]], \"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}, \"Tokenizer:LegalReadersDecide\": {\"fields\": \"text\", \"filename_embd_weight\": 0.1, \"search_method\": [\"embedding\", \"full_text\"]}, \"Extractor:AutoExtractDefault\": {\"field_name\": \"\", \"auto_keywords\": 0, \"auto_questions\": 0, \"llm_id\": \"\", \"auto_tags\": 0}}", - "laws": "{\"File\": {}, \"Parser:HipSignsRhyme\": {\"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\"], \"remove_toc\": true, \"suffix\": [\"pdf\"], \"vlm\": {}}, \"text&code\": {\"output_format\": \"json\", \"preprocess\": [\"main_content\"], \"suffix\": [\"txt\"]}}, \"TitleChunker:SpicyKeysKick\": {\"hierarchy\": 2, \"include_heading_content\": false, \"levels\": [[\"^#[^#]\", \"^##[^#]\", \"^###[^#]\", \"^####[^#]\"], [\"第[零一二三四五六七八九十百0-9]+(分?编|部分)\", \"第[零一二三四五六七八九十百0-9]+章\", \"第[零一二三四五六七八九十百0-9]+节\", \"第[零一二三四五六七八九十百0-9]+条\", \"[\\\\(（][零一二三四五六七八九十百]+[\\\\)）]\"], [\"第[0-9]+章\", \"第[0-9]+节\", \"[0-9]{1,2}[\\\\. 、]\", \"[0-9]{1,2}\\\\.[0-9]{1,2}($|[^a-zA-Z/%~.-])\", \"[0-9]{1,2}\\\\.[0-9]{1,2}\\\\.[0-9]{1,2}\"], [\"第[零一二三四五六七八九十百0-9]+章\", \"第[零一二三四五六七八九十百0-9]+节\", \"[零一二三四五六七八九十百]+[ 、]\", \"[\\\\(（][零一二三四五六七八九十百]+[\\\\)）]\", \"[\\\\(（][0-9]{,2}[\\\\)）]\"], [\"PART (ONE|TWO|THREE|FOUR|FIVE|SIX|SEVEN|EIGHT|NINE|TEN)\", \"Chapter (I+V?|VI*|XI|IX|X)\", \"Section [0-9]+\", \"Article [0-9]+\"]], \"method\": \"hierarchy\"}, \"Tokenizer:PublicJobsTake\": {\"fields\": \"text\", \"filename_embd_weight\": 0.1, \"search_method\": [\"embedding\", \"full_text\"]}, \"Extractor:AutoExtractDefault\": {\"field_name\": \"\", \"auto_keywords\": 0, \"auto_questions\": 0, \"llm_id\": \"\", \"auto_tags\": 0}}", - "manual": "{\"File\": {}, \"Parser:HipSignsRhyme\": {\"doc\": {\"output_format\": \"json\", \"preprocess\": [\"main_content\"], \"suffix\": [\"doc\"]}, \"docx\": {\"flatten_media_to_text\": false, \"output_format\": \"json\", \"preprocess\": [\"main_content\"], \"suffix\": [\"docx\"], \"vlm\": {}}, \"pdf\": {\"flatten_media_to_text\": false, \"output_format\": \"json\", \"parse_method\": \"DeepDOC\", \"preprocess\": [\"main_content\"], \"suffix\": [\"pdf\"], \"vlm\": {}}}, \"TitleChunker:NineInsectsFind\": {\"hierarchy\": 0, \"include_heading_content\": false, \"levels\": [[\"^#[^#]\", \"^##[^#]\", \"^###[^#]\", \"^####[^#]\"], [\"第[零一二三四五六七八九十百0-9]+(分?编|部分)\", \"第[零一二三四五六七八九十百0-9]+章\", \"第[零一二三四五六七八九十百0-9]+节\", \"第[零一二三四五六七八九十百0-9]+条\", \"[\\\\(（][零一二三四五六七八九十百]+[\\\\)）]\"], [\"第[0-9]+章\", \"第[0-9]+节\", \"[0-9]{1,2}[\\\\. 、]\", \"[0-9]{1,2}\\\\.[0-9]{1,2}($|[^a-zA-Z/%~.-])\", \"[0-9]{1,2}\\\\.[0-9]{1,2}\\\\.[0-9]{1,2}\"], [\"第[零一二三四五六七八九十百0-9]+章\", \"第[零一二三四五六七八九十百0-9]+节\", \"[零一二三四五六七八九十百]+[ 、]\", \"[\\\\(（][零一二三四五六七八九十百]+[\\\\)）]\", \"[\\\\(（][0-9]{,2}[\\\\)）]\"], [\"PART (ONE|TWO|THREE|FOUR|FIVE|SIX|SEVEN|EIGHT|NINE|TEN)\", \"Chapter (I+V?|VI*|XI|IX|X)\", \"Section [0-9]+\", \"Article [0-9]+\"]], \"method\": \"group\"}, \"Tokenizer:FunnyBalloonsGrin\": {\"fields\": \"text\", \"filename_embd_weight\": 0.1, \"search_method\": [\"embedding\", \"full_text\"]}, \"Extractor:AutoExtractDefault\": {\"field_name\": \"\", \"auto_keywords\": 0, \"auto_questions\": 0, \"llm_id\": \"\", \"auto_tags\": 0}}", - "one": "{\"File\": {}, \"Parser:HipSignsRhyme\": {\"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\"], \"vlm\": {}}, \"text&code\": {\"output_format\": \"json\", \"preprocess\": [\"main_content\"], \"suffix\": [\"txt\"]}}, \"OneChunker:DryDrinksVisit\": {}, \"Tokenizer:FrankWeeksListen\": {\"fields\": \"text\", \"filename_embd_weight\": 0.1, \"search_method\": [\"embedding\", \"full_text\"]}, \"Extractor:AutoExtractDefault\": {\"field_name\": \"\", \"auto_keywords\": 0, \"auto_questions\": 0, \"llm_id\": \"\", \"auto_tags\": 0}}", - "paper": "{\"File\": {}, \"Parser:HipSignsRhyme\": {\"pdf\": {\"enable_multi_column\": true, \"flatten_media_to_text\": false, \"output_format\": \"json\", \"parse_method\": \"DeepDOC\", \"preprocess\": [\"main_content\"], \"suffix\": [\"pdf\"], \"vlm\": {}}}, \"TitleChunker:SparklySchoolsTravel\": {\"hierarchy\": 0, \"include_heading_content\": false, \"levels\": [[\"^#[^#]\", \"^##[^#]\", \"^###[^#]\", \"^####[^#]\"], [\"第[零一二三四五六七八九十百0-9]+(分?编|部分)\", \"第[零一二三四五六七八九十百0-9]+章\", \"第[零一二三四五六七八九十百0-9]+节\", \"第[零一二三四五六七八九十百0-9]+条\", \"[\\\\(（][零一二三四五六七八九十百]+[\\\\)）]\"], [\"第[0-9]+章\", \"第[0-9]+节\", \"[0-9]{1,2}[\\\\. 、]\", \"[0-9]{1,2}\\\\.[0-9]{1,2}($|[^a-zA-Z/%~.-])\", \"[0-9]{1,2}\\\\.[0-9]{1,2}\\\\.[0-9]{1,2}\"], [\"第[零一二三四五六七八九十百0-9]+章\", \"第[零一二三四五六七八九十百0-9]+节\", \"[零一二三四五六七八九十百]+[ 、]\", \"[\\\\(（][零一二三四五六七八九十百]+[\\\\)）]\", \"[\\\\(（][0-9]{,2}[\\\\)）]\"], [\"PART (ONE|TWO|THREE|FOUR|FIVE|SIX|SEVEN|EIGHT|NINE|TEN)\", \"Chapter (I+V?|VI*|XI|IX|X)\", \"Section [0-9]+\", \"Article [0-9]+\"]], \"method\": \"group\"}, \"Tokenizer:GreatCarsWash\": {\"fields\": \"text\", \"filename_embd_weight\": 0.1, \"search_method\": [\"embedding\", \"full_text\"]}, \"Extractor:AutoExtractDefault\": {\"field_name\": \"\", \"auto_keywords\": 0, \"auto_questions\": 0, \"llm_id\": \"\", \"auto_tags\": 0}}", - "picture": "{\"File\": {}, \"Parser:ViewsCaptureLight\": {\"image\": {\"output_format\": \"json\", \"parse_method\": \"ocr\", \"preprocess\": [\"main_content\"], \"suffix\": [\"bmp\", \"gif\", \"jpeg\", \"jpg\", \"png\", \"svg\", \"tif\", \"tiff\", \"webp\"]}, \"video\": {\"output_format\": \"text\", \"preprocess\": [\"main_content\"], \"suffix\": [\"3gp\", \"3gpp\", \"avi\", \"flv\", \"mkv\", \"mov\", \"mp4\", \"mpeg\", \"mpg\", \"webm\", \"wmv\"]}}, \"TokenChunker:BrightColorsGlow\": {}, \"Tokenizer:SharpLensFocus\": {\"fields\": \"text\", \"filename_embd_weight\": 0.1, \"search_method\": [\"embedding\", \"full_text\"]}, \"Extractor:AutoExtractDefault\": {\"field_name\": \"\", \"auto_keywords\": 0, \"auto_questions\": 0, \"llm_id\": \"\", \"auto_tags\": 0}}", - "presentation": "{\"File\": {}, \"Parser:HipSignsRhyme\": {\"pdf\": {\"flatten_media_to_text\": false, \"output_format\": \"json\", \"parse_method\": \"DeepDOC\", \"preprocess\": [\"main_content\"], \"suffix\": [\"pdf\"], \"vlm\": {}}, \"slides\": {\"output_format\": \"json\", \"parse_method\": \"DeepDOC\", \"preprocess\": [\"main_content\"], \"suffix\": [\"pptx\", \"ppt\"]}}, \"Tokenizer:TallTreesDance\": {\"fields\": \"text\", \"filename_embd_weight\": 0.1, \"search_method\": [\"embedding\", \"full_text\"]}, \"PresentationChunker:HappyHillsGlow\": {}, \"Extractor:AutoExtractDefault\": {\"field_name\": \"\", \"auto_keywords\": 0, \"auto_questions\": 0, \"llm_id\": \"\", \"auto_tags\": 0}}", - "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}, \"File\": {}, \"Parser:HipSignsRhyme\": {\"docx\": {\"flatten_media_to_text\": true, \"output_format\": \"json\", \"preprocess\": [\"main_content\"], \"suffix\": [\"docx\"], \"vlm\": {}}, \"pdf\": {\"flatten_media_to_text\": true, \"output_format\": \"json\", \"parse_method\": \"DeepDOC\", \"preprocess\": [\"main_content\"], \"suffix\": [\"pdf\"], \"vlm\": {}}, \"text&code\": {\"output_format\": \"json\", \"preprocess\": [\"main_content\"], \"suffix\": [\"txt\"]}}, \"TitleChunker:FlatMiceFix\": {\"hierarchy\": 1, \"include_heading_content\": false, \"levels\": [[\"^\\\\s*(?i:(?:\\\\d+[\\\\.\\\\)]\\\\s*)?(?:EDUCATION|ACADEMIC\\\\s*BACKGROUND|ACADEMIC\\\\s*HISTORY|EDUCATIONAL\\\\s*BACKGROUND|RELEVANT\\\\s*COURSEWORK|COURSEWORK|EXPERIENCE|WORK\\\\s*EXPERIENCE|PROFESSIONAL\\\\s*EXPERIENCE|RELEVANT\\\\s*EXPERIENCE|EMPLOYMENT\\\\s*HISTORY|CAREER\\\\s*HISTORY|INTERNSHIP\\\\s*EXPERIENCE|PROJECTS|PROJECT\\\\s*EXPERIENCE|ACADEMIC\\\\s*PROJECTS|PROFESSIONAL\\\\s*PROJECTS|SKILLS|TECHNICAL\\\\s*SKILLS|CORE\\\\s*COMPETENCIES|COMPETENCIES|QUALIFICATIONS|SUMMARY\\\\s*OF\\\\s*QUALIFICATIONS|CERTIFICATIONS|LICENSES|CERTIFICATES|AWARDS|HONORS|HONOURS|ACHIEVEMENTS|PUBLICATIONS|RESEARCH|RESEARCH\\\\s*EXPERIENCE|LEADERSHIP|LEADERSHIP\\\\s*EXPERIENCE|ACTIVITIES|EXTRACURRICULAR\\\\s*ACTIVITIES|ACTIVITIES\\\\s*(?:&|AND)\\\\s*SKILLS|INVOLVEMENT|CAMPUS\\\\s*INVOLVEMENT|VOLUNTEER\\\\s*EXPERIENCE|VOLUNTEERING|COMMUNITY\\\\s*SERVICE|LANGUAGES|INTERESTS|HOBBIES|PROFILE|PROFESSIONAL\\\\s*PROFILE|SUMMARY|PROFESSIONAL\\\\s*SUMMARY|CAREER\\\\s*SUMMARY|OBJECTIVE|CAREER\\\\s*OBJECTIVE|PERSONAL\\\\s*INFORMATION|CONTACT\\\\s*INFORMATION|ADDITIONAL\\\\s*INFORMATION|TRAINING))\\\\s*[:：]?\\\\s*$\"], [\"^\\\\s*(?:\\\\d+[\\\\.、\\\\)]\\\\s*)?(?:教育背景|教育经历|学历背景|学术背景|技术背景|工作经历|工作经验|实习经历|项目经历|项目经验|科研经历|研究经历|校园经历|实践经历|专业经历|职业经历|技能|专业技能|技能特长|核心技能|技术栈|个人技能|工作技能|职业技能|技能与评价|技能与自我评价|工作技能与自我评价|职业技能与自我评价|证书|资格证书|职业资格|资质证书|获奖情况|获奖经历|荣誉|荣誉奖项|奖项|科研成果|论文发表|发表论文|领导经历|学生工作|校园活动|社团经历|活动经历|志愿经历|志愿服务|社会实践|语言能力|语言|自我评价|个人评价|自我总结|个人总结|个人优势|个人简介|个人信息|基本信息|联系方式|求职意向|应聘意向|职业目标|求职目标|兴趣爱好|兴趣特长|培训经历|其他信息|附加信息)\\\\s*[:：]?\\\\s*$\"]], \"method\": \"hierarchy\"}, \"Tokenizer:KindHandsWin\": {\"fields\": \"text\", \"filename_embd_weight\": 0.1, \"search_method\": [\"embedding\", \"full_text\"]}}", - "table": "{\"File\": {}, \"Parser:HipSignsRhyme\": {\"spreadsheet\": {\"flatten_media_to_text\": false, \"output_format\": \"html\", \"parse_method\": \"DeepDOC\", \"preprocess\": [\"main_content\"], \"suffix\": [\"xls\", \"xlsx\", \"csv\"], \"vlm\": {}, \"column_mode\": \"auto\", \"column_roles\": {}, \"column_names\": []}, \"text&code\": {\"output_format\": \"json\", \"preprocess\": [\"main_content\"], \"suffix\": [\"txt\"]}}, \"TableChunker:FastFoxesJump\": {}, \"Tokenizer:DeepLakesShine\": {\"fields\": \"text\", \"filename_embd_weight\": 0.1, \"search_method\": [\"embedding\", \"full_text\"]}}", + "audio": "{\"File\": {}, \"Parser:SongsFillAir\": {\"audio\": {\"output_format\": \"text\", \"preprocess\": [\"main_content\"], \"suffix\": [\"aac\", \"aiff\", \"ape\", \"au\", \"da\", \"flac\", \"midi\", \"mp3\", \"ogg\", \"oggvorbis\", \"realaudio\", \"vqf\", \"wav\", \"wave\", \"wma\"]}}, \"TokenChunker:BlueSkiesLaugh\": {}, \"Tokenizer:KindEyesWatch\": {\"fields\": \"text\", \"filename_embd_weight\": 0.1, \"search_method\": [\"embedding\", \"full_text\"]}, \"Extractor:AutoExtractDefault\": {\"field_name\": \"\", \"auto_keywords\": 0, \"auto_questions\": 0, \"llm_id\": \"\", \"auto_tags\": 0, \"enable_metadata\": 0, \"metadata\": []}}", + "book": "{\"File\": {}, \"Parser:HipSignsRhyme\": {\"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\"]}, \"pdf\": {\"flatten_media_to_text\": false, \"output_format\": \"json\", \"parse_method\": \"DeepDOC\", \"preprocess\": [\"main_content\"], \"remove_toc\": true, \"suffix\": [\"pdf\"], \"vlm\": {}}, \"text&code\": {\"output_format\": \"json\", \"preprocess\": [\"main_content\"], \"suffix\": [\"txt\"]}}, \"TitleChunker:GrumpyGarlicsBake\": {\"hierarchy\": 5, \"include_heading_content\": true, \"levels\": [[\"^#[^#]\", \"^##[^#]\", \"^###[^#]\", \"^####[^#]\"], [\"第[零一二三四五六七八九十百0-9]+(分?编|部分)\", \"第[零一二三四五六七八九十百0-9]+章\", \"第[零一二三四五六七八九十百0-9]+节\", \"第[零一二三四五六七八九十百0-9]+条\", \"[\\\\(（][零一二三四五六七八九十百]+[\\\\)）]\"], [\"第[0-9]+章\", \"第[0-9]+节\", \"[0-9]{1,2}[\\\\. 、]\", \"[0-9]{1,2}\\\\.[0-9]{1,2}($|[^a-zA-Z/%~.-])\", \"[0-9]{1,2}\\\\.[0-9]{1,2}\\\\.[0-9]{1,2}\"], [\"第[零一二三四五六七八九十百0-9]+章\", \"第[零一二三四五六七八九十百0-9]+节\", \"[零一二三四五六七八九十百]+[ 、]\", \"[\\\\(（][零一二三四五六七八九十百]+[\\\\)）]\", \"[\\\\(（][0-9]{,2}[\\\\)）]\"], [\"PART (ONE|TWO|THREE|FOUR|FIVE|SIX|SEVEN|EIGHT|NINE|TEN)\", \"Chapter (I+V?|VI*|XI|IX|X)\", \"Section [0-9]+\", \"Article [0-9]+\"]], \"method\": \"hierarchy\"}, \"Tokenizer:HotDonutsRing\": {\"fields\": \"text\", \"filename_embd_weight\": 0.1, \"search_method\": [\"embedding\", \"full_text\"]}, \"Extractor:AutoExtractDefault\": {\"field_name\": \"\", \"auto_keywords\": 0, \"auto_questions\": 0, \"llm_id\": \"\", \"auto_tags\": 0, \"enable_metadata\": 0, \"metadata\": []}}", + "email": "{\"File\": {}, \"Parser:BirdsFlutterHigh\": {\"email\": {\"fields\": [\"from\", \"to\", \"cc\", \"bcc\", \"date\", \"subject\", \"body\", \"attachments\"], \"output_format\": \"text\", \"preprocess\": [\"main_content\"], \"suffix\": [\"eml\"]}}, \"TokenChunker:WarmBreadSmells\": {\"children_delimiters\": [], \"chunk_token_size\": 512, \"delimiter_mode\": \"token_size\", \"delimiters\": [\"\\n\", \"!\", \"?\", \"。\", \"；\", \"！\", \"？\"], \"image_context_size\": 0, \"overlapped_percent\": 0, \"table_context_size\": 0}, \"Tokenizer:NiceWordsSpoken\": {\"fields\": \"text\", \"filename_embd_weight\": 0.1, \"search_method\": [\"embedding\", \"full_text\"]}, \"Extractor:AutoExtractDefault\": {\"field_name\": \"\", \"auto_keywords\": 0, \"auto_questions\": 0, \"llm_id\": \"\", \"auto_tags\": 0, \"enable_metadata\": 0, \"metadata\": []}}", + "general": "{\"File\": {}, \"Parser:HipSignsRhyme\": {\"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\", \"pages\": [[1, 100000]], \"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}, \"Tokenizer:LegalReadersDecide\": {\"fields\": \"text\", \"filename_embd_weight\": 0.1, \"search_method\": [\"embedding\", \"full_text\"]}, \"Extractor:AutoExtractDefault\": {\"field_name\": \"\", \"auto_keywords\": 0, \"auto_questions\": 0, \"llm_id\": \"\", \"auto_tags\": 0, \"enable_metadata\": 0, \"metadata\": []}}", + "laws": "{\"File\": {}, \"Parser:HipSignsRhyme\": {\"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\"], \"remove_toc\": true, \"suffix\": [\"pdf\"], \"vlm\": {}}, \"text&code\": {\"output_format\": \"json\", \"preprocess\": [\"main_content\"], \"suffix\": [\"txt\"]}}, \"TitleChunker:SpicyKeysKick\": {\"hierarchy\": 2, \"include_heading_content\": false, \"levels\": [[\"^#[^#]\", \"^##[^#]\", \"^###[^#]\", \"^####[^#]\"], [\"第[零一二三四五六七八九十百0-9]+(分?编|部分)\", \"第[零一二三四五六七八九十百0-9]+章\", \"第[零一二三四五六七八九十百0-9]+节\", \"第[零一二三四五六七八九十百0-9]+条\", \"[\\\\(（][零一二三四五六七八九十百]+[\\\\)）]\"], [\"第[0-9]+章\", \"第[0-9]+节\", \"[0-9]{1,2}[\\\\. 、]\", \"[0-9]{1,2}\\\\.[0-9]{1,2}($|[^a-zA-Z/%~.-])\", \"[0-9]{1,2}\\\\.[0-9]{1,2}\\\\.[0-9]{1,2}\"], [\"第[零一二三四五六七八九十百0-9]+章\", \"第[零一二三四五六七八九十百0-9]+节\", \"[零一二三四五六七八九十百]+[ 、]\", \"[\\\\(（][零一二三四五六七八九十百]+[\\\\)）]\", \"[\\\\(（][0-9]{,2}[\\\\)）]\"], [\"PART (ONE|TWO|THREE|FOUR|FIVE|SIX|SEVEN|EIGHT|NINE|TEN)\", \"Chapter (I+V?|VI*|XI|IX|X)\", \"Section [0-9]+\", \"Article [0-9]+\"]], \"method\": \"hierarchy\"}, \"Tokenizer:PublicJobsTake\": {\"fields\": \"text\", \"filename_embd_weight\": 0.1, \"search_method\": [\"embedding\", \"full_text\"]}, \"Extractor:AutoExtractDefault\": {\"field_name\": \"\", \"auto_keywords\": 0, \"auto_questions\": 0, \"llm_id\": \"\", \"auto_tags\": 0, \"enable_metadata\": 0, \"metadata\": []}}", + "manual": "{\"File\": {}, \"Parser:HipSignsRhyme\": {\"doc\": {\"output_format\": \"json\", \"preprocess\": [\"main_content\"], \"suffix\": [\"doc\"]}, \"docx\": {\"flatten_media_to_text\": false, \"output_format\": \"json\", \"preprocess\": [\"main_content\"], \"suffix\": [\"docx\"], \"vlm\": {}}, \"pdf\": {\"flatten_media_to_text\": false, \"output_format\": \"json\", \"parse_method\": \"DeepDOC\", \"preprocess\": [\"main_content\"], \"suffix\": [\"pdf\"], \"vlm\": {}}}, \"TitleChunker:NineInsectsFind\": {\"hierarchy\": 0, \"include_heading_content\": false, \"levels\": [[\"^#[^#]\", \"^##[^#]\", \"^###[^#]\", \"^####[^#]\"], [\"第[零一二三四五六七八九十百0-9]+(分?编|部分)\", \"第[零一二三四五六七八九十百0-9]+章\", \"第[零一二三四五六七八九十百0-9]+节\", \"第[零一二三四五六七八九十百0-9]+条\", \"[\\\\(（][零一二三四五六七八九十百]+[\\\\)）]\"], [\"第[0-9]+章\", \"第[0-9]+节\", \"[0-9]{1,2}[\\\\. 、]\", \"[0-9]{1,2}\\\\.[0-9]{1,2}($|[^a-zA-Z/%~.-])\", \"[0-9]{1,2}\\\\.[0-9]{1,2}\\\\.[0-9]{1,2}\"], [\"第[零一二三四五六七八九十百0-9]+章\", \"第[零一二三四五六七八九十百0-9]+节\", \"[零一二三四五六七八九十百]+[ 、]\", \"[\\\\(（][零一二三四五六七八九十百]+[\\\\)）]\", \"[\\\\(（][0-9]{,2}[\\\\)）]\"], [\"PART (ONE|TWO|THREE|FOUR|FIVE|SIX|SEVEN|EIGHT|NINE|TEN)\", \"Chapter (I+V?|VI*|XI|IX|X)\", \"Section [0-9]+\", \"Article [0-9]+\"]], \"method\": \"group\"}, \"Tokenizer:FunnyBalloonsGrin\": {\"fields\": \"text\", \"filename_embd_weight\": 0.1, \"search_method\": [\"embedding\", \"full_text\"]}, \"Extractor:AutoExtractDefault\": {\"field_name\": \"\", \"auto_keywords\": 0, \"auto_questions\": 0, \"llm_id\": \"\", \"auto_tags\": 0, \"enable_metadata\": 0, \"metadata\": []}}", + "one": "{\"File\": {}, \"Parser:HipSignsRhyme\": {\"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\"], \"vlm\": {}}, \"text&code\": {\"output_format\": \"json\", \"preprocess\": [\"main_content\"], \"suffix\": [\"txt\"]}}, \"OneChunker:DryDrinksVisit\": {}, \"Tokenizer:FrankWeeksListen\": {\"fields\": \"text\", \"filename_embd_weight\": 0.1, \"search_method\": [\"embedding\", \"full_text\"]}, \"Extractor:AutoExtractDefault\": {\"field_name\": \"\", \"auto_keywords\": 0, \"auto_questions\": 0, \"llm_id\": \"\", \"auto_tags\": 0, \"enable_metadata\": 0, \"metadata\": []}}", + "paper": "{\"File\": {}, \"Parser:HipSignsRhyme\": {\"pdf\": {\"enable_multi_column\": true, \"flatten_media_to_text\": false, \"output_format\": \"json\", \"parse_method\": \"DeepDOC\", \"preprocess\": [\"main_content\"], \"suffix\": [\"pdf\"], \"vlm\": {}}}, \"TitleChunker:SparklySchoolsTravel\": {\"hierarchy\": 0, \"include_heading_content\": false, \"levels\": [[\"^#[^#]\", \"^##[^#]\", \"^###[^#]\", \"^####[^#]\"], [\"第[零一二三四五六七八九十百0-9]+(分?编|部分)\", \"第[零一二三四五六七八九十百0-9]+章\", \"第[零一二三四五六七八九十百0-9]+节\", \"第[零一二三四五六七八九十百0-9]+条\", \"[\\\\(（][零一二三四五六七八九十百]+[\\\\)）]\"], [\"第[0-9]+章\", \"第[0-9]+节\", \"[0-9]{1,2}[\\\\. 、]\", \"[0-9]{1,2}\\\\.[0-9]{1,2}($|[^a-zA-Z/%~.-])\", \"[0-9]{1,2}\\\\.[0-9]{1,2}\\\\.[0-9]{1,2}\"], [\"第[零一二三四五六七八九十百0-9]+章\", \"第[零一二三四五六七八九十百0-9]+节\", \"[零一二三四五六七八九十百]+[ 、]\", \"[\\\\(（][零一二三四五六七八九十百]+[\\\\)）]\", \"[\\\\(（][0-9]{,2}[\\\\)）]\"], [\"PART (ONE|TWO|THREE|FOUR|FIVE|SIX|SEVEN|EIGHT|NINE|TEN)\", \"Chapter (I+V?|VI*|XI|IX|X)\", \"Section [0-9]+\", \"Article [0-9]+\"]], \"method\": \"group\"}, \"Tokenizer:GreatCarsWash\": {\"fields\": \"text\", \"filename_embd_weight\": 0.1, \"search_method\": [\"embedding\", \"full_text\"]}, \"Extractor:AutoExtractDefault\": {\"field_name\": \"\", \"auto_keywords\": 0, \"auto_questions\": 0, \"llm_id\": \"\", \"auto_tags\": 0, \"enable_metadata\": 0, \"metadata\": []}}", + "picture": "{\"File\": {}, \"Parser:ViewsCaptureLight\": {\"image\": {\"output_format\": \"json\", \"parse_method\": \"ocr\", \"preprocess\": [\"main_content\"], \"suffix\": [\"bmp\", \"gif\", \"jpeg\", \"jpg\", \"png\", \"svg\", \"tif\", \"tiff\", \"webp\"]}, \"video\": {\"output_format\": \"text\", \"preprocess\": [\"main_content\"], \"suffix\": [\"3gp\", \"3gpp\", \"avi\", \"flv\", \"mkv\", \"mov\", \"mp4\", \"mpeg\", \"mpg\", \"webm\", \"wmv\"]}}, \"TokenChunker:BrightColorsGlow\": {}, \"Tokenizer:SharpLensFocus\": {\"fields\": \"text\", \"filename_embd_weight\": 0.1, \"search_method\": [\"embedding\", \"full_text\"]}, \"Extractor:AutoExtractDefault\": {\"field_name\": \"\", \"auto_keywords\": 0, \"auto_questions\": 0, \"llm_id\": \"\", \"auto_tags\": 0, \"enable_metadata\": 0, \"metadata\": []}}", + "presentation": "{\"File\": {}, \"Parser:HipSignsRhyme\": {\"pdf\": {\"flatten_media_to_text\": false, \"output_format\": \"json\", \"parse_method\": \"DeepDOC\", \"preprocess\": [\"main_content\"], \"suffix\": [\"pdf\"], \"vlm\": {}}, \"slides\": {\"output_format\": \"json\", \"parse_method\": \"DeepDOC\", \"preprocess\": [\"main_content\"], \"suffix\": [\"pptx\", \"ppt\"]}}, \"Tokenizer:TallTreesDance\": {\"fields\": \"text\", \"filename_embd_weight\": 0.1, \"search_method\": [\"embedding\", \"full_text\"]}, \"PresentationChunker:HappyHillsGlow\": {}, \"Extractor:AutoExtractDefault\": {\"field_name\": \"\", \"auto_keywords\": 0, \"auto_questions\": 0, \"llm_id\": \"\", \"auto_tags\": 0, \"enable_metadata\": 0, \"metadata\": []}}", + "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\": []}, \"File\": {}, \"Parser:HipSignsRhyme\": {\"docx\": {\"flatten_media_to_text\": true, \"output_format\": \"json\", \"preprocess\": [\"main_content\"], \"suffix\": [\"docx\"], \"vlm\": {}}, \"pdf\": {\"flatten_media_to_text\": true, \"output_format\": \"json\", \"parse_method\": \"DeepDOC\", \"preprocess\": [\"main_content\"], \"suffix\": [\"pdf\"], \"vlm\": {}}, \"text&code\": {\"output_format\": \"json\", \"preprocess\": [\"main_content\"], \"suffix\": [\"txt\"]}}, \"TitleChunker:FlatMiceFix\": {\"hierarchy\": 1, \"include_heading_content\": false, \"levels\": [[\"^\\\\s*(?i:(?:\\\\d+[\\\\.\\\\)]\\\\s*)?(?:EDUCATION|ACADEMIC\\\\s*BACKGROUND|ACADEMIC\\\\s*HISTORY|EDUCATIONAL\\\\s*BACKGROUND|RELEVANT\\\\s*COURSEWORK|COURSEWORK|EXPERIENCE|WORK\\\\s*EXPERIENCE|PROFESSIONAL\\\\s*EXPERIENCE|RELEVANT\\\\s*EXPERIENCE|EMPLOYMENT\\\\s*HISTORY|CAREER\\\\s*HISTORY|INTERNSHIP\\\\s*EXPERIENCE|PROJECTS|PROJECT\\\\s*EXPERIENCE|ACADEMIC\\\\s*PROJECTS|PROFESSIONAL\\\\s*PROJECTS|SKILLS|TECHNICAL\\\\s*SKILLS|CORE\\\\s*COMPETENCIES|COMPETENCIES|QUALIFICATIONS|SUMMARY\\\\s*OF\\\\s*QUALIFICATIONS|CERTIFICATIONS|LICENSES|CERTIFICATES|AWARDS|HONORS|HONOURS|ACHIEVEMENTS|PUBLICATIONS|RESEARCH|RESEARCH\\\\s*EXPERIENCE|LEADERSHIP|LEADERSHIP\\\\s*EXPERIENCE|ACTIVITIES|EXTRACURRICULAR\\\\s*ACTIVITIES|ACTIVITIES\\\\s*(?:&|AND)\\\\s*SKILLS|INVOLVEMENT|CAMPUS\\\\s*INVOLVEMENT|VOLUNTEER\\\\s*EXPERIENCE|VOLUNTEERING|COMMUNITY\\\\s*SERVICE|LANGUAGES|INTERESTS|HOBBIES|PROFILE|PROFESSIONAL\\\\s*PROFILE|SUMMARY|PROFESSIONAL\\\\s*SUMMARY|CAREER\\\\s*SUMMARY|OBJECTIVE|CAREER\\\\s*OBJECTIVE|PERSONAL\\\\s*INFORMATION|CONTACT\\\\s*INFORMATION|ADDITIONAL\\\\s*INFORMATION|TRAINING))\\\\s*[:：]?\\\\s*$\"], [\"^\\\\s*(?:\\\\d+[\\\\.、\\\\)]\\\\s*)?(?:教育背景|教育经历|学历背景|学术背景|技术背景|工作经历|工作经验|实习经历|项目经历|项目经验|科研经历|研究经历|校园经历|实践经历|专业经历|职业经历|技能|专业技能|技能特长|核心技能|技术栈|个人技能|工作技能|职业技能|技能与评价|技能与自我评价|工作技能与自我评价|职业技能与自我评价|证书|资格证书|职业资格|资质证书|获奖情况|获奖经历|荣誉|荣誉奖项|奖项|科研成果|论文发表|发表论文|领导经历|学生工作|校园活动|社团经历|活动经历|志愿经历|志愿服务|社会实践|语言能力|语言|自我评价|个人评价|自我总结|个人总结|个人优势|个人简介|个人信息|基本信息|联系方式|求职意向|应聘意向|职业目标|求职目标|兴趣爱好|兴趣特长|培训经历|其他信息|附加信息)\\\\s*[:：]?\\\\s*$\"]], \"method\": \"hierarchy\"}, \"Tokenizer:KindHandsWin\": {\"fields\": \"text\", \"filename_embd_weight\": 0.1, \"search_method\": [\"embedding\", \"full_text\"]}}", + "table": "{\"File\": {}, \"Parser:HipSignsRhyme\": {\"spreadsheet\": {\"flatten_media_to_text\": false, \"output_format\": \"html\", \"parse_method\": \"DeepDOC\", \"preprocess\": [\"main_content\"], \"suffix\": [\"xls\", \"xlsx\", \"csv\"], \"vlm\": {}, \"column_mode\": \"auto\", \"column_roles\": {}, \"column_names\": []}, \"text&code\": {\"output_format\": \"json\", \"preprocess\": [\"main_content\"], \"suffix\": [\"txt\"]}}, \"TableChunker:FastFoxesJump\": {}, \"Tokenizer:DeepLakesShine\": {\"fields\": \"text\", \"filename_embd_weight\": 0.1, \"search_method\": [\"embedding\", \"full_text\"]}}", + "knowledge_compiler": "{\"File\": {}, \"KnowledgeCompiler:KnownSwiftLions\": {\"language\": \"English\", \"variant\": \"structure\"}, \"Parser:HipSignsRhyme\": {\"setups\": {\"doc\": {\"output_format\": \"json\", \"preprocess\": [\"main_content\"], \"suffix\": [\"doc\"]}, \"docx\": {\"flatten_media_to_text\": false, \"output_format\": \"json\", \"preprocess\": [\"main_content\"], \"suffix\": [\"docx\"], \"vlm\": {}}, \"html\": {\"output_format\": \"json\", \"preprocess\": [\"main_content\"], \"suffix\": [\"htm\", \"html\"]}, \"markdown\": {\"flatten_media_to_text\": false, \"output_format\": \"json\", \"preprocess\": [\"main_content\"], \"suffix\": [\"md\", \"markdown\", \"mdx\"], \"vlm\": {}}, \"pdf\": {\"flatten_media_to_text\": false, \"output_format\": \"json\", \"parse_method\": \"DeepDOC\", \"preprocess\": [\"main_content\"], \"suffix\": [\"pdf\"], \"vlm\": {}}, \"spreadsheet\": {\"flatten_media_to_text\": false, \"output_format\": \"html\", \"parse_method\": \"DeepDOC\", \"preprocess\": [\"main_content\"], \"suffix\": [\"xls\", \"xlsx\", \"csv\"], \"vlm\": {}}, \"text&code\": {\"output_format\": \"json\", \"preprocess\": [\"main_content\"], \"suffix\": [\"txt\", \"py\", \"js\", \"java\", \"c\", \"cpp\", \"h\", \"php\", \"go\", \"ts\", \"sh\", \"cs\", \"kt\", \"sql\"]}}}, \"TokenChunker:SixApplesFall\": {\"children_delimiters\": [], \"chunk_token_size\": 512, \"delimiter_mode\": \"token_size\", \"delimiters\": [\"\\n\", \"!\", \"?\", \"。\", \"；\", \"！\", \"？\"], \"image_context_size\": 0, \"overlapped_percent\": 0, \"table_context_size\": 0}}", } // Per-template test methods. Each resolves default component params from a @@ -60,6 +61,9 @@ func TestBuildComponentParams_Audio(t *testing.T) { assertTemplateComponentPar func TestBuildComponentParams_Book(t *testing.T) { assertTemplateComponentParams(t, "book") } func TestBuildComponentParams_Email(t *testing.T) { assertTemplateComponentParams(t, "email") } func TestBuildComponentParams_General(t *testing.T) { assertTemplateComponentParams(t, "general") } +func TestBuildComponentParams_KnowledgeCompiler(t *testing.T) { + assertTemplateComponentParams(t, "knowledge_compiler") +} func TestBuildComponentParams_Laws(t *testing.T) { assertTemplateComponentParams(t, "laws") } func TestBuildComponentParams_Manual(t *testing.T) { assertTemplateComponentParams(t, "manual") } func TestBuildComponentParams_One(t *testing.T) { assertTemplateComponentParams(t, "one") } diff --git a/internal/ingestion/task/pipeline_executor_test.go b/internal/ingestion/task/pipeline_executor_test.go index 67692f9ad7..296063fc01 100644 --- a/internal/ingestion/task/pipeline_executor_test.go +++ b/internal/ingestion/task/pipeline_executor_test.go @@ -22,6 +22,34 @@ import ( func strPtr(s string) *string { return &s } +// TestMarkCompiledProductsHidden verifies the pipeline caller hides +// per-document compiled knowledge products (compile_kwd present) as +// available_int=0 while leaving ordinary source chunks searchable +// (available_int=1, the index default). Merged dataset-level products are written +// by the consumer and never reach this path, so they are never double-marked. +func TestMarkCompiledProductsHidden(t *testing.T) { + chunks := []map[string]any{ + {"id": "src-1", "content_with_weight": "ordinary source chunk"}, + {"id": "struct-1", "compile_kwd": "structure", "content_with_weight": "entity A"}, + {"id": "wiki-1", "compile_kwd": "artifact_page", "content_with_weight": "page X"}, + {"id": "src-2", "content_with_weight": "another source chunk"}, + } + markCompiledProductsHidden(chunks) + + if v, ok := chunks[0]["available_int"]; ok { + t.Fatalf("ordinary source chunk should keep default available_int, got %v", v) + } + if chunks[1]["available_int"] != 0 { + t.Fatalf("compiled structure chunk should be available_int=0, got %v", chunks[1]["available_int"]) + } + if chunks[2]["available_int"] != 0 { + t.Fatalf("compiled wiki chunk should be available_int=0, got %v", chunks[2]["available_int"]) + } + if v, ok := chunks[3]["available_int"]; ok { + t.Fatalf("source chunk without compile_kwd should keep default available_int, got %v", v) + } +} + func makeTaskCtx() *TaskContext { return &TaskContext{ IngestionTask: &entity.IngestionTask{ diff --git a/internal/service/document/document_crud.go b/internal/service/document/document_crud.go index c0a16f5ad0..0034e2c12d 100644 --- a/internal/service/document/document_crud.go +++ b/internal/service/document/document_crud.go @@ -4,10 +4,12 @@ import ( "context" "errors" "fmt" + "time" "ragflow/internal/common" "ragflow/internal/dao" "ragflow/internal/entity" + "ragflow/internal/ingestion/knowledge_compile" "ragflow/internal/storage" "gorm.io/gorm" @@ -357,6 +359,18 @@ func (s *DocumentService) deleteDocEngineData(docID, tenantID, kbID string) { if _, delErr := s.docEngine.DeleteChunks(ctx, map[string]interface{}{"doc_id": docID}, indexName, kbID); delErr != nil { common.Logger.Warn(fmt.Sprintf("deleteDocEngineData: failed to delete chunks for %s: %v", docID, delErr)) } + // Notify the dataset-level post-processing consumer (§11) that this document's + // source + per-doc compiled chunks are gone. The consumer removes the + // dataset-scoped merged products and triggers incremental re-dedup. Best- + // effort, non-fatal: the synchronous deletion above already removed the + // source/per-doc chunks, so the consumer only owns merged-product cleanup. + // Bound the publish with a timeout so a stalled scheduler (MySQL/NATS) can + // never block the document delete, which already succeeded above. + pubCtx, cancel := context.WithTimeout(context.Background(), 15*time.Second) + defer cancel() + if err := knowledge_compile.PublishDeleted(pubCtx, tenantID, kbID, docID, 0); err != nil { + common.Logger.Warn(fmt.Sprintf("deleteDocEngineData: publish doc_deleted for %s failed: %v", docID, err)) + } if s.metadataSvc != nil { _ = s.DeleteDocumentAllMetadata(ctx, docID) // logs internally } diff --git a/internal/utility/metadata.go b/internal/utility/metadata.go index 2e68dfea47..49a034a7f3 100644 --- a/internal/utility/metadata.go +++ b/internal/utility/metadata.go @@ -45,22 +45,24 @@ func UpdateMetadataTo(target map[string]any, meta any) map[string]any { continue } - // Merge with existing value - existStr, existIsStr := existing.(string) - newStr, newIsStr := normVal.(string) - existList, existIsList := existing.([]string) - newList, newIsList := normVal.([]string) - - if existIsStr && newIsStr { - // Both strings: convert to list, append - target[k] = dedupeStrings(append([]string{existStr}, newStr)) - } else if existIsStr && newIsList { - target[k] = dedupeStrings(append([]string{existStr}, newList...)) - } else if existIsList && newIsStr { - target[k] = dedupeStrings(append(existList, newStr)) - } else if existIsList && newIsList { - target[k] = dedupeStrings(append(existList, newList...)) + // Merge with existing value, mirroring Python + // common.metadata_utils.update_metadata_to exactly: + // - both lists: extend + dedupe + // - target is list, incoming is scalar: append + dedupe + // - target is scalar (or any non-list): overwrite with the + // incoming value (the stored/merged-in side wins), NOT a list. + targetList, targetIsList := existing.([]string) + normList, normIsList := normVal.([]string) + if targetIsList { + if normIsList { + target[k] = dedupeStrings(append(targetList, normList...)) + } else if s, ok := normVal.(string); ok { + target[k] = dedupeStrings(append(targetList, s)) + } + continue } + // target is a scalar: overwrite with the incoming value. + target[k] = normVal } return target diff --git a/internal/utility/metadata_test.go b/internal/utility/metadata_test.go index 9d96f88082..fa6179886e 100644 --- a/internal/utility/metadata_test.go +++ b/internal/utility/metadata_test.go @@ -25,14 +25,12 @@ func TestUpdateMetadataTo_NewKeysAdded(t *testing.T) { } } -func TestUpdateMetadataTo_ExistingKeyMergedToList(t *testing.T) { +func TestUpdateMetadataTo_ExistingKeyScalarOverwrite(t *testing.T) { + // Python update_metadata_to overwrites a scalar target with the incoming + // (merged-in) value rather than list-ifying it. result := UpdateMetadataTo(map[string]any{"author": "Alice"}, map[string]any{"author": "Bob"}) - tags, ok := result["author"].([]string) - if !ok { - t.Fatalf("author should be []string, got %T", result["author"]) - } - if len(tags) != 2 || tags[0] != "Alice" || tags[1] != "Bob" { - t.Errorf("author = %v, want [Alice Bob]", tags) + if result["author"] != "Bob" { + t.Errorf("author = %v, want Bob (incoming scalar overwrites)", result["author"]) } } @@ -57,25 +55,25 @@ func TestUpdateMetadataTo_ListAppend(t *testing.T) { } } -func TestUpdateMetadataTo_StringToListMerge(t *testing.T) { +func TestUpdateMetadataTo_StringScalarOverwrite(t *testing.T) { + // Scalar target is overwritten by the incoming scalar (mirrors Python). result := UpdateMetadataTo( map[string]any{"tags": "a"}, map[string]any{"tags": "b"}, ) - tags := result["tags"].([]string) - if len(tags) != 2 || tags[0] != "a" || tags[1] != "b" { - t.Errorf("tags = %v, want [a b]", tags) + if result["tags"] != "b" { + t.Errorf("tags = %v, want b", result["tags"]) } } func TestUpdateMetadataTo_DeduplicateList(t *testing.T) { result := UpdateMetadataTo( - map[string]any{"tags": "a"}, - map[string]any{"tags": "a"}, + map[string]any{"tags": []string{"a", "b"}}, + map[string]any{"tags": []string{"b", "c"}}, ) tags := result["tags"].([]string) - if len(tags) != 1 { - t.Errorf("tags should be deduplicated: got %v", tags) + if len(tags) != 3 { + t.Errorf("tags should be deduplicated union: got %v", tags) } } @@ -113,3 +111,21 @@ func TestUpdateMetadataTo_SkipOnlyEmptyStringsList(t *testing.T) { t.Error("all-empty list should be skipped") } } + +// TestUpdateMetadataTo_ScalarScalarKeepsExistingPythonParity pins the exact +// Python update_metadata_to (common/metadata_utils.py:301) scalar semantics for +// the "re-ingest + same key + both sides scalar" scenario: the loop iterates +// the second argument (existing_meta), so when the target already holds a +// scalar the EXISTING (stored) value wins — NOT a [old,new] list. This guards +// against a regression where scalar+scalar would be list-ified. +func TestUpdateMetadataTo_ScalarScalarKeepsExistingPythonParity(t *testing.T) { + // mergeDocMetadata calls UpdateMetadataTo(new, existing), so existing is + // the second argument and must win for a shared scalar key. + result := UpdateMetadataTo( + map[string]any{"author": "NEWLY_EXTRACTED"}, + map[string]any{"author": "STORED"}, + ) + if result["author"] != "STORED" { + t.Errorf("author = %v, want STORED (existing scalar wins, Python parity)", result["author"]) + } +} diff --git a/ragflow_deps/download_go_deps.py b/ragflow_deps/download_go_deps.py index 3d43018a87..6f6e484ee8 100644 --- a/ragflow_deps/download_go_deps.py +++ b/ragflow_deps/download_go_deps.py @@ -5,7 +5,8 @@ # requires-python = ">=3.10" # dependencies = [ # "nltk", -# "huggingface-hub" +# "huggingface-hub", +# "requests", # ] # ///