From 75ac8cec2e6c33062bb6e60445fca855a62bd26b Mon Sep 17 00:00:00 2001 From: Zhichang Yu Date: Thu, 30 Jul 2026 19:56:17 +0800 Subject: [PATCH] Align Go knowledge compiler RAPTOR with Python and drop guardrails (#17573) Port RAPTOR knowledge compilation to match Python: remove Go-only capacity guardrails, fix prompts, add token truncation, response post-processing, retries, replace AHC with watershed. --- .../knowledge_compiler/common/common_test.go | 100 --- .../knowledge_compiler/common/deps.go | 16 +- .../knowledge_compiler/common/types.go | 240 ++----- .../component/knowledge_compiler/component.go | 61 +- .../knowledge_compiler/component_test.go | 145 +--- .../datasetnav/datasetnav.go | 32 +- .../golden/fixtures/raptor_baseline.json | 16 +- .../knowledge_compiler/golden_test.go | 46 +- .../knowledge_compiler/mindmap/mindmap.go | 35 +- .../mindmap/mindmap_test.go | 167 +---- .../knowledge_compiler/raptor/clustering.go | 680 ------------------ .../raptor/clustering_test.go | 84 --- .../knowledge_compiler/raptor/psi.go | 119 --- .../knowledge_compiler/raptor/raptor.go | 504 +++++++++---- .../knowledge_compiler/raptor/raptor_test.go | 148 ++++ .../knowledge_compiler/raptor/watershed.go | 142 ++++ .../raptor/watershed_test.go | 127 ++++ .../knowledge_compiler/structure/structure.go | 26 +- .../component/knowledge_compiler/wiki/wiki.go | 17 +- .../ingestion/knowledge_compile/reader.go | 31 +- .../task/knowledge_compiler_wiring.go | 26 +- .../markdown_to_json.go} | 176 +++-- internal/utility/markdown_to_json_test.go | 177 +++++ 23 files changed, 1311 insertions(+), 1804 deletions(-) delete mode 100644 internal/ingestion/component/knowledge_compiler/raptor/clustering.go delete mode 100644 internal/ingestion/component/knowledge_compiler/raptor/clustering_test.go delete mode 100644 internal/ingestion/component/knowledge_compiler/raptor/psi.go create mode 100644 internal/ingestion/component/knowledge_compiler/raptor/raptor_test.go create mode 100644 internal/ingestion/component/knowledge_compiler/raptor/watershed.go create mode 100644 internal/ingestion/component/knowledge_compiler/raptor/watershed_test.go rename internal/{ingestion/component/knowledge_compiler/mindmap/dictify.go => utility/markdown_to_json.go} (61%) create mode 100644 internal/utility/markdown_to_json_test.go diff --git a/internal/ingestion/component/knowledge_compiler/common/common_test.go b/internal/ingestion/component/knowledge_compiler/common/common_test.go index 5076846fd8..ac9c03d365 100644 --- a/internal/ingestion/component/knowledge_compiler/common/common_test.go +++ b/internal/ingestion/component/knowledge_compiler/common/common_test.go @@ -1,7 +1,6 @@ package common import ( - "context" "testing" ) @@ -97,93 +96,6 @@ func TestLocalize(t *testing.T) { } } -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"}, @@ -205,15 +117,3 @@ func TestPackBatches(t *testing.T) { 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 index f02f4f7ce4..b730f74975 100644 --- a/internal/ingestion/component/knowledge_compiler/common/deps.go +++ b/internal/ingestion/component/knowledge_compiler/common/deps.go @@ -6,6 +6,11 @@ import ( "sync" ) +// DefaultLLMContextLength is the fallback chat-model context window (tokens) +// used for per-cluster text truncation when the model's real max length is +// unknown (e.g. in tests). Mirrors Python's default llm_model.max_length. +const DefaultLLMContextLength = 8192 + // ChatRequest is the minimal surface a variant needs to dispatch one LLM call. type ChatRequest struct { LLMID string @@ -16,8 +21,11 @@ type ChatRequest struct { // 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 + // MaxTokens caps the generated summary length (mirrors Python's + // {"max_tokens": max(self._max_token, 512)} config, issue #10235). + MaxTokens *int + APIKey string + BaseURL string } // ChatResponse holds the LLM's text answer. @@ -96,6 +104,10 @@ type Deps struct { Redis RedisClient // optional (datasetnav) TenantID string DatasetID string + // LLMMaxLength is the chat model's context window in tokens. RAPTOR uses it + // to truncate each cluster's texts so the summary prompt fits the window + // (mirrors Python self._llm_model.max_length). + LLMMaxLength int } // DepsResolver resolves the per-run Deps from a tenant/llm/embedding triple. diff --git a/internal/ingestion/component/knowledge_compiler/common/types.go b/internal/ingestion/component/knowledge_compiler/common/types.go index 1ff645d788..cb8fe3bf4d 100644 --- a/internal/ingestion/component/knowledge_compiler/common/types.go +++ b/internal/ingestion/component/knowledge_compiler/common/types.go @@ -1,7 +1,7 @@ // 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. +// KnowledgeCompiler component: parameter/IO types, 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 @@ -9,7 +9,6 @@ package common import ( - "context" "errors" "fmt" "log" @@ -29,9 +28,8 @@ const ( // 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") + 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 @@ -47,7 +45,6 @@ type Param struct { MaxWorkers int EnableHistoricalDedup bool Extra map[string]any - Guardrails CapacityGuardrails } // Defaults returns a Param populated with safe production fallbacks. @@ -55,12 +52,6 @@ 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 } @@ -75,7 +66,11 @@ type Chunk struct { ID string Text string // "text" field from upstream Content string // "content_with_weight" field from upstream - Meta map[string]any + // Vector is the pre-computed embedding of the chunk, when the caller has + // already embedded it. Variants reuse it instead of re-embedding; a nil + // Vector means "not embedded yet" and the variant computes one on demand. + Vector []float32 + Meta map[string]any } // Inputs is the resolved input set passed to a variant Run. @@ -86,7 +81,6 @@ type Inputs struct { 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. @@ -94,12 +88,6 @@ 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 @@ -112,158 +100,13 @@ type Product struct { Meta map[string]any } -// Outputs is the result of a variant Run. +// Outputs is the result of a variant Run. All compiled products are buffered +// here; the component merges them into the upstream chunk stream. 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) { @@ -390,3 +233,62 @@ func parseStringList(m map[string]any, keys ...string) []string { } return nil } + +// VectorFromChunkMap extracts a pre-computed embedding from a chunk map. The +// upstream pipeline (and the knowledge-compile store) carry the vector under a +// dimension-specific column key "q__vec" — set by the tokenizer via +// SetExtraValue and flattened into the map by ChunkDoc.ToMap. The value may +// arrive as []float32 (store read) or []float64 (pipeline map); both are +// normalised to []float32. A nil result means the chunk has no embedding yet. +// +// dim selects the exact "q__vec" key when the embedding dimension is known +// (dim > 0); otherwise every "q_*_vec" key is considered. A chunk must carry at +// most one embedding vector: multiple q_*_vec fields signal mixed embedding +// models, and because Go map iteration order is nondeterministic, picking among +// them would be unstable across runs — such chunks are rejected with an error. +func VectorFromChunkMap(m map[string]any, dim int) ([]float32, error) { + if dim > 0 { + if v, ok := m[fmt.Sprintf("q_%d_vec", dim)]; ok { + return toFloat32Slice(v), nil + } + } + var keys []string + for k := range m { + if strings.HasPrefix(k, "q_") && strings.HasSuffix(k, "_vec") { + keys = append(keys, k) + } + } + switch len(keys) { + case 0: + return nil, nil + case 1: + return toFloat32Slice(m[keys[0]]), nil + default: + return nil, fmt.Errorf("knowledge_compiler: chunk carries %d embedding vectors %v; expected exactly one", len(keys), keys) + } +} + +// toFloat32Slice normalises a numeric slice (any of []float32, []float64, or +// []any of float64) to []float32, returning nil when the value is not a usable +// vector. +func toFloat32Slice(v any) []float32 { + switch arr := v.(type) { + case []float32: + return arr + case []float64: + out := make([]float32, len(arr)) + for i, x := range arr { + out[i] = float32(x) + } + return out + case []any: + 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/component/knowledge_compiler/component.go b/internal/ingestion/component/knowledge_compiler/component.go index a0b2130750..668908a276 100644 --- a/internal/ingestion/component/knowledge_compiler/component.go +++ b/internal/ingestion/component/knowledge_compiler/component.go @@ -124,19 +124,9 @@ func (c *KnowledgeCompilerComponent) Invoke(ctx context.Context, db *gorm.DB, in 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} - } - + // Stamp the resolved template ids onto every compiled product after the + // variant returns. All products are buffered in out.Products (there is no + // streaming sink path), so the post-Run loop below covers every row (M1). var out common.Outputs switch param.Variant { case common.VariantStructure: @@ -565,6 +555,13 @@ func mergeChunks(inputs map[string]any, compiled []schema.ChunkDoc) map[string]a } // buildInputs converts the runtime inputs map into a typed common.Inputs. +// It is necessary because the pipeline passes components a generic +// map[string]any contract while every variant Run consumes a strongly-typed +// common.Inputs: this function is the single translation seam that decouples +// the dependency-light common package (and thus the variants) from the raw +// serialization shape, and the one place where inputs are validated, defaulted, +// and enriched (e.g. extracting each chunk's pre-computed embedding) before any +// LLM/embedding work begins. func buildInputs(inputs map[string]any, param common.Param) (common.Inputs, error) { in := common.Inputs{ LLMID: param.LLMID, @@ -590,17 +587,22 @@ func buildInputs(inputs map[string]any, param common.Param) (common.Inputs, erro if cw, ok := m["content_with_weight"].(string); ok { ch.Content = cw } + // Reuse the embedding the upstream pipeline already computed on the + // chunk (stored under q__vec); variants fall back to embedding + // on demand when it is absent. A chunk must carry exactly one vector. + vec, err := common.VectorFromChunkMap(m, 0) + if err != nil { + return in, err + } + ch.Vector = vec in.Chunks = append(in.Chunks, ch) } } 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, + "doc_id": true, "chunks": true, "historical_candidates": true, "llm_id": true, "embedding_model": true, "tenant_id": true, "dataset_id": true, } for k, v := range inputs { @@ -611,31 +613,6 @@ func buildInputs(inputs map[string]any, param common.Param) (common.Inputs, erro 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{ diff --git a/internal/ingestion/component/knowledge_compiler/component_test.go b/internal/ingestion/component/knowledge_compiler/component_test.go index 582eb41f9c..f7107cf6d0 100644 --- a/internal/ingestion/component/knowledge_compiler/component_test.go +++ b/internal/ingestion/component/knowledge_compiler/component_test.go @@ -336,8 +336,8 @@ func TestKnowledgeCompiler_Wiki_EndToEnd(t *testing.T) { 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"}}) + // watershed (default tree_order): zero external clustering dependency. + chunks := runVariant(t, "raptor", nil) foundRoot := false for _, c := range chunks { if kind, _ := c["kc_kind"].(string); kind == "root" { @@ -345,29 +345,23 @@ func TestKnowledgeCompiler_Raptor_EndToEnd(t *testing.T) { } } if !foundRoot { - t.Fatalf("raptor(PSI): no 'root' chunk; got %d chunks", len(chunks)) + t.Fatalf("raptor(default): 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") + // A smaller tree_order (more, smaller clusters) must also run and still + // produce a well-formed tree (root present, chunks non-empty). + chunksCoarse := runVariant(t, "raptor", map[string]any{"extra": map[string]any{"tree_order": 2}}) + if len(chunksCoarse) == 0 { + t.Fatalf("raptor(tree_order=2): 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 { + foundRootCoarse := false + for _, c := range chunksCoarse { if kind, _ := c["kc_kind"].(string); kind == "root" { - foundRootGMM = true + foundRootCoarse = true } } - if !foundRootGMM { - t.Fatalf("raptor(GMM): no 'root' chunk; got %d chunks", len(chunksGMM)) + if !foundRootCoarse { + t.Fatalf("raptor(tree_order=2): no 'root' chunk; got %d chunks", len(chunksCoarse)) } } @@ -540,35 +534,20 @@ func (m constEmbedder) Encode(_ context.Context, texts []string) ([][]float32, e 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.) +// vectors, so under the watershed default (tree_order=4, ratio 25) every adjacent pair +// has cosine >= 0 and a pathological input can collapse into one cluster; with +// more than one point 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"}, + "extra": map[string]any{"tree_order": 4}, }) if err != nil { t.Fatalf("NewKnowledgeCompilerComponent: %v", err) @@ -722,63 +701,6 @@ func TestKnowledgeCompiler_Wiki_UpdateMergesExistingPage(t *testing.T) { } } -// 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. @@ -933,15 +855,14 @@ func (r *recordingNavChat) Chat(_ context.Context, req common.ChatRequest) (*com 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) { +// TestKnowledgeCompiler_Datasetnav_RootIncludesAllSummaries is a regression for +// root-overview completeness: the root synthesis must be built from ALL child +// summaries, not from a partial buffer. The Go implementation buffers every +// product in Outputs.Products (no streaming sink), so the root always sees the +// full child set even when the result set is large. Here nav_radius>1 forces +// each chunk into its own nav node, and we assert every child marker appears in +// the recorded root-synthesis prompt. +func TestKnowledgeCompiler_Datasetnav_RootIncludesAllSummaries(t *testing.T) { chat := &recordingNavChat{} common.SetDepsResolver(func(tenantID, llmID, embeddingModel string) (common.Deps, error) { return common.Deps{ @@ -961,13 +882,6 @@ func TestKnowledgeCompiler_Datasetnav_RootIncludesFlushedSummaries(t *testing.T) 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)) @@ -978,7 +892,6 @@ func TestKnowledgeCompiler_Datasetnav_RootIncludesFlushedSummaries(t *testing.T) "chunks": chunks, "doc_id": "d1", "tenant_id": "t1", - "sink": &sinkRecorder{}, }); err != nil { t.Fatalf("Invoke: %v", err) } @@ -991,7 +904,7 @@ func TestKnowledgeCompiler_Datasetnav_RootIncludesFlushedSummaries(t *testing.T) } 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) + t.Fatalf("root overview is missing child summary %q; the root must be built from ALL child summaries.\nroot prompt:\n%s", m, root) } } } diff --git a/internal/ingestion/component/knowledge_compiler/datasetnav/datasetnav.go b/internal/ingestion/component/knowledge_compiler/datasetnav/datasetnav.go index f893ae92ab..f604e7be4b 100644 --- a/internal/ingestion/component/knowledge_compiler/datasetnav/datasetnav.go +++ b/internal/ingestion/component/knowledge_compiler/datasetnav/datasetnav.go @@ -145,7 +145,7 @@ func Run(ctx context.Context, deps common.Deps, param common.Param, inputs commo // 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) + var products []common.Product clusterProductID := map[string]string{} // Pre-compute every cluster id (deterministic from StableRowID) before // resolving parent edges, so a child emitted before its (later-created) @@ -162,7 +162,7 @@ func Run(ctx context.Context, deps common.Deps, param common.Param, inputs commo pid = clusterProductID[c.Parent] } id := clusterProductID[c.Name] - if err := sink.Add(common.Product{ + products = append(products, common.Product{ ID: id, DocID: docID, TenantID: tenantID, @@ -179,12 +179,10 @@ func Run(ctx context.Context, deps common.Deps, param common.Param, inputs commo "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{ + products = append(products, 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). @@ -203,18 +201,16 @@ func Run(ctx context.Context, deps common.Deps, param common.Param, inputs commo "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). + // Root overview built from EVERY cluster description (all products are + // buffered in the same slice, so every node contributes 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{ + products = append(products, common.Product{ ID: common.StableRowID(tenantID, docID, string(common.VariantDatasetnav), "root"), DocID: docID, TenantID: tenantID, @@ -228,20 +224,12 @@ func Run(ctx context.Context, deps common.Deps, param common.Param, inputs commo "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 + Products: products, } return out, nil } diff --git a/internal/ingestion/component/knowledge_compiler/golden/fixtures/raptor_baseline.json b/internal/ingestion/component/knowledge_compiler/golden/fixtures/raptor_baseline.json index 472c0c3fdd..d1a7527493 100644 --- a/internal/ingestion/component/knowledge_compiler/golden/fixtures/raptor_baseline.json +++ b/internal/ingestion/component/knowledge_compiler/golden/fixtures/raptor_baseline.json @@ -1,12 +1,10 @@ { - "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, + "default_leaf_clusters_min": 4, + "default_leaf_clusters_max": 4, + "fine_leaf_clusters_min": 6, + "fine_leaf_clusters_max": 6, + "coarse_leaf_clusters_min": 2, + "coarse_leaf_clusters_max": 2, "leaf_coverage_min": 1.0, - "max_depth_max": 6 + "max_depth_max": 5 } diff --git a/internal/ingestion/component/knowledge_compiler/golden_test.go b/internal/ingestion/component/knowledge_compiler/golden_test.go index 75df1a964e..281d5089e7 100644 --- a/internal/ingestion/component/knowledge_compiler/golden_test.go +++ b/internal/ingestion/component/knowledge_compiler/golden_test.go @@ -225,53 +225,59 @@ func TestGolden_Wiki_ProductCount(t *testing.T) { } // 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. +// fixed corpus, the watershed builder (raptor.go::watershed) must produce a +// well-formed tree (single root, fully parented, vectors + schema intact, full +// chunk coverage, bounded depth/cluster count) across several tree_order +// settings. The structural contract is what a regression would violate; the +// exact cluster counts are locked in raptor_baseline.json. func TestGolden_Raptor_Structure(t *testing.T) { installSignedProseDeps(t) baseline := loadBaseline(t, "raptor_baseline.json") nChunks := len(golden.FixedCorpus()) cases := []struct { - method string + name string + order int // <=0 means default (omit extra) 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"}, + {"default", 0, "default_leaf_clusters_min", "default_leaf_clusters_max"}, + {"fine", 2, "fine_leaf_clusters_min", "fine_leaf_clusters_max"}, + {"coarse", 10, "coarse_leaf_clusters_min", "coarse_leaf_clusters_max"}, } for _, tc := range cases { - prods := runVariantChunks(t, "raptor", map[string]any{ - "extra": map[string]any{"clustering_method": tc.method}, - }) + var prods []schema.ChunkDoc + if tc.order <= 0 { + prods = runVariantChunks(t, "raptor", nil) + } else { + prods = runVariantChunks(t, "raptor", map[string]any{ + "extra": map[string]any{"tree_order": tc.order}, + }) + } 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)) + tc.name, 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) + t.Errorf("raptor(%s): rootCount=%d, want 1", tc.name, m.RootCount) } if !m.AllParented { - t.Errorf("raptor(%s): tree has dangling parent_id references", tc.method) + t.Errorf("raptor(%s): tree has dangling parent_id references", tc.name) } if !m.VectorOK { - t.Errorf("raptor(%s): some products missing vectors", tc.method) + t.Errorf("raptor(%s): some products missing vectors", tc.name) } if !m.SchemaOK { - t.Errorf("raptor(%s): some products missing schema fields", tc.method) + t.Errorf("raptor(%s): some products missing schema fields", tc.name) } 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")) + t.Errorf("raptor(%s): leaf coverage %.2f < baseline %.2f", tc.name, 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)) + t.Errorf("raptor(%s): leafClusters=%d outside [%d,%d]", tc.name, 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")) + t.Errorf("raptor(%s): maxDepth=%d > baseline %d", tc.name, m.MaxDepth, num(baseline, "max_depth_max")) } } } diff --git a/internal/ingestion/component/knowledge_compiler/mindmap/mindmap.go b/internal/ingestion/component/knowledge_compiler/mindmap/mindmap.go index 78c02e4c1d..5a78d0e5b2 100644 --- a/internal/ingestion/component/knowledge_compiler/mindmap/mindmap.go +++ b/internal/ingestion/component/knowledge_compiler/mindmap/mindmap.go @@ -39,7 +39,7 @@ func Run(ctx context.Context, deps common.Deps, param common.Param, inputs commo // One LLM task per token-budget batch (mirrors __call__'s task fan-out). batches := packSections(sections, deps.Tokenizer) - results := make([]omap, len(batches)) + results := make([]utility.OMap, len(batches)) // One LLM task per token-budget batch (mirrors __call__'s task fan-out). n := param.MaxWorkers if n <= 0 { @@ -59,7 +59,7 @@ func Run(ctx context.Context, deps common.Deps, param common.Param, inputs commo if err != nil { return err } - results[i] = todict(dictify(StripFences(resp.Content))) + results[i] = utility.Todict(utility.Dictify(utility.StripFences(resp.Content))) return nil }) if err != nil { @@ -77,14 +77,14 @@ func Run(ctx context.Context, deps common.Deps, param common.Param, inputs commo // 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 + var merged utility.OMap if len(results) > 0 { merged = results[0] for _, r := range results[1:] { - merged = mergeDicts(merged, r) + merged = utility.MergeDicts(merged, r) } } - root := shapeTree(merged) + root := utility.ShapeTree(merged) products := treeToProducts(tenantID, docID, root) @@ -105,24 +105,13 @@ func Run(ctx context.Context, deps common.Deps, param common.Param, inputs commo } } - // 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 - } - } + // Buffer every tree node in one slice; the component merges them into the + // upstream chunk stream (matching Python, which appends compiled units onto + // the chunk list). out := common.Outputs{ - Products: sink.Products(), - VectorBytes: sink.Bytes(), - Items: sink.TotalItems(), - Flushed: sink.Flushed(), + Products: products, } - if err := out.EnforceGuardrails(param.Guardrails, inputs.Sink, ctx); err != nil { - return common.Outputs{}, err - } return out, nil } @@ -130,7 +119,7 @@ func Run(ctx context.Context, deps common.Deps, param common.Param, inputs commo // 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 { +func treeToProducts(tenantID, docID string, root *utility.Node) []common.Product { var out []common.Product rootID := common.StableRowID(tenantID, docID, string(common.VariantMindmap), "root") out = append(out, common.Product{ @@ -147,7 +136,7 @@ func treeToProducts(tenantID, docID string, root *Node) []common.Product { }) type pending struct { - node *Node + node *utility.Node parentID string level int } @@ -192,7 +181,7 @@ func treeToProducts(tenantID, docID string, root *Node) []common.Product { // serializeNode renders the tree in Python's MindMapResult.output shape: // {"id": ..., "children": [...]}. -func serializeNode(node *Node) string { +func serializeNode(node *utility.Node) string { var b strings.Builder b.WriteString(`{"id":`) b.WriteString(quoteJSON(node.ID)) diff --git a/internal/ingestion/component/knowledge_compiler/mindmap/mindmap_test.go b/internal/ingestion/component/knowledge_compiler/mindmap/mindmap_test.go index ea78616a64..b6eb5bc594 100644 --- a/internal/ingestion/component/knowledge_compiler/mindmap/mindmap_test.go +++ b/internal/ingestion/component/knowledge_compiler/mindmap/mindmap_test.go @@ -4,166 +4,9 @@ import ( "strings" "testing" - "ragflow/internal/ingestion/component/knowledge_compiler/common" + "ragflow/internal/utility" ) -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") { @@ -199,8 +42,8 @@ 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"}}}, + root := &utility.Node{ID: "root", Children: []*utility.Node{ + {ID: "A", Children: []*utility.Node{{ID: "A1"}, {ID: "A2"}}}, {ID: "B"}, }} products := treeToProducts("t1", "d1", root) @@ -229,11 +72,9 @@ func TestTreeToProducts_ParentLinks(t *testing.T) { } func TestSerializeNode_Shape(t *testing.T) { - root := &Node{ID: "Top \"quoted\"", Children: []*Node{{ID: "child"}}} + root := &utility.Node{ID: "Top \"quoted\"", Children: []*utility.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/raptor/clustering.go b/internal/ingestion/component/knowledge_compiler/raptor/clustering.go deleted file mode 100644 index 993e30d0ef..0000000000 --- a/internal/ingestion/component/knowledge_compiler/raptor/clustering.go +++ /dev/null @@ -1,680 +0,0 @@ -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 deleted file mode 100644 index 78d250f272..0000000000 --- a/internal/ingestion/component/knowledge_compiler/raptor/clustering_test.go +++ /dev/null @@ -1,84 +0,0 @@ -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 deleted file mode 100644 index d623729174..0000000000 --- a/internal/ingestion/component/knowledge_compiler/raptor/psi.go +++ /dev/null @@ -1,119 +0,0 @@ -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 index e69eb34ff0..3dde45ddf1 100644 --- a/internal/ingestion/component/knowledge_compiler/raptor/raptor.go +++ b/internal/ingestion/component/knowledge_compiler/raptor/raptor.go @@ -1,33 +1,30 @@ // 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: +// until a single root summary remains. // -// - 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). +// Clustering uses a single method, watershed (raptor.go::watershed): a 1D +// watershed over the document-ordered embeddings. The granularity of each +// level's clustering is controlled by the "tree_order" integer parameter +// (raptor.go::resolveTreeOrder) — the branching factor of the RAPTOR tree +// (the B+ tree "order"): the average chunk count per level-0 cluster. It +// defaults to DefaultTreeOrder. // -// 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. +// See PORT_PLAN.md §3.3. package raptor import ( "context" "fmt" + "log" + "regexp" "strings" + "time" "ragflow/internal/ingestion/component/knowledge_compiler/common" + "ragflow/internal/tokenizer" ) -// 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 { @@ -43,82 +40,130 @@ func Run(ctx context.Context, deps common.Deps, param common.Param, inputs commo 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) + treeOrder := resolveTreeOrder(param) + taskPrompt := resolveRaptorPrompt(param) // 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 { + // collecting every node into a plain slice. The component merges these into + // the upstream chunk stream (matching Python, which appends summaries onto the + // chunk list). buildTree reads the texts, ids, and embeddings directly from + // the source chunks. + var products []common.Product + if err := buildTree(ctx, deps, llmID, tenantID, docID, inputs.Chunks, treeOrder, taskPrompt, param, &products); err != nil { return common.Outputs{}, err } out := common.Outputs{ - Products: sink.Products(), - VectorBytes: sink.Bytes(), - Items: sink.TotalItems(), - Flushed: sink.Flushed(), + Products: products, } - 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, +// resolveTreeOrder reads the watershed "tree_order" parameter (an integer +// branching factor, the B+ tree "order", in [MinTreeOrder, MaxTreeOrder]) from +// the component extra map, defaulting to DefaultTreeOrder when absent or invalid. +func resolveTreeOrder(param common.Param) int { + raw, ok := param.Extra["tree_order"] + if !ok { + return DefaultTreeOrder } - if m, ok := param.Extra["clustering_method"].(string); ok && m != "" { - spec.Method = ClusteringMethod(strings.ToUpper(m)) + v, ok := toInt(raw) + if !ok || v < MinTreeOrder || v > MaxTreeOrder { + return DefaultTreeOrder } - if t, ok := param.Extra["clustering_threshold"].(float64); ok { - spec.Threshold = t + return v +} + +// toInt coerces the common param-extra value types to int. +func toInt(v any) (int, bool) { + switch x := v.(type) { + case int: + return x, true + case int64: + return int(x), true + case float64: + return int(x), true + case float32: + return int(x), true + default: + return 0, false } - if v, ok := param.Extra["min_clusters"].(float64); ok { - spec.MinClusters = int(v) +} + +const ( + raptorMaxRetries = 3 + raptorDefaultMaxToken = 512 + raptorDefaultMaxErrors = 3 +) + +// raptorTruncationMarkerRE strips model truncation notices that some LLMs emit +// when a response is cut short. Mirrors Python _summarize_texts (raptor.py:405). +var raptorTruncationMarkerRE = regexp.MustCompile( + strings.Repeat("\u00b7", 6) + "\n由于长度的原因，回答被截断了，要继续吗？|For the content length reason, it stopped, continue?", +) + +// resolveMaxToken returns the generation cap (max_tokens) for summaries. Python +// uses max(self._max_token, 512) (raptor.py:403, issue #10235); we honour an +// extra["max_token"] override with the same 512 floor. +func resolveMaxToken(param common.Param) int { + if v, ok := param.Extra["max_token"]; ok { + if n, ok := toInt(v); ok && n > 0 { + if n < raptorDefaultMaxToken { + return raptorDefaultMaxToken + } + return n + } } - if v, ok := param.Extra["max_clusters"].(float64); ok { - spec.MaxClusters = int(v) + return raptorDefaultMaxToken +} + +// resolveMaxErrors returns the per-tree error ceiling before RAPTOR aborts. +// Python defaults _max_errors to 3 and skips individual failed clusters below +// that threshold (raptor.py:_summarize_texts). +func resolveMaxErrors(param common.Param) int { + if v, ok := param.Extra["max_errors"]; ok { + if n, ok := toInt(v); ok && n > 0 { + return n + } } - 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 + return raptorDefaultMaxErrors } // 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 { +// "raptor" product with a stable id. Every node is appended to products (a +// caller-owned slice), and only the deepest-level summaries are retained for the +// root synthesis. buildTree reads the source texts, ids, and embeddings directly +// from the passed chunks; when the embeddings are missing or incomplete, it +// embeds the texts itself. Every summary node is embedded here too, so the caller +// never calls the embedder directly. treeOrder drives both the top-level +// clustering (performed inside buildTree via watershed) 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, chunks []common.Chunk, treeOrder int, taskPrompt string, param common.Param, products *[]common.Product) error { + maxToken := resolveMaxToken(param) + maxErrors := resolveMaxErrors(param) + errorCount := 0 + texts := chunkTexts(chunks) + if len(texts) == 0 { + return nil + } + chunkIDs := chunkIDsOf(chunks) + // Reuse the embeddings already carried on the source chunks when they are + // complete; only re-embed the texts as a fallback (e.g. the caller has not + // pre-embedded, or some rows lack a vector). + embeddings := chunkVectors(chunks) + if !embeddingsReady(embeddings, len(texts)) { + vectors, err := deps.Embed.Encode(ctx, texts) + if err != nil { + return err + } + embeddings = toFloat64Matrix(vectors) + } + labels, err := watershed(embeddings, treeOrder) + if err != nil { + return err + } n := len(labels) // Group point indices by cluster label. groups := map[int][]int{} @@ -146,15 +191,20 @@ func buildTree(ctx context.Context, deps common.Deps, llmID, tenantID, docID str 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()) + // Gather the source text for this cluster's points and truncate each + // text to a per-chunk token budget so the cluster fits the LLM context + // (Python: len_per_chunk = (max_length - max_token) / len(texts); + // truncate(t, len_per_chunk), raptor.py:389-390). + content := buildClusterContent(texts, task.pointIdxs, deps.LLMMaxLength, maxToken) + system := raptorSystemHelper + strings.Replace(taskPrompt, "{cluster_content}", content, 1) + summary, err := summarizeTexts(ctx, deps, llmID, system, raptorTitleInstruction, maxToken) if err != nil { - return err + errorCount++ + if errorCount >= maxErrors { + return fmt.Errorf("raptor: aborted after %d summarization errors: %w", errorCount, err) + } + log.Printf("raptor: skipping cluster due to summarization error (continuing): %v", err) + continue } nodeID := common.StableRowID(tenantID, docID, string(common.VariantRaptor), fmt.Sprintf("L%d", task.level), summary) @@ -166,7 +216,7 @@ func buildTree(ctx context.Context, deps common.Deps, llmID, tenantID, docID str if len(embedding) > 0 { vec = embedding[0] } - if err := sink.Add(common.Product{ + *products = append(*products, common.Product{ ID: nodeID, DocID: docID, TenantID: tenantID, @@ -175,13 +225,12 @@ func buildTree(ctx context.Context, deps common.Deps, llmID, tenantID, docID str Vector: vec, ParentID: task.parentID, Meta: map[string]any{ + "title": titleOf(summary), "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. @@ -193,57 +242,76 @@ func buildTree(ctx context.Context, deps common.Deps, llmID, tenantID, docID str 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}) + // Recurse: re-cluster this cluster's points and enqueue sub-clusters + // under this node. The sub-clustering uses the same watershed pass as the + // top-level (treeOrder is carried through), so every level of the tree is + // built with the same method. Recursion terminates because the + // sub-clustering is only re-enqueued when it actually splits the cluster + // into more than one group (the degenerate single-label guard after the + // watershed call); once a cluster cannot be split further it stops, and a + // cluster is always driven down to single-point leaves, so depth is + // naturally bounded by log_2(len(embeddings)). + subEmb := make([][]float64, 0, len(task.pointIdxs)) + for _, pi := range task.pointIdxs { + subEmb = append(subEmb, embeddings[pi]) + } + subLabels, err := watershed(subEmb, treeOrder) + if err != nil { + return err + } + 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. a corpus whose vectors all share one sign + // so the watershed never cuts), 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")) + // Root product ties the tree together. The root summary is the synthesis of + // the highest-level summaries (the leaves of the upper tree) using the same + // standard task prompt Python applies to every cluster — NOT a separate + // "Synthesize the overall theme" prompt, and NOT the union of all leaf + // summaries (which would just re-summarize the same text repeatedly). + if len(topLevelTexts) == 0 { + // No summaries survived (e.g. every deepest cluster failed while the + // error budget was not yet exhausted). Return the partial tree without a + // root node — Python drops the root in this case rather than crashing. + log.Printf("raptor: no top-level summaries produced, skipping root node") + return nil + } + rootContent := buildClusterContent(topLevelTexts, allIndices(len(topLevelTexts)), deps.LLMMaxLength, maxToken) + rootSummary, err := summarizeTexts(ctx, deps, llmID, + raptorSystemHelper+strings.Replace(taskPrompt, "{cluster_content}", rootContent, 1), + raptorTitleInstruction, maxToken) if err != nil { - return fmt.Errorf("raptor: root summarization failed: %w", err) + // A failed root must not abort the whole tree: Python drops the root + // node and returns the rest of the tree. + log.Printf("raptor: root synthesis failed, skipping root node: %v", err) + return nil } embedding, err := deps.Embed.Encode(ctx, []string{rootSummary}) if err != nil { - return fmt.Errorf("raptor: root embedding failed: %w", err) + log.Printf("raptor: root embedding failed, skipping root node: %v", err) + return nil } if len(embedding) == 0 { - return fmt.Errorf("raptor: root embedding returned no vectors") + log.Printf("raptor: root embedding returned no vectors, skipping root node") + return nil } - return sink.Add(common.Product{ + *products = append(*products, common.Product{ ID: rootID, DocID: docID, TenantID: tenantID, @@ -251,30 +319,123 @@ func buildTree(ctx context.Context, deps common.Deps, llmID, tenantID, docID str Content: rootSummary, Vector: embedding[0], ParentID: "", - Meta: map[string]any{"kind": "root", "level": -1}, + Meta: map[string]any{"title": titleOf(rootSummary), "kind": "root", "level": -1}, }) + return nil } -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 +// summarizeTexts asks the LLM for a summary and normalises the response to match +// Python _summarize_texts / _chat: +// - retries up to raptorMaxRetries times with linear backoff (Python: 3 attempts); +// - strips a reasoning-model preamble up to the final / tag; +// - treats a "**ERROR**" marker as a transient failure (Python raises on it, so +// it is retried like any other LLM error); +// - strips model truncation notices (raptor.py:405); +// - trims surrounding whitespace. +// +// systemText is the fully-built system prompt (helper + filled task template) and +// userText is the user turn; the title instruction is passed as the user turn so +// the model emits a one-line title on the first line of the summary. +func summarizeTexts(ctx context.Context, deps common.Deps, llmID, systemText, userText string, maxToken int) (string, error) { + mt := maxToken + for attempt := 0; attempt < raptorMaxRetries; attempt++ { + if attempt > 0 { + select { + case <-ctx.Done(): + return "", ctx.Err() + case <-time.After(time.Duration(1+attempt) * time.Second): + } + } + resp, err := deps.Chat.Chat(ctx, common.ChatRequest{ + LLMID: llmID, + SystemPrompt: systemText, + UserPrompt: userText, + MaxTokens: &mt, + }) + if err != nil { + continue + } + content := resp.Content + // Strip reasoning preamble up to the final thinking close tag + // (Python: re.sub(r"^.*", "", response, DOTALL)). + if i := strings.LastIndex(content, ""); i >= 0 { + content = content[i+len(""):] + } else if i := strings.LastIndex(content, ""); i >= 0 { + content = content[i+len(""):] + } + // Python raises on the "**ERROR**" marker; treat it as a retryable error. + if strings.Contains(content, "**ERROR**") { + continue + } + // Strip model truncation notices (Python _summarize_texts). + content = raptorTruncationMarkerRE.ReplaceAllString(content, "") + return strings.TrimSpace(content), nil } - return strings.TrimSpace(resp.Content), nil + return "", fmt.Errorf("raptor: summarization failed after %d attempts", raptorMaxRetries) +} + +// titleOf returns the first line of a summary completion (Python stores +// summary_ti = cnt.split("\n")[0]). +func titleOf(summary string) string { + if i := strings.IndexByte(summary, '\n'); i >= 0 { + return summary[:i] + } + return summary +} + +// buildClusterContent joins the selected texts with a single newline (matching +// Python's "\n".join) and truncates each text to a per-chunk token budget so the +// whole cluster fits the LLM context window (Python: len_per_chunk = +// (max_length - max_token) / len(texts); truncate(t, len_per_chunk)). The token +// budget uses the cl100k_base encoder, mirroring Python's truncate (token-level, +// not character-level). +func buildClusterContent(texts []string, idxs []int, llmMaxLength, maxToken int) string { + if len(idxs) == 0 { + return "" + } + if llmMaxLength <= 0 { + llmMaxLength = common.DefaultLLMContextLength + } + per := (llmMaxLength - maxToken) / len(idxs) + if per < 1 { + per = 1 + } + parts := make([]string, 0, len(idxs)) + for _, i := range idxs { + parts = append(parts, tokenizer.TrimContentToTokenLimit(texts[i], per)) + } + return strings.Join(parts, "\n") +} + +// allIndices returns [0, 1, ..., n-1] so a whole slice can be fed to +// buildClusterContent without copying. +func allIndices(n int) []int { + idxs := make([]int, n) + for i := range idxs { + idxs[i] = i + } + return idxs +} + +// chunkTextOf returns the text of a chunk, preferring Text over Content. +func chunkTextOf(c common.Chunk) string { + return firstNonEmpty(c.Text, c.Content) +} + +// chunkHasText reports whether a chunk carries non-blank text. Empty chunks +// contribute nothing to clustering or summarization and are skipped uniformly +// across chunkTexts, chunkIDsOf, and chunkVectors. +func chunkHasText(c common.Chunk) bool { + return strings.TrimSpace(chunkTextOf(c)) != "" } func chunkTexts(chunks []common.Chunk) []string { var out []string for _, c := range chunks { - t := firstNonEmpty(c.Text, c.Content) - if strings.TrimSpace(t) == "" { + if !chunkHasText(c) { continue } - out = append(out, t) + out = append(out, chunkTextOf(c)) } return out } @@ -285,8 +446,7 @@ func chunkTexts(chunks []common.Chunk) []string { func chunkIDsOf(chunks []common.Chunk) []string { var out []string for i, c := range chunks { - t := firstNonEmpty(c.Text, c.Content) - if strings.TrimSpace(t) == "" { + if !chunkHasText(c) { continue } id := c.ID @@ -298,6 +458,48 @@ func chunkIDsOf(chunks []common.Chunk) []string { return out } +// chunkVectors returns the pre-computed embedding of every non-empty chunk, in +// document order and parallel to chunkTexts/chunkIDsOf. A chunk without a +// Vector yields a nil entry; buildTree re-embeds the texts when any entry is +// missing so the caller can choose whether to pre-embed. +func chunkVectors(chunks []common.Chunk) [][]float64 { + out := make([][]float64, 0, len(chunks)) + for _, c := range chunks { + if !chunkHasText(c) { + continue + } + out = append(out, toFloat64Slice(c.Vector)) + } + return out +} + +// toFloat64Slice converts a float32 vector to float64, returning nil for an +// empty input (so a missing embedding is distinguishable from a zero vector). +func toFloat64Slice(v []float32) []float64 { + if len(v) == 0 { + return nil + } + out := make([]float64, len(v)) + for i, x := range v { + out[i] = float64(x) + } + return out +} + +// embeddingsReady reports whether emb is a complete, usable embedding set for +// n points: exactly n rows, none empty. +func embeddingsReady(emb [][]float64, n int) bool { + if len(emb) != n { + return false + } + for _, e := range emb { + if len(e) == 0 { + return false + } + } + return true +} + // collectIDs returns the chunk ids at the given point indices, de-duplicated // and in stable order. func collectIDs(chunkIDs []string, pointIdxs []int) []string { @@ -338,4 +540,30 @@ func toFloat64Matrix(vecs [][]float32) [][]float64 { 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.` +// defaultRaptorPrompt is the summary task template. It mirrors Python +// rag/advanced_rag/knowlege_compile/raptor.py, whose _summarize_texts uses +// "Please write a concise summary of the following texts:\n{cluster_content}". +// The {cluster_content} placeholder is filled with the joined cluster text. A +// caller may override it via extra["prompt"] (Python: raptor_cfg["prompt"]). +const defaultRaptorPrompt = "Please write a concise summary of the following texts:\n{cluster_content}" + +// raptorSystemHelper mirrors the leading "You're a helpful assistant.\n\nHelp me +// with the following task.\n\n" wrapper Python prepends to the task prompt. +const raptorSystemHelper = "You're a helpful assistant.\n\nHelp me with the following task.\n\n" + +// raptorTitleInstruction is sent as the user turn so the model also emits a +// one-line title on the first line of the summary (Python: "Beside the +// summarization, give a title at the first line of your summarization. Must be +// in the same language as the paragraphs."). +const raptorTitleInstruction = "Beside the summarization, give a title at the first line of your summarization. Must be in the same language as the paragraphs." + +// resolveRaptorPrompt returns the summary task template, honouring an +// extra["prompt"] override; falls back to defaultRaptorPrompt. +func resolveRaptorPrompt(param common.Param) string { + if raw, ok := param.Extra["prompt"]; ok { + if s, ok := raw.(string); ok && strings.TrimSpace(s) != "" { + return s + } + } + return defaultRaptorPrompt +} diff --git a/internal/ingestion/component/knowledge_compiler/raptor/raptor_test.go b/internal/ingestion/component/knowledge_compiler/raptor/raptor_test.go new file mode 100644 index 0000000000..eb8d0c0f1a --- /dev/null +++ b/internal/ingestion/component/knowledge_compiler/raptor/raptor_test.go @@ -0,0 +1,148 @@ +package raptor + +import ( + "context" + "strings" + "testing" + + "ragflow/internal/ingestion/component/knowledge_compiler/common" + "ragflow/internal/tokenizer" +) + +// fakeChat records the last request and returns scripted responses. +type fakeChat struct { + calls int + lastReq common.ChatRequest + // responses[i] is returned on the i-th call; nil entry means an error. + responses []*common.ChatResponse + errs []error +} + +func (f *fakeChat) Chat(_ context.Context, req common.ChatRequest) (*common.ChatResponse, error) { + f.lastReq = req + i := f.calls + f.calls++ + if i < len(f.errs) && f.errs[i] != nil { + return nil, f.errs[i] + } + if i < len(f.responses) { + return f.responses[i], nil + } + return &common.ChatResponse{Content: "ok"}, nil +} + +func depsWithChat(c common.ChatInvoker) common.Deps { + return common.Deps{Chat: c, Embed: nil, TenantID: "t"} +} + +func TestSummarizeTextsStripsThinkPreamble(t *testing.T) { + f := &fakeChat{responses: []*common.ChatResponse{{Content: "let me think...\n\nFinal summary title\nbody"}}} + got, err := summarizeTexts(context.Background(), depsWithChat(f), "llm", "sys", "user", 512) + if err != nil { + t.Fatalf("unexpected err: %v", err) + } + if strings.Contains(got, "") || strings.Contains(got, "") { + t.Fatalf("think preamble not stripped: %q", got) + } + if !strings.Contains(got, "Final summary title") { + t.Fatalf("body lost: %q", got) + } +} + +func TestSummarizeTextsStripsTruncationMarker(t *testing.T) { + marker := strings.Repeat("·", 6) + "\n由于长度的原因，回答被截断了，要继续吗？" + f := &fakeChat{responses: []*common.ChatResponse{{Content: "title\nbody " + marker}}} + got, err := summarizeTexts(context.Background(), depsWithChat(f), "llm", "sys", "user", 512) + if err != nil { + t.Fatalf("unexpected err: %v", err) + } + if strings.Contains(got, "回答被截断了") { + t.Fatalf("truncation marker not stripped: %q", got) + } +} + +func TestSummarizeTextsRetriesOnErrorMarker(t *testing.T) { + f := &fakeChat{ + responses: []*common.ChatResponse{ + {Content: "**ERROR** something broke"}, + {Content: "title\nclean summary"}, + }, + } + got, err := summarizeTexts(context.Background(), depsWithChat(f), "llm", "sys", "user", 512) + if err != nil { + t.Fatalf("unexpected err after retry: %v", err) + } + if got != "title\nclean summary" { + t.Fatalf("expected clean summary after retry, got %q", got) + } + if f.calls != 2 { + t.Fatalf("expected 2 calls (1 retry), got %d", f.calls) + } +} + +func TestSummarizeTextsFailsAfterMaxRetries(t *testing.T) { + f := &fakeChat{responses: []*common.ChatResponse{ + {Content: "**ERROR** 1"}, {Content: "**ERROR** 2"}, {Content: "**ERROR** 3"}, + }} + if _, err := summarizeTexts(context.Background(), depsWithChat(f), "llm", "sys", "user", 512); err == nil { + t.Fatal("expected error after exhausting retries") + } + if f.calls != raptorMaxRetries { + t.Fatalf("expected %d attempts, got %d", raptorMaxRetries, f.calls) + } +} + +func TestSummarizeTextsPassesMaxTokens(t *testing.T) { + f := &fakeChat{responses: []*common.ChatResponse{{Content: "title\nbody"}}} + if _, err := summarizeTexts(context.Background(), depsWithChat(f), "llm", "sys", "user", 1024); err != nil { + t.Fatalf("unexpected err: %v", err) + } + if f.lastReq.MaxTokens == nil || *f.lastReq.MaxTokens != 1024 { + t.Fatalf("MaxTokens not passed through: %v", f.lastReq.MaxTokens) + } +} + +func TestBuildClusterContentJoinsWithSingleNewline(t *testing.T) { + // delimiter must be "\n" to match Python's "\n".join, not "\n\n". + out := buildClusterContent([]string{"a", "b", "c"}, []int{0, 1, 2}, common.DefaultLLMContextLength, 512) + if out != "a\nb\nc" { + t.Fatalf("expected single-newline join, got %q", out) + } +} + +func TestBuildClusterContentTruncatesPerChunk(t *testing.T) { + // A long text must be truncated to the per-chunk token budget so the + // cluster fits the LLM context window (Python len_per_chunk). + long := strings.Repeat("hello world ", 200) + out := buildClusterContent([]string{long}, []int{0}, common.DefaultLLMContextLength, 512) + per := (common.DefaultLLMContextLength - 512) / 1 + if tokenizer.NumTokensFromString(out) > per { + t.Fatalf("output exceeded per-chunk budget: %d > %d", tokenizer.NumTokensFromString(out), per) + } +} + +// TestBuildTreeNoPanicWhenAllSummariesFail guards the divide-by-zero that +// occurred when every deepest cluster failed: buildClusterContent divides by +// len(idxs), and the root synthesis built a cluster from allIndices(0) when +// topLevelTexts was empty. The root is now skipped and the partial tree is +// returned without error. +func TestBuildTreeNoPanicWhenAllSummariesFail(t *testing.T) { + errs := make([]error, 16) + for i := range errs { + errs[i] = context.DeadlineExceeded + } + f := &fakeChat{errs: errs} + deps := common.Deps{Chat: f, Embed: nil, TenantID: "t"} + // Pre-computed vectors so the tree never needs to call the embedder. + chunks := []common.Chunk{ + {Text: "alpha", Vector: []float32{1, 0, 0, 0}}, + {Text: "beta", Vector: []float32{0, 1, 0, 0}}, + } + var products []common.Product + if err := buildTree(context.Background(), deps, "llm", "t", "d", chunks, 4, "", common.Param{}, &products); err != nil { + t.Fatalf("buildTree returned unexpected error: %v", err) + } + if len(products) != 0 { + t.Fatalf("expected no products when every summary fails, got %d", len(products)) + } +} diff --git a/internal/ingestion/component/knowledge_compiler/raptor/watershed.go b/internal/ingestion/component/knowledge_compiler/raptor/watershed.go new file mode 100644 index 0000000000..c392ef27ac --- /dev/null +++ b/internal/ingestion/component/knowledge_compiler/raptor/watershed.go @@ -0,0 +1,142 @@ +// watershed is the sole clustering method used by the RAPTOR tree. +// +// It is the algorithm historically named "AHC" in the Python raptor +// implementation (raptor.py::_get_clusters_ahc, logged as "RAPTOR seq-clus"); +// the name "watershed" reflects what the code actually does, which is not +// agglomerative hierarchical clustering. +// +// See PORT_PLAN.md §3.3 and the M6 validation gate (缺口 C). +package raptor + +import ( + "errors" + "math" + "sort" +) + +// DefaultTreeOrder is the default branching factor (the B+ tree "order") used by +// watershed clustering. It is the average chunk count per level-0 cluster: the +// expected number of level-0 clusters is close to len(embeddings)/DefaultTreeOrder, +// so the recursion depth is roughly log_DefaultTreeOrder(len(embeddings)). 4 +// (ratio 25) is close to the Python raptor "cluster_percentile" default of 30. +const DefaultTreeOrder = 4 + +// MinTreeOrder and MaxTreeOrder bound the integer tree_order parameter. +const ( + MinTreeOrder = 2 + MaxTreeOrder = 16 +) + +// watershed performs a 1D watershed over the sequence of embeddings: it scans +// adjacent pairs only and starts a new cluster whenever the cosine similarity +// between two consecutive embeddings falls below a percentile threshold. +// +// The cut points are contiguous segments of the input in its given order. +// Callers pass document-ordered embeddings, so the boundaries correspond to +// topic shifts in the text. This is a cheap, deterministic text-segmentation +// heuristic (not a semantic clustering), and is O(N) in the number of chunks. +// +// treeOrder is the branching factor of the RAPTOR tree (the B+ tree "order"): +// the average chunk count per level-0 cluster. It maps to the percentile +// threshold as 100/treeOrder, so the expected level-0 cluster count is +// approximately len(embeddings)/treeOrder. Larger treeOrder => fewer, bigger +// clusters and a shallower tree; smaller treeOrder => more, smaller clusters +// and a deeper tree. treeOrder is clamped to [MinTreeOrder, MaxTreeOrder]. +func watershed(embeddings [][]float64, treeOrder int) ([]int, error) { + n := len(embeddings) + switch { + case n == 0: + return nil, errors.New("raptor: watershed requires at least one embedding") + case n == 1: + return []int{0}, nil + case n == 2: + // One adjacent pair: the Python reference always splits two points into + // two clusters, so the recursion terminates with one-point leaves. + return []int{0, 1}, nil + } + + if treeOrder < MinTreeOrder { + treeOrder = MinTreeOrder + } else if treeOrder > MaxTreeOrder { + treeOrder = MaxTreeOrder + } + + // L2-normalize so the dot product equals cosine similarity. + normalized := normalizeRows(embeddings) + + // Adjacent cosine similarities (n-1 pairs). + adj := make([]float64, n-1) + for i := 0; i < n-1; i++ { + adj[i] = dot(normalized[i], normalized[i+1]) + } + + // The percentile threshold is 100/treeOrder: a larger treeOrder lowers the + // threshold percentile, producing fewer cuts (fewer, bigger clusters). + clusterRatio := 100.0 / float64(treeOrder) + threshold := percentile(adj, clusterRatio) + labels := make([]int, n) + clusterID := 0 + for i := 1; i < n; i++ { + if adj[i-1] >= threshold { + labels[i] = clusterID + } else { + clusterID++ + labels[i] = clusterID + } + } + return labels, nil +} + +// normalizeRows returns a copy of rows with each row L2-normalized. Zero rows +// are left unchanged (their dot products are zero) so they cannot produce NaN. +func normalizeRows(rows [][]float64) [][]float64 { + out := make([][]float64, len(rows)) + for i, r := range rows { + var norm float64 + for _, v := range r { + norm += v * v + } + norm = math.Sqrt(norm) + if norm == 0 { + norm = 1 + } + row := make([]float64, len(r)) + for j, v := range r { + row[j] = v / norm + } + out[i] = row + } + return out +} + +func dot(a, b []float64) float64 { + var s float64 + for i := range a { + s += a[i] * b[i] + } + return s +} + +// percentile returns the p-th percentile of values using linear interpolation, +// matching numpy's default. p is clamped to [0,100]. +func percentile(values []float64, p float64) float64 { + if len(values) == 0 { + return 0 + } + cp := append([]float64(nil), values...) + sort.Float64s(cp) + if p < 0 { + p = 0 + } + if p > 100 { + p = 100 + } + rank := (p / 100.0) * float64(len(cp)-1) + lo := int(math.Floor(rank)) + hi := int(math.Ceil(rank)) + if lo == hi { + return cp[lo] + } + frac := rank - float64(lo) + return cp[lo]*(1-frac) + cp[hi]*frac +} diff --git a/internal/ingestion/component/knowledge_compiler/raptor/watershed_test.go b/internal/ingestion/component/knowledge_compiler/raptor/watershed_test.go new file mode 100644 index 0000000000..f6b9735eba --- /dev/null +++ b/internal/ingestion/component/knowledge_compiler/raptor/watershed_test.go @@ -0,0 +1,127 @@ +package raptor + +import ( + "math" + "testing" +) + +// cosEmbeddings places points on the unit circle so that the cosine similarity +// between consecutive points equals sims[i]. It lets tests assert exact +// watershed cut points from a known similarity profile. +func cosEmbeddings(sims []float64) [][]float64 { + angles := make([]float64, len(sims)+1) + for i, c := range sims { + if c > 1 { + c = 1 + } + if c < -1 { + c = -1 + } + angles[i+1] = angles[i] + math.Acos(c) + } + emb := make([][]float64, len(angles)) + for i, a := range angles { + emb[i] = []float64{math.Cos(a), math.Sin(a)} + } + return emb +} + +func uniqueCount(labels []int) int { + seen := map[int]struct{}{} + for _, l := range labels { + seen[l] = struct{}{} + } + return len(seen) +} + +func TestWatershedEdgeCases(t *testing.T) { + if _, err := watershed(nil, 4); err == nil { + t.Fatalf("expected error for empty input") + } + labels, err := watershed([][]float64{{1, 0}}, 4) + if err != nil || len(labels) != 1 || labels[0] != 0 { + t.Fatalf("n=1: got %v err %v", labels, err) + } + // n=2: the Python reference always splits two points into two clusters. + labels, err = watershed([][]float64{{1, 0}, {0, 1}}, 4) + if err != nil || len(labels) != 2 || labels[0] != 0 || labels[1] != 1 { + t.Fatalf("n=2: got %v err %v", labels, err) + } +} + +// TestWatershedTreeOrderControlsCount verifies that the tree_order branching +// factor drives the expected cluster count: for a monotone similarity drop the +// number of clusters is non-increasing as the order grows (larger order => fewer, +// bigger clusters). The expected count is approximately 1 + (N-1)/order. +func TestWatershedTreeOrderControlsCount(t *testing.T) { + emb := cosEmbeddings([]float64{0.2, 0.4, 0.6, 0.8}) // N = 5 points + cases := []struct { + order int + want int + }{ + {2, 3}, + {3, 2}, + {4, 2}, + {10, 2}, + } + prev := 1 << 30 + for _, c := range cases { + labels, err := watershed(emb, c.order) + if err != nil { + t.Fatalf("order %v: %v", c.order, err) + } + got := uniqueCount(labels) + if got != c.want { + t.Fatalf("order %v: got %d clusters, want %d", c.order, got, c.want) + } + if got > prev { + t.Fatalf("order %v: clusters %d > previous %d (not non-increasing)", c.order, got, prev) + } + prev = got + } +} + +func TestWatershedDeterministic(t *testing.T) { + emb := cosEmbeddings([]float64{0.1, 0.9, 0.2, 0.8, 0.3}) + a, err := watershed(emb, 4) + if err != nil { + t.Fatal(err) + } + b, err := watershed(emb, 4) + if err != nil { + t.Fatal(err) + } + for i := range a { + if a[i] != b[i] { + t.Fatalf("non-deterministic at %d: %v vs %v", i, a, b) + } + } +} + +func TestWatershedScaleInvariant(t *testing.T) { + emb := cosEmbeddings([]float64{0.2, 0.4, 0.6, 0.8}) + scaled := make([][]float64, len(emb)) + for i, r := range emb { + scaled[i] = []float64{r[0] * 5, r[1] * 5} + } + a, _ := watershed(emb, 2) + b, _ := watershed(scaled, 2) + for i := range a { + if a[i] != b[i] { + t.Fatalf("scale changed labels at %d: %v vs %v", i, a, b) + } + } +} + +// TestWatershedTreeOrderDefault verifies the exported default is honored when +// the parameter is absent. +func TestWatershedTreeOrderDefault(t *testing.T) { + emb := cosEmbeddings([]float64{0.2, 0.4, 0.6, 0.8}) + a, _ := watershed(emb, DefaultTreeOrder) + b, _ := watershed(emb, 4) + for i := range a { + if a[i] != b[i] { + t.Fatalf("default order != 4 at %d: %v vs %v", i, a, b) + } + } +} diff --git a/internal/ingestion/component/knowledge_compiler/structure/structure.go b/internal/ingestion/component/knowledge_compiler/structure/structure.go index a36d64391e..788abaed2c 100644 --- a/internal/ingestion/component/knowledge_compiler/structure/structure.go +++ b/internal/ingestion/component/knowledge_compiler/structure/structure.go @@ -141,30 +141,16 @@ func Run(ctx context.Context, deps common.Deps, param common.Param, inputs commo 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 - } + // Buffer every product (plus the graph) in one slice; the component merges + // them into the upstream chunk stream (matching Python, which appends + // compiled units onto the chunk list). + products := append([]common.Product{}, prods...) + products = append(products, graphProduct) out := common.Outputs{ - Products: sink.Products(), - VectorBytes: sink.Bytes(), - Items: sink.TotalItems(), - Flushed: sink.Flushed(), + Products: products, 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 } diff --git a/internal/ingestion/component/knowledge_compiler/wiki/wiki.go b/internal/ingestion/component/knowledge_compiler/wiki/wiki.go index 8ffe6c3e33..2cb95570af 100644 --- a/internal/ingestion/component/knowledge_compiler/wiki/wiki.go +++ b/internal/ingestion/component/knowledge_compiler/wiki/wiki.go @@ -181,22 +181,13 @@ func Run(ctx context.Context, deps common.Deps, param common.Param, inputs commo 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 - } - } + // Buffer every product in one slice; the component merges them into the + // upstream chunk stream (matching Python, which appends compiled units onto + // the chunk list). out := common.Outputs{ - Products: sink.Products(), - VectorBytes: sink.Bytes(), - Items: sink.TotalItems(), - Flushed: sink.Flushed(), + Products: prod, DuplicatesDropped: stats.DuplicatesDropped, } - if err := out.EnforceGuardrails(param.Guardrails, inputs.Sink, ctx); err != nil { - return common.Outputs{}, err - } return out, nil } diff --git a/internal/ingestion/knowledge_compile/reader.go b/internal/ingestion/knowledge_compile/reader.go index 1b6d9d43bb..93c0a54104 100644 --- a/internal/ingestion/knowledge_compile/reader.go +++ b/internal/ingestion/knowledge_compile/reader.go @@ -18,7 +18,6 @@ package knowledge_compile import ( "context" "fmt" - "strings" "ragflow/internal/engine" "ragflow/internal/engine/types" @@ -129,40 +128,14 @@ func productFromChunkMap(c map[string]interface{}, tenant string) (kccommon.Prod meta["source_doc_ids"] = v } + vec, _ := kccommon.VectorFromChunkMap(c, 0) return kccommon.Product{ ID: id, DocID: docID, TenantID: tenant, Variant: kccommon.Variant(variant), Content: content, - Vector: vectorFromChunkMap(c), + Vector: vec, Meta: meta, }, true } - -// vectorFromChunkMap extracts the embedding from the q__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/task/knowledge_compiler_wiring.go b/internal/ingestion/task/knowledge_compiler_wiring.go index 244bef3e3e..06a651986b 100644 --- a/internal/ingestion/task/knowledge_compiler_wiring.go +++ b/internal/ingestion/task/knowledge_compiler_wiring.go @@ -21,10 +21,12 @@ import ( "fmt" "strings" "sync/atomic" + "time" "ragflow/internal/dao" "ragflow/internal/engine" enginetypes "ragflow/internal/engine/types" + "ragflow/internal/entity" "ragflow/internal/entity/models" kc "ragflow/internal/ingestion/component/knowledge_compiler/common" "ragflow/internal/service" @@ -70,6 +72,17 @@ func newKnowledgeCompilerDepsResolver() kc.DepsResolver { if strings.TrimSpace(llmID) == "" { return kc.Deps{}, fmt.Errorf("knowledge_compiler: llm_id is required for production deps resolution") } + // Resolve the chat model's context window so RAPTOR can truncate each + // cluster's texts to fit the LLM context (mirrors Python self._llm_model.max_length). + llmMax := kc.DefaultLLMContextLength + // Bound the model-config lookup so a stalled provider/instance DB read + // cannot block document ingestion indefinitely. + ctx, cancel := context.WithTimeout(context.Background(), 30*time.Second) + defer cancel() + if _, _, _, ml, merr := svc.ResolveModelConfig(ctx, tenantID, entity.ModelTypeChat, llmID); merr == nil && ml > 0 { + llmMax = ml + } + return kc.Deps{ Chat: &kcChatInvoker{svc: svc, tenantID: tenantID, llmID: llmID}, Embed: &kcEmbedder{svc: svc, tenantID: tenantID, embdID: embeddingModel}, @@ -77,6 +90,7 @@ func newKnowledgeCompilerDepsResolver() kc.DepsResolver { // HistoricalKNN / Redis are optional (wiki historical dedup, // datasetnav lock). They are wired separately when the // surrounding pipeline supplies the backing services. + LLMMaxLength: llmMax, }, nil } } @@ -101,8 +115,16 @@ func (c *kcChatInvoker) Chat(ctx context.Context, req kc.ChatRequest) (*kc.ChatR // 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} + if req.Temperature != nil || req.MaxTokens != nil { + config = &models.ChatConfig{} + if req.Temperature != nil { + config.Temperature = req.Temperature + } + // MaxTokens caps the generated summary length (mirrors Python's + // {"max_tokens": max(self._max_token, 512)}, issue #10235). + if req.MaxTokens != nil { + config.MaxTokens = req.MaxTokens + } } resp, err := c.svc.Chat(ctx, c.tenantID, llmID, msgs, config) if err != nil { diff --git a/internal/ingestion/component/knowledge_compiler/mindmap/dictify.go b/internal/utility/markdown_to_json.go similarity index 61% rename from internal/ingestion/component/knowledge_compiler/mindmap/dictify.go rename to internal/utility/markdown_to_json.go index 5c55266319..7ea56c7ed0 100644 --- a/internal/ingestion/component/knowledge_compiler/mindmap/dictify.go +++ b/internal/utility/markdown_to_json.go @@ -1,4 +1,75 @@ -package mindmap +// Package utility provides a Go port of the Python `markdown_to_json` +// library's markdown-outline → nested-dict pipeline, plus the +// list-to-key/value, merge, and tree-shaping steps that RAGFlow's +// MindMapExtractor layers on top. +// +// # Provenance +// +// This file reproduces two Python sources, kept byte-for-behavior where the +// port is faithful: +// +// 1. The `markdown_to_json` package (vendored in this repo at +// .venv/lib/python3.13/site-packages/markdown_to_json/markdown_to_json.py): +// - dictify -> Dictify +// - CMarkASTNester.nest / _dictify_blocks / _ensure_list_singleton +// + dictify_list_by -> dictifyBlocks +// - Renderer.stringify_dict / _valuify / _render_block / +// _render_generic_block / _render_List / _render_FencedCode +// -> valuify / renderBlockText / renderBlockAny / renderList / toBlock +// +// 2. RAGFlow's MindMapExtractor +// (rag/advanced_rag/knowlege_compile/mind_map_extractor.py): +// - _todict (+ _list_to_kv) -> Todict +// - _merge -> MergeDicts +// - _be_children -> BeChildren +// - _key -> StripStars +// - __call__'s final shaping -> ShapeTree +// - re.sub(r"```[^\n]*", "") -> StripFences +// +// The public entry points used by the knowledge-compiler mindmap component are +// Dictify, Todict, MergeDicts, ShapeTree, StripFences, and the Node / OMap +// types. The remaining helpers are unexported. +// +// # Deviations from the Python sources +// +// The port is behaviorally faithful for the common mindmap output shape, but +// the following intentional or incidental differences exist (all documented so +// the divergences are auditable against the Python golden behavior): +// +// 1. (Fixed here) Closed ATX headings: Python's CommonMark strips BOTH the +// leading and the trailing '#' run ("## Title ##" -> "Title"). This port +// now strips the trailing '#' too, in headingRawText. +// +// 2. Heading text is taken from the raw source line (goldmark's +// Heading.Lines) instead of goldmark's inline-parsed text, so inline +// emphasis markers (asterisks) survive exactly as Python's block.strings +// would keep them. +// +// 3. A list that is NOT the first child of a heading is rendered as readable +// flattened item text joined by "\n"; Python's _valuify falls back to +// str() of the nested list (a Python list literal). This is a deliberate +// divergence for the rare "list mid-children" case. +// +// 4. Fenced code blocks: goldmark keeps each source line's trailing newline, +// so the rendered block carries one extra trailing "\n" compared with +// Python's string_content. Minor. +// +// 5. Todict term keying types the preceding item as a string; if it is not a +// string the key becomes "" and is later dropped by BeChildren. Python +// would raise TypeError (unhashable) on a list term. The port is more +// robust. +// +// 6. BeChildren skips non-string list items instead of raising on unhashable +// list members (same defensive rationale as #5). +// +// 7. The fence-strip-before-parse behavior is inherited from Python +// (re.sub(r"```[^\n]*","") runs before parsing), so any fenced code in an +// LLM reply is stripped before it is parsed. Faithful to the original. +// +// 8. An empty merged dict degrades to a bare root node (ShapeTree returns +// &Node{ID:"root"}); Python's __call__ would raise IndexError on +// keys()[0]. The port keeps the pipeline alive. +package utility import ( "regexp" @@ -9,19 +80,13 @@ import ( "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 +// 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 +type OMap []kv // Node is one shaped mind-map node (Python's {"id", "children"} shape). type Node struct { @@ -57,11 +122,11 @@ func StripFences(s string) string { return fenceStripRE.ReplaceAllString(s, "") } -// dictify mirrors markdown_to_json.dictify: nest top-level blocks by the +// 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 { +func Dictify(md string) OMap { src := []byte(md) doc := goldmark.DefaultParser().Parse(text.NewReader(src)) var blocks []mdBlock @@ -69,7 +134,7 @@ func dictify(md string) omap { blocks = append(blocks, toBlock(n, src)) } if len(blocks) == 0 { - return omap{} + return OMap{} } min := 100000 for _, b := range blocks { @@ -82,8 +147,8 @@ func dictify(md string) omap { } } switch t := dictifyBlocks(blocks, min).(type) { - case omap: - out := omap{} + case OMap: + out := OMap{} for _, e := range t { out = append(out, kv{e.k, valuify(e.v)}) } @@ -95,13 +160,13 @@ func dictify(md string) omap { for _, b := range t { rendered = append(rendered, renderBlockAny(b)) } - return omap{{"root", rendered}} + return OMap{{"root", rendered}} } - return omap{} + return OMap{} } // dictifyBlocks mirrors CMarkASTNester._dictify_blocks: split blocks by -// headings at exactly `level`, recursing at level+1. Returns omap or the raw +// 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 @@ -114,7 +179,7 @@ func dictifyBlocks(blocks []mdBlock, level int) any { if !has { return blocks } - out := omap{} + out := OMap{} var cur string var children []mdBlock started := false @@ -145,8 +210,8 @@ func dictifyBlocks(blocks []mdBlock, level int) any { // a list → its items; otherwise the blocks joined with "\n\n". func valuify(v any) any { switch t := v.(type) { - case omap: - out := omap{} + case OMap: + out := OMap{} for _, e := range t { out = append(out, kv{e.k, valuify(e.v)}) } @@ -169,7 +234,8 @@ func valuify(v any) any { // 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). +// (Python str()'s the nested list — a pathological case we render readably; +// see deviation #3). func renderBlockText(b mdBlock) string { switch b.typ { case fencedBlock: @@ -194,17 +260,17 @@ func renderBlockAny(b mdBlock) any { return renderBlockText(b) } -// todict mirrors _todict + _list_to_kv: lists under dict keys convert to +// 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 { +func Todict(m OMap) OMap { for i := range m { switch val := m[i].v.(type) { - case omap: - m[i].v = todict(val) + case OMap: + m[i].v = Todict(val) case []any: - nv := omap{} + nv := OMap{} for j, item := range val { lst, isList := item.([]any) if !isList || j == 0 { @@ -223,10 +289,10 @@ func todict(m omap) omap { return m } -// mergeDicts mirrors _merge: d1 merges INTO d2 (ordered). Both dicts recurse; +// 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 { +func MergeDicts(d1, d2 OMap) OMap { for _, e := range d1 { idx := indexOMap(d2, e.k) if idx < 0 { @@ -234,13 +300,13 @@ func mergeDicts(d1, d2 omap) omap { continue } dv := d2[idx].v - d1Dict, d1IsDict := e.v.(omap) - d2Dict, d2IsDict := dv.(omap) + 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) + d2[idx].v = MergeDicts(d1Dict, d2Dict) case d1IsList && d2IsList: d2[idx].v = append(d2List, d1List...) default: @@ -250,50 +316,51 @@ func mergeDicts(d1, d2 omap) omap { return d2 } -// shapeTree mirrors __call__'s final shaping: >1 top-level keys → synthetic +// 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 { +// keys()[0]; the Go port keeps the pipeline alive — deviation #8). +func ShapeTree(merged OMap) *Node { if len(merged) > 1 { keyset := map[string]bool{} for _, e := range merged { - if _, ok := e.v.(omap); !ok { + if _, ok := e.v.(OMap); !ok { continue } - if k := stripStars(e.k); k != "" { + if k := StripStars(e.k); k != "" { keyset[k] = true } } var children []*Node for _, e := range merged { - if _, ok := e.v.(omap); !ok { + if _, ok := e.v.(OMap); !ok { continue } - k := stripStars(e.k) + k := StripStars(e.k) if k == "" { continue } - children = append(children, &Node{ID: k, Children: beChildren(e.v, keyset)}) + 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})} + 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 +// 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 { +// Non-string list items are skipped (deviation #6) rather than raising. +func BeChildren(v any, keyset map[string]bool) []*Node { switch t := v.(type) { case string: keyset[t] = true - if id := stripStars(t); id != "" { + if id := StripStars(t); id != "" { return []*Node{{ID: id}} } return nil @@ -309,20 +376,20 @@ func beChildren(v any, keyset map[string]bool) []*Node { if !ok { continue } - if id := stripStars(s); id != "" { + if id := StripStars(s); id != "" { out = append(out, &Node{ID: id}) } } return out - case omap: + case OMap: var out []*Node for _, e := range t { - k := stripStars(e.k) + k := StripStars(e.k) if k == "" || keyset[k] { continue } keyset[k] = true - out = append(out, &Node{ID: k, Children: beChildren(e.v, keyset)}) + out = append(out, &Node{ID: k, Children: BeChildren(e.v, keyset)}) } return out } @@ -332,8 +399,8 @@ func beChildren(v any, keyset map[string]bool) []*Node { // starsRE mirrors re.compile(r"\*+") used by _key. var starsRE = regexp.MustCompile(`\*+`) -// stripStars mirrors _key (re.sub(r"\*+", "", k)). -func stripStars(s string) string { +// StripStars mirrors _key (re.sub(r"\*+", "", k)). +func StripStars(s string) string { return starsRE.ReplaceAllString(s, "") } @@ -356,10 +423,13 @@ func toBlock(n ast.Node, src []byte) mdBlock { // 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). +// block.strings (goldmark's inline text would strip emphasis). Both leading +// and trailing '#' runs are removed (deviation #1 fixed: Python's CommonMark +// strips the closing '#' sequence too). func headingRawText(n *ast.Heading, src []byte) string { raw := strings.TrimSpace(rawLines(n, src, "\n")) raw = strings.TrimLeft(raw, "#") + raw = strings.TrimRight(raw, "#") return strings.TrimSpace(raw) } @@ -399,7 +469,7 @@ func renderList(list *ast.List, src []byte) []any { // 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 { +func setOMap(m OMap, k string, v any) OMap { if idx := indexOMap(m, k); idx >= 0 { m[idx].v = v return m @@ -407,7 +477,7 @@ func setOMap(m omap, k string, v any) omap { return append(m, kv{k, v}) } -func indexOMap(m omap, k string) int { +func indexOMap(m OMap, k string) int { for i, e := range m { if e.k == k { return i diff --git a/internal/utility/markdown_to_json_test.go b/internal/utility/markdown_to_json_test.go new file mode 100644 index 0000000000..cc06f912a3 --- /dev/null +++ b/internal/utility/markdown_to_json_test.go @@ -0,0 +1,177 @@ +package utility + +import ( + "strings" + "testing" +) + +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 TestDictify_ClosedHeadingStripsTrailingHashes(t *testing.T) { + // Deviation #1 fix: "## Title ##" must yield "Title" like Python's + // CommonMark (leading AND trailing '#' runs removed). + md := "# Root ##\n\n## Title ##\n\nbody\n" + got := Dictify(md) + if len(got) != 1 || got[0].k != "Root" { + t.Fatalf("closed root heading = %+v, want key \"Root\"", got) + } + children, ok := got[0].v.(OMap) + if !ok || len(children) != 1 || children[0].k != "Title" { + t.Fatalf("closed child heading = %+v, want key \"Title\"", got[0].v) + } +} + +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) + } +}