Refactor Search Module and Update Documentation
- 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
parent
d3865d782b
commit
534f4d6ed5
11 changed files with 1192 additions and 159 deletions
|
|
@ -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", "<assistant-id>", "mcp:<server>.<tool>"
|
||||
Web 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
|
||||
// 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", "<assistant-id>", "mcp:<server>.<tool>"
|
||||
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`:
|
||||
|
||||
|
|
|
|||
|
|
@ -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)
|
||||
}
|
||||
|
|
|
|||
232
agent/search/rerank/agent.go
Normal file
232
agent/search/rerank/agent.go
Normal 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
123
agent/search/rerank/agent_test.go
Normal 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
62
agent/search/rerank/builtin.go
Normal 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
124
agent/search/rerank/builtin_test.go
Normal 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
171
agent/search/rerank/mcp.go
Normal 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
148
agent/search/rerank/mcp_test.go
Normal 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
99
agent/search/rerank/reranker.go
Normal 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
99
agent/search/rerank/reranker_test.go
Normal 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)
|
||||
}
|
||||
|
|
@ -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
|
||||
}
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue