534f4d6ed5
- Introduced a new TODO.md file to outline the implementation plan and progress for the search module. - Updated DESIGN.md to reflect changes in the directory structure and clarify the roles of various components, including the new Handler + Registry pattern for reranking and keyword extraction. - Refactored the Searcher struct to utilize a direct reference to the rerank package, enhancing modularity and clarity in the search process. - Modified the Search and SearchMultiple methods to include context parameters, improving flexibility for agent mode operations. - Revised the Reranker interface to require context for Agent and MCP modes, ensuring compatibility with different reranking strategies. - Enhanced documentation to provide comprehensive guidance on the updated search architecture and its components.
124 lines
3.8 KiB
Go
124 lines
3.8 KiB
Go
package rerank
|
|
|
|
import (
|
|
"testing"
|
|
|
|
"github.com/stretchr/testify/assert"
|
|
"github.com/yaoapp/yao/agent/search/types"
|
|
)
|
|
|
|
func TestBuiltinReranker_EmptyItems(t *testing.T) {
|
|
reranker := NewBuiltinReranker()
|
|
result, err := reranker.Rerank("test query", []*types.ResultItem{}, &types.RerankOptions{TopN: 5})
|
|
|
|
assert.NoError(t, err)
|
|
assert.Empty(t, result)
|
|
}
|
|
|
|
func TestBuiltinReranker_SortByWeightedScore(t *testing.T) {
|
|
reranker := NewBuiltinReranker()
|
|
|
|
items := []*types.ResultItem{
|
|
{CitationID: "ref_001", Score: 0.8, Weight: 0.6}, // weighted: 0.48
|
|
{CitationID: "ref_002", Score: 0.6, Weight: 1.0}, // weighted: 0.60
|
|
{CitationID: "ref_003", Score: 0.9, Weight: 0.8}, // weighted: 0.72
|
|
{CitationID: "ref_004", Score: 0.5, Weight: 1.0}, // weighted: 0.50
|
|
}
|
|
|
|
result, err := reranker.Rerank("test query", items, &types.RerankOptions{TopN: 10})
|
|
|
|
assert.NoError(t, err)
|
|
assert.Len(t, result, 4)
|
|
|
|
// Should be sorted by weighted score: ref_003 (0.72) > ref_002 (0.60) > ref_004 (0.50) > ref_001 (0.48)
|
|
assert.Equal(t, "ref_003", result[0].CitationID)
|
|
assert.Equal(t, "ref_002", result[1].CitationID)
|
|
assert.Equal(t, "ref_004", result[2].CitationID)
|
|
assert.Equal(t, "ref_001", result[3].CitationID)
|
|
}
|
|
|
|
func TestBuiltinReranker_TopN(t *testing.T) {
|
|
reranker := NewBuiltinReranker()
|
|
|
|
items := []*types.ResultItem{
|
|
{CitationID: "ref_001", Score: 0.9, Weight: 1.0},
|
|
{CitationID: "ref_002", Score: 0.8, Weight: 1.0},
|
|
{CitationID: "ref_003", Score: 0.7, Weight: 1.0},
|
|
{CitationID: "ref_004", Score: 0.6, Weight: 1.0},
|
|
{CitationID: "ref_005", Score: 0.5, Weight: 1.0},
|
|
}
|
|
|
|
result, err := reranker.Rerank("test query", items, &types.RerankOptions{TopN: 3})
|
|
|
|
assert.NoError(t, err)
|
|
assert.Len(t, result, 3)
|
|
assert.Equal(t, "ref_001", result[0].CitationID)
|
|
assert.Equal(t, "ref_002", result[1].CitationID)
|
|
assert.Equal(t, "ref_003", result[2].CitationID)
|
|
}
|
|
|
|
func TestBuiltinReranker_DefaultWeight(t *testing.T) {
|
|
reranker := NewBuiltinReranker()
|
|
|
|
// Items without weight should use default 0.6
|
|
items := []*types.ResultItem{
|
|
{CitationID: "ref_001", Score: 0.9, Weight: 0}, // weighted: 0.9 * 0.6 = 0.54
|
|
{CitationID: "ref_002", Score: 0.5, Weight: 1.0}, // weighted: 0.5 * 1.0 = 0.50
|
|
}
|
|
|
|
result, err := reranker.Rerank("test query", items, &types.RerankOptions{TopN: 10})
|
|
|
|
assert.NoError(t, err)
|
|
assert.Len(t, result, 2)
|
|
// ref_001 (0.54) > ref_002 (0.50)
|
|
assert.Equal(t, "ref_001", result[0].CitationID)
|
|
assert.Equal(t, "ref_002", result[1].CitationID)
|
|
}
|
|
|
|
func TestBuiltinReranker_TopNLargerThanItems(t *testing.T) {
|
|
reranker := NewBuiltinReranker()
|
|
|
|
items := []*types.ResultItem{
|
|
{CitationID: "ref_001", Score: 0.9, Weight: 1.0},
|
|
{CitationID: "ref_002", Score: 0.8, Weight: 1.0},
|
|
}
|
|
|
|
// TopN > len(items) should return all items
|
|
result, err := reranker.Rerank("test query", items, &types.RerankOptions{TopN: 10})
|
|
|
|
assert.NoError(t, err)
|
|
assert.Len(t, result, 2)
|
|
}
|
|
|
|
func TestBuiltinReranker_ZeroTopN(t *testing.T) {
|
|
reranker := NewBuiltinReranker()
|
|
|
|
items := []*types.ResultItem{
|
|
{CitationID: "ref_001", Score: 0.9, Weight: 1.0},
|
|
{CitationID: "ref_002", Score: 0.8, Weight: 1.0},
|
|
}
|
|
|
|
// TopN = 0 should return all items
|
|
result, err := reranker.Rerank("test query", items, &types.RerankOptions{TopN: 0})
|
|
|
|
assert.NoError(t, err)
|
|
assert.Len(t, result, 2)
|
|
}
|
|
|
|
func TestBuiltinReranker_SameWeightedScore(t *testing.T) {
|
|
reranker := NewBuiltinReranker()
|
|
|
|
// Items with same weighted score - order should be stable
|
|
items := []*types.ResultItem{
|
|
{CitationID: "ref_001", Score: 0.8, Weight: 1.0}, // weighted: 0.80
|
|
{CitationID: "ref_002", Score: 0.8, Weight: 1.0}, // weighted: 0.80
|
|
{CitationID: "ref_003", Score: 0.4, Weight: 1.0}, // weighted: 0.40
|
|
}
|
|
|
|
result, err := reranker.Rerank("test query", items, &types.RerankOptions{TopN: 10})
|
|
|
|
assert.NoError(t, err)
|
|
assert.Len(t, result, 3)
|
|
// ref_003 should be last
|
|
assert.Equal(t, "ref_003", result[2].CitationID)
|
|
}
|