50 lines
1.1 KiB
Go
50 lines
1.1 KiB
Go
package rerank
|
|
|
|
import (
|
|
"context"
|
|
"fmt"
|
|
"strings"
|
|
|
|
"knowlege-lsxd/internal/types"
|
|
)
|
|
|
|
// Reranker defines the interface for document reranking
|
|
type Reranker interface {
|
|
// Rerank reranks documents based on relevance to the query
|
|
Rerank(ctx context.Context, query string, documents []string) ([]RankResult, error)
|
|
|
|
// GetModelName returns the model name
|
|
GetModelName() string
|
|
|
|
// GetModelID returns the model ID
|
|
GetModelID() string
|
|
}
|
|
|
|
type RankResult struct {
|
|
Index int `json:"index"`
|
|
Document DocumentInfo `json:"document"`
|
|
RelevanceScore float64 `json:"relevance_score"`
|
|
}
|
|
|
|
type DocumentInfo struct {
|
|
Text string `json:"text"`
|
|
}
|
|
|
|
type RerankerConfig struct {
|
|
APIKey string
|
|
BaseURL string
|
|
ModelName string
|
|
Source types.ModelSource
|
|
ModelID string
|
|
}
|
|
|
|
// NewReranker creates a reranker
|
|
func NewReranker(config *RerankerConfig) (Reranker, error) {
|
|
switch strings.ToLower(string(config.Source)) {
|
|
case string(types.ModelSourceRemote):
|
|
return NewOpenAIReranker(config)
|
|
default:
|
|
return nil, fmt.Errorf("unsupported rerank model source: %s", config.Source)
|
|
}
|
|
}
|