258 lines
9.4 KiB
Go
258 lines
9.4 KiB
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")
|
||
|
|
}
|
||
|
|
}
|