ragflow/internal/common/kg_scoring_test.go

//
//  Copyright 2026 The InfiniFlow Authors. All Rights Reserved.
//
//  Licensed under the Apache License, Version 2.0 (the "License");
//  you may not use this file except in compliance with the License.
//  You may obtain a copy of the License at
//
//      http://www.apache.org/licenses/LICENSE-2.0
//
//  Unless required by applicable law or agreed to in writing, software
//  distributed under the License is distributed on an "AS IS" BASIS,
//  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
//  See the License for the specific language governing permissions and
//  limitations under the License.
//

package common

import (
	"strings"
	"testing"
)

// --- AnalyzeNHopPaths ---

func TestAnalyzeNHopPaths_Basic(t *testing.T) {
	ents := map[string]*KGEntity{
		"A": {
			Sim: 0.9,
			NhopEnts: []NhopEntity{
				{Path: []string{"A", "B", "C"}, Weights: []float64{0.8, 0.5}},
			},
		},
	}
	result := AnalyzeNHopPaths(ents)
	// A→B: 0.9 / (2+0) = 0.45
	// B→C: 0.9 / (2+1) = 0.3
	if len(result) != 2 {
		t.Fatalf("expected 2 edges, got %d", len(result))
	}
	if result[Edge{"A", "B"}].Sim != 0.45 {
		t.Errorf("expected A→B sim=0.45, got %f", result[Edge{"A", "B"}].Sim)
	}
	if result[Edge{"B", "C"}].Sim != 0.3 {
		t.Errorf("expected B→C sim=0.3, got %f", result[Edge{"B", "C"}].Sim)
	}
}

func TestAnalyzeNHopPaths_MultipleContributors(t *testing.T) {
	ents := map[string]*KGEntity{
		"A": {
			Sim: 0.8,
			NhopEnts: []NhopEntity{
				{Path: []string{"A", "B"}, Weights: []float64{0.7}},
			},
		},
		"X": {
			Sim: 0.6,
			NhopEnts: []NhopEntity{
				{Path: []string{"X", "B"}, Weights: []float64{0.5}},
			},
		},
	}
	result := AnalyzeNHopPaths(ents)
	// A→B: 0.8 / 2 = 0.4
	// X→B: 0.6 / 2 = 0.3
	if result[Edge{"A", "B"}].Sim != 0.4 {
		t.Errorf("expected A→B sim=0.4, got %f", result[Edge{"A", "B"}].Sim)
	}
	if result[Edge{"X", "B"}].Sim != 0.3 {
		t.Errorf("expected X→B sim=0.3, got %f", result[Edge{"X", "B"}].Sim)
	}
}

func TestAnalyzeNHopPaths_Empty(t *testing.T) {
	result := AnalyzeNHopPaths(nil)
	if len(result) != 0 {
		t.Errorf("expected empty, got %d", len(result))
	}
}

// --- DoubleHitBoost ---

func TestDoubleHitBoost(t *testing.T) {
	ents := map[string]*KGEntity{
		"A": {Sim: 0.5},
		"B": {Sim: 0.3},
	}
	types := map[string]struct{}{"A": {}}
	DoubleHitBoost(ents, types)
	if ents["A"].Sim != 1.0 {
		t.Errorf("expected A sim=1.0 after boost, got %f", ents["A"].Sim)
	}
	if ents["B"].Sim != 0.3 {
		t.Errorf("expected B sim unchanged at 0.3, got %f", ents["B"].Sim)
	}
}

func TestDoubleHitBoost_Empty(t *testing.T) {
	ents := map[string]*KGEntity{"A": {Sim: 0.5}}
	DoubleHitBoost(ents, map[string]struct{}{})
	if ents["A"].Sim != 0.5 {
		t.Errorf("expected unchanged, got %f", ents["A"].Sim)
	}
}

// --- FuseRelationScores ---

func TestFuseRelationScores_NhopContribution(t *testing.T) {
	rels := map[Edge]*KGRelation{
		{"A", "B"}: {Sim: 0.5, PageRank: 0.8},
	}
	types := map[string]struct{}{}
	nhop := map[Edge]EdgeScore{
		{"A", "B"}: {Sim: 0.3},
	}
	FuseRelationScores(rels, types, nhop)
	// sim = 0.5 * (0.3 + 1) = 0.65
	if rels[Edge{"A", "B"}].Sim != 0.65 {
		t.Errorf("expected 0.65, got %f", rels[Edge{"A", "B"}].Sim)
	}
}

