xinmi/internal/models/embedding/request_budget_test.go

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")
	}
}