50 lines
1.5 KiB
Go
50 lines
1.5 KiB
Go
|
package embedding
|
||
|
|
|
||
|
|
import (
|
||
|
|
"context"
|
||
|
|
"errors"
|
||
|
|
|
||
|
|
"github.com/Tencent/WeKnora/internal/models/requestbudget"
|
||
|
|
)
|
||
|
|
|
||
|
|
func stopProviderRetry(ctx context.Context, err error) bool {
|
||
|
|
return err != nil && (ctx.Err() != nil || errors.Is(err, requestbudget.ErrStopped))
|
||
|
|
}
|
||
|
|
|
||
|
|
type requestBudgetEmbedder struct {
|
||
|
|
Embedder
|
||
|
|
supported bool
|
||
|
|
}
|
||
|
|
|
||
|
|
func wrapEmbeddingRequestBudget(inner Embedder) Embedder {
|
||
|
|
supported := false
|
||
|
|
switch inner.(type) {
|
||
|
|
case *OpenAIEmbedder, *AzureOpenAIEmbedder, *AliyunEmbedder, *VolcengineEmbedder,
|
||
|
|
*JinaEmbedder, *NvidiaEmbedder, *GeminiEmbedder, *ZhipuEmbedder, *WeKnoraCloudEmbedder:
|
||
|
|
supported = true
|
||
|
|
}
|
||
|
|
return &requestBudgetEmbedder{Embedder: inner, supported: supported}
|
||
|
|
}
|
||
|
|
|
||
|
|
func (w *requestBudgetEmbedder) Embed(ctx context.Context, text string) ([]float32, error) {
|
||
|
|
if err := requestbudget.RequireSupported(ctx, w.supported); err != nil {
|
||
|
|
return nil, err
|
||
|
|
}
|
||
|
|
return w.Embedder.Embed(ctx, text)
|
||
|
|
}
|
||
|
|
|
||
|
|
func (w *requestBudgetEmbedder) BatchEmbed(ctx context.Context, texts []string) ([][]float32, error) {
|
||
|
|
if err := requestbudget.RequireSupported(ctx, w.supported); err != nil {
|
||
|
|
return nil, err
|
||
|
|
}
|
||
|
|
return w.Embedder.BatchEmbed(ctx, texts)
|
||
|
|
}
|
||
|
|
|
||
|
|
func (w *requestBudgetEmbedder) BatchEmbedWithPool(ctx context.Context, model Embedder, texts []string) ([][]float32, error) {
|
||
|
|
if err := requestbudget.RequireSupported(ctx, w.supported); err != nil {
|
||
|
|
return nil, err
|
||
|
|
}
|
||
|
|
// Preserve the concurrency wrapper supplied by the factory for each batch.
|
||
|
|
return w.Embedder.BatchEmbedWithPool(ctx, model, texts)
|
||
|
|
}
|