diff --git a/agent/search/DESIGN.md b/agent/search/DESIGN.md index c97d650c..b14377c3 100644 --- a/agent/search/DESIGN.md +++ b/agent/search/DESIGN.md @@ -124,11 +124,10 @@ sequenceDiagram ``` agent/search/ ├── DESIGN.md # This document +├── TODO.md # Implementation plan and progress ├── search.go # Main Searcher implementation and public API ├── registry.go # Handler registry (manages web/kb/db handlers) -├── jsapi.go # JavaScript API bindings for hooks -├── trace.go # Trace node creation and management -├── output.go # Real-time output/streaming to client +├── jsapi.go # JavaScript API bindings for hooks (skeleton) ├── citation.go # Citation ID generation and tracking ├── reference.go # Reference building and LLM context formatting │ @@ -144,19 +143,19 @@ agent/search/ │ ├── reranker.go # Reranker interface │ └── nlp.go # NLP interfaces (KeywordExtractor, QueryDSLGenerator) │ -├── rerank/ # Result reranking implementations -│ ├── rerank.go # Reranker factory and common logic -│ ├── builtin.go # Built-in score-based reranking (default) -│ ├── agent.go # Agent-based reranking (delegate to another assistant) -│ └── mcp.go # MCP-based reranking (call MCP server tool) +├── rerank/ # Result reranking implementations (Handler + Registry pattern) ✅ +│ ├── reranker.go # Main entry point (mode dispatch) +│ ├── builtin.go # Builtin: weighted score sorting +│ ├── agent.go # Agent mode (delegate to LLM assistant) +│ └── mcp.go # MCP mode (external service) │ ├── nlp/ # Natural language processing for search -│ ├── keyword/ # Keyword extraction (Handler + Registry pattern) +│ ├── keyword/ # Keyword extraction (Handler + Registry pattern) ✅ │ │ ├── extractor.go # Main extractor (mode dispatch) │ │ ├── builtin.go # Builtin frequency-based extraction │ │ ├── agent.go # Agent mode (LLM-powered) │ │ └── mcp.go # MCP mode (external service) -│ └── querydsl/ # QueryDSL generation for DB search +│ └── querydsl/ # QueryDSL generation for DB search (待实现) │ ├── generator.go # Main generator (mode dispatch) │ ├── builtin.go # Builtin template-based generation │ ├── agent.go # Agent mode (LLM-powered) @@ -164,7 +163,7 @@ agent/search/ │ # Note: Embedding follows KB collection config, not in this package │ ├── handlers/ # Search handler implementations -│ ├── web/ # Web search +│ ├── web/ # Web search ✅ │ │ ├── handler.go # Web search handler (mode dispatch) │ │ ├── tavily.go # Tavily provider (builtin) │ │ ├── serper.go # Serper provider (serper.dev, builtin) @@ -172,18 +171,22 @@ agent/search/ │ │ ├── agent.go # Agent mode (AI Search) │ │ └── mcp.go # MCP mode (external service) │ │ -│ ├── kb/ # Knowledge base search +│ ├── kb/ # Knowledge base search (骨架) │ │ ├── handler.go # KB search handler -│ │ ├── vector.go # Vector similarity search -│ │ └── graph.go # Graph-based association (GraphRAG) +│ │ ├── vector.go # Vector similarity search (待实现) +│ │ └── graph.go # Graph-based association (待实现) │ │ -│ └── db/ # Database search (Yao Model/QueryDSL) +│ └── db/ # Database search (骨架) │ ├── handler.go # DB search handler -│ ├── query.go # QueryDSL builder -│ └── schema.go # Model schema introspection +│ ├── query.go # QueryDSL builder (待实现) +│ └── schema.go # Model schema introspection (待实现) │ └── defaults/ # Default configuration values └── defaults.go # System built-in defaults (used by agent/load.go) + +# 待实现文件: +# - trace.go # Trace node creation and management +# - output.go # Real-time output/streaming to client ``` ### Dependency Graph @@ -249,13 +252,13 @@ import ( type Searcher struct { config *types.Config // Merged config (global + assistant) handlers map[types.SearchType]interfaces.Handler - reranker interfaces.Reranker + reranker *rerank.Reranker // Uses rerank package directly citation *CitationGenerator } -// SearchUses contains the search-specific uses configuration +// Uses contains the search-specific uses configuration // These are extracted from context.Uses and search config -type SearchUses struct { +type Uses struct { Search string // "builtin", "disabled", "", "mcp:." Web string // "builtin", "", "mcp:." Keyword string // "builtin", "", "mcp:." @@ -266,7 +269,7 @@ type SearchUses struct { // New creates a new Searcher instance // cfg: merged config from agent/load.go + assistant config // uses: merged uses configuration (global → assistant → hook) -func New(cfg *types.Config, uses *SearchUses) *Searcher { +func New(cfg *types.Config, uses *Uses) *Searcher { return &Searcher{ config: cfg, handlers: map[types.SearchType]interfaces.Handler{ @@ -274,7 +277,7 @@ func New(cfg *types.Config, uses *SearchUses) *Searcher { types.SearchTypeKB: kb.NewHandler(cfg.KB), // KB always builtin types.SearchTypeDB: db.NewHandler(uses.QueryDSL, cfg.DB), }, - reranker: rerank.NewReranker(uses.Rerank), + reranker: rerank.NewReranker(uses.Rerank, cfg.Rerank), citation: NewCitationGenerator(), } } @@ -286,8 +289,8 @@ func (s *Searcher) Search(ctx *context.Context, req *types.Request) (*types.Resu return &types.Result{Error: "unsupported search type"}, nil } - // Execute search - result, err := handler.Search(ctx, req) + // Execute search (handler doesn't need ctx) + result, err := handler.Search(req) if err != nil { return &types.Result{Error: err.Error()}, nil } @@ -297,8 +300,8 @@ func (s *Searcher) Search(ctx *context.Context, req *types.Request) (*types.Resu item.Weight = s.config.GetWeight(req.Source) } - // Rerank if requested - if req.Rerank != nil { + // Rerank if requested (reranker needs ctx for Agent/MCP modes) + if req.Rerank != nil && s.reranker != nil { result.Items, _ = s.reranker.Rerank(ctx, req.Query, result.Items, req.Rerank) } @@ -396,7 +399,6 @@ All interfaces are defined in `search/interfaces/` package to prevent circular d package interfaces import ( - "github.com/yaoapp/yao/agent/context" "github.com/yaoapp/yao/agent/search/types" ) @@ -405,11 +407,8 @@ type Handler interface { // Type returns the search type this handler supports Type() types.SearchType - // CanHandle checks if this handler can process the given request - CanHandle(ctx *context.Context, req *types.Request) bool - // Search executes the search and returns results - Search(ctx *context.Context, req *types.Request) (*types.Result, error) + Search(req *types.Request) (*types.Result, error) } ``` @@ -419,29 +418,32 @@ type Handler interface { package interfaces import ( - "github.com/yaoapp/yao/agent/context" "github.com/yaoapp/yao/agent/search/types" ) // Searcher is the main interface exposed to external callers type Searcher interface { // Search executes a single search request - Search(ctx *context.Context, req *types.Request) (*types.Result, error) + Search(req *types.Request) (*types.Result, error) // SearchMultiple executes multiple searches (potentially in parallel) - SearchMultiple(ctx *context.Context, reqs []*types.Request) ([]*types.Result, error) + SearchMultiple(reqs []*types.Request) ([]*types.Result, error) // BuildReferences converts search results to unified Reference format for LLM BuildReferences(results []*types.Result) []*types.Reference } ``` +> **Note**: The actual `Searcher` struct in `search.go` has `Search(ctx, req)` and `SearchMultiple(ctx, reqs)` signatures that include context for reranking support. The interface is kept minimal for flexibility. + ### NLP Interfaces (`interfaces/nlp.go`) ```go package interfaces import ( + "github.com/yaoapp/gou/model" + "github.com/yaoapp/gou/query/gou" "github.com/yaoapp/yao/agent/context" "github.com/yaoapp/yao/agent/search/types" ) @@ -449,6 +451,7 @@ import ( // KeywordExtractor extracts keywords for web search type KeywordExtractor interface { // Extract extracts search keywords from user message + // ctx is required for Agent and MCP modes, can be nil for builtin mode Extract(ctx *context.Context, content string, opts *types.KeywordOptions) ([]string, error) } @@ -1768,7 +1771,7 @@ package web import ( "strings" - "github.com/yaoapp/yao/agent/context" + agentContext "github.com/yaoapp/yao/agent/context" "github.com/yaoapp/yao/agent/search/types" ) @@ -1779,61 +1782,28 @@ type Handler struct { } // NewHandler creates a new web search handler -func NewHandler(usesWeb string, cfg *types.WebConfig) *Handler { - return &Handler{usesWeb: usesWeb, config: cfg} -} +func NewHandler(usesWeb string, cfg *types.WebConfig) *Handler -// Search executes web search based on uses.web mode -func (h *Handler) Search(ctx *context.Context, req *types.Request) (*types.Result, error) { - switch { - case h.usesWeb == "builtin" || h.usesWeb == "": - return h.builtinSearch(ctx, req) - case strings.HasPrefix(h.usesWeb, "mcp:"): - return h.mcpSearch(ctx, req) - default: - // Agent mode: delegate to assistant for AI-powered search - return h.agentSearch(ctx, req) - } -} +// Type returns the search type this handler supports +func (h *Handler) Type() types.SearchType -// builtinSearch uses Tavily/Serper/SerpAPI directly -func (h *Handler) builtinSearch(ctx *context.Context, req *types.Request) (*types.Result, error) { - switch h.config.Provider { - case "tavily": - return NewTavilyProvider(h.config).Search(req) - case "serper": - // Serper (serper.dev) - POST request with X-API-KEY header - return NewSerperProvider(h.config).Search(req) - case "serpapi": - // SerpAPI (serpapi.com) - GET request with api_key parameter - // Supports multiple engines: google, bing, baidu, yandex, etc. - return NewSerpAPIProvider(h.config).Search(req) - default: - return nil, fmt.Errorf("unknown provider: %s", h.config.Provider) - } -} +// Search implements interfaces.Handler (without context) +func (h *Handler) Search(req *types.Request) (*types.Result, error) -// agentSearch delegates to an assistant for AI-powered search -func (h *Handler) agentSearch(ctx *context.Context, req *types.Request) (*types.Result, error) { - // 1. Call assistant with search request - // 2. Assistant understands intent, generates optimized queries - // 3. Assistant executes searches (may call builtin internally) - // 4. Assistant analyzes and returns structured results - return nil, nil -} +// SearchWithContext executes web search with context (for Agent/MCP modes) +func (h *Handler) SearchWithContext(ctx *agentContext.Context, req *types.Request) (*types.Result, error) +``` -// mcpSearch calls external MCP tool -func (h *Handler) mcpSearch(ctx *context.Context, req *types.Request) (*types.Result, error) { - // Parse "mcp:server.tool" - mcpRef := strings.TrimPrefix(h.usesWeb, "mcp:") - parts := strings.SplitN(mcpRef, ".", 2) - if len(parts) != 2 { - return nil, fmt.Errorf("invalid MCP format, expected 'mcp:server.tool', got '%s'", h.usesWeb) - } - serverID, toolName := parts[0], parts[1] - // Call MCP tool - return nil, nil -} +**Directory Structure:** + +``` +handlers/web/ +├── handler.go # Main entry point (mode dispatch) +├── tavily.go # Tavily provider (builtin) +├── serper.go # Serper provider (serper.dev) +├── serpapi.go # SerpAPI provider (serpapi.com, multi-engine) +├── agent.go # Agent mode (AI Search) +└── mcp.go # MCP mode (external service) ``` **Built-in Providers (when `uses.web = "builtin"`):** @@ -1942,8 +1912,6 @@ function Create(ctx, messages, options) { package kb import ( - "github.com/yaoapp/yao/agent/context" - "github.com/yaoapp/yao/agent/search/interfaces" "github.com/yaoapp/yao/agent/search/types" ) @@ -1953,18 +1921,14 @@ type Handler struct { } // NewHandler creates a new KB search handler -func NewHandler(cfg *types.KBConfig) *Handler { - return &Handler{config: cfg} -} +func NewHandler(cfg *types.KBConfig) *Handler + +// Type returns the search type this handler supports +func (h *Handler) Type() types.SearchType // Search executes vector search and optional graph association -func (h *Handler) Search(ctx *context.Context, req *types.Request) (*types.Result, error) { - // 1. Generate embedding via query processor - // 2. Vector search in collections - // 3. Optional: Graph association (if req.Graph) - // 4. Return results - return nil, nil -} +// TODO: Implement actual search logic +func (h *Handler) Search(req *types.Request) (*types.Result, error) ``` | File | Description | @@ -1980,8 +1944,6 @@ func (h *Handler) Search(ctx *context.Context, req *types.Request) (*types.Resul package db import ( - "github.com/yaoapp/yao/agent/context" - "github.com/yaoapp/yao/agent/search/interfaces" "github.com/yaoapp/yao/agent/search/types" ) @@ -1992,18 +1954,14 @@ type Handler struct { } // NewHandler creates a new DB search handler -func NewHandler(usesQueryDSL string, cfg *types.DBConfig) *Handler { - return &Handler{usesQueryDSL: usesQueryDSL, config: cfg} -} +func NewHandler(usesQueryDSL string, cfg *types.DBConfig) *Handler + +// Type returns the search type this handler supports +func (h *Handler) Type() types.SearchType // Search converts NL to QueryDSL and executes -func (h *Handler) Search(ctx *context.Context, req *types.Request) (*types.Result, error) { - // 1. Get model schemas - // 2. Generate QueryDSL via query processor - // 3. Execute queries on models - // 4. Return results - return nil, nil -} +// TODO: Implement actual search logic +func (h *Handler) Search(req *types.Request) (*types.Result, error) ``` | File | Description | @@ -2023,42 +1981,71 @@ Integrates with Yao's Model/QueryDSL system: ### Reranking (`rerank/`) +The rerank module follows the Handler + Registry pattern, consistent with `keyword/` and `web/`. + ```go -// rerank/rerank.go +// rerank/reranker.go package rerank import ( - "github.com/yaoapp/yao/agent/search/interfaces" + "strings" + + "github.com/yaoapp/yao/agent/context" "github.com/yaoapp/yao/agent/search/types" ) -// NewReranker creates a reranker based on uses.rerank config -func NewReranker(usesRerank string) interfaces.Reranker { - switch { - case usesRerank == "builtin" || usesRerank == "": - return NewBuiltinReranker() - case strings.HasPrefix(usesRerank, "mcp:"): - // Parse "mcp:server.tool" - mcpRef := strings.TrimPrefix(usesRerank, "mcp:") - parts := strings.SplitN(mcpRef, ".", 2) - if len(parts) != 2 { - // Invalid format, fallback to builtin - return NewBuiltinReranker() - } - return NewMCPReranker(parts[0], parts[1]) // serverID, toolName - default: - // Assume it's an assistant ID - return NewAgentReranker(usesRerank) - } +// Reranker reorders search results by relevance +// Mode is determined by uses.rerank configuration +type Reranker struct { + usesRerank string // "builtin", "", "mcp:." + config *types.RerankConfig } + +// NewReranker creates a new reranker +func NewReranker(usesRerank string, cfg *types.RerankConfig) *Reranker + +// Rerank reorders results based on configured mode +func (r *Reranker) Rerank(ctx *context.Context, query string, items []*types.ResultItem, opts *types.RerankOptions) ([]*types.ResultItem, error) ``` -| File | Description | -| ------------ | -------------------------------- | -| `rerank.go` | Factory and common logic | -| `builtin.go` | Simple score sorting (default) | -| `agent.go` | Delegate to an assistant (Agent) | -| `mcp.go` | Call MCP tool for reranking | +**Directory Structure:** + +``` +rerank/ +├── reranker.go # Main entry point (mode dispatch) +├── builtin.go # Builtin: weighted score sorting (score * weight) +├── agent.go # Agent mode (delegate to LLM assistant) +└── mcp.go # MCP mode (external service) +``` + +**Builtin Implementation:** + +The builtin reranker uses weighted score sorting: + +- Calculate `weightedScore = score * weight` +- Sort items by weighted score descending +- Return top N items + +> **Note**: For production use cases requiring semantic understanding, use Agent or MCP mode. + +**Agent Response Format:** + +The agent should return reordered items in one of these formats: + +```json +// Format 1: Order list (recommended) +{ "order": ["ref_003", "ref_001", "ref_002"] } + +// Format 2: Items list with citation_id +{ "items": [{ "citation_id": "ref_003" }, { "citation_id": "ref_001" }] } +``` + +| File | Description | +| ------------- | ---------------------------------------- | +| `reranker.go` | Main entry point and mode dispatch | +| `builtin.go` | Weighted score sorting (score \* weight) | +| `agent.go` | Delegate to LLM assistant for reranking | +| `mcp.go` | Call external MCP tool for reranking | Configure via `uses.rerank` in `agent/agent.yml`: diff --git a/agent/search/interfaces/reranker.go b/agent/search/interfaces/reranker.go index 58f3dd3e..770272ff 100644 --- a/agent/search/interfaces/reranker.go +++ b/agent/search/interfaces/reranker.go @@ -1,11 +1,13 @@ package interfaces import ( + "github.com/yaoapp/yao/agent/context" "github.com/yaoapp/yao/agent/search/types" ) // Reranker reorders search results by relevance type Reranker interface { // Rerank reorders results based on query relevance - Rerank(query string, items []*types.ResultItem, opts *types.RerankOptions) ([]*types.ResultItem, error) + // ctx is required for Agent and MCP modes, can be nil for builtin mode + Rerank(ctx *context.Context, query string, items []*types.ResultItem, opts *types.RerankOptions) ([]*types.ResultItem, error) } diff --git a/agent/search/rerank/agent.go b/agent/search/rerank/agent.go new file mode 100644 index 00000000..0a5137e6 --- /dev/null +++ b/agent/search/rerank/agent.go @@ -0,0 +1,232 @@ +package rerank + +import ( + "encoding/json" + "fmt" + "strings" + + "github.com/yaoapp/yao/agent/caller" + "github.com/yaoapp/yao/agent/context" + "github.com/yaoapp/yao/agent/search/types" +) + +// AgentProvider implements reranking by delegating to another agent +// The agent should have a Next Hook that accepts rerank request and returns reordered items +type AgentProvider struct { + agentID string // Assistant ID to delegate to +} + +// NewAgentProvider creates a new agent reranker +func NewAgentProvider(agentID string) *AgentProvider { + return &AgentProvider{agentID: agentID} +} + +// Rerank delegates reranking to an LLM-powered assistant +// The assistant receives items and query, returns reordered item IDs or items +func (p *AgentProvider) Rerank(ctx *context.Context, query string, items []*types.ResultItem, opts *types.RerankOptions) ([]*types.ResultItem, error) { + if ctx == nil { + return nil, fmt.Errorf("context is required for agent rerank") + } + + // Get agent via caller interface (avoids circular dependency) + agent, err := caller.AgentGetterFunc(p.agentID) + if err != nil { + return nil, fmt.Errorf("failed to get agent %s: %w", p.agentID, err) + } + + // Build request message with items to rerank + requestData := map[string]interface{}{ + "query": query, + "items": items, + "top_n": opts.TopN, + "action": "rerank", + } + requestJSON, _ := json.Marshal(requestData) + + // Create messages for agent + messages := []context.Message{ + { + Role: "user", + Content: string(requestJSON), + }, + } + + // Call agent's Stream method with skip options (no history, no output) + options := &context.Options{ + Skip: &context.Skip{ + History: true, + Output: true, + }, + } + + result, err := agent.Stream(ctx, messages, options) + if err != nil { + return nil, fmt.Errorf("agent stream failed: %w", err) + } + + // Parse response + return p.parseResponse(result, items, opts) +} + +// parseResponse extracts reranked items from agent response +// The response format from agent.Stream is typically: +// { "next": { "data": { "order": [...] } } } +func (p *AgentProvider) parseResponse(result interface{}, originalItems []*types.ResultItem, opts *types.RerankOptions) ([]*types.ResultItem, error) { + if result == nil { + return originalItems, nil + } + + // Build index map for quick lookup + itemMap := make(map[string]*types.ResultItem) + for _, item := range originalItems { + if item.CitationID != "" { + itemMap[item.CitationID] = item + } + } + + // Extract response data + response := extractResponseData(result) + if response == nil { + return originalItems, nil + } + + // Try to get reranked order from response + // Expected format: { "order": ["ref_001", "ref_003", "ref_002"] } + // Or: { "items": [{ "citation_id": "ref_001", ... }, ...] } + + var reranked []*types.ResultItem + + // Try "order" field (list of citation IDs) + if order, ok := response["order"]; ok { + if orderList := toStringSlice(order); len(orderList) > 0 { + for _, id := range orderList { + if item, exists := itemMap[id]; exists { + reranked = append(reranked, item) + delete(itemMap, id) // Avoid duplicates + } + } + // Append remaining items not in order + for _, item := range originalItems { + if _, exists := itemMap[item.CitationID]; exists { + reranked = append(reranked, item) + } + } + } + } + + // Try "items" field (full items or items with citation_id) + if len(reranked) == 0 { + if items, ok := response["items"]; ok { + if itemsList := toItemsList(items); len(itemsList) > 0 { + for _, respItem := range itemsList { + // Check if it's just a reference or full item + if citationID, ok := respItem["citation_id"].(string); ok { + if item, exists := itemMap[citationID]; exists { + reranked = append(reranked, item) + delete(itemMap, citationID) + } + } + } + // Append remaining items + for _, item := range originalItems { + if _, exists := itemMap[item.CitationID]; exists { + reranked = append(reranked, item) + } + } + } + } + } + + // If no valid response, return original items + if len(reranked) == 0 { + reranked = originalItems + } + + // Apply top N + if opts.TopN > 0 && opts.TopN < len(reranked) { + reranked = reranked[:opts.TopN] + } + + return reranked, nil +} + +// extractResponseData extracts the actual response data from agent.Stream result +// Handles nested structures like { "next": { "data": { ... } } } +func extractResponseData(result interface{}) map[string]interface{} { + switch v := result.(type) { + case map[string]interface{}: + // Check for "next" wrapper (from NextHookResponse) + if next, ok := v["next"].(map[string]interface{}); ok { + // Check for "data" inside next + if data, ok := next["data"].(map[string]interface{}); ok { + return data + } + return next + } + // Check for direct "data" wrapper + if data, ok := v["data"].(map[string]interface{}); ok { + return data + } + return v + case string: + // Try to parse as JSON + var data map[string]interface{} + if err := json.Unmarshal([]byte(v), &data); err == nil { + return extractResponseData(data) + } + } + // Try to handle other types by converting to JSON and back + if result != nil { + if bytes, err := json.Marshal(result); err == nil { + var data map[string]interface{} + if err := json.Unmarshal(bytes, &data); err == nil { + return extractResponseData(data) + } + } + } + return nil +} + +// toStringSlice converts interface to string slice +func toStringSlice(v interface{}) []string { + switch val := v.(type) { + case []string: + return val + case []interface{}: + result := make([]string, 0, len(val)) + for _, item := range val { + if s, ok := item.(string); ok { + result = append(result, s) + } + } + return result + } + return nil +} + +// toItemsList converts interface to list of maps +func toItemsList(v interface{}) []map[string]interface{} { + switch val := v.(type) { + case []map[string]interface{}: + return val + case []interface{}: + result := make([]map[string]interface{}, 0, len(val)) + for _, item := range val { + if m, ok := item.(map[string]interface{}); ok { + result = append(result, m) + } + } + return result + } + return nil +} + +// extractAgentID extracts assistant ID from uses.rerank value +// For backward compatibility, strips any prefix if present +func extractAgentID(usesRerank string) string { + // Remove any prefix like "agent:" if present + if strings.HasPrefix(usesRerank, "agent:") { + return strings.TrimPrefix(usesRerank, "agent:") + } + return usesRerank +} diff --git a/agent/search/rerank/agent_test.go b/agent/search/rerank/agent_test.go new file mode 100644 index 00000000..1f430ab2 --- /dev/null +++ b/agent/search/rerank/agent_test.go @@ -0,0 +1,123 @@ +package rerank_test + +import ( + "testing" + + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" + "github.com/yaoapp/yao/agent/assistant" + "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 TestAgentProviderWithAssistantConfig(t *testing.T) { + testutils.Prepare(t) + defer testutils.Clean(t) + + // Load the rerank-agent assistant + ast, err := assistant.Get("tests.rerank-agent") + require.NoError(t, err) + require.NotNil(t, ast) + + // Create test context + ctx := newTestContext(t) + + // Create provider with test assistant + provider := rerank.NewAgentProvider("tests.rerank-agent") + + 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 agent 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 TestAgentProviderWithTopN(t *testing.T) { + testutils.Prepare(t) + defer testutils.Clean(t) + + ctx := newTestContext(t) + provider := rerank.NewAgentProvider("tests.rerank-agent") + + 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 TestAgentProviderWithoutContext(t *testing.T) { + provider := rerank.NewAgentProvider("tests.rerank-agent") + + 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 TestAgentProviderAgentNotFound(t *testing.T) { + testutils.Prepare(t) + defer testutils.Clean(t) + + ctx := newTestContext(t) + provider := rerank.NewAgentProvider("non-existent-agent") + + 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(), "failed to get agent") +} + +func TestAgentProviderEmptyItems(t *testing.T) { + testutils.Prepare(t) + defer testutils.Clean(t) + + ctx := newTestContext(t) + provider := rerank.NewAgentProvider("tests.rerank-agent") + + result, err := provider.Rerank(ctx, "test query", []*types.ResultItem{}, &types.RerankOptions{TopN: 10}) + + require.NoError(t, err) + assert.Empty(t, result) +} + +// newTestContext creates a test context with required fields +func newTestContext(t *testing.T) *context.Context { + t.Helper() + authorized := &oauthTypes.AuthorizedInfo{ + UserID: "test-user", + } + chatID := "test-chat-rerank" + return context.New(t.Context(), authorized, chatID) +} diff --git a/agent/search/rerank/builtin.go b/agent/search/rerank/builtin.go new file mode 100644 index 00000000..94923886 --- /dev/null +++ b/agent/search/rerank/builtin.go @@ -0,0 +1,62 @@ +package rerank + +import ( + "sort" + + "github.com/yaoapp/yao/agent/search/types" +) + +// BuiltinReranker implements simple score-based reranking +// For production use cases requiring semantic understanding, use Agent or MCP mode. +type BuiltinReranker struct{} + +// NewBuiltinReranker creates a new builtin reranker +func NewBuiltinReranker() *BuiltinReranker { + return &BuiltinReranker{} +} + +// Rerank sorts items by weighted score (score * weight) and returns top N +// This is a simple implementation without semantic understanding. +func (r *BuiltinReranker) Rerank(query string, items []*types.ResultItem, opts *types.RerankOptions) ([]*types.ResultItem, error) { + if len(items) == 0 { + return items, nil + } + + // Calculate weighted scores + type scoredItem struct { + item *types.ResultItem + weightedScore float64 + } + + scored := make([]scoredItem, len(items)) + for i, item := range items { + // Weighted score = base score * source weight + // Higher weight sources (user=1.0) get priority over lower (auto=0.6) + weight := item.Weight + if weight == 0 { + weight = 0.6 // Default weight for items without weight + } + scored[i] = scoredItem{ + item: item, + weightedScore: item.Score * weight, + } + } + + // Sort by weighted score descending + sort.Slice(scored, func(i, j int) bool { + return scored[i].weightedScore > scored[j].weightedScore + }) + + // Get top N + topN := opts.TopN + if topN <= 0 || topN > len(scored) { + topN = len(scored) + } + + result := make([]*types.ResultItem, topN) + for i := 0; i < topN; i++ { + result[i] = scored[i].item + } + + return result, nil +} diff --git a/agent/search/rerank/builtin_test.go b/agent/search/rerank/builtin_test.go new file mode 100644 index 00000000..01b47ea0 --- /dev/null +++ b/agent/search/rerank/builtin_test.go @@ -0,0 +1,124 @@ +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) +} diff --git a/agent/search/rerank/mcp.go b/agent/search/rerank/mcp.go new file mode 100644 index 00000000..6bc86108 --- /dev/null +++ b/agent/search/rerank/mcp.go @@ -0,0 +1,171 @@ +package rerank + +import ( + "encoding/json" + "fmt" + "strings" + + "github.com/yaoapp/gou/mcp" + gouMCPTypes "github.com/yaoapp/gou/mcp/types" + "github.com/yaoapp/yao/agent/context" + "github.com/yaoapp/yao/agent/search/types" +) + +// MCPProvider implements reranking by calling an MCP tool +type MCPProvider struct { + serverID string // MCP server ID + toolName string // Tool name +} + +// NewMCPProvider creates a new MCP reranker +// mcpRef format: "server_id.tool_name" +func NewMCPProvider(mcpRef string) (*MCPProvider, error) { + parts := strings.SplitN(mcpRef, ".", 2) + if len(parts) != 2 { + return nil, fmt.Errorf("invalid MCP format, expected 'server.tool', got '%s'", mcpRef) + } + return &MCPProvider{ + serverID: parts[0], + toolName: parts[1], + }, nil +} + +// Rerank calls MCP tool to rerank items +func (p *MCPProvider) Rerank(ctx *context.Context, query string, items []*types.ResultItem, opts *types.RerankOptions) ([]*types.ResultItem, error) { + if ctx == nil { + return nil, fmt.Errorf("context is required for MCP rerank") + } + + // Get MCP client + client, err := mcp.Select(p.serverID) + if err != nil { + return nil, fmt.Errorf("MCP server %s not found: %w", p.serverID, err) + } + + // Build arguments for MCP tool + args := map[string]interface{}{ + "query": query, + "items": items, + "top_n": opts.TopN, + } + + // Call MCP tool + result, err := client.CallTool(ctx.Context, p.toolName, args) + if err != nil { + return nil, fmt.Errorf("MCP tool call failed: %w", err) + } + + // Parse result + return p.parseResult(result, items, opts) +} + +// parseResult extracts reranked items from MCP response +func (p *MCPProvider) parseResult(result *gouMCPTypes.CallToolResponse, originalItems []*types.ResultItem, opts *types.RerankOptions) ([]*types.ResultItem, error) { + if result == nil || len(result.Content) == 0 { + return originalItems, nil + } + + // Build index map for quick lookup + itemMap := make(map[string]*types.ResultItem) + for _, item := range originalItems { + if item.CitationID != "" { + itemMap[item.CitationID] = item + } + } + + // Extract text content from MCP response + var textContent string + for _, content := range result.Content { + if content.Type == gouMCPTypes.ToolContentTypeText && content.Text != "" { + textContent = content.Text + break + } + } + + if textContent == "" { + return originalItems, nil + } + + // Parse JSON response + var response map[string]interface{} + if err := json.Unmarshal([]byte(textContent), &response); err != nil { + // Try parsing as array of IDs + var orderList []string + if err := json.Unmarshal([]byte(textContent), &orderList); err == nil { + return p.reorderByIDs(orderList, itemMap, originalItems, opts) + } + return originalItems, nil + } + + // Try "order" field (list of citation IDs) + if order, ok := response["order"]; ok { + if orderList := toStringSlice(order); len(orderList) > 0 { + return p.reorderByIDs(orderList, itemMap, originalItems, opts) + } + } + + // Try "items" field + if items, ok := response["items"]; ok { + if itemsList := toItemsList(items); len(itemsList) > 0 { + return p.reorderByItems(itemsList, itemMap, originalItems, opts) + } + } + + return originalItems, nil +} + +// reorderByIDs reorders items based on list of citation IDs +func (p *MCPProvider) reorderByIDs(order []string, itemMap map[string]*types.ResultItem, originalItems []*types.ResultItem, opts *types.RerankOptions) ([]*types.ResultItem, error) { + var result []*types.ResultItem + + // Add items in specified order + for _, id := range order { + if item, exists := itemMap[id]; exists { + result = append(result, item) + delete(itemMap, id) + } + } + + // Append remaining items + for _, item := range originalItems { + if _, exists := itemMap[item.CitationID]; exists { + result = append(result, item) + } + } + + // Apply top N + if opts.TopN > 0 && opts.TopN < len(result) { + result = result[:opts.TopN] + } + + return result, nil +} + +// reorderByItems reorders items based on list of item references +func (p *MCPProvider) reorderByItems(itemsList []map[string]interface{}, itemMap map[string]*types.ResultItem, originalItems []*types.ResultItem, opts *types.RerankOptions) ([]*types.ResultItem, error) { + var result []*types.ResultItem + + // Add items in specified order + for _, respItem := range itemsList { + if citationID, ok := respItem["citation_id"].(string); ok { + if item, exists := itemMap[citationID]; exists { + result = append(result, item) + delete(itemMap, citationID) + } + } + } + + // Append remaining items + for _, item := range originalItems { + if _, exists := itemMap[item.CitationID]; exists { + result = append(result, item) + } + } + + // Apply top N + if opts.TopN > 0 && opts.TopN < len(result) { + result = result[:opts.TopN] + } + + return result, nil +} diff --git a/agent/search/rerank/mcp_test.go b/agent/search/rerank/mcp_test.go new file mode 100644 index 00000000..9518432b --- /dev/null +++ b/agent/search/rerank/mcp_test.go @@ -0,0 +1,148 @@ +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) +} diff --git a/agent/search/rerank/reranker.go b/agent/search/rerank/reranker.go new file mode 100644 index 00000000..34b9e46f --- /dev/null +++ b/agent/search/rerank/reranker.go @@ -0,0 +1,99 @@ +// Package rerank provides result reranking for search module +// Supports three modes via uses.rerank configuration: +// - "builtin": Simple score-based sorting (no external dependencies) +// - "": Delegate to an LLM-powered assistant for semantic reranking +// - "mcp:.": Call external MCP tool +// +// For production use cases requiring high accuracy, use Agent or MCP mode. +package rerank + +import ( + "strings" + + "github.com/yaoapp/yao/agent/context" + "github.com/yaoapp/yao/agent/search/types" +) + +// Reranker reorders search results by relevance +// Mode is determined by uses.rerank configuration +type Reranker struct { + usesRerank string // "builtin", "", "mcp:." + config *types.RerankConfig // Rerank options +} + +// NewReranker creates a new reranker +// usesRerank: value from uses.rerank config +// cfg: rerank options from search config +func NewReranker(usesRerank string, cfg *types.RerankConfig) *Reranker { + return &Reranker{ + usesRerank: usesRerank, + config: cfg, + } +} + +// Rerank reorders results based on configured mode +// Returns reordered items, potentially truncated to top N +func (r *Reranker) Rerank(ctx *context.Context, query string, items []*types.ResultItem, opts *types.RerankOptions) ([]*types.ResultItem, error) { + if len(items) == 0 { + return items, nil + } + + // Merge options with config defaults + mergedOpts := r.mergeOptions(opts) + + switch { + case r.usesRerank == "builtin" || r.usesRerank == "": + return r.builtinRerank(query, items, mergedOpts) + case strings.HasPrefix(r.usesRerank, "mcp:"): + return r.mcpRerank(ctx, query, items, mergedOpts) + default: + // Assume it's an assistant ID for Agent mode + return r.agentRerank(ctx, query, items, mergedOpts) + } +} + +// mergeOptions merges runtime options with config defaults +func (r *Reranker) mergeOptions(opts *types.RerankOptions) *types.RerankOptions { + result := &types.RerankOptions{ + TopN: 10, // default + } + + // Apply config defaults + if r.config != nil { + if r.config.TopN > 0 { + result.TopN = r.config.TopN + } + } + + // Apply runtime options (highest priority) + if opts != nil { + if opts.TopN > 0 { + result.TopN = opts.TopN + } + } + + return result +} + +// builtinRerank uses simple score-based sorting +func (r *Reranker) builtinRerank(query string, items []*types.ResultItem, opts *types.RerankOptions) ([]*types.ResultItem, error) { + reranker := NewBuiltinReranker() + return reranker.Rerank(query, items, opts) +} + +// agentRerank delegates to an LLM-powered assistant +func (r *Reranker) agentRerank(ctx *context.Context, query string, items []*types.ResultItem, opts *types.RerankOptions) ([]*types.ResultItem, error) { + provider := NewAgentProvider(r.usesRerank) + return provider.Rerank(ctx, query, items, opts) +} + +// mcpRerank calls an external MCP tool +func (r *Reranker) mcpRerank(ctx *context.Context, query string, items []*types.ResultItem, opts *types.RerankOptions) ([]*types.ResultItem, error) { + mcpRef := strings.TrimPrefix(r.usesRerank, "mcp:") + provider, err := NewMCPProvider(mcpRef) + if err != nil { + // Fallback to builtin on invalid MCP format + return r.builtinRerank(query, items, opts) + } + return provider.Rerank(ctx, query, items, opts) +} diff --git a/agent/search/rerank/reranker_test.go b/agent/search/rerank/reranker_test.go new file mode 100644 index 00000000..9b58e2b9 --- /dev/null +++ b/agent/search/rerank/reranker_test.go @@ -0,0 +1,99 @@ +package rerank + +import ( + "testing" + + "github.com/stretchr/testify/assert" + "github.com/yaoapp/yao/agent/search/types" +) + +func TestReranker_BuiltinMode(t *testing.T) { + reranker := NewReranker("builtin", &types.RerankConfig{TopN: 5}) + + 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}, + } + + result, err := reranker.Rerank(nil, "test query", items, nil) + + assert.NoError(t, err) + assert.Len(t, result, 3) + assert.Equal(t, "ref_001", result[0].CitationID) +} + +func TestReranker_EmptyUsesRerank(t *testing.T) { + // Empty usesRerank should use builtin + reranker := NewReranker("", &types.RerankConfig{TopN: 5}) + + items := []*types.ResultItem{ + {CitationID: "ref_001", Score: 0.9, Weight: 1.0}, + } + + result, err := reranker.Rerank(nil, "test query", items, nil) + + assert.NoError(t, err) + assert.Len(t, result, 1) +} + +func TestReranker_MergeOptions(t *testing.T) { + // Config sets TopN = 5 + reranker := NewReranker("builtin", &types.RerankConfig{TopN: 5}) + + 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}, + {CitationID: "ref_006", Score: 0.4, Weight: 1.0}, + } + + // Runtime opts override config + result, err := reranker.Rerank(nil, "test query", items, &types.RerankOptions{TopN: 3}) + + assert.NoError(t, err) + assert.Len(t, result, 3) +} + +func TestReranker_ConfigTopN(t *testing.T) { + // Config sets TopN = 3 + reranker := NewReranker("builtin", &types.RerankConfig{TopN: 3}) + + 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}, + } + + // No runtime opts, should use config TopN + result, err := reranker.Rerank(nil, "test query", items, nil) + + assert.NoError(t, err) + assert.Len(t, result, 3) +} + +func TestReranker_NilConfig(t *testing.T) { + reranker := NewReranker("builtin", nil) + + items := []*types.ResultItem{ + {CitationID: "ref_001", Score: 0.9, Weight: 1.0}, + } + + result, err := reranker.Rerank(nil, "test query", items, nil) + + assert.NoError(t, err) + assert.Len(t, result, 1) +} + +func TestReranker_EmptyItems(t *testing.T) { + reranker := NewReranker("builtin", &types.RerankConfig{TopN: 5}) + + result, err := reranker.Rerank(nil, "test query", []*types.ResultItem{}, nil) + + assert.NoError(t, err) + assert.Empty(t, result) +} diff --git a/agent/search/search.go b/agent/search/search.go index 23896756..2332f852 100644 --- a/agent/search/search.go +++ b/agent/search/search.go @@ -3,10 +3,12 @@ package search import ( "sync" + "github.com/yaoapp/yao/agent/context" "github.com/yaoapp/yao/agent/search/handlers/db" "github.com/yaoapp/yao/agent/search/handlers/kb" "github.com/yaoapp/yao/agent/search/handlers/web" "github.com/yaoapp/yao/agent/search/interfaces" + "github.com/yaoapp/yao/agent/search/rerank" "github.com/yaoapp/yao/agent/search/types" ) @@ -14,7 +16,7 @@ import ( type Searcher struct { config *types.Config // Merged config (global + assistant) handlers map[types.SearchType]interfaces.Handler - reranker interfaces.Reranker + reranker *rerank.Reranker citation *CitationGenerator } @@ -46,13 +48,13 @@ func New(cfg *types.Config, uses *Uses) *Searcher { types.SearchTypeKB: kb.NewHandler(cfg.KB), types.SearchTypeDB: db.NewHandler(uses.QueryDSL, cfg.DB), }, - reranker: newBuiltinReranker(), // TODO: use uses.Rerank to select reranker + reranker: rerank.NewReranker(uses.Rerank, cfg.Rerank), citation: NewCitationGenerator(), } } // Search executes a single search request -func (s *Searcher) Search(req *types.Request) (*types.Result, error) { +func (s *Searcher) Search(ctx *context.Context, req *types.Request) (*types.Result, error) { handler, ok := s.handlers[req.Type] if !ok { return &types.Result{Error: "unsupported search type"}, nil @@ -71,7 +73,7 @@ func (s *Searcher) Search(req *types.Request) (*types.Result, error) { // Rerank if requested if req.Rerank != nil && s.reranker != nil { - result.Items, _ = s.reranker.Rerank(req.Query, result.Items, req.Rerank) + result.Items, _ = s.reranker.Rerank(ctx, req.Query, result.Items, req.Rerank) } // Generate citation IDs @@ -83,7 +85,7 @@ func (s *Searcher) Search(req *types.Request) (*types.Result, error) { } // SearchMultiple executes multiple searches in parallel -func (s *Searcher) SearchMultiple(reqs []*types.Request) ([]*types.Result, error) { +func (s *Searcher) SearchMultiple(ctx *context.Context, reqs []*types.Request) ([]*types.Result, error) { results := make([]*types.Result, len(reqs)) var wg sync.WaitGroup var mu sync.Mutex @@ -92,7 +94,7 @@ func (s *Searcher) SearchMultiple(reqs []*types.Request) ([]*types.Result, error wg.Add(1) go func(idx int, r *types.Request) { defer wg.Done() - result, _ := s.Search(r) + result, _ := s.Search(ctx, r) mu.Lock() results[idx] = result mu.Unlock() @@ -107,19 +109,3 @@ func (s *Searcher) SearchMultiple(reqs []*types.Request) ([]*types.Result, error func (s *Searcher) BuildReferences(results []*types.Result) []*types.Reference { return BuildReferences(results) } - -// builtinReranker is a simple score-based reranker -type builtinReranker struct{} - -func newBuiltinReranker() *builtinReranker { - return &builtinReranker{} -} - -func (r *builtinReranker) Rerank(query string, items []*types.ResultItem, opts *types.RerankOptions) ([]*types.ResultItem, error) { - // Simple implementation: sort by score (already sorted in most cases) - // TODO: Implement proper reranking logic - if opts != nil && opts.TopN > 0 && opts.TopN < len(items) { - return items[:opts.TopN], nil - } - return items, nil -}