zlei6/ragflow
1
0
Fork 0
mirror of https://github.com/infiniflow/ragflow.git synced 2026-07-31 21:13:49 +08:00
Code Issues Packages Projects Releases Wiki Activity

Align Go knowledge compiler RAPTOR with Python and drop guardrails (#17573)

Browse Source 
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.
This commit is contained in:
Zhichang Yu
2026-07-30 19:56:17 +08:00
committed by GitHub
parent 38cb3f13de
commit 75ac8cec2e
23 changed files with 1311 additions and 1804 deletions
Download Patch File Download Diff File Expand all files Collapse all files

100
internal/ingestion/component/knowledge_compiler/common/common_test.go
View File

@@ -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) }

16
internal/ingestion/component/knowledge_compiler/common/deps.go
View File

@@ -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.

240
internal/ingestion/component/knowledge_compiler/common/types.go
View File

@@ -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_<dim>_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_<dim>_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
}

61
internal/ingestion/component/knowledge_compiler/component.go
View File

@@ -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_<dim>_vec); variants fall back to embedding
// on demand when it is absent. A chunk must carry exactly one vector.
vec, err := common.VectorFromChunkMap(m, 0)
if err != nil {
return in, err
}
ch.Vector = vec
in.Chunks = append(in.Chunks, ch)
}
}
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{

145
internal/ingestion/component/knowledge_compiler/component_test.go
View File

@@ -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)
}
}
}

32
internal/ingestion/component/knowledge_compiler/datasetnav/datasetnav.go
View File

@@ -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
}

16
internal/ingestion/component/knowledge_compiler/golden/fixtures/raptor_baseline.json
View File

@@ -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
}

46
internal/ingestion/component/knowledge_compiler/golden_test.go
View File

@@ -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"))
}
}
}

35
internal/ingestion/component/knowledge_compiler/mindmap/mindmap.go
View File

@@ -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))

167
internal/ingestion/component/knowledge_compiler/mindmap/mindmap_test.go
View File

@@ -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

680
internal/ingestion/component/knowledge_compiler/raptor/clustering.go
View File

@@ -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
}

84
internal/ingestion/component/knowledge_compiler/raptor/clustering_test.go
View File

@@ -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))
}
}

119
internal/ingestion/component/knowledge_compiler/raptor/psi.go
View File

@@ -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
}

504
internal/ingestion/component/knowledge_compiler/raptor/raptor.go
View File

@@ -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 </think> / </think:6124c78e> 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"^.*</think>", "", response, DOTALL)).
if i := strings.LastIndex(content, "</think>"); i >= 0 {
content = content[i+len("</think>"):]
} else if i := strings.LastIndex(content, "</think:6124c78e>"); i >= 0 {
content = content[i+len("</think:6124c78e>"):]
}
// 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
}

148
internal/ingestion/component/knowledge_compiler/raptor/raptor_test.go Normal file
View File

@@ -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: "<think>let me think...\n\n</think>Final 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, "<think>") || strings.Contains(got, "</think>") {
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))
}
}

142
internal/ingestion/component/knowledge_compiler/raptor/watershed.go Normal file
View File

@@ -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
}

127
internal/ingestion/component/knowledge_compiler/raptor/watershed_test.go Normal file
View File

@@ -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)
}
}
}

26
internal/ingestion/component/knowledge_compiler/structure/structure.go
View File

@@ -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
}

17
internal/ingestion/component/knowledge_compiler/wiki/wiki.go
View File

@@ -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
}

31
internal/ingestion/knowledge_compile/reader.go
View File

@@ -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_<dim>_vec column, whose
// exact name depends on the embedding dimension.
func vectorFromChunkMap(c map[string]interface{}) []float32 {
for k, v := range c {
if strings.HasPrefix(k, "q_") && strings.HasSuffix(k, "_vec") {
return toFloat32Slice(v)
}
}
return nil
}
func toFloat32Slice(v interface{}) []float32 {
switch arr := v.(type) {
case []float32:
return arr
case []interface{}:
out := make([]float32, 0, len(arr))
for _, e := range arr {
if f, ok := e.(float64); ok {
out = append(out, float32(f))
}
}
return out
}
return nil
}

26
internal/ingestion/task/knowledge_compiler_wiring.go
View File

@@ -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 {

176
internal/ingestion/component/knowledge_compiler/mindmap/dictify.go → internal/utility/markdown_to_json.go
View File

@@ -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

177
internal/utility/markdown_to_json_test.go Normal file
View File

@@ -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)
}
}