ragflow/internal/agent/tool/retrieval_nlp_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.
//

// retrieval_nlp_test.go — NLPRetrievalAdapter tests.
//
// The adapter's Search method calls into the real
// nlp.RetrievalService, which requires a doc engine + document
// DAO. Those require network access and aren't worth standing up
// in a unit test. Instead we exercise the translation layer
// (translateChunk + helpers) directly — that's the surface that
// differs between the nlp chunk shape and the agent-tool chunk
// shape, and it's the only piece the adapter owns.

package tool

import (
	"context"
	"math"
	"testing"

	"ragflow/internal/entity"
)

// floatEqual compares two floats with a small epsilon so
// arithmetic like 0.4+0.8/2 == 0.6000000000000001 doesn't fail.
func floatEqual(a, b float64) bool { return math.Abs(a-b) < 1e-9 }

// TestTranslateChunk_FullFields: every chunk field is present and
// maps to the expected RetrievalChunk field.
func TestTranslateChunk_FullFields(t *testing.T) {
	raw := map[string]any{
		"chunk_id":            "ck-42",
		"doc_id":              "doc-7",
		"docnm_kwd":           "report.pdf",
		"kb_id":               "kb-1",
		"image_id":            "img-1",
		"position_int":        [][]float64{{1, 2, 3, 4}},
		"content_with_weight": "the answer is 42",
		"content_ltks":        "answer 42",
		"similarity":          0.87,
		"term_similarity":     0.5,
		"vector_similarity":   0.9,
	}
	got := translateChunk(raw)
	if got.ID != "ck-42" {
		t.Errorf("ID = %q, want \"ck-42\"", got.ID)
	}
	if got.Content != "the answer is 42" {
		t.Errorf("Content = %q, want content_with_weight", got.Content)
	}
	if got.DocumentID != "doc-7" {
		t.Errorf("DocumentID = %q, want \"doc-7\"", got.DocumentID)
	}
	if got.DocumentName != "report.pdf" {
		t.Errorf("DocumentName = %q, want \"report.pdf\"", got.DocumentName)
	}
	if got.DatasetID != "kb-1" {
		t.Errorf("DatasetID = %q, want \"kb-1\"", got.DatasetID)
	}
	if got.ImageID != "img-1" {
		t.Errorf("ImageID = %q, want \"img-1\"", got.ImageID)
	}
	if got.Positions == nil {
		t.Errorf("Positions is nil, want position_int payload")
	}
	if got.Score != 0.87 {
		t.Errorf("Score = %v, want 0.87 (similarity preferred)", got.Score)
	}
	if got.TermSimilarity != 0.5 {
		t.Errorf("TermSimilarity = %v, want 0.5", got.TermSimilarity)
	}
	if got.VectorSimilarity != 0.9 {
		t.Errorf("VectorSimilarity = %v, want 0.9", got.VectorSimilarity)
	}
}

// TestTranslateChunk_ContentFallback: when content_with_weight is
// empty, content_ltks is used.
func TestTranslateChunk_ContentFallback(t *testing.T) {
	raw := map[string]any{
		"chunk_id":     "ck-1",
		"content_ltks": "answer 42",
		"doc_id":       "doc-1",
		"similarity":   0.5,
	}
	got := translateChunk(raw)
	if got.Content != "answer 42" {
		t.Errorf("Content = %q, want content_ltks fallback", got.Content)
	}
}

// TestTranslateChunk_EmptyContent: both content fields empty →
// empty string (don't crash, don't synthesise).
func TestTranslateChunk_EmptyContent(t *testing.T) {
	raw := map[string]any{
		"chunk_id":   "ck-1",
		"doc_id":     "doc-1",
		"similarity": 0.5,
	}
	got := translateChunk(raw)
	if got.Content != "" {
		t.Errorf("Content = %q, want \"\"", got.Content)
	}
}

// TestTranslateChunk_ScoreFallback: when "similarity" is missing,
// the average of term_similarity + vector_similarity is used.
func TestTranslateChunk_ScoreFallback(t *testing.T) {
	raw := map[string]any{
		"chunk_id":          "ck-1",
		"term_similarity":   0.4,
		"vector_similarity": 0.8,
	}
	got := translateChunk(raw)
	if !floatEqual(got.Score, 0.6) {
		t.Errorf("Score = %v, want ~0.6 (avg of term+vec)", got.Score)
	}
}

