mirror of
https://github.com/infiniflow/ragflow.git
synced 2026-08-16 21:50:58 +08:00
0784bef5b0
Ports Python dataset-level structure aggregation (timeline, graph, mindmap) to Go. Mindmap emits entity/relation rows and merges like graph. Adds dataset_merge guard, engine gate, resolveDatasetStructureKind, kind-required structure graph GET/DELETE API, per-index task-id fields.
97 lines
2.9 KiB
Go
97 lines
2.9 KiB
Go
package mindmap
|
|
|
|
import (
|
|
"strings"
|
|
"testing"
|
|
|
|
"ragflow/internal/utility"
|
|
)
|
|
|
|
func TestRenderPrompt(t *testing.T) {
|
|
got := renderPrompt("THE TEXT")
|
|
if !strings.Contains(got, "-TEXT-\nTHE TEXT") {
|
|
t.Fatalf("input_text not substituted: %q", got)
|
|
}
|
|
if strings.Contains(got, "{input_text}") {
|
|
t.Fatalf("placeholder survived: %q", got)
|
|
}
|
|
// The verbatim prompt body must be present.
|
|
if !strings.Contains(got, "Generate a title for user's 'TEXT'。") {
|
|
t.Fatalf("prompt body drifted")
|
|
}
|
|
}
|
|
|
|
func TestPackSections_Budget(t *testing.T) {
|
|
// budget = max(4096*0.8, 4096-512) = 3584; with the len-based fake
|
|
// tokenizer each section of 2000 tokens packs one per batch.
|
|
tok := fakeTok{}
|
|
sections := []string{strings.Repeat("a", 8000), strings.Repeat("b", 8000), "short"}
|
|
got := packSections(sections, tok)
|
|
if len(got) != 2 {
|
|
t.Fatalf("batches = %d, want 2 (2000+2000 > 3584 splits)", len(got))
|
|
}
|
|
// An oversized section is never split: it forms its own batch.
|
|
big := packSections([]string{strings.Repeat("x", 20000)}, tok)
|
|
if len(big) != 1 {
|
|
t.Fatalf("oversized section must stay whole: %d batches", len(big))
|
|
}
|
|
}
|
|
|
|
type fakeTok struct{}
|
|
|
|
func (fakeTok) NumTokens(s string) int { return len(s) / 4 }
|
|
|
|
func TestTreeToProducts_ParentLinks(t *testing.T) {
|
|
root := &utility.Node{ID: "root", Children: []*utility.Node{
|
|
{ID: "A", Children: []*utility.Node{{ID: "A1"}, {ID: "A2"}}},
|
|
{ID: "B"},
|
|
}}
|
|
products := treeToProducts("t1", "d1", root)
|
|
// 5 entities (root, A, A1, A2, B) + 4 relations (root→A, A→A1, A→A2, root→B).
|
|
if len(products) != 9 {
|
|
t.Fatalf("products = %d, want 9 (5 entities + 4 relations)", len(products))
|
|
}
|
|
entCount, relCount := 0, 0
|
|
fromTo := map[string]bool{}
|
|
for _, p := range products {
|
|
kind, _ := p.Meta["kind"].(string)
|
|
switch kind {
|
|
case "entity":
|
|
entCount++
|
|
if p.Meta["entity_type"] != "mindmap" {
|
|
t.Errorf("entity %v type = %v, want mindmap", p.Meta["name"], p.Meta["entity_type"])
|
|
}
|
|
if p.Meta["compile_kwd"] != "mindmap" {
|
|
t.Errorf("entity %v compile_kwd = %v, want mindmap", p.Meta["name"], p.Meta["compile_kwd"])
|
|
}
|
|
case "relation":
|
|
relCount++
|
|
from, _ := p.Meta["from"].(string)
|
|
to, _ := p.Meta["to"].(string)
|
|
fromTo[from+"->"+to] = true
|
|
if p.Meta["relation_type"] != "related" {
|
|
t.Errorf("relation %v->%v type = %v, want related", from, to, p.Meta["relation_type"])
|
|
}
|
|
default:
|
|
t.Errorf("unexpected kind %q", kind)
|
|
}
|
|
}
|
|
if entCount != 5 || relCount != 4 {
|
|
t.Errorf("entities=%d relations=%d, want 5/4", entCount, relCount)
|
|
}
|
|
for _, edge := range []string{"root->A", "A->A1", "A->A2", "root->B"} {
|
|
if !fromTo[edge] {
|
|
t.Errorf("missing relation %s", edge)
|
|
}
|
|
}
|
|
}
|
|
|
|
func TestTreeToProducts_EmptyAndNil(t *testing.T) {
|
|
if got := treeToProducts("t1", "d1", nil); len(got) != 0 {
|
|
t.Errorf("nil root produced %d products", len(got))
|
|
}
|
|
if got := treeToProducts("t1", "d1", &utility.Node{ID: ""}); len(got) != 0 {
|
|
t.Errorf("empty root id produced %d products", len(got))
|
|
}
|
|
}
|