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.
148 lines
3.9 KiB
Go
148 lines
3.9 KiB
Go
package rerank_test
|
|
|
|
import (
|
|
"testing"
|
|
|
|
"github.com/stretchr/testify/assert"
|
|
"github.com/stretchr/testify/require"
|
|
"github.com/yaoapp/yao/agent/context"
|
|
"github.com/yaoapp/yao/agent/search/rerank"
|
|
"github.com/yaoapp/yao/agent/search/types"
|
|
"github.com/yaoapp/yao/agent/testutils"
|
|
oauthTypes "github.com/yaoapp/yao/openapi/oauth/types"
|
|
)
|
|
|
|
func TestMCPProviderWithSearchRerank(t *testing.T) {
|
|
testutils.Prepare(t)
|
|
defer testutils.Clean(t)
|
|
|
|
ctx := newMCPTestContext(t)
|
|
|
|
provider, err := rerank.NewMCPProvider("search.rerank")
|
|
require.NoError(t, err)
|
|
|
|
items := []*types.ResultItem{
|
|
{CitationID: "ref_001", Score: 0.9, Weight: 1.0, Title: "First"},
|
|
{CitationID: "ref_002", Score: 0.8, Weight: 1.0, Title: "Second"},
|
|
{CitationID: "ref_003", Score: 0.7, Weight: 1.0, Title: "Third"},
|
|
}
|
|
|
|
result, err := provider.Rerank(ctx, "test query", items, &types.RerankOptions{TopN: 10})
|
|
|
|
require.NoError(t, err)
|
|
assert.NotEmpty(t, result)
|
|
|
|
// The mock MCP reverses the order
|
|
// So we expect: ref_003, ref_002, ref_001
|
|
assert.Len(t, result, 3)
|
|
assert.Equal(t, "ref_003", result[0].CitationID)
|
|
assert.Equal(t, "ref_002", result[1].CitationID)
|
|
assert.Equal(t, "ref_001", result[2].CitationID)
|
|
}
|
|
|
|
func TestMCPProviderWithTopN(t *testing.T) {
|
|
testutils.Prepare(t)
|
|
defer testutils.Clean(t)
|
|
|
|
ctx := newMCPTestContext(t)
|
|
|
|
provider, err := rerank.NewMCPProvider("search.rerank")
|
|
require.NoError(t, err)
|
|
|
|
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},
|
|
}
|
|
|
|
// Request top 2 only
|
|
result, err := provider.Rerank(ctx, "test query", items, &types.RerankOptions{TopN: 2})
|
|
|
|
require.NoError(t, err)
|
|
assert.Len(t, result, 2)
|
|
}
|
|
|
|
func TestMCPProviderInvalidFormat(t *testing.T) {
|
|
_, err := rerank.NewMCPProvider("invalid-format")
|
|
|
|
assert.Error(t, err)
|
|
assert.Contains(t, err.Error(), "invalid MCP format")
|
|
}
|
|
|
|
func TestMCPProviderServerNotFound(t *testing.T) {
|
|
testutils.Prepare(t)
|
|
defer testutils.Clean(t)
|
|
|
|
ctx := newMCPTestContext(t)
|
|
|
|
provider, err := rerank.NewMCPProvider("nonexistent.rerank")
|
|
require.NoError(t, err)
|
|
|
|
items := []*types.ResultItem{
|
|
{CitationID: "ref_001", Score: 0.9, Weight: 1.0},
|
|
}
|
|
|
|
_, err = provider.Rerank(ctx, "test query", items, &types.RerankOptions{TopN: 10})
|
|
|
|
assert.Error(t, err)
|
|
assert.Contains(t, err.Error(), "not found")
|
|
}
|
|
|
|
func TestMCPProviderToolNotFound(t *testing.T) {
|
|
testutils.Prepare(t)
|
|
defer testutils.Clean(t)
|
|
|
|
ctx := newMCPTestContext(t)
|
|
|
|
provider, err := rerank.NewMCPProvider("search.nonexistent_tool")
|
|
require.NoError(t, err)
|
|
|
|
items := []*types.ResultItem{
|
|
{CitationID: "ref_001", Score: 0.9, Weight: 1.0},
|
|
}
|
|
|
|
_, err = provider.Rerank(ctx, "test query", items, &types.RerankOptions{TopN: 10})
|
|
|
|
assert.Error(t, err)
|
|
}
|
|
|
|
func TestMCPProviderWithoutContext(t *testing.T) {
|
|
provider, err := rerank.NewMCPProvider("search.rerank")
|
|
require.NoError(t, err)
|
|
|
|
items := []*types.ResultItem{
|
|
{CitationID: "ref_001", Score: 0.9, Weight: 1.0},
|
|
}
|
|
|
|
_, err = provider.Rerank(nil, "test query", items, &types.RerankOptions{TopN: 10})
|
|
|
|
assert.Error(t, err)
|
|
assert.Contains(t, err.Error(), "context is required")
|
|
}
|
|
|
|
func TestMCPProviderEmptyItems(t *testing.T) {
|
|
testutils.Prepare(t)
|
|
defer testutils.Clean(t)
|
|
|
|
ctx := newMCPTestContext(t)
|
|
|
|
provider, err := rerank.NewMCPProvider("search.rerank")
|
|
require.NoError(t, err)
|
|
|
|
result, err := provider.Rerank(ctx, "test query", []*types.ResultItem{}, &types.RerankOptions{TopN: 10})
|
|
|
|
require.NoError(t, err)
|
|
assert.Empty(t, result)
|
|
}
|
|
|
|
// newMCPTestContext creates a test context with required fields
|
|
func newMCPTestContext(t *testing.T) *context.Context {
|
|
t.Helper()
|
|
authorized := &oauthTypes.AuthorizedInfo{
|
|
UserID: "test-user",
|
|
}
|
|
chatID := "test-chat-rerank-mcp"
|
|
return context.New(t.Context(), authorized, chatID)
|
|
}
|