func TestFuseRelationScores_TypeBoost(t *testing.T) {
	rels := map[Edge]*KGRelation{
		{"A", "B"}: {Sim: 0.5},
	}
	types := map[string]struct{}{"A": {}, "B": {}}
	nhop := map[Edge]EdgeScore{}
	FuseRelationScores(rels, types, nhop)
	// Both endpoints in types: s=2, sim = 0.5 * (2+1) = 1.5
	if rels[Edge{"A", "B"}].Sim != 1.5 {
		t.Errorf("expected 1.5, got %f", rels[Edge{"A", "B"}].Sim)
	}
}

func TestFuseRelationScores_NhopNewEdge(t *testing.T) {
	rels := map[Edge]*KGRelation{}
	types := map[string]struct{}{}
	nhop := map[Edge]EdgeScore{
		{"A", "B"}: {Sim: 0.4, PageRank: 0.7},
	}
	FuseRelationScores(rels, types, nhop)
	if _, ok := rels[Edge{"A", "B"}]; !ok {
		t.Fatal("expected new edge from N-hop")
	}
	if rels[Edge{"A", "B"}].Sim != 0.4 {
		t.Errorf("expected sim=0.4, got %f", rels[Edge{"A", "B"}].Sim)
	}
}

// --- SortAndTrim ---

func TestSortAndTrimEntities(t *testing.T) {
	ents := map[string]*KGEntity{
		"A": {Sim: 0.5, PageRank: 0.9},
		"B": {Sim: 0.8, PageRank: 0.3},
		"C": {Sim: 0.9, PageRank: 0.1},
	}
	result := SortAndTrimEntities(ents, 2)
	if len(result) != 2 {
		t.Fatalf("expected 2, got %d", len(result))
	}
	// A: 0.45, B: 0.24, C: 0.09 → top 2 should be A, B
	if result[0].Entity != "A" {
		t.Errorf("expected A first (0.45), got %s (%f)", result[0].Entity, result[0].Score)
	}
}

func TestSortAndTrimEntities_DefaultTopN(t *testing.T) {
	ents := map[string]*KGEntity{
		"A": {Sim: 0.5, PageRank: 0.9},
		"B": {Sim: 0.8, PageRank: 0.3},
	}
	result := SortAndTrimEntities(ents, 0)
	if len(result) != 2 {
		t.Errorf("expected default topN to include all, got %d", len(result))
	}
}

func TestSortAndTrimRelations(t *testing.T) {
	rels := map[Edge]*KGRelation{
		{"A", "B"}: {Sim: 0.9, PageRank: 0.1},
		{"C", "D"}: {Sim: 0.3, PageRank: 0.8},
	}
	result := SortAndTrimRelations(rels, 1)
	if len(result) != 1 {
		t.Fatalf("expected 1, got %d", len(result))
	}
	// A→B: 0.09, C→D: 0.24 → C→D should be first
	if result[0].From != "C" {
		t.Errorf("expected C first (0.24), got %s (%f)", result[0].From, result[0].Score)
	}
}

// --- Format and Build ---

func TestBuildKGContent_Basic(t *testing.T) {
	entities := []ScoredEntity{
		{Entity: "A", Score: 0.45, Description: `{"description": "Entity A desc"}`},
	}
	relations := []ScoredRelation{
		{From: "A", To: "B", Score: 0.3, Description: `{"description": "rel A-B"}`},
	}
	result := BuildKGContent(entities, relations, 10000)
	if !contains(result, "Entity A desc") {
		t.Errorf("expected entity description in output, got: %s", result)
	}
	if !contains(result, "rel A-B") {
		t.Errorf("expected relation description in output, got: %s", result)
	}
}

func TestBuildKGContent_TokenBudget(t *testing.T) {
	longDesc := strings.Repeat("This is a very long description. ", 50)
	entities := []ScoredEntity{
		{Entity: "LongEntityName", Score: 1.0, Description: longDesc},
	}
	relations := []ScoredRelation{
		{From: "X", To: "Y", Score: 1.0, Description: "relation desc"},
	}
	result := BuildKGContent(entities, relations, 50)
	// Token budget is very small, should truncate and not include relations
	if contains(result, "relation desc") {
		t.Log("Note: relations included despite small budget (depending on token count)")
	}
}

func TestExtractDescription_JSON(t *testing.T) {
	result := extractDescription(`{"description": "Entity A description", "other": "value"}`)
	if result != "Entity A description" {
		t.Errorf("expected 'Entity A description', got %q", result)
	}
}

func TestExtractDescription_Plain(t *testing.T) {
	result := extractDescription("plain description")
	if result != "plain description" {
		t.Errorf("expected 'plain description', got %q", result)
	}
}

func TestNumTokensFromString(t *testing.T) {
	s := "This is a test string with multiple words"
	tokens := NumTokensFromString(s)
	if tokens <= 0 {
		t.Errorf("expected positive token count, got %d", tokens)
	}
}

// contains checks if a string contains a substring.
func contains(s, substr string) bool {
	return len(s) >= len(substr) && containsStr(s, substr)
}

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