53 lines
1.8 KiB
Go
53 lines
1.8 KiB
Go
|
package embeddingprovider
|
||
|
|
|
||
|
|
import (
|
||
|
|
"context"
|
||
|
|
"encoding/json"
|
||
|
|
"net/http"
|
||
|
|
"net/http/httptest"
|
||
|
|
"testing"
|
||
|
|
"time"
|
||
|
|
)
|
||
|
|
|
||
|
|
func TestArkEmbeddingRequestAndResponse(t *testing.T) {
|
||
|
|
vector := make([]float64, Dimensions)
|
||
|
|
vector[0] = 0.5
|
||
|
|
server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||
|
|
if r.URL.Path != "/api/v3/embeddings" || r.Header.Get("Authorization") != "Bearer test-key" {
|
||
|
|
t.Fatalf("unexpected embedding request")
|
||
|
|
}
|
||
|
|
var body struct {
|
||
|
|
Model string `json:"model"`
|
||
|
|
Input string `json:"input"`
|
||
|
|
Dimensions int `json:"dimensions"`
|
||
|
|
}
|
||
|
|
if err := json.NewDecoder(r.Body).Decode(&body); err != nil || body.Model != "embedding-test" || body.Input != "河边小学" || body.Dimensions != Dimensions {
|
||
|
|
t.Fatalf("body=%+v err=%v", body, err)
|
||
|
|
}
|
||
|
|
_ = json.NewEncoder(w).Encode(map[string]any{"data": []any{map[string]any{"index": 0, "embedding": vector}}})
|
||
|
|
}))
|
||
|
|
defer server.Close()
|
||
|
|
provider, err := NewArk(server.URL+"/api/v3", "test-key", "embedding-test", time.Second)
|
||
|
|
if err != nil {
|
||
|
|
t.Fatal(err)
|
||
|
|
}
|
||
|
|
result, err := provider.Embed(context.Background(), "河边小学")
|
||
|
|
if err != nil || len(result.Vector) != Dimensions || result.ProviderRef != "ark:embedding-test" || result.Vector[0] != 0.5 {
|
||
|
|
t.Fatalf("result=%+v err=%v", result, err)
|
||
|
|
}
|
||
|
|
}
|
||
|
|
|
||
|
|
func TestArkEmbeddingRejectsWrongDimensions(t *testing.T) {
|
||
|
|
server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) {
|
||
|
|
_ = json.NewEncoder(w).Encode(map[string]any{"data": []any{map[string]any{"index": 0, "embedding": []float64{1, 2}}}})
|
||
|
|
}))
|
||
|
|
defer server.Close()
|
||
|
|
provider, err := NewArk(server.URL, "test-key", "embedding-test", time.Second)
|
||
|
|
if err != nil {
|
||
|
|
t.Fatal(err)
|
||
|
|
}
|
||
|
|
if _, err = provider.Embed(context.Background(), "测试"); err == nil {
|
||
|
|
t.Fatal("wrong vector dimensions were accepted")
|
||
|
|
}
|
||
|
|
}
|