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

internal/models/embedding/request_budget_test.go

@@ -0,0 +1,257 @@
+package embedding
+
+import (
+	"context"
+	"encoding/json"
+	"errors"
+	"fmt"
+	"net/http"
+	"net/http/httptest"
+	"sync"
+	"sync/atomic"
+	"testing"
+	"time"
+
+	"github.com/Tencent/WeKnora/internal/models/requestbudget"
+	"github.com/Tencent/WeKnora/internal/types"
+	"github.com/panjf2000/ants/v2"
+)
+
+type embeddingBudgetGuard struct {
+	mu             sync.Mutex
+	limit          int
+	models, phases []string
+}
+
+func (g *embeddingBudgetGuard) Key() string                     { return "local-embedding-test" }
+func (g *embeddingBudgetGuard) Check(ctx context.Context) error { return ctx.Err() }
+func (g *embeddingBudgetGuard) Reserve(ctx context.Context, model, phase string) error {
+	g.mu.Lock()
+	defer g.mu.Unlock()
+	if err := ctx.Err(); err != nil {
+		return err
+	}
+	if len(g.models) >= g.limit {
+		return requestbudget.ErrStopped
+	}
+	g.models, g.phases = append(g.models, model), append(g.phases, phase)
+	return nil
+}
+func embeddingBudgetContext(g *embeddingBudgetGuard) context.Context {
+	return requestbudget.WithModel(requestbudget.WithPhase(requestbudget.WithGuard(context.Background(), g), "wiki_embedding"), "caller-forged-model")
+}
+func assertEmbeddingBudget(t *testing.T, g *embeddingBudgetGuard, sends int) {
+	t.Helper()
+	g.mu.Lock()
+	defer g.mu.Unlock()
+	if len(g.models) != sends {
+		t.Fatalf("sends=%d reservations=%d", sends, len(g.models))
+	}
+	for i := range g.models {
+		if g.models[i] != "concrete-embedding" || g.phases[i] != "wiki_embedding" {
+			t.Fatalf("lost adapter identity/context: %v %v", g.models, g.phases)
+		}
+	}
+}
+
+func TestRequestBudgetEmbeddingProviderTransports(t *testing.T) {
+	t.Setenv("SSRF_WHITELIST", "127.0.0.1")
+	var sends atomic.Int32
+	server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+		sends.Add(1)
+		w.Header().Set("Content-Type", "application/json")
+		fmt.Fprint(w, `{"data":[{"index":0,"embedding":[0.1]}]}`)
+	}))
+	defer server.Close()
+	client := newEmbeddingHTTPClient(3 * time.Second)
+	calls := map[string]func(context.Context, []byte) (*http.Response, error){
+		"openai":     (&OpenAIEmbedder{baseURL: server.URL, modelID: "concrete-embedding", httpClient: client}).doRequestWithRetry,
+		"azure":      (&AzureOpenAIEmbedder{baseURL: server.URL, modelID: "concrete-embedding", httpClient: client}).doRequestWithRetry,
+		"aliyun":     (&AliyunEmbedder{baseURL: server.URL, modelID: "concrete-embedding", httpClient: client}).doRequestWithRetry,
+		"volcengine": (&VolcengineEmbedder{baseURL: server.URL, modelID: "concrete-embedding", httpClient: client}).doRequestWithRetry,
+		"jina":       (&JinaEmbedder{baseURL: server.URL, modelID: "concrete-embedding", httpClient: client}).doRequestWithRetry,
+		"nvidia":     (&NvidiaEmbedder{baseURL: server.URL, modelID: "concrete-embedding", httpClient: client}).doRequestWithRetry,
+		"gemini":     (&GeminiEmbedder{baseURL: server.URL, modelID: "concrete-embedding", httpClient: client}).doRequestWithRetry,
+		"zhipu":      (&ZhipuEmbedder{baseURL: server.URL, modelID: "concrete-embedding", httpClient: client}).doRequestWithRetry,
+	}
+	for name, call := range calls {
+		for _, limit := range []int{0, 1} {
+			t.Run(fmt.Sprintf("%s/limit=%d", name, limit), func(t *testing.T) {
+				sends.Store(0)
+				g := &embeddingBudgetGuard{limit: limit}
+				resp, err := call(embeddingBudgetContext(g), []byte(`{}`))
+				if resp != nil {
+					resp.Body.Close()
+				}
+				if limit == 0 && !errors.Is(err, requestbudget.ErrStopped) {
+					t.Fatalf("guard not applied: %v", err)
+				}
+				if limit == 1 && err != nil {
+					t.Fatal(err)
+				}
+				if int(sends.Load()) != limit {
+					t.Fatalf("sends=%d want=%d", sends.Load(), limit)
+				}
+				assertEmbeddingBudget(t, g, limit)
+			})
+		}
+	}
+	for _, limit := range []int{0, 1} {
+		t.Run(fmt.Sprintf("weknoracloud/limit=%d", limit), func(t *testing.T) {
+			sends.Store(0)
+			e := &WeKnoraCloudEmbedder{baseURL: server.URL, modelName: "local-test", modelID: "concrete-embedding", client: client}
+			g := &embeddingBudgetGuard{limit: limit}
+			_, err := e.BatchEmbed(embeddingBudgetContext(g), []string{"local test"})
+			if limit == 0 && !errors.Is(err, requestbudget.ErrStopped) {
+				t.Fatalf("guard not applied: %v", err)
+			}
+			if limit == 1 && err != nil {
+				t.Fatal(err)
+			}
+			if int(sends.Load()) != limit {
+				t.Fatalf("sends=%d want=%d", sends.Load(), limit)
+			}
+			assertEmbeddingBudget(t, g, limit)
+		})
+	}
+}
+
+func TestRequestBudgetEmbeddingPoolCountsEachBatch(t *testing.T) {
+	t.Setenv("SSRF_WHITELIST", "127.0.0.1")
+	t.Setenv("BATCH_EMBED_SIZE", "2")
+	for _, limit := range []int{0, 2, 4} {
+		t.Run(fmt.Sprint(limit), func(t *testing.T) {
+			var sends atomic.Int32
+			server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+				sends.Add(1)
+				var body struct {
+					Input []string `json:"input"`
+				}
+				if err := json.NewDecoder(r.Body).Decode(&body); err != nil {
+					http.Error(w, "bad request", http.StatusBadRequest)
+					return
+				}
+				data := make([]map[string]any, len(body.Input))
+				for i := range body.Input {
+					data[i] = map[string]any{"index": i, "embedding": []float32{0.1}}
+				}
+				w.Header().Set("Content-Type", "application/json")
+				json.NewEncoder(w).Encode(map[string]any{"data": data})
+			}))
+			defer server.Close()
+			pool, err := ants.NewPool(3)
+			if err != nil {
+				t.Fatal(err)
+			}
+			defer pool.Release()
+			e, err := NewEmbedder(Config{Source: types.ModelSourceRemote, Provider: "openai", BaseURL: server.URL, APIKey: "local-only", ModelID: "concrete-embedding", ModelName: "local-test"}, NewBatchEmbedder(pool), nil)
+			if err != nil {
+				t.Fatal(err)
+			}
+			g := &embeddingBudgetGuard{limit: limit}
+			ctx, cancel := context.WithTimeout(embeddingBudgetContext(g), 3*time.Second)
+			defer cancel()
+			vectors, err := e.BatchEmbedWithPool(ctx, e, []string{"a", "b", "c", "d", "e", "f", "g"})
+			if limit < 4 && !errors.Is(err, requestbudget.ErrStopped) {
+				t.Fatalf("batch escaped budget: %v", err)
+			}
+			if limit == 4 && (err != nil || len(vectors) != 7) {
+				t.Fatalf("batch failed: %d %v", len(vectors), err)
+			}
+			if int(sends.Load()) != limit {
+				t.Fatalf("batch sends=%d want=%d", sends.Load(), limit)
+			}
+			assertEmbeddingBudget(t, g, limit)
+		})
+	}
+}
+
+func TestRequestBudgetEmbeddingRetriesConsumeReservations(t *testing.T) {
+	t.Setenv("SSRF_WHITELIST", "127.0.0.1")
+	for _, limit := range []int{1, 2} {
+		t.Run(fmt.Sprint(limit), func(t *testing.T) {
+			var sends atomic.Int32
+			server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+				if sends.Add(1) == 1 {
+					connection, _, err := w.(http.Hijacker).Hijack()
+					if err == nil {
+						connection.Close()
+					}
+					return
+				}
+				w.Header().Set("Content-Type", "application/json")
+				fmt.Fprint(w, `{"data":[{"index":0,"embedding":[0.1]}]}`)
+			}))
+			defer server.Close()
+			e, err := NewOpenAIEmbedder("local-only", server.URL, "local-test", 0, 0, "concrete-embedding", nil)
+			if err != nil {
+				t.Fatal(err)
+			}
+			e.maxRetries = 1
+			g := &embeddingBudgetGuard{limit: limit}
+			_, err = e.BatchEmbed(embeddingBudgetContext(g), []string{"test"})
+			if limit == 1 && !errors.Is(err, requestbudget.ErrStopped) {
+				t.Fatalf("retry escaped budget: %v", err)
+			}
+			if limit == 2 && err != nil {
+				t.Fatal(err)
+			}
+			if int(sends.Load()) != limit {
+				t.Fatalf("retry sends=%d want=%d", sends.Load(), limit)
+			}
+			assertEmbeddingBudget(t, g, limit)
+		})
+	}
+}
+
+func TestRequestBudgetVolcengineCountsEachText(t *testing.T) {
+	t.Setenv("SSRF_WHITELIST", "127.0.0.1")
+	var sends atomic.Int32
+	server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+		sends.Add(1)
+		w.Header().Set("Content-Type", "application/json")
+		fmt.Fprint(w, `{"data":{"embedding":[0.1]}}`)
+	}))
+	defer server.Close()
+	e := &VolcengineEmbedder{baseURL: server.URL, modelID: "concrete-embedding", modelName: "local-test", httpClient: newEmbeddingHTTPClient(time.Second)}
+	g := &embeddingBudgetGuard{limit: 2}
+	_, err := e.BatchEmbed(embeddingBudgetContext(g), []string{"a", "b", "c"})
+	if !errors.Is(err, requestbudget.ErrStopped) || sends.Load() != 2 {
+		t.Fatalf("per-text budget not enforced: sends=%d err=%v", sends.Load(), err)
+	}
+	assertEmbeddingBudget(t, g, 2)
+}
+
+func TestRequestBudgetEmbeddingFactoryFailsClosed(t *testing.T) {
+	g := &embeddingBudgetGuard{limit: 1}
+	ctx := embeddingBudgetContext(g)
+	e, err := NewEmbedder(Config{Source: types.ModelSourceLocal, ModelID: "local-model"}, nil, nil)
+	if err != nil {
+		t.Fatal(err)
+	}
+	if _, err := e.Embed(ctx, "test"); !errors.Is(err, requestbudget.ErrStopped) {
+		t.Fatalf("local embed invoked: %v", err)
+	}
+	if _, err := e.BatchEmbed(ctx, []string{"test"}); !errors.Is(err, requestbudget.ErrStopped) {
+		t.Fatalf("local batch invoked: %v", err)
+	}
+	if _, err := e.BatchEmbedWithPool(ctx, e, []string{"test"}); !errors.Is(err, requestbudget.ErrStopped) {
+		t.Fatalf("local pool invoked: %v", err)
+	}
+	assertEmbeddingBudget(t, g, 0)
+	fake := newFakeEmbedder("unknown")
+	close(fake.release)
+	unknown := wrapEmbeddingRequestBudget(fake)
+	if _, err := unknown.Embed(ctx, "test"); !errors.Is(err, requestbudget.ErrStopped) {
+		t.Fatalf("unknown embed adapter allowed: %v", err)
+	}
+	if _, err := unknown.BatchEmbed(ctx, []string{"test"}); !errors.Is(err, requestbudget.ErrStopped) {
+		t.Fatalf("unknown batch adapter allowed: %v", err)
+	}
+	if _, err := unknown.BatchEmbedWithPool(ctx, unknown, []string{"test"}); !errors.Is(err, requestbudget.ErrStopped) {
+		t.Fatalf("unknown pool adapter allowed: %v", err)
+	}
+	if len(fake.enter) != 0 {
+		t.Fatal("unsupported adapter was invoked")
+	}
+}