mirror of
https://github.com/infiniflow/ragflow.git
synced 2026-07-11 22:25:41 +08:00
428 lines
11 KiB
Go
428 lines
11 KiB
Go
package task
|
|
|
|
import (
|
|
"testing"
|
|
)
|
|
|
|
// =============================================================================
|
|
// NormalizeChunks
|
|
// =============================================================================
|
|
|
|
func TestNormalizeChunks_ChunksFormat(t *testing.T) {
|
|
input := map[string]any{
|
|
"chunks": []map[string]any{
|
|
{"text": "hello", "doc_type_kwd": "text"},
|
|
{"text": "world", "doc_type_kwd": "text"},
|
|
},
|
|
}
|
|
result := NormalizeChunks(input)
|
|
if len(result) != 2 {
|
|
t.Fatalf("len = %d, want 2", len(result))
|
|
}
|
|
if result[0]["text"] != "hello" {
|
|
t.Errorf("result[0][\"text\"] = %q, want \"hello\"", result[0]["text"])
|
|
}
|
|
}
|
|
|
|
func TestNormalizeChunks_JSONFormat(t *testing.T) {
|
|
input := map[string]any{
|
|
"json": []map[string]any{
|
|
{"text": "section 1", "doc_type_kwd": "text"},
|
|
},
|
|
}
|
|
result := NormalizeChunks(input)
|
|
if len(result) != 1 {
|
|
t.Fatalf("len = %d, want 1", len(result))
|
|
}
|
|
if result[0]["text"] != "section 1" {
|
|
t.Errorf("result[0][\"text\"] = %q, want \"section 1\"", result[0]["text"])
|
|
}
|
|
}
|
|
|
|
func TestNormalizeChunks_JSONFormatFromGenericSlice(t *testing.T) {
|
|
input := map[string]any{
|
|
"json": []any{
|
|
map[string]any{"text": "section 1", "doc_type_kwd": "text"},
|
|
},
|
|
}
|
|
result := NormalizeChunks(input)
|
|
if len(result) != 1 {
|
|
t.Fatalf("len = %d, want 1", len(result))
|
|
}
|
|
if result[0]["text"] != "section 1" {
|
|
t.Errorf("result[0][\"text\"] = %q, want \"section 1\"", result[0]["text"])
|
|
}
|
|
}
|
|
|
|
func TestNormalizeChunks_MarkdownFormat(t *testing.T) {
|
|
input := map[string]any{
|
|
"markdown": "# Title\n\nContent",
|
|
}
|
|
result := NormalizeChunks(input)
|
|
if len(result) != 1 {
|
|
t.Fatalf("len = %d, want 1", len(result))
|
|
}
|
|
text, ok := result[0]["text"].(string)
|
|
if !ok {
|
|
t.Fatalf("text should be string for markdown format, got %T", result[0]["text"])
|
|
}
|
|
if text != "# Title\n\nContent" {
|
|
t.Errorf("text = %q, want \"# Title\\n\\nContent\"", text)
|
|
}
|
|
}
|
|
|
|
func TestNormalizeChunks_TextFormat(t *testing.T) {
|
|
input := map[string]any{
|
|
"text": "plain text",
|
|
}
|
|
result := NormalizeChunks(input)
|
|
if len(result) != 1 {
|
|
t.Fatalf("len = %d, want 1", len(result))
|
|
}
|
|
text, ok := result[0]["text"].(string)
|
|
if !ok {
|
|
t.Fatalf("text should be string for text format, got %T", result[0]["text"])
|
|
}
|
|
if text != "plain text" {
|
|
t.Errorf("text = %q, want \"plain text\"", text)
|
|
}
|
|
}
|
|
|
|
func TestNormalizeChunks_HTMLFormat(t *testing.T) {
|
|
input := map[string]any{
|
|
"html": "<p>Hello</p>",
|
|
}
|
|
result := NormalizeChunks(input)
|
|
if len(result) != 1 {
|
|
t.Fatalf("len = %d, want 1", len(result))
|
|
}
|
|
text, ok := result[0]["text"].(string)
|
|
if !ok {
|
|
t.Fatalf("text should be string for html format, got %T", result[0]["text"])
|
|
}
|
|
if text != "<p>Hello</p>" {
|
|
t.Errorf("text = %q, want \"<p>Hello</p>\"", text)
|
|
}
|
|
}
|
|
|
|
func TestNormalizeChunks_EmptyOutput(t *testing.T) {
|
|
result := NormalizeChunks(map[string]any{})
|
|
if result != nil {
|
|
t.Errorf("expected nil for empty input, got %v", result)
|
|
}
|
|
}
|
|
|
|
func TestNormalizeChunks_EmptyMarkdown(t *testing.T) {
|
|
result := NormalizeChunks(map[string]any{"markdown": ""})
|
|
if result != nil {
|
|
t.Errorf("expected nil for empty markdown, got %v", result)
|
|
}
|
|
}
|
|
|
|
func TestNormalizeChunks_EmptyText(t *testing.T) {
|
|
result := NormalizeChunks(map[string]any{"text": ""})
|
|
if result != nil {
|
|
t.Errorf("expected nil for empty text, got %v", result)
|
|
}
|
|
}
|
|
|
|
func TestNormalizeChunks_EmptyHTML(t *testing.T) {
|
|
result := NormalizeChunks(map[string]any{"html": ""})
|
|
if result != nil {
|
|
t.Errorf("expected nil for empty html, got %v", result)
|
|
}
|
|
}
|
|
|
|
func TestNormalizeChunks_EmptyChunksList(t *testing.T) {
|
|
result := NormalizeChunks(map[string]any{"chunks": []map[string]any{}})
|
|
if len(result) != 0 {
|
|
t.Errorf("expected empty slice, got len=%d", len(result))
|
|
}
|
|
}
|
|
|
|
func TestNormalizeChunks_EmptyJSONList(t *testing.T) {
|
|
result := NormalizeChunks(map[string]any{"json": []map[string]any{}})
|
|
if len(result) != 0 {
|
|
t.Errorf("expected empty slice, got len=%d", len(result))
|
|
}
|
|
}
|
|
|
|
func TestNormalizeChunks_Priority(t *testing.T) {
|
|
t.Run("chunks over json", func(t *testing.T) {
|
|
input := map[string]any{
|
|
"chunks": []map[string]any{{"text": "from chunks"}},
|
|
"json": []map[string]any{{"text": "from json"}},
|
|
}
|
|
result := NormalizeChunks(input)
|
|
if result[0]["text"] != "from chunks" {
|
|
t.Errorf("chunks should win: got %q", result[0]["text"])
|
|
}
|
|
})
|
|
|
|
t.Run("json over markdown", func(t *testing.T) {
|
|
input := map[string]any{
|
|
"json": []map[string]any{{"text": "from json"}},
|
|
"markdown": "from markdown",
|
|
}
|
|
result := NormalizeChunks(input)
|
|
if result[0]["text"] != "from json" {
|
|
t.Errorf("json should win: got %q", result[0]["text"])
|
|
}
|
|
})
|
|
|
|
t.Run("markdown over text", func(t *testing.T) {
|
|
input := map[string]any{
|
|
"markdown": "from markdown",
|
|
"text": "from text",
|
|
}
|
|
result := NormalizeChunks(input)
|
|
text, ok := result[0]["text"].(string)
|
|
if !ok {
|
|
t.Fatalf("text should be string, got %T", result[0]["text"])
|
|
}
|
|
if text != "from markdown" {
|
|
t.Errorf("markdown should win: got %q", text)
|
|
}
|
|
})
|
|
}
|
|
|
|
func TestNormalizeChunks_DoesNotMutateInput(t *testing.T) {
|
|
original := []map[string]any{{"text": "original"}}
|
|
input := map[string]any{"chunks": original}
|
|
result := NormalizeChunks(input)
|
|
result[0]["text"] = "modified"
|
|
if original[0]["text"] != "original" {
|
|
t.Error("should deep copy, not mutate input")
|
|
}
|
|
}
|
|
|
|
func TestNormalizeChunks_DeepCopyVectors(t *testing.T) {
|
|
// Python: copy.deepcopy creates fully independent copies.
|
|
// Mutating a slice element in the result must NOT affect the original.
|
|
original_vec := []float64{0.1, 0.2, 0.3}
|
|
original := []map[string]any{{"text": "hello", "q_3_vec": original_vec}}
|
|
input := map[string]any{"chunks": original}
|
|
result := NormalizeChunks(input)
|
|
// Mutate the slice *element* in-place (not replace the slice)
|
|
result[0]["q_3_vec"].([]float64)[0] = 0.9
|
|
// Original must be unchanged
|
|
if original[0]["q_3_vec"].([]float64)[0] != 0.1 {
|
|
t.Error("mutating result vector element should not affect original")
|
|
}
|
|
}
|
|
|
|
func TestNormalizeChunks_NilInput(t *testing.T) {
|
|
result := NormalizeChunks(nil)
|
|
if result != nil {
|
|
t.Errorf("expected nil for nil input, got %v", result)
|
|
}
|
|
}
|
|
|
|
// =============================================================================
|
|
// PrepareTextsForDataflowEmbedding
|
|
// =============================================================================
|
|
|
|
func TestPrepareTexts_QuestionsPriority(t *testing.T) {
|
|
chunks := []map[string]any{
|
|
{"questions": "Q1\nQ2", "summary": "a summary", "text": "plain text"},
|
|
}
|
|
result := PrepareTextsForDataflowEmbedding(chunks)
|
|
if len(result) != 1 {
|
|
t.Fatalf("len = %d, want 1", len(result))
|
|
}
|
|
if result[0] != "Q1\nQ2" {
|
|
t.Errorf("questions should take priority: got %q", result[0])
|
|
}
|
|
}
|
|
|
|
func TestPrepareTexts_SummaryFallback(t *testing.T) {
|
|
chunks := []map[string]any{
|
|
{"summary": "a summary", "text": "plain text"},
|
|
}
|
|
result := PrepareTextsForDataflowEmbedding(chunks)
|
|
if result[0] != "a summary" {
|
|
t.Errorf("summary should be used when no questions: got %q", result[0])
|
|
}
|
|
}
|
|
|
|
func TestPrepareTexts_TextFallback(t *testing.T) {
|
|
chunks := []map[string]any{
|
|
{"text": "plain text"},
|
|
}
|
|
result := PrepareTextsForDataflowEmbedding(chunks)
|
|
if result[0] != "plain text" {
|
|
t.Errorf("text should be used when no questions/summary: got %q", result[0])
|
|
}
|
|
}
|
|
|
|
func TestPrepareTexts_EmptyStringFallback(t *testing.T) {
|
|
chunks := []map[string]any{
|
|
{"text": ""},
|
|
}
|
|
result := PrepareTextsForDataflowEmbedding(chunks)
|
|
if result[0] != "" {
|
|
t.Errorf("expected empty string, got %q", result[0])
|
|
}
|
|
}
|
|
|
|
func TestPrepareTexts_MultipleChunks(t *testing.T) {
|
|
chunks := []map[string]any{
|
|
{"questions": "Q1", "text": "t1"},
|
|
{"summary": "S2", "text": "t2"},
|
|
{"text": "t3"},
|
|
}
|
|
result := PrepareTextsForDataflowEmbedding(chunks)
|
|
if len(result) != 3 {
|
|
t.Fatalf("len = %d, want 3", len(result))
|
|
}
|
|
if result[0] != "Q1" {
|
|
t.Errorf("result[0] = %q, want \"Q1\"", result[0])
|
|
}
|
|
if result[1] != "S2" {
|
|
t.Errorf("result[1] = %q, want \"S2\"", result[1])
|
|
}
|
|
if result[2] != "t3" {
|
|
t.Errorf("result[2] = %q, want \"t3\"", result[2])
|
|
}
|
|
}
|
|
|
|
func TestPrepareTexts_NilChunks(t *testing.T) {
|
|
result := PrepareTextsForDataflowEmbedding(nil)
|
|
if result != nil {
|
|
t.Errorf("expected nil for nil chunks, got %v", result)
|
|
}
|
|
}
|
|
|
|
func TestPrepareTexts_EmptyChunks(t *testing.T) {
|
|
result := PrepareTextsForDataflowEmbedding([]map[string]any{})
|
|
if len(result) != 0 {
|
|
t.Errorf("expected empty slice, got len=%d", len(result))
|
|
}
|
|
}
|
|
|
|
func TestPrepareTexts_MissingTextKey(t *testing.T) {
|
|
chunks := []map[string]any{
|
|
{"other_key": "value"},
|
|
}
|
|
result := PrepareTextsForDataflowEmbedding(chunks)
|
|
if result[0] != "" {
|
|
t.Errorf("expected empty string for missing text key, got %q", result[0])
|
|
}
|
|
}
|
|
|
|
func TestPrepareTexts_PanicOnListText(t *testing.T) {
|
|
chunks := []map[string]any{
|
|
{"text": []any{"bad-shape"}},
|
|
}
|
|
defer func() {
|
|
if r := recover(); r == nil {
|
|
t.Error("expected panic for list-shaped text")
|
|
}
|
|
}()
|
|
_ = PrepareTextsForDataflowEmbedding(chunks)
|
|
}
|
|
|
|
func TestMustGetChunkTextString_PanicOnStringSlice(t *testing.T) {
|
|
chunk := map[string]any{"text": []string{"bad-shape"}}
|
|
defer func() {
|
|
if r := recover(); r == nil {
|
|
t.Error("expected panic for []string text")
|
|
}
|
|
}()
|
|
_ = MustGetChunkTextString(chunk, "unit-test")
|
|
}
|
|
|
|
// =============================================================================
|
|
// AttachVectors
|
|
// =============================================================================
|
|
|
|
func TestAttachVectors_Basic(t *testing.T) {
|
|
chunks := []map[string]any{
|
|
{"text": "hello"},
|
|
{"text": "world"},
|
|
}
|
|
vectors := [][]float64{
|
|
{0.1, 0.2, 0.3},
|
|
{0.4, 0.5, 0.6},
|
|
}
|
|
dim := AttachVectors(chunks, vectors)
|
|
if dim != 3 {
|
|
t.Errorf("dim = %d, want 3", dim)
|
|
}
|
|
if len(chunks[0]) == 0 || chunks[0]["q_3_vec"] == nil {
|
|
t.Errorf("expected q_3_vec in chunk[0], got %v", chunks[0])
|
|
}
|
|
vec0 := chunks[0]["q_3_vec"].([]float64)
|
|
if len(vec0) != 3 || vec0[0] != 0.1 {
|
|
t.Errorf("chunk[0] vector = %v, want [0.1 0.2 0.3]", vec0)
|
|
}
|
|
vec1 := chunks[1]["q_3_vec"].([]float64)
|
|
if len(vec1) != 3 || vec1[2] != 0.6 {
|
|
t.Errorf("chunk[1] vector = %v, want [0.4 0.5 0.6]", vec1)
|
|
}
|
|
}
|
|
|
|
func TestAttachVectors_EmptyChunks(t *testing.T) {
|
|
result := AttachVectors(nil, nil)
|
|
if result != 0 {
|
|
t.Errorf("expected 0 for empty chunks, got %d", result)
|
|
}
|
|
}
|
|
|
|
func TestAttachVectors_DifferentDimensions(t *testing.T) {
|
|
chunks := []map[string]any{
|
|
{"text": "a"},
|
|
{"text": "b"},
|
|
}
|
|
vectors := [][]float64{
|
|
{0.1, 0.2},
|
|
{0.3, 0.4, 0.5},
|
|
}
|
|
// Each chunk gets its own vector with key based on its dimension
|
|
AttachVectors(chunks, vectors)
|
|
if chunks[0]["q_2_vec"] == nil {
|
|
t.Errorf("chunk[0] should have q_2_vec, got keys: %v", chunkKeys(chunks[0]))
|
|
}
|
|
if chunks[1]["q_3_vec"] == nil {
|
|
t.Errorf("chunk[1] should have q_3_vec, got keys: %v", chunkKeys(chunks[1]))
|
|
}
|
|
}
|
|
|
|
func TestAttachVectors_MismatchedCount(t *testing.T) {
|
|
chunks := []map[string]any{
|
|
{"text": "hello"},
|
|
}
|
|
vectors := [][]float64{
|
|
{0.1},
|
|
{0.2},
|
|
}
|
|
defer func() {
|
|
if r := recover(); r == nil {
|
|
t.Error("expected panic for mismatched count")
|
|
}
|
|
}()
|
|
AttachVectors(chunks, vectors)
|
|
}
|
|
|
|
func TestAttachVectors_KeyFormat(t *testing.T) {
|
|
chunks := []map[string]any{
|
|
{"text": "hello"},
|
|
}
|
|
vectors := [][]float64{
|
|
{0.1, 0.2, 0.3, 0.4, 0.5},
|
|
}
|
|
AttachVectors(chunks, vectors)
|
|
key := "q_5_vec"
|
|
if chunks[0][key] == nil {
|
|
t.Errorf("expected key %q in chunk, got keys: %v", key, chunkKeys(chunks[0]))
|
|
}
|
|
}
|
|
|
|
func chunkKeys(m map[string]any) []string {
|
|
keys := make([]string, 0, len(m))
|
|
for k := range m {
|
|
keys = append(keys, k)
|
|
}
|
|
return keys
|
|
}
|