xinmi/internal/models/embedding/request_budget.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)
}