WeKnora deployment source with the reviewed wiki identity, summary and recovery fixes

internal/models/embedding/request_budget.go

@@ -0,0 +1,49 @@
+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)
+}