ragflow/internal/service/kb_prompt_test.go

package service

import (
	"testing"

	"ragflow/internal/tokenizer"
)

func TestKbPrompt_Empty(t *testing.T) {
	if got := KbPrompt(nil, 100); got != "" {
		t.Errorf("expected empty for nil chunks")
	}
	if got := KbPrompt([]SourcedChunk{}, 100); got != "" {
		t.Errorf("expected empty for empty chunks")
	}
	if got := KbPrompt([]SourcedChunk{{Content: "x"}}, 0); got != "" {
		t.Errorf("expected empty for maxTokens=0")
	}
	if got := KbPrompt([]SourcedChunk{{Content: "x"}}, -1); got != "" {
		t.Errorf("expected empty for maxTokens=-1")
	}
}

func TestKbPrompt_Format(t *testing.T) {
	chunks := []SourcedChunk{{
		ID:      "abc",
		Content: "chunk content here",
		DocName: "Test Document",
		URL:     "http://example.com",
	}}
	result := KbPrompt(chunks, 10000)
	if result == "" {
		t.Fatal("expected non-empty prompt")
	}
	// Verify ID appears
	if !contains(result, "ID: abc") {
		t.Errorf("missing ID line: %s", result)
	}
	// Verify title
	if !contains(result, "Title: Test Document") {
		t.Errorf("missing title: %s", result)
	}
	// Verify URL
	if !contains(result, "URL: http://example.com") {
		t.Errorf("missing URL: %s", result)
	}
	// Verify content
	if !contains(result, "chunk content here") {
		t.Errorf("missing content: %s", result)
	}
	// Verify unicode box-drawing chars
	if !contains(result, "├──") {
		t.Errorf("missing tree drawing: %s", result)
	}
}

func TestKbPrompt_TokenLimit(t *testing.T) {
	chunks := []SourcedChunk{
		{ID: "1", Content: "a very long content that takes many tokens "},
		{ID: "2", Content: "second chunk content here"},
	}
	// Compute limit dynamically so the test works with both the C++
	// tokenizer and the rune-based fallback.
	entryTokens := tokenizer.NumTokensFromString(formatChunkEntry(chunks[0]))
	maxToks := int(float64(entryTokens+1) / 0.97) // just enough for first
	result := KbPrompt(chunks, maxToks)
	if !contains(result, "ID: 1") {
		t.Error("first chunk should be included")
	}
	if contains(result, "ID: 2") {
		t.Error("second chunk should be excluded under tight limit")
	}
}

func TestKbPrompt_DocMetadata(t *testing.T) {
	chunks := []SourcedChunk{{
		ID:      "abc",
		Content: "content",
		DocumentMetadata: map[string]interface{}{
			"author": "test author",
			"year":   "2024",
		},
	}}
	result := KbPrompt(chunks, 10000)
	if !contains(result, "author: test author") {
		t.Errorf("missing metadata author: %s", result)
	}
	if !contains(result, "year: 2024") {
		t.Errorf("missing metadata year: %s", result)
	}
}

func TestKbPrompt_NoDocNameOrURL(t *testing.T) {
	chunks := []SourcedChunk{{
		ID:      "simple",
		Content: "plain content",
	}}
	result := KbPrompt(chunks, 10000)
	if contains(result, "Title:") {
		t.Error("should not have title when empty")
	}
	if contains(result, "URL:") {
		t.Error("should not have URL when empty")
	}
}

func contains(s, substr string) bool {
	for i := 0; i <= len(s)-len(substr); i++ {
		if s[i:i+len(substr)] == substr {
			return true
		}
	}
	return false
}

func TestKbPrompt_TokenLimitAccurate(t *testing.T) {
	// Verify truncation uses tokenizer.NumTokensFromString by computing the limit
	// dynamically from the actual token count (works in both fallback
	// and C++ tokenizer environments).
	chunks := []SourcedChunk{
		{ID: "1", Content: "hello"},
		{ID: "2", Content: "world"},
	}
	entryTokens := tokenizer.NumTokensFromString(formatChunkEntry(chunks[0]))
	maxToks := int(float64(entryTokens+1) / 0.97) // just enough for first entry
	result := KbPrompt(chunks, maxToks)
	if !contains(result, "ID: 1") {
		t.Error("first chunk should fit")
	}
	if contains(result, "ID: 2") {
		t.Errorf("second chunk should be excluded: result = %q", result)
	}
}

func TestKbPrompt_AllFit(t *testing.T) {
	chunks := []SourcedChunk{
		{ID: "1", Content: "a"},
		{ID: "2", Content: "b"},
	}
	result := KbPrompt(chunks, 1000)
	if !contains(result, "ID: 1") || !contains(result, "ID: 2") {
		t.Error("both chunks should fit under generous limit")
	}
}