mirror of
https://github.com/infiniflow/ragflow.git
synced 2026-07-11 14:15:40 +08:00
401 lines
14 KiB
Go
401 lines
14 KiB
Go
package task
|
||
|
||
import (
|
||
"math"
|
||
"testing"
|
||
"time"
|
||
)
|
||
|
||
// =============================================================================
|
||
// TruncateTexts — Python: tiktoken.cl100k_base token-level truncation
|
||
// =============================================================================
|
||
|
||
func TestTruncateTexts_TokenLevel(t *testing.T) {
|
||
// Python: enc.encode("hello world") → [15339, 1917] (2 tokens)
|
||
// With maxLength=12, safeMax=2, both tokens fit → unchanged
|
||
result := TruncateTexts([]string{"hello world"}, 12)
|
||
if len(result) != 1 {
|
||
t.Fatalf("len = %d, want 1", len(result))
|
||
}
|
||
if result[0] != "hello world" {
|
||
t.Errorf("with safeMax=2, 'hello world' (2 tokens) should fit, got %q", result[0])
|
||
}
|
||
}
|
||
|
||
func TestTruncateTexts_TruncatesByToken(t *testing.T) {
|
||
// Python: with safeMax=1, keep only 1 token → shorter than original
|
||
result := TruncateTexts([]string{"hello world"}, 11)
|
||
// safeMax = 11-10 = 1, keeps only 1 token
|
||
if len(result[0]) >= len("hello world") || len(result[0]) == 0 {
|
||
t.Errorf("safeMax=1 should produce shorter output, got %q", result[0])
|
||
}
|
||
}
|
||
|
||
func TestTruncateTexts_EmptyText(t *testing.T) {
|
||
result := TruncateTexts([]string{""}, 100)
|
||
if len(result) != 1 || result[0] != "" {
|
||
t.Errorf("empty text should remain empty, got %q", result[0])
|
||
}
|
||
}
|
||
|
||
func TestTruncateTexts_NilInput(t *testing.T) {
|
||
result := TruncateTexts(nil, 100)
|
||
if result != nil {
|
||
t.Errorf("expected nil for nil input, got %v", result)
|
||
}
|
||
}
|
||
|
||
func TestTruncateTexts_EmptySlice(t *testing.T) {
|
||
result := TruncateTexts([]string{}, 100)
|
||
if len(result) != 0 {
|
||
t.Errorf("expected empty slice, got len=%d", len(result))
|
||
}
|
||
}
|
||
|
||
func TestTruncateTexts_MultipleTexts(t *testing.T) {
|
||
texts := []string{"hello world", "short"}
|
||
result := TruncateTexts(texts, 100)
|
||
if len(result) != 2 {
|
||
t.Fatalf("len = %d, want 2", len(result))
|
||
}
|
||
if result[0] != "hello world" {
|
||
t.Errorf("first text should be unchanged")
|
||
}
|
||
if result[1] != "short" {
|
||
t.Errorf("second text should be unchanged")
|
||
}
|
||
}
|
||
|
||
// =============================================================================
|
||
// SplitQuestions — Python: str.split("\n"), keeps empty strings
|
||
// =============================================================================
|
||
|
||
func TestSplitQuestions_Basic(t *testing.T) {
|
||
result := SplitQuestions("Q1\nQ2\nQ3")
|
||
if len(result) != 3 {
|
||
t.Fatalf("len = %d, want 3", len(result))
|
||
}
|
||
if result[0] != "Q1" || result[1] != "Q2" || result[2] != "Q3" {
|
||
t.Errorf("got %v, want [Q1 Q2 Q3]", result)
|
||
}
|
||
}
|
||
|
||
func TestSplitQuestions_Empty(t *testing.T) {
|
||
result := SplitQuestions("")
|
||
// Python: "".split("\n") → [""]
|
||
if len(result) != 1 || result[0] != "" {
|
||
t.Errorf("got %v, want [\"\"]", result)
|
||
}
|
||
}
|
||
|
||
func TestSplitQuestions_TrailingNewline(t *testing.T) {
|
||
result := SplitQuestions("Q1\nQ2\n")
|
||
// Python: "Q1\nQ2\n".split("\n") → ["Q1", "Q2", ""]
|
||
if len(result) != 3 {
|
||
t.Fatalf("len = %d, want 3, got %v", len(result), result)
|
||
}
|
||
if result[2] != "" {
|
||
t.Errorf("last element should be empty string, got %q", result[2])
|
||
}
|
||
}
|
||
|
||
func TestSplitQuestions_NoNewline(t *testing.T) {
|
||
result := SplitQuestions("single")
|
||
if len(result) != 1 || result[0] != "single" {
|
||
t.Errorf("got %v, want [single]", result)
|
||
}
|
||
}
|
||
|
||
// =============================================================================
|
||
// SplitKeywords — Python: re.split(r"[,，;；、\r\n]+", ...)
|
||
// =============================================================================
|
||
|
||
func TestSplitKeywords_Comma(t *testing.T) {
|
||
result := SplitKeywords("kw1,kw2,kw3")
|
||
if len(result) != 3 {
|
||
t.Fatalf("len = %d, want 3", len(result))
|
||
}
|
||
}
|
||
|
||
func TestSplitKeywords_ChineseComma(t *testing.T) {
|
||
result := SplitKeywords("kw1，kw2，kw3")
|
||
if len(result) != 3 {
|
||
t.Fatalf("len = %d, want 3", len(result))
|
||
}
|
||
}
|
||
|
||
func TestSplitKeywords_MixedDelimiters(t *testing.T) {
|
||
result := SplitKeywords("kw1,kw2；kw3")
|
||
if len(result) != 3 {
|
||
t.Fatalf("len = %d, want 3, got %v", len(result), result)
|
||
}
|
||
}
|
||
|
||
func TestSplitKeywords_FiltersEmptyStrings(t *testing.T) {
|
||
result := SplitKeywords("kw1,,kw2")
|
||
// Python: re.split → ["kw1", "", "kw2"] → after filter: ["kw1", "kw2"]
|
||
if len(result) != 2 {
|
||
t.Errorf("empty strings should be filtered: got %v", result)
|
||
}
|
||
}
|
||
|
||
func TestSplitKeywords_Empty(t *testing.T) {
|
||
result := SplitKeywords("")
|
||
// Python: re.split on "" → [""], then filtered by if k.strip()
|
||
if len(result) != 0 {
|
||
t.Errorf("got %v, want empty", result)
|
||
}
|
||
}
|
||
|
||
// =============================================================================
|
||
// CreateChunkTime — Python: datetime.now().timestamp() has sub-second precision
|
||
// =============================================================================
|
||
|
||
func TestCreateChunkTime_Format(t *testing.T) {
|
||
timeStr, ts := CreateChunkTime()
|
||
if timeStr == "" {
|
||
t.Error("create_time should not be empty")
|
||
}
|
||
if len(timeStr) != 19 {
|
||
t.Errorf("expected 19 chars (2006-01-02 15:04:05), got %q", timeStr)
|
||
}
|
||
// Python: datetime.now().timestamp() returns float with fractional part
|
||
frac := ts - math.Floor(ts)
|
||
if frac <= 0 {
|
||
t.Errorf("timestamp should have sub-second precision, got frac=%f", frac)
|
||
}
|
||
}
|
||
|
||
// =============================================================================
|
||
// RenameTextToContentWithWeight — Python processChunks logic
|
||
// =============================================================================
|
||
|
||
func TestRenameTextToContentWithWeight_Basic(t *testing.T) {
|
||
chunk := map[string]any{"text": "hello world"}
|
||
RenameTextToContentWithWeight(chunk)
|
||
if _, exists := chunk["text"]; exists {
|
||
t.Error("text key should be removed")
|
||
}
|
||
if chunk["content_with_weight"] != "hello world" {
|
||
t.Errorf("content_with_weight = %q, want \"hello world\"", chunk["content_with_weight"])
|
||
}
|
||
}
|
||
|
||
func TestRenameTextToContentWithWeight_PreservesExisting(t *testing.T) {
|
||
chunk := map[string]any{"content_with_weight": "already set", "text": "hello"}
|
||
RenameTextToContentWithWeight(chunk)
|
||
if chunk["content_with_weight"] != "already set" {
|
||
t.Errorf("preserved value should not be overwritten")
|
||
}
|
||
if _, exists := chunk["text"]; exists {
|
||
t.Error("text should still be removed")
|
||
}
|
||
}
|
||
|
||
func TestRenameTextToContentWithWeight_NoTextKey(t *testing.T) {
|
||
chunk := map[string]any{"other": "value"}
|
||
RenameTextToContentWithWeight(chunk)
|
||
if _, exists := chunk["content_with_weight"]; exists {
|
||
t.Error("should not add content_with_weight when no text key")
|
||
}
|
||
}
|
||
|
||
// =============================================================================
|
||
// ProcessChunksForDataflow — Python: processChunks()
|
||
// =============================================================================
|
||
|
||
func TestProcessChunksForDataflow_SetsDocIDAndKBID(t *testing.T) {
|
||
chunks := []map[string]any{{"text": "hello world"}}
|
||
_ = ProcessChunksForDataflow(chunks, "doc-1", "kb-1", "test-doc.pdf", time.Now())
|
||
|
||
if chunks[0]["doc_id"] != "doc-1" {
|
||
t.Errorf("doc_id = %q, want \"doc-1\"", chunks[0]["doc_id"])
|
||
}
|
||
if kbIDs, ok := chunks[0]["kb_id"].([]string); ok {
|
||
if len(kbIDs) != 1 || kbIDs[0] != "kb-1" {
|
||
t.Errorf("kb_id = %v, want [\"kb-1\"]", chunks[0]["kb_id"])
|
||
}
|
||
} else {
|
||
t.Errorf("kb_id should be []string, got %T", chunks[0]["kb_id"])
|
||
}
|
||
}
|
||
|
||
func TestProcessChunksForDataflow_SetsDocNameKwd(t *testing.T) {
|
||
chunks := []map[string]any{{"text": "hello"}}
|
||
_ = ProcessChunksForDataflow(chunks, "doc-1", "kb-1", "test-doc.pdf", time.Now())
|
||
if chunks[0]["docnm_kwd"] != "test-doc.pdf" {
|
||
t.Errorf("docnm_kwd = %q, want \"test-doc.pdf\"", chunks[0]["docnm_kwd"])
|
||
}
|
||
}
|
||
|
||
func TestProcessChunksForDataflow_SetsTimeFields(t *testing.T) {
|
||
now := time.Now()
|
||
chunks := []map[string]any{{"text": "hello"}}
|
||
_ = ProcessChunksForDataflow(chunks, "doc-1", "kb-1", "test-doc.pdf", now)
|
||
|
||
if timeStr, ok := chunks[0]["create_time"].(string); ok {
|
||
if timeStr != now.Format("2006-01-02 15:04:05") {
|
||
t.Errorf("create_time = %q, want %q", timeStr, now.Format("2006-01-02 15:04:05"))
|
||
}
|
||
} else {
|
||
t.Errorf("create_time should be string, got %T", chunks[0]["create_time"])
|
||
}
|
||
if ts, ok := chunks[0]["create_timestamp_flt"].(float64); ok {
|
||
expected := float64(now.UnixMicro()) / 1e6
|
||
if ts != expected {
|
||
t.Errorf("create_timestamp_flt = %f, want %f", ts, expected)
|
||
}
|
||
} else {
|
||
t.Errorf("create_timestamp_flt should be float64, got %T", chunks[0]["create_timestamp_flt"])
|
||
}
|
||
}
|
||
|
||
func TestProcessChunksForDataflow_GeneratesID(t *testing.T) {
|
||
chunks := []map[string]any{{"text": "hello"}}
|
||
_ = ProcessChunksForDataflow(chunks, "doc-1", "kb-1", "test-doc.pdf", time.Now())
|
||
id, ok := chunks[0]["id"].(string)
|
||
if !ok || id == "" {
|
||
t.Errorf("id should be non-empty string, got %v", chunks[0]["id"])
|
||
}
|
||
}
|
||
|
||
func TestProcessChunksForDataflow_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")
|
||
}
|
||
}()
|
||
_ = ProcessChunksForDataflow(chunks, "doc-1", "kb-1", "test-doc.pdf", time.Now())
|
||
}
|
||
|
||
func TestProcessChunksForDataflow_RemovesInternalPipelineFields(t *testing.T) {
|
||
chunks := []map[string]any{{
|
||
"text": "hello",
|
||
"_pdf_positions": []any{[]any{0, 1, 2, 3, 4}},
|
||
"image": "data:image/png;base64,abc",
|
||
}}
|
||
|
||
_ = ProcessChunksForDataflow(chunks, "doc-1", "kb-1", "test-doc.pdf", time.Now())
|
||
if _, exists := chunks[0]["_pdf_positions"]; exists {
|
||
t.Fatalf("_pdf_positions should be removed before indexing: %v", chunks[0]["_pdf_positions"])
|
||
}
|
||
if _, exists := chunks[0]["image"]; exists {
|
||
t.Fatalf("image should be removed before indexing: %v", chunks[0]["image"])
|
||
}
|
||
}
|
||
|
||
func TestProcessChunksForDataflow_PreservesExistingID(t *testing.T) {
|
||
chunks := []map[string]any{{"text": "hello", "id": "existing-id"}}
|
||
_ = ProcessChunksForDataflow(chunks, "doc-1", "kb-1", "test-doc.pdf", time.Now())
|
||
if chunks[0]["id"] != "existing-id" {
|
||
t.Errorf("existing id should be preserved, got %q", chunks[0]["id"])
|
||
}
|
||
}
|
||
|
||
func TestProcessChunksForDataflow_QuestionsProcessing(t *testing.T) {
|
||
chunks := []map[string]any{{"text": "hello", "questions": "Q1\nQ2\nQ3"}}
|
||
_ = ProcessChunksForDataflow(chunks, "doc-1", "kb-1", "test-doc.pdf", time.Now())
|
||
|
||
if _, exists := chunks[0]["questions"]; exists {
|
||
t.Error("questions key should be removed")
|
||
}
|
||
kwd, ok := chunks[0]["question_kwd"].([]string)
|
||
if !ok {
|
||
t.Fatalf("question_kwd should be []string, got %T", chunks[0]["question_kwd"])
|
||
}
|
||
if len(kwd) != 3 {
|
||
t.Errorf("question_kwd len = %d, want 3", len(kwd))
|
||
}
|
||
if _, ok := chunks[0]["question_tks"].(string); !ok {
|
||
t.Errorf("question_tks should be string, got %T", chunks[0]["question_tks"])
|
||
}
|
||
}
|
||
|
||
func TestProcessChunksForDataflow_KeywordsProcessing(t *testing.T) {
|
||
chunks := []map[string]any{{"text": "hello", "keywords": "kw1,kw2;kw3"}}
|
||
_ = ProcessChunksForDataflow(chunks, "doc-1", "kb-1", "test-doc.pdf", time.Now())
|
||
|
||
if _, exists := chunks[0]["keywords"]; exists {
|
||
t.Error("keywords key should be removed")
|
||
}
|
||
kwd, ok := chunks[0]["important_kwd"].([]string)
|
||
if !ok || len(kwd) == 0 {
|
||
t.Errorf("important_kwd should be non-empty []string, got %v", chunks[0]["important_kwd"])
|
||
}
|
||
if _, ok := chunks[0]["important_tks"].(string); !ok {
|
||
t.Errorf("important_tks should be string, got %T", chunks[0]["important_tks"])
|
||
}
|
||
}
|
||
|
||
func TestProcessChunksForDataflow_SummaryProcessing(t *testing.T) {
|
||
chunks := []map[string]any{{"text": "hello", "summary": "This is a summary."}}
|
||
_ = ProcessChunksForDataflow(chunks, "doc-1", "kb-1", "test-doc.pdf", time.Now())
|
||
|
||
if _, exists := chunks[0]["summary"]; exists {
|
||
t.Error("summary key should be removed")
|
||
}
|
||
if _, ok := chunks[0]["content_ltks"].(string); !ok {
|
||
t.Errorf("content_ltks should be string, got %T", chunks[0]["content_ltks"])
|
||
}
|
||
if _, ok := chunks[0]["content_sm_ltks"].(string); !ok {
|
||
t.Errorf("content_sm_ltks should be string, got %T", chunks[0]["content_sm_ltks"])
|
||
}
|
||
}
|
||
|
||
func TestProcessChunksForDataflow_TextRenamed(t *testing.T) {
|
||
chunks := []map[string]any{{"text": "hello world"}}
|
||
_ = ProcessChunksForDataflow(chunks, "doc-1", "kb-1", "test-doc.pdf", time.Now())
|
||
|
||
if _, exists := chunks[0]["text"]; exists {
|
||
t.Error("text key should be removed")
|
||
}
|
||
if chunks[0]["content_with_weight"] != "hello world" {
|
||
t.Errorf("content_with_weight = %q, want \"hello world\"", chunks[0]["content_with_weight"])
|
||
}
|
||
}
|
||
|
||
func TestProcessChunksForDataflow_PreservesContentWithWeight(t *testing.T) {
|
||
chunks := []map[string]any{{"content_with_weight": "already set", "text": "hello"}}
|
||
_ = ProcessChunksForDataflow(chunks, "doc-1", "kb-1", "test-doc.pdf", time.Now())
|
||
if chunks[0]["content_with_weight"] != "already set" {
|
||
t.Errorf("content_with_weight = %q, want \"already set\"", chunks[0]["content_with_weight"])
|
||
}
|
||
}
|
||
|
||
func TestProcessChunksForDataflow_MetadataExtraction(t *testing.T) {
|
||
chunks := []map[string]any{
|
||
{"text": "hello", "metadata": map[string]any{"author": "Alice", "year": "2024"}},
|
||
}
|
||
meta := ProcessChunksForDataflow(chunks, "doc-1", "kb-1", "test-doc.pdf", time.Now())
|
||
|
||
if _, exists := chunks[0]["metadata"]; exists {
|
||
t.Error("metadata key should be removed from chunk")
|
||
}
|
||
if meta["author"] != "Alice" {
|
||
t.Errorf("metadata[\"author\"] = %q, want \"Alice\"", meta["author"])
|
||
}
|
||
}
|
||
|
||
func TestProcessChunksForDataflow_MetadataMergeAcrossChunks(t *testing.T) {
|
||
chunks := []map[string]any{
|
||
{"text": "chunk1", "metadata": map[string]any{"author": "Alice"}},
|
||
{"text": "chunk2", "metadata": map[string]any{"year": "2024"}},
|
||
}
|
||
meta := ProcessChunksForDataflow(chunks, "doc-1", "kb-1", "test-doc.pdf", time.Now())
|
||
|
||
if meta["author"] != "Alice" {
|
||
t.Errorf("author = %q, want \"Alice\"", meta["author"])
|
||
}
|
||
if meta["year"] != "2024" {
|
||
t.Errorf("year = %v, want \"2024\"", meta["year"])
|
||
}
|
||
}
|
||
|
||
func TestProcessChunksForDataflow_NilChunks(t *testing.T) {
|
||
meta := ProcessChunksForDataflow(nil, "doc-1", "kb-1", "test-doc.pdf", time.Now())
|
||
if len(meta) != 0 {
|
||
t.Errorf("metadata should be empty for nil chunks, got %v", meta)
|
||
}
|
||
}
|