// TestTranslateChunk_ScoreOnlyOneSub: only one of term_similarity
// or vector_similarity present → that one wins (not averaged
// against zero, which would be misleading).
func TestTranslateChunk_ScoreOnlyOneSub(t *testing.T) {
	cases := []struct {
		name string
		raw  map[string]any
		want float64
	}{
		{"term only", map[string]any{"term_similarity": 0.7}, 0.7},
		{"vector only", map[string]any{"vector_similarity": 0.3}, 0.3},
	}
	for _, tc := range cases {
		got := translateChunk(tc.raw)
		if !floatEqual(got.Score, tc.want) {
			t.Errorf("%s: Score = %v, want %v", tc.name, got.Score, tc.want)
		}
	}
}

// TestTranslateChunk_NumberTypesTolerated: numeric similarity
// fields may come back as float64, float32, int, or int64 depending
// on how the upstream serialised them. All should coerce.
func TestTranslateChunk_NumberTypesTolerated(t *testing.T) {
	cases := []struct {
		name string
		val  any
		want float64
	}{
		{"float64", float64(0.5), 0.5},
		{"float32", float32(0.5), 0.5},
		{"int", int(1), 1},
		{"int64", int64(1), 1},
	}
	for _, tc := range cases {
		raw := map[string]any{"similarity": tc.val}
		got := translateChunk(raw)
		if !floatEqual(got.Score, tc.want) {
			t.Errorf("%s: Score = %v, want %v", tc.name, got.Score, tc.want)
		}
	}
}

// TestTranslateChunk_WrongTypesIgnored: a similarity field that's
// a string or nil must not crash. We fall back to the sub-scores
// (which are also missing in this test → zero).
func TestTranslateChunk_WrongTypesIgnored(t *testing.T) {
	cases := []map[string]any{
		{"similarity": "0.5"},
		{"similarity": nil},
		{"similarity": []any{0.5}},
	}
	for _, raw := range cases {
		got := translateChunk(raw)
		if got.Score != 0 {
			t.Errorf("wrong-type similarity: Score = %v, want 0", got.Score)
		}
	}
}

// TestTranslateChunk_WrongStringTypesIgnored: a string field
// that's actually a number must not crash.
func TestTranslateChunk_WrongStringTypesIgnored(t *testing.T) {
	raw := map[string]any{
		"chunk_id": 42,  // int, not string
		"doc_id":   nil, // nil
	}
	got := translateChunk(raw)
	if got.ID != "" {
		t.Errorf("ID = %q, want \"\"", got.ID)
	}
	if got.DocumentID != "" {
		t.Errorf("DocumentID = %q, want \"\"", got.DocumentID)
	}
}

// TestTranslateChunk_MissingAllScores: a chunk with no score
// fields at all → score 0 (don't panic).
func TestTranslateChunk_MissingAllScores(t *testing.T) {
	raw := map[string]any{"chunk_id": "ck-1"}
	got := translateChunk(raw)
	if got.Score != 0 {
		t.Errorf("Score = %v, want 0", got.Score)
	}
}

// TestNewNLPRetrievalAdapter_NilService: nil constructor inputs
// must produce an adapter whose Search returns the missing-service
// error, not a panic.
func TestNewNLPRetrievalAdapter_NilService(t *testing.T) {
	a := NewNLPRetrievalAdapter(nil)
	_, err := a.Search(context.TODO(), RetrievalRequest{Query: "hi"})
	if err == nil {
		t.Fatal("expected error from nil-service adapter")
	}
	if err != ErrRetrievalServiceMissing {
		t.Errorf("err = %v, want ErrRetrievalServiceMissing", err)
	}
}

func TestNLPRequestFromRetrieval_ThreadsSearchControls(t *testing.T) {
	keywordWeight := 0.25
	got := nlpRequestFromRetrieval(RetrievalRequest{
		Query:                    "hi",
		DatasetIDs:               []string{"kb-1"},
		TopN:                     3,
		TopK:                     99,
		KeywordsSimilarityWeight: &keywordWeight,
		SimilarityThreshold:      0.42,
	}, []string{"tenant-a"}, 3)

	if got.Question != "hi" {
		t.Fatalf("Question=%q want hi", got.Question)
	}
	if len(got.TenantIDs) != 1 || got.TenantIDs[0] != "tenant-a" {
		t.Fatalf("TenantIDs=%v want [tenant-a]", got.TenantIDs)
	}
	if len(got.KbIDs) != 1 || got.KbIDs[0] != "kb-1" {
		t.Fatalf("KbIDs=%v want [kb-1]", got.KbIDs)
	}
	if got.Page != 1 || got.PageSize != 3 {
		t.Fatalf("Page/PageSize=%d/%d want 1/3", got.Page, got.PageSize)
	}
	if got.Top == nil || *got.Top != 99 {
		t.Fatalf("Top=%v want 99", got.Top)
	}
	if got.SimilarityThreshold == nil || *got.SimilarityThreshold != 0.42 {
		t.Fatalf("SimilarityThreshold=%v want 0.42", got.SimilarityThreshold)
	}
	if got.VectorSimilarityWeight == nil || !floatEqual(*got.VectorSimilarityWeight, 0.75) {
		t.Fatalf("VectorSimilarityWeight=%v want 0.75", got.VectorSimilarityWeight)
	}
}

func TestNLPRequestFromRetrieval_FallsBackToTopNHeadroom(t *testing.T) {
	got := nlpRequestFromRetrieval(RetrievalRequest{
		Query:      "hi",
		DatasetIDs: []string{"kb-1"},
		TopN:       3,
	}, []string{"tenant-a"}, 3)

	if got.Top == nil || *got.Top != 12 {
		t.Fatalf("Top=%v want 12", got.Top)
	}
	if got.VectorSimilarityWeight != nil {
		t.Fatalf("VectorSimilarityWeight=%v want nil", got.VectorSimilarityWeight)
	}
}

func TestNLPRetrievalAdapter_ResolveTenantIDsStaysWithinRequestTenant(t *testing.T) {
	a := &NLPRetrievalAdapter{}
	got, err := a.resolveTenantIDs(RetrievalRequest{
		TenantID:   "tenant-a",
		DatasetIDs: []string{"kb-1", "kb-2", "kb-missing"},
	})
	if err != nil {
		t.Fatalf("resolveTenantIDs: %v", err)
	}

	if len(got) != 1 {
		t.Fatalf("tenantIDs len=%d want 1, got=%v", len(got), got)
	}
	if got[0] != "tenant-a" {
		t.Fatalf("tenantIDs=%v want [tenant-a]", got)
	}
}

func TestNLPRetrievalAdapter_ResolveTenantIDsFromDatasetIDs(t *testing.T) {
	a := &NLPRetrievalAdapter{
		kbDAO: fakeKnowledgebaseLookup{
			kbs: []*entity.Knowledgebase{
				{ID: "kb-1", TenantID: "tenant-a"},
				{ID: "kb-2", TenantID: "tenant-b"},
				{ID: "kb-3", TenantID: "tenant-a"},
			},
		},
	}
	got, err := a.resolveTenantIDs(RetrievalRequest{
		DatasetIDs: []string{"kb-1", "kb-2", "kb-3", " "},
	})
	if err != nil {
		t.Fatalf("resolveTenantIDs: %v", err)
	}
	if len(got) != 2 || got[0] != "tenant-a" || got[1] != "tenant-b" {
		t.Fatalf("tenantIDs=%v want [tenant-a tenant-b]", got)
	}
}

func TestNLPRetrievalAdapter_ResolveTenantIDsKeepsRequestTenantFirst(t *testing.T) {
	a := &NLPRetrievalAdapter{
		kbDAO: fakeKnowledgebaseLookup{
			kbs: []*entity.Knowledgebase{
				{ID: "kb-1", TenantID: "tenant-b"},
			},
		},
	}
	got, err := a.resolveTenantIDs(RetrievalRequest{
		TenantID:   "tenant-a",
		DatasetIDs: []string{"kb-1"},
	})
	if err != nil {
		t.Fatalf("resolveTenantIDs: %v", err)
	}
	if len(got) != 2 || got[0] != "tenant-a" || got[1] != "tenant-b" {
		t.Fatalf("tenantIDs=%v want [tenant-a tenant-b]", got)
	}
}

type fakeKnowledgebaseLookup struct {
	kbs []*entity.Knowledgebase
	err error
}

func (f fakeKnowledgebaseLookup) GetByIDs(_ []string) ([]*entity.Knowledgebase, error) {
	return f.kbs, f.err
}