ai-agent/internal/ai/rag/retrieve_test.go

package rag

import (
	"context"
	"errors"
	"testing"

	"code.tczkiot.com/wlw/ai-agent/internal/models"
)

func TestResolveKnowledgeBaseSearchOptionsUsesKnowledgeBaseDefaults(t *testing.T) {
	topK, scoreThreshold := resolveKnowledgeBaseSearchOptions(RetrieveRequest{}, &models.KnowledgeBase{
		DefaultTopK:           6,
		DefaultScoreThreshold: 0.42,
	})

	if topK != 6 {
		t.Fatalf("expected topK 6, got %d", topK)
	}
	if scoreThreshold != float32(0.42) {
		t.Fatalf("expected score threshold 0.42, got %v", scoreThreshold)
	}
}

func TestApplyRerankFallsBackWithoutRetrievingAgain(t *testing.T) {
	calls := 0
	retriever := &retrieve{rerankResults: func(context.Context, string, []RetrieveResult, int) ([]RetrieveResult, error) {
		calls++
		return nil, errors.New("platform rerank is unavailable")
	}}
	results := []RetrieveResult{{ChunkID: 1, Score: 0.9}, {ChunkID: 2, Score: 0.8}, {ChunkID: 3, Score: 0.7}}

	got, err := retriever.ApplyRerank(context.Background(), "refund", results, 2)
	if err != nil {
		t.Fatalf("ApplyRerank() error = %v", err)
	}
	if calls != 1 {
		t.Fatalf("rerank calls = %d, want 1", calls)
	}
	if len(got) != 2 || got[0].ChunkID != 1 || got[1].ChunkID != 2 {
		t.Fatalf("ApplyRerank() fallback = %+v", got)
	}
}

func TestResolveKnowledgeBaseSearchOptionsRequestOverridesKnowledgeBaseDefaults(t *testing.T) {
	topK, scoreThreshold := resolveKnowledgeBaseSearchOptions(RetrieveRequest{
		TopK:           9,
		ScoreThreshold: 0.55,
	}, &models.KnowledgeBase{
		DefaultTopK:           6,
		DefaultScoreThreshold: 0.42,
	})

	if topK != 9 {
		t.Fatalf("expected request topK 9, got %d", topK)
	}
	if scoreThreshold != float32(0.55) {
		t.Fatalf("expected request score threshold 0.55, got %v", scoreThreshold)
	}
}

func TestResolveKnowledgeBaseSearchOptionsUsesSystemDefaults(t *testing.T) {
	topK, scoreThreshold := resolveKnowledgeBaseSearchOptions(RetrieveRequest{}, nil)

	if topK != 8 {
		t.Fatalf("expected fallback topK 8, got %d", topK)
	}
	if scoreThreshold != float32(0.3) {
		t.Fatalf("expected fallback score threshold 0.3, got %v", scoreThreshold)
	}
}