opencode/yao
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions

Refactor Search Module and Update Documentation

Browse source 
- 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.
This commit is contained in:
Max 2025-12-13 15:01:53 +08:00
parent d3865d782b
commit 534f4d6ed5
11 changed files with 1192 additions and 159 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

259
agent/search/DESIGN.md
View file

@ -124,11 +124,10 @@ sequenceDiagram
``` ```
agent/search/ agent/search/
├── DESIGN.md # This document ├── DESIGN.md # This document
├── TODO.md # Implementation plan and progress
├── search.go # Main Searcher implementation and public API ├── search.go # Main Searcher implementation and public API
├── registry.go # Handler registry (manages web/kb/db handlers) ├── registry.go # Handler registry (manages web/kb/db handlers)
├── jsapi.go # JavaScript API bindings for hooks ├── jsapi.go # JavaScript API bindings for hooks (skeleton)
├── trace.go # Trace node creation and management
├── output.go # Real-time output/streaming to client
├── citation.go # Citation ID generation and tracking ├── citation.go # Citation ID generation and tracking
├── reference.go # Reference building and LLM context formatting ├── reference.go # Reference building and LLM context formatting
│ │
@ -144,19 +143,19 @@ agent/search/
│ ├── reranker.go # Reranker interface │ ├── reranker.go # Reranker interface
│ └── nlp.go # NLP interfaces (KeywordExtractor, QueryDSLGenerator) │ └── nlp.go # NLP interfaces (KeywordExtractor, QueryDSLGenerator)
│ │
├── rerank/ # Result reranking implementations ├── rerank/ # Result reranking implementations (Handler + Registry pattern) ✅
│ ├── rerank.go # Reranker factory and common logic │ ├── reranker.go # Main entry point (mode dispatch)
│ ├── builtin.go # Built-in score-based reranking (default) │ ├── builtin.go # Builtin: weighted score sorting
│ ├── agent.go # Agent-based reranking (delegate to another assistant) │ ├── agent.go # Agent mode (delegate to LLM assistant)
│ └── mcp.go # MCP-based reranking (call MCP server tool) │ └── mcp.go # MCP mode (external service)
│ │
├── nlp/ # Natural language processing for search ├── nlp/ # Natural language processing for search
│ ├── keyword/ # Keyword extraction (Handler + Registry pattern) │ ├── keyword/ # Keyword extraction (Handler + Registry pattern) ✅
│ │ ├── extractor.go # Main extractor (mode dispatch) │ │ ├── extractor.go # Main extractor (mode dispatch)
│ │ ├── builtin.go # Builtin frequency-based extraction │ │ ├── builtin.go # Builtin frequency-based extraction
│ │ ├── agent.go # Agent mode (LLM-powered) │ │ ├── agent.go # Agent mode (LLM-powered)
│ │ └── mcp.go # MCP mode (external service) │ │ └── mcp.go # MCP mode (external service)
│ └── querydsl/ # QueryDSL generation for DB search │ └── querydsl/ # QueryDSL generation for DB search (待实现)
│ ├── generator.go # Main generator (mode dispatch) │ ├── generator.go # Main generator (mode dispatch)
│ ├── builtin.go # Builtin template-based generation │ ├── builtin.go # Builtin template-based generation
│ ├── agent.go # Agent mode (LLM-powered) │ ├── agent.go # Agent mode (LLM-powered)
@ -164,7 +163,7 @@ agent/search/
│ # Note: Embedding follows KB collection config, not in this package │ # Note: Embedding follows KB collection config, not in this package
│ │
├── handlers/ # Search handler implementations ├── handlers/ # Search handler implementations
│ ├── web/ # Web search │ ├── web/ # Web search ✅
│ │ ├── handler.go # Web search handler (mode dispatch) │ │ ├── handler.go # Web search handler (mode dispatch)
│ │ ├── tavily.go # Tavily provider (builtin) │ │ ├── tavily.go # Tavily provider (builtin)
│ │ ├── serper.go # Serper provider (serper.dev, builtin) │ │ ├── serper.go # Serper provider (serper.dev, builtin)
@ -172,18 +171,22 @@ agent/search/
│ │ ├── agent.go # Agent mode (AI Search) │ │ ├── agent.go # Agent mode (AI Search)
│ │ └── mcp.go # MCP mode (external service) │ │ └── mcp.go # MCP mode (external service)
│ │ │ │
│ ├── kb/ # Knowledge base search │ ├── kb/ # Knowledge base search (骨架)
│ │ ├── handler.go # KB search handler │ │ ├── handler.go # KB search handler
│ │ ├── vector.go # Vector similarity search │ │ ├── vector.go # Vector similarity search (待实现)
│ │ └── graph.go # Graph-based association (GraphRAG) │ │ └── graph.go # Graph-based association (待实现)
│ │ │ │
│ └── db/ # Database search (Yao Model/QueryDSL) │ └── db/ # Database search (骨架)
│ ├── handler.go # DB search handler │ ├── handler.go # DB search handler
│ ├── query.go # QueryDSL builder │ ├── query.go # QueryDSL builder (待实现)
│ └── schema.go # Model schema introspection │ └── schema.go # Model schema introspection (待实现)
│ │
└── defaults/ # Default configuration values └── defaults/ # Default configuration values
└── defaults.go # System built-in defaults (used by agent/load.go) └── 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 ### Dependency Graph
@ -249,13 +252,13 @@ import (
type Searcher struct { type Searcher struct {
config *types.Config // Merged config (global + assistant) config *types.Config // Merged config (global + assistant)
handlers map[types.SearchType]interfaces.Handler handlers map[types.SearchType]interfaces.Handler
reranker interfaces.Reranker reranker *rerank.Reranker // Uses rerank package directly
citation *CitationGenerator 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 // These are extracted from context.Uses and search config
type SearchUses struct { type Uses struct {
Search string // "builtin", "disabled", "<assistant-id>", "mcp:<server>.<tool>" Search string // "builtin", "disabled", "<assistant-id>", "mcp:<server>.<tool>"
Web string // "builtin", "<assistant-id>", "mcp:<server>.<tool>" Web string // "builtin", "<assistant-id>", "mcp:<server>.<tool>"
Keyword string // "builtin", "<assistant-id>", "mcp:<server>.<tool>" Keyword string // "builtin", "<assistant-id>", "mcp:<server>.<tool>"
@ -266,7 +269,7 @@ type SearchUses struct {
// New creates a new Searcher instance // New creates a new Searcher instance
// cfg: merged config from agent/load.go + assistant config // cfg: merged config from agent/load.go + assistant config
// uses: merged uses configuration (global → assistant → hook) // 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{ return &Searcher{
config: cfg, config: cfg,
handlers: map[types.SearchType]interfaces.Handler{ 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.SearchTypeKB: kb.NewHandler(cfg.KB), // KB always builtin
types.SearchTypeDB: db.NewHandler(uses.QueryDSL, cfg.DB), types.SearchTypeDB: db.NewHandler(uses.QueryDSL, cfg.DB),
}, },
reranker: rerank.NewReranker(uses.Rerank), reranker: rerank.NewReranker(uses.Rerank, cfg.Rerank),
citation: NewCitationGenerator(), 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 return &types.Result{Error: "unsupported search type"}, nil
} }
// Execute search // Execute search (handler doesn't need ctx)
result, err := handler.Search(ctx, req) result, err := handler.Search(req)
if err != nil { if err != nil {
return &types.Result{Error: err.Error()}, 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) item.Weight = s.config.GetWeight(req.Source)
} }
// Rerank if requested // Rerank if requested (reranker needs ctx for Agent/MCP modes)
if req.Rerank != nil { if req.Rerank != nil && s.reranker != nil {
result.Items, _ = s.reranker.Rerank(ctx, req.Query, result.Items, req.Rerank) 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 package interfaces
import ( import (
"github.com/yaoapp/yao/agent/context"
"github.com/yaoapp/yao/agent/search/types" "github.com/yaoapp/yao/agent/search/types"
) )
@ -405,11 +407,8 @@ type Handler interface {
// Type returns the search type this handler supports // Type returns the search type this handler supports
Type() types.SearchType 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 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 package interfaces
import ( import (
"github.com/yaoapp/yao/agent/context"
"github.com/yaoapp/yao/agent/search/types" "github.com/yaoapp/yao/agent/search/types"
) )
// Searcher is the main interface exposed to external callers // Searcher is the main interface exposed to external callers
type Searcher interface { type Searcher interface {
// Search executes a single search request // 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 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 converts search results to unified Reference format for LLM
BuildReferences(results []*types.Result) []*types.Reference 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`) ### NLP Interfaces (`interfaces/nlp.go`)
```go ```go
package interfaces package interfaces
import ( import (
"github.com/yaoapp/gou/model"
"github.com/yaoapp/gou/query/gou"
"github.com/yaoapp/yao/agent/context" "github.com/yaoapp/yao/agent/context"
"github.com/yaoapp/yao/agent/search/types" "github.com/yaoapp/yao/agent/search/types"
) )
@ -449,6 +451,7 @@ import (
// KeywordExtractor extracts keywords for web search // KeywordExtractor extracts keywords for web search
type KeywordExtractor interface { type KeywordExtractor interface {
// Extract extracts search keywords from user message // 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) Extract(ctx *context.Context, content string, opts *types.KeywordOptions) ([]string, error)
} }
@ -1768,7 +1771,7 @@ package web
import ( import (
"strings" "strings"
"github.com/yaoapp/yao/agent/context" agentContext "github.com/yaoapp/yao/agent/context"
"github.com/yaoapp/yao/agent/search/types" "github.com/yaoapp/yao/agent/search/types"
) )
@ -1779,61 +1782,28 @@ type Handler struct {
} }
// NewHandler creates a new web search handler // NewHandler creates a new web search handler
func NewHandler(usesWeb string, cfg *types.WebConfig) *Handler { func NewHandler(usesWeb string, cfg *types.WebConfig) *Handler
return &Handler{usesWeb: usesWeb, config: cfg}
}
// Search executes web search based on uses.web mode // Type returns the search type this handler supports
func (h *Handler) Search(ctx *context.Context, req *types.Request) (*types.Result, error) { func (h *Handler) Type() types.SearchType
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)
}
}
// builtinSearch uses Tavily/Serper/SerpAPI directly // Search implements interfaces.Handler (without context)
func (h *Handler) builtinSearch(ctx *context.Context, req *types.Request) (*types.Result, error) { func (h *Handler) Search(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)
}
}
// agentSearch delegates to an assistant for AI-powered search // SearchWithContext executes web search with context (for Agent/MCP modes)
func (h *Handler) agentSearch(ctx *context.Context, req *types.Request) (*types.Result, error) { func (h *Handler) SearchWithContext(ctx *agentContext.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
}
// mcpSearch calls external MCP tool **Directory Structure:**
func (h *Handler) mcpSearch(ctx *context.Context, req *types.Request) (*types.Result, error) {
// Parse "mcp:server.tool" ```
mcpRef := strings.TrimPrefix(h.usesWeb, "mcp:") handlers/web/
parts := strings.SplitN(mcpRef, ".", 2) ├── handler.go # Main entry point (mode dispatch)
if len(parts) != 2 { ├── tavily.go # Tavily provider (builtin)
return nil, fmt.Errorf("invalid MCP format, expected 'mcp:server.tool', got '%s'", h.usesWeb) ├── serper.go # Serper provider (serper.dev)
} ├── serpapi.go # SerpAPI provider (serpapi.com, multi-engine)
serverID, toolName := parts[0], parts[1] ├── agent.go # Agent mode (AI Search)
// Call MCP tool └── mcp.go # MCP mode (external service)
return nil, nil
}
``` ```
**Built-in Providers (when `uses.web = "builtin"`):** **Built-in Providers (when `uses.web = "builtin"`):**
@ -1942,8 +1912,6 @@ function Create(ctx, messages, options) {
package kb package kb
import ( import (
"github.com/yaoapp/yao/agent/context"
"github.com/yaoapp/yao/agent/search/interfaces"
"github.com/yaoapp/yao/agent/search/types" "github.com/yaoapp/yao/agent/search/types"
) )
@ -1953,18 +1921,14 @@ type Handler struct {
} }
// NewHandler creates a new KB search handler // NewHandler creates a new KB search handler
func NewHandler(cfg *types.KBConfig) *Handler { func NewHandler(cfg *types.KBConfig) *Handler
return &Handler{config: cfg}
} // Type returns the search type this handler supports
func (h *Handler) Type() types.SearchType
// Search executes vector search and optional graph association // Search executes vector search and optional graph association
func (h *Handler) Search(ctx *context.Context, req *types.Request) (*types.Result, error) { // TODO: Implement actual search logic
// 1. Generate embedding via query processor func (h *Handler) Search(req *types.Request) (*types.Result, error)
// 2. Vector search in collections
// 3. Optional: Graph association (if req.Graph)
// 4. Return results
return nil, nil
}
``` ```
| File | Description | | File | Description |
@ -1980,8 +1944,6 @@ func (h *Handler) Search(ctx *context.Context, req *types.Request) (*types.Resul
package db package db
import ( import (
"github.com/yaoapp/yao/agent/context"
"github.com/yaoapp/yao/agent/search/interfaces"
"github.com/yaoapp/yao/agent/search/types" "github.com/yaoapp/yao/agent/search/types"
) )
@ -1992,18 +1954,14 @@ type Handler struct {
} }
// NewHandler creates a new DB search handler // NewHandler creates a new DB search handler
func NewHandler(usesQueryDSL string, cfg *types.DBConfig) *Handler { func NewHandler(usesQueryDSL string, cfg *types.DBConfig) *Handler
return &Handler{usesQueryDSL: usesQueryDSL, config: cfg}
} // Type returns the search type this handler supports
func (h *Handler) Type() types.SearchType
// Search converts NL to QueryDSL and executes // Search converts NL to QueryDSL and executes
func (h *Handler) Search(ctx *context.Context, req *types.Request) (*types.Result, error) { // TODO: Implement actual search logic
// 1. Get model schemas func (h *Handler) Search(req *types.Request) (*types.Result, error)
// 2. Generate QueryDSL via query processor
// 3. Execute queries on models
// 4. Return results
return nil, nil
}
``` ```
| File | Description | | File | Description |
@ -2023,42 +1981,71 @@ Integrates with Yao's Model/QueryDSL system:
### Reranking (`rerank/`) ### Reranking (`rerank/`)
The rerank module follows the Handler + Registry pattern, consistent with `keyword/` and `web/`.
```go ```go
// rerank/rerank.go // rerank/reranker.go
package rerank package rerank
import ( import (
"github.com/yaoapp/yao/agent/search/interfaces" "strings"
"github.com/yaoapp/yao/agent/context"
"github.com/yaoapp/yao/agent/search/types" "github.com/yaoapp/yao/agent/search/types"
) )
// NewReranker creates a reranker based on uses.rerank config // Reranker reorders search results by relevance
func NewReranker(usesRerank string) interfaces.Reranker { // Mode is determined by uses.rerank configuration
switch { type Reranker struct {
case usesRerank == "builtin" || usesRerank == "": usesRerank string // "builtin", "<assistant-id>", "mcp:<server>.<tool>"
return NewBuiltinReranker() config *types.RerankConfig
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)
}
} }
// 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 | **Directory Structure:**
| ------------ | -------------------------------- |
| `rerank.go` | Factory and common logic | ```
| `builtin.go` | Simple score sorting (default) | rerank/
| `agent.go` | Delegate to an assistant (Agent) | ├── reranker.go # Main entry point (mode dispatch)
| `mcp.go` | Call MCP tool for reranking | ├── 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`: Configure via `uses.rerank` in `agent/agent.yml`:

4
agent/search/interfaces/reranker.go
View file

@ -1,11 +1,13 @@
package interfaces package interfaces
import ( import (
"github.com/yaoapp/yao/agent/context"
"github.com/yaoapp/yao/agent/search/types" "github.com/yaoapp/yao/agent/search/types"
) )
// Reranker reorders search results by relevance // Reranker reorders search results by relevance
type Reranker interface { type Reranker interface {
// Rerank reorders results based on query relevance // 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)
} }

232
agent/search/rerank/agent.go Normal file
View file

@ -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
}

123
agent/search/rerank/agent_test.go Normal file
View file

@ -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)
}

62
agent/search/rerank/builtin.go Normal file
View file

@ -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
}

124
agent/search/rerank/builtin_test.go Normal file
View file

@ -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)
}

171
agent/search/rerank/mcp.go Normal file
View file

@ -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
}

148
agent/search/rerank/mcp_test.go Normal file
View file

@ -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)
}

99
agent/search/rerank/reranker.go Normal file
View file

@ -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)
// - "<assistant-id>": Delegate to an LLM-powered assistant for semantic reranking
// - "mcp:<server>.<tool>": 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", "<assistant-id>", "mcp:<server>.<tool>"
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)
}

99
agent/search/rerank/reranker_test.go Normal file
View file

@ -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)
}

30
agent/search/search.go
View file

@ -3,10 +3,12 @@ package search
import ( import (
"sync" "sync"
"github.com/yaoapp/yao/agent/context"
"github.com/yaoapp/yao/agent/search/handlers/db" "github.com/yaoapp/yao/agent/search/handlers/db"
"github.com/yaoapp/yao/agent/search/handlers/kb" "github.com/yaoapp/yao/agent/search/handlers/kb"
"github.com/yaoapp/yao/agent/search/handlers/web" "github.com/yaoapp/yao/agent/search/handlers/web"
"github.com/yaoapp/yao/agent/search/interfaces" "github.com/yaoapp/yao/agent/search/interfaces"
"github.com/yaoapp/yao/agent/search/rerank"
"github.com/yaoapp/yao/agent/search/types" "github.com/yaoapp/yao/agent/search/types"
) )
@ -14,7 +16,7 @@ import (
type Searcher struct { type Searcher struct {
config *types.Config // Merged config (global + assistant) config *types.Config // Merged config (global + assistant)
handlers map[types.SearchType]interfaces.Handler handlers map[types.SearchType]interfaces.Handler
reranker interfaces.Reranker reranker *rerank.Reranker
citation *CitationGenerator citation *CitationGenerator
} }
@ -46,13 +48,13 @@ func New(cfg *types.Config, uses *Uses) *Searcher {
types.SearchTypeKB: kb.NewHandler(cfg.KB), types.SearchTypeKB: kb.NewHandler(cfg.KB),
types.SearchTypeDB: db.NewHandler(uses.QueryDSL, cfg.DB), 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(), citation: NewCitationGenerator(),
} }
} }
// Search executes a single search request // 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] handler, ok := s.handlers[req.Type]
if !ok { if !ok {
return &types.Result{Error: "unsupported search type"}, nil 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 // Rerank if requested
if req.Rerank != nil && s.reranker != nil { 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 // Generate citation IDs
@ -83,7 +85,7 @@ func (s *Searcher) Search(req *types.Request) (*types.Result, error) {
} }
// SearchMultiple executes multiple searches in parallel // 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)) results := make([]*types.Result, len(reqs))
var wg sync.WaitGroup var wg sync.WaitGroup
var mu sync.Mutex var mu sync.Mutex
@ -92,7 +94,7 @@ func (s *Searcher) SearchMultiple(reqs []*types.Request) ([]*types.Result, error
wg.Add(1) wg.Add(1)
go func(idx int, r *types.Request) { go func(idx int, r *types.Request) {
defer wg.Done() defer wg.Done()
result, _ := s.Search(r) result, _ := s.Search(ctx, r)
mu.Lock() mu.Lock()
results[idx] = result results[idx] = result
mu.Unlock() 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 { func (s *Searcher) BuildReferences(results []*types.Result) []*types.Reference {
return BuildReferences(results) 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
}