rerank

type CohereConfig struct {
	APIKey  string        `json:"api_key" yaml:"api_key"`
	BaseURL string        `json:"base_url" yaml:"base_url"`
	Model   string        `json:"model,omitempty" yaml:"model,omitempty"` // rerank-v3.5
	Timeout time.Duration `json:"timeout,omitempty" yaml:"timeout,omitempty"`
}

type CohereProvider struct {
	// contains filtered or unexported fields
}

type Document struct {
	Text  string `json:"text"`
	ID    string `json:"id,omitempty"`
	Title string `json:"title,omitempty"`
}

type JinaConfig struct {
	APIKey  string        `json:"api_key" yaml:"api_key"`
	BaseURL string        `json:"base_url" yaml:"base_url"`
	Model   string        `json:"model,omitempty" yaml:"model,omitempty"` // jina-reranker-v2-base-multilingual
	Timeout time.Duration `json:"timeout,omitempty" yaml:"timeout,omitempty"`
}

type JinaProvider struct {
	// contains filtered or unexported fields
}

type Provider interface {
	// Rerank 根据查询的相关性重排序文档.
	Rerank(ctx context.Context, req *RerankRequest) (*RerankResponse, error)

	// RerankSimple 是简单重排序的便捷方法.
	RerankSimple(ctx context.Context, query string, documents []string, topN int) ([]RerankResult, error)

	// Name 返回提供者名称.
	Name() string

	// MaxDocuments 返回支持的最大文档数量.
	MaxDocuments() int
}

type RerankRequest struct {
	Query           string            `json:"query"`
	Documents       []Document        `json:"documents"`
	Model           string            `json:"model,omitempty"`
	TopN            int               `json:"top_n,omitempty"`              // Return top N results
	ReturnDocuments bool              `json:"return_documents,omitempty"`   // Include document text in response
	MaxChunksPerDoc int               `json:"max_chunks_per_doc,omitempty"` // For long documents
	Metadata        map[string]string `json:"metadata,omitempty"`
}

type RerankResponse struct {
	ID        string         `json:"id,omitempty"`
	Provider  string         `json:"provider"`
	Model     string         `json:"model"`
	Results   []RerankResult `json:"results"`
	Usage     RerankUsage    `json:"usage"`
	CreatedAt time.Time      `json:"created_at,omitempty"`
}

type RerankResult struct {
	Index          int      `json:"index"`           // Original index in input
	RelevanceScore float64  `json:"relevance_score"` // 0-1 normalized score
	Document       Document `json:"document,omitempty"`
}

type RerankUsage struct {
	SearchUnits int     `json:"search_units,omitempty"`
	TotalTokens int     `json:"total_tokens,omitempty"`
	Cost        float64 `json:"cost,omitempty"`
}

type VoyageConfig struct {
	APIKey  string        `json:"api_key" yaml:"api_key"`
	BaseURL string        `json:"base_url" yaml:"base_url"`
	Model   string        `json:"model,omitempty" yaml:"model,omitempty"` // rerank-2
	Timeout time.Duration `json:"timeout,omitempty" yaml:"timeout,omitempty"`
}

type VoyageProvider struct {
	// contains filtered or unexported fields
}