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

Port RAPTOR knowledge compilation to match Python: remove Go-only capacity guardrails, fix prompts, add token truncation, response post-processing, retries, replace AHC with watershed.
2026-08-01 21:37:33 +08:00 · 2026-07-30 19:56:17 +08:00
parent 38cb3f13de
commit 75ac8cec2e
23 changed files with 1311 additions and 1804 deletions
--- a/internal/utility/markdown_to_json.go
+++ b/internal/utility/markdown_to_json.go
@@ -0,0 +1,487 @@
+// 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"
+	"strings"
+
+	"github.com/yuin/goldmark"
+	"github.com/yuin/goldmark/ast"
+	"github.com/yuin/goldmark/text"
+)
+
+// kv is one ordered key/value pair; OMap mirrors Python's OrderedDict (key
+// order is load-bearing for tree shape and merge order).
+type kv struct {
+	k string
+	v any
+}
+type OMap []kv
+
+// Node is one shaped mind-map node (Python's {"id", "children"} shape).
+type Node struct {
+	ID       string
+	Children []*Node
+}
+
+type blockType int
+
+const (
+	headingBlock blockType = iota
+	listBlock
+	paraBlock
+	fencedBlock
+	otherBlock
+)
+
+// mdBlock is one markdown block in dictify terms: heading (level+text), list
+// (flattened items), paragraph/fenced/other (text).
+type mdBlock struct {
+	typ   blockType
+	level int
+	text  string
+	items []any
+}
+
+// fenceStripRE mirrors re.sub(r"```[^\n]*", "", response): whole fence marker
+// lines (``` or ```lang) are removed before parsing.
+var fenceStripRE = regexp.MustCompile("```[^\\n]*")
+
+// StripFences removes markdown code-fence marker lines from an LLM reply.
+func StripFences(s string) string {
+	return fenceStripRE.ReplaceAllString(s, "")
+}
+
+// Dictify mirrors markdown_to_json.dictify: nest top-level blocks by the
+// MINIMUM top heading level; content before the first heading at each level
+// is dropped; when a level has no headings the blocks become a list (the top
+// level then renders as {"root": [...]}).
+func Dictify(md string) OMap {
+	src := []byte(md)
+	doc := goldmark.DefaultParser().Parse(text.NewReader(src))
+	var blocks []mdBlock
+	for n := doc.FirstChild(); n != nil; n = n.NextSibling() {
+		blocks = append(blocks, toBlock(n, src))
+	}
+	if len(blocks) == 0 {
+		return OMap{}
+	}
+	min := 100000
+	for _, b := range blocks {
+		lvl := 100000
+		if b.typ == headingBlock {
+			lvl = b.level
+		}
+		if lvl < min {
+			min = lvl
+		}
+	}
+	switch t := dictifyBlocks(blocks, min).(type) {
+	case OMap:
+		out := OMap{}
+		for _, e := range t {
+			out = append(out, kv{e.k, valuify(e.v)})
+		}
+		return out
+	case []mdBlock:
+		// No headings anywhere: {"root": [rendered blocks]} (Renderer's
+		// non-dict branch).
+		rendered := make([]any, 0, len(t))
+		for _, b := range t {
+			rendered = append(rendered, renderBlockAny(b))
+		}
+		return OMap{{"root", rendered}}
+	}
+	return OMap{}
+}
+
+// dictifyBlocks mirrors CMarkASTNester._dictify_blocks: split blocks by
+// headings at exactly `level`, recursing at level+1. Returns OMap or the raw
+// []mdBlock (no headings at this level).
+func dictifyBlocks(blocks []mdBlock, level int) any {
+	has := false
+	for _, b := range blocks {
+		if b.typ == headingBlock && b.level == level {
+			has = true
+			break
+		}
+	}
+	if !has {
+		return blocks
+	}
+	out := OMap{}
+	var cur string
+	var children []mdBlock
+	started := false
+	for _, b := range blocks {
+		if b.typ == headingBlock && b.level == level {
+			if started {
+				out = setOMap(out, cur, children)
+			}
+			cur = b.text
+			children = nil
+			started = true
+			continue
+		}
+		if started {
+			children = append(children, b)
+		}
+	}
+	if started {
+		out = setOMap(out, cur, children)
+	}
+	for i := range out {
+		out[i].v = dictifyBlocks(out[i].v.([]mdBlock), level+1)
+	}
+	return out
+}
+
+// valuify mirrors Renderer._valuify: dict → recurse; empty → ""; first child
+// a list → its items; otherwise the blocks joined with "\n\n".
+func valuify(v any) any {
+	switch t := v.(type) {
+	case OMap:
+		out := OMap{}
+		for _, e := range t {
+			out = append(out, kv{e.k, valuify(e.v)})
+		}
+		return out
+	case []mdBlock:
+		if len(t) == 0 {
+			return ""
+		}
+		if t[0].typ == listBlock {
+			return t[0].items
+		}
+		parts := make([]string, 0, len(t))
+		for _, b := range t {
+			parts = append(parts, renderBlockText(b))
+		}
+		return strings.Join(parts, "\n\n")
+	}
+	return ""
+}
+
+// renderBlockText renders one block to a string (Renderer._render_*).
+// A list appearing mid-children is rendered as its flattened item texts
+// (Python str()'s the nested list — a pathological case we render readably;
+// see deviation #3).
+func renderBlockText(b mdBlock) string {
+	switch b.typ {
+	case fencedBlock:
+		return "```\n" + b.text + "```"
+	case listBlock:
+		var parts []string
+		for _, item := range b.items {
+			if s, ok := item.(string); ok {
+				parts = append(parts, s)
+			}
+		}
+		return strings.Join(parts, "\n")
+	default:
+		return b.text
+	}
+}
+
+func renderBlockAny(b mdBlock) any {
+	if b.typ == listBlock {
+		return b.items
+	}
+	return renderBlockText(b)
+}
+
+// Todict mirrors _todict + _list_to_kv: lists under dict keys convert to
+// key→definition pairs ([prev scalar] → sublist[0]); a list with no
+// definition sub-lists collapses to an empty dict (the bullets are dropped —
+// faithful to Python). Recurses into nested dicts.
+func Todict(m OMap) OMap {
+	for i := range m {
+		switch val := m[i].v.(type) {
+		case OMap:
+			m[i].v = Todict(val)
+		case []any:
+			nv := OMap{}
+			for j, item := range val {
+				lst, isList := item.([]any)
+				if !isList || j == 0 {
+					continue
+				}
+				key, _ := val[j-1].(string)
+				var first any
+				if len(lst) > 0 {
+					first = lst[0]
+				}
+				nv = append(nv, kv{key, first})
+			}
+			m[i].v = nv
+		}
+	}
+	return m
+}
+
+// MergeDicts mirrors _merge: d1 merges INTO d2 (ordered). Both dicts recurse;
+// both lists concatenate d2's then d1's items; otherwise d1 wins. Keys absent
+// from d2 append in d1's order.
+func MergeDicts(d1, d2 OMap) OMap {
+	for _, e := range d1 {
+		idx := indexOMap(d2, e.k)
+		if idx < 0 {
+			d2 = append(d2, e)
+			continue
+		}
+		dv := d2[idx].v
+		d1Dict, d1IsDict := e.v.(OMap)
+		d2Dict, d2IsDict := dv.(OMap)
+		d1List, d1IsList := e.v.([]any)
+		d2List, d2IsList := dv.([]any)
+		switch {
+		case d1IsDict && d2IsDict:
+			d2[idx].v = MergeDicts(d1Dict, d2Dict)
+		case d1IsList && d2IsList:
+			d2[idx].v = append(d2List, d1List...)
+		default:
+			d2[idx].v = e.v
+		}
+	}
+	return d2
+}
+
+// ShapeTree mirrors __call__'s final shaping: >1 top-level keys → synthetic
+// "root" node whose children are the dict-valued keys (with a keyset seeded
+// by those keys); exactly 1 → that key becomes the root id. An empty merged
+// dict degrades to a bare root (Python would raise IndexError on
+// keys()[0]; the Go port keeps the pipeline alive — 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 {
+				continue
+			}
+			if k := StripStars(e.k); k != "" {
+				keyset[k] = true
+			}
+		}
+		var children []*Node
+		for _, e := range merged {
+			if _, ok := e.v.(OMap); !ok {
+				continue
+			}
+			k := StripStars(e.k)
+			if k == "" {
+				continue
+			}
+			children = append(children, &Node{ID: k, Children: BeChildren(e.v, keyset)})
+		}
+		return &Node{ID: "root", Children: children}
+	}
+	if len(merged) == 1 {
+		k := StripStars(merged[0].k)
+		return &Node{ID: k, Children: BeChildren(merged[0].v, map[string]bool{k: true})}
+	}
+	return &Node{ID: "root"}
+}
+
+// BeChildren mirrors _be_children: strings/lists become leaf nodes; dict keys
+// become nested nodes. keyset is SHARED across the whole subtree walk — a key
+// already seen anywhere below is skipped (silent dedup, as in Python).
+// 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 != "" {
+			return []*Node{{ID: id}}
+		}
+		return nil
+	case []any:
+		for _, item := range t {
+			if s, ok := item.(string); ok {
+				keyset[s] = true
+			}
+		}
+		var out []*Node
+		for _, item := range t {
+			s, ok := item.(string)
+			if !ok {
+				continue
+			}
+			if id := StripStars(s); id != "" {
+				out = append(out, &Node{ID: id})
+			}
+		}
+		return out
+	case OMap:
+		var out []*Node
+		for _, e := range t {
+			k := StripStars(e.k)
+			if k == "" || keyset[k] {
+				continue
+			}
+			keyset[k] = true
+			out = append(out, &Node{ID: k, Children: BeChildren(e.v, keyset)})
+		}
+		return out
+	}
+	return nil
+}
+
+// starsRE mirrors re.compile(r"\*+") used by _key.
+var starsRE = regexp.MustCompile(`\*+`)
+
+// StripStars mirrors _key (re.sub(r"\*+", "", k)).
+func StripStars(s string) string {
+	return starsRE.ReplaceAllString(s, "")
+}
+
+// ---- goldmark AST → mdBlock ----
+
+func toBlock(n ast.Node, src []byte) mdBlock {
+	switch node := n.(type) {
+	case *ast.Heading:
+		return mdBlock{typ: headingBlock, level: node.Level, text: headingRawText(node, src)}
+	case *ast.List:
+		return mdBlock{typ: listBlock, items: renderList(node, src)}
+	case *ast.FencedCodeBlock:
+		return mdBlock{typ: fencedBlock, text: rawLines(n, src, "")}
+	case *ast.Paragraph:
+		return mdBlock{typ: paraBlock, text: rawLines(n, src, "\n")}
+	default:
+		return mdBlock{typ: otherBlock, text: rawLines(n, src, "\n")}
+	}
+}
+
+// headingRawText renders the heading's raw source line without the ATX
+// markers, keeping inline markdown (e.g. asterisks) intact like CommonMark's
+// block.strings (goldmark's inline text would strip emphasis). 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)
+}
+
+// rawLines joins a node's source lines (sep between lines). For fenced code
+// the lines already carry their newlines (sep="").
+func rawLines(n ast.Node, src []byte, sep string) string {
+	lines := n.Lines()
+	parts := make([]string, 0, lines.Len())
+	for i := 0; i < lines.Len(); i++ {
+		seg := lines.At(i)
+		v := string(seg.Value(src))
+		if sep != "" {
+			v = strings.TrimRight(v, "\n")
+		}
+		parts = append(parts, v)
+	}
+	return strings.Join(parts, sep)
+}
+
+// renderList mirrors Renderer._render_List: each ListItem's children render
+// (paragraph text / nested list) and the results flatten one level, so an
+// item contributes its text and (if any) a nested item list.
+func renderList(list *ast.List, src []byte) []any {
+	var out []any
+	for item := list.FirstChild(); item != nil; item = item.NextSibling() {
+		for c := item.FirstChild(); c != nil; c = c.NextSibling() {
+			switch n := c.(type) {
+			case *ast.List:
+				out = append(out, renderList(n, src))
+			default:
+				out = append(out, rawLines(n, src, "\n"))
+			}
+		}
+	}
+	return out
+}
+
+// setOMap sets key k, replacing the value in place when the key repeats
+// (OrderedDict semantics: position kept), else appending.
+func setOMap(m OMap, k string, v any) OMap {
+	if idx := indexOMap(m, k); idx >= 0 {
+		m[idx].v = v
+		return m
+	}
+	return append(m, kv{k, v})
+}
+
+func indexOMap(m OMap, k string) int {
+	for i, e := range m {
+		if e.k == k {
+			return i
+		}
+	}
+	return -1
+}
--- a/internal/utility/markdown_to_json_test.go
+++ b/internal/utility/markdown_to_json_test.go
@@ -0,0 +1,177 @@
+package utility
+
+import (
+	"strings"
+	"testing"
+)
+
+func TestDictify_HeadingsNest(t *testing.T) {
+	md := "# Top\n\n## Mid A\n\n## Mid B\n\n### Deep\n"
+	got := Dictify(md)
+	if len(got) != 1 || got[0].k != "Top" {
+		t.Fatalf("top-level = %+v, want single key Top", got)
+	}
+	mid, ok := got[0].v.(OMap)
+	if !ok || len(mid) != 2 {
+		t.Fatalf("Top value = %+v, want {Mid A, Mid B}", got[0].v)
+	}
+	if mid[0].k != "Mid A" || mid[0].v != "" {
+		t.Errorf("Mid A = %+v, want empty string value (no children)", mid[0])
+	}
+	deep, ok := mid[1].v.(OMap)
+	if !ok || len(deep) != 1 || deep[0].k != "Deep" {
+		t.Errorf("Mid B value = %+v, want {Deep}", mid[1].v)
+	}
+}
+
+func TestDictify_BulletPairsViaTodict(t *testing.T) {
+	// "- term\n  - def" pairs become {term: def}; the unpaired bullet
+	// ("term2") is dropped — faithful to _list_to_kv.
+	md := "# T\n\n## S\n- term1\n  - def1\n- term2\n"
+	got := Todict(Dictify(md))
+	if len(got) != 1 || got[0].k != "T" {
+		t.Fatalf("dictify = %+v", got)
+	}
+	s, ok := got[0].v.(OMap)
+	if !ok || len(s) != 1 || s[0].k != "S" {
+		t.Fatalf("T value = %+v", got[0].v)
+	}
+	pairs, ok := s[0].v.(OMap)
+	if !ok || len(pairs) != 1 || pairs[0].k != "term1" || pairs[0].v != "def1" {
+		t.Fatalf("S value = %+v, want {term1: def1} (term2 dropped)", s[0].v)
+	}
+}
+
+func TestDictify_PlainListVanishes(t *testing.T) {
+	// Bullets without definition sub-lists collapse to an empty dict
+	// (Python _list_to_kv parity).
+	md := "# T\n\n- a\n- b\n"
+	got := Todict(Dictify(md))
+	if len(got) != 1 {
+		t.Fatalf("dictify = %+v", got)
+	}
+	if v, ok := got[0].v.(OMap); !ok || len(v) != 0 {
+		t.Fatalf("plain bullets must collapse to empty dict, got %+v", got[0].v)
+	}
+}
+
+func TestDictify_PlainProse(t *testing.T) {
+	got := Todict(Dictify("hello world, no headings here"))
+	if len(got) != 1 || got[0].k != "root" {
+		t.Fatalf("prose dictify = %+v, want single root key", got)
+	}
+	// The root list has no definition pairs → collapses to empty dict.
+	if v, ok := got[0].v.(OMap); !ok || len(v) != 0 {
+		t.Fatalf("root value = %+v, want empty dict", got[0].v)
+	}
+}
+
+func TestDictify_JumpedHeadingBecomesText(t *testing.T) {
+	// H1 → H3 (no H2): the jumped heading is not treated as a key.
+	md := "# H1\n\n### H3\n\ntext\n"
+	got := Dictify(md)
+	if len(got) != 1 || got[0].k != "H1" {
+		t.Fatalf("dictify = %+v", got)
+	}
+	if got[0].v != "H3\n\ntext" {
+		t.Fatalf("H1 value = %q, want jumped heading rendered as text", got[0].v)
+	}
+}
+
+func TestDictify_LeadingContentDropped(t *testing.T) {
+	// Content before the first heading at the minimum level is discarded
+	// (dictify_list_by parity).
+	md := "preamble\n\n# H\n\nbody\n"
+	got := Dictify(md)
+	if len(got) != 1 || got[0].k != "H" || got[0].v != "body" {
+		t.Fatalf("leading content must be dropped: %+v", got)
+	}
+}
+
+func TestDictify_ClosedHeadingStripsTrailingHashes(t *testing.T) {
+	// Deviation #1 fix: "## Title ##" must yield "Title" like Python's
+	// CommonMark (leading AND trailing '#' runs removed).
+	md := "# Root ##\n\n## Title ##\n\nbody\n"
+	got := Dictify(md)
+	if len(got) != 1 || got[0].k != "Root" {
+		t.Fatalf("closed root heading = %+v, want key \"Root\"", got)
+	}
+	children, ok := got[0].v.(OMap)
+	if !ok || len(children) != 1 || children[0].k != "Title" {
+		t.Fatalf("closed child heading = %+v, want key \"Title\"", got[0].v)
+	}
+}
+
+func TestMergeDicts_Order(t *testing.T) {
+	d1 := OMap{{"A", []any{"x"}}}
+	d2 := OMap{{"A", []any{"y"}}, {"B", "b"}}
+	got := MergeDicts(d1, d2)
+	if len(got) != 2 || got[0].k != "A" || got[1].k != "B" {
+		t.Fatalf("merge = %+v", got)
+	}
+	list, _ := got[0].v.([]any)
+	if len(list) != 2 || list[0] != "y" || list[1] != "x" {
+		t.Fatalf("list merge = %v, want [y x] (d2 items before d1)", list)
+	}
+}
+
+func TestShapeTree_MultiTopKeys(t *testing.T) {
+	merged := OMap{
+		{"A", OMap{{"shared", "x"}}},
+		{"B", OMap{{"shared", "y"}}},
+	}
+	root := ShapeTree(merged)
+	if root.ID != "root" || len(root.Children) != 2 {
+		t.Fatalf("shape = %+v, want root with A,B", root)
+	}
+	// The second "shared" key must survive (keyset dedups within a subtree
+	// walk, and A's subtree consumes "shared" first — Python parity: B's
+	// "shared" is dropped because the keyset is shared across children).
+	seen := map[string]int{}
+	var walk func(n *Node)
+	walk = func(n *Node) {
+		seen[n.ID]++
+		for _, c := range n.Children {
+			walk(c)
+		}
+	}
+	walk(root)
+	if seen["shared"] != 1 {
+		t.Fatalf("shared key count = %d, want 1 (global keyset dedup)", seen["shared"])
+	}
+}
+
+func TestShapeTree_SingleKey(t *testing.T) {
+	merged := OMap{{"Only", OMap{{"child", "text"}}}}
+	root := ShapeTree(merged)
+	if root.ID != "Only" {
+		t.Fatalf("single-key root id = %q, want Only", root.ID)
+	}
+	if len(root.Children) != 1 || root.Children[0].ID != "child" {
+		t.Fatalf("root children = %+v", root.Children)
+	}
+}
+
+func TestBeChildren_StringLeaf(t *testing.T) {
+	keyset := map[string]bool{}
+	out := BeChildren("plain string", keyset)
+	if len(out) != 1 || out[0].ID != "plain string" {
+		t.Fatalf("string leaf = %+v", out)
+	}
+	if !keyset["plain string"] {
+		t.Fatalf("raw string must enter the keyset")
+	}
+}
+
+func TestStripStars(t *testing.T) {
+	if got := StripStars("**Bold** title"); got != "Bold title" {
+		t.Fatalf("StripStars = %q", got)
+	}
+}
+
+func TestStripFences(t *testing.T) {
+	in := "```json\n{\"a\": 1}\n```\n"
+	if got := StripFences(in); strings.Contains(got, "```") {
+		t.Fatalf("fences not stripped: %q", got)
+	}
+}