fix(embedding): validate WeKnoraCloud response indexes

Reject out-of-range and duplicate response indexes before mapping them to batch inputs. Return an error when a requested input has no response item instead of silently passing an incomplete batch downstream.

Add focused coverage for malformed responses and valid out-of-order results.

Signed-off-by: PP1 <74917296+pengpengyi92@users.noreply.github.com>
This commit is contained in:
PP1
2026-08-11 16:59:48 +08:00
committed by lyingbug
parent 9bf3dee28a
commit 06eeb95b4e
2 changed files with 117 additions and 2 deletions
+17 -2
View File
@@ -130,9 +130,24 @@ func (e *WeKnoraCloudEmbedder) BatchEmbed(ctx context.Context, texts []string) (
}
result := make([][]float32, len(texts))
seen := make([]bool, len(texts))
for _, item := range embedResp.Data {
if item.Index < len(result) {
result[item.Index] = item.Embedding
if item.Index < 0 || item.Index >= len(result) {
return nil, fmt.Errorf(
"weknoracloud embedder: response index %d out of range for %d inputs",
item.Index,
len(texts),
)
}
if seen[item.Index] {
return nil, fmt.Errorf("weknoracloud embedder: duplicate response index %d", item.Index)
}
result[item.Index] = item.Embedding
seen[item.Index] = true
}
for index, found := range seen {
if !found {
return nil, fmt.Errorf("weknoracloud embedder: missing embedding for input index %d", index)
}
}
return result, nil
@@ -0,0 +1,100 @@
package embedding
import (
"context"
"fmt"
"net/http"
"net/http/httptest"
"reflect"
"strings"
"testing"
)
func newWeKnoraCloudEmbedderTestServer(t *testing.T, response string) *WeKnoraCloudEmbedder {
t.Helper()
server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
if r.URL.Path != weKnoraCloudEmbedPath {
http.Error(w, fmt.Sprintf("request path = %q, want %q", r.URL.Path, weKnoraCloudEmbedPath), http.StatusNotFound)
return
}
w.Header().Set("Content-Type", "application/json")
_, _ = fmt.Fprint(w, response)
}))
t.Cleanup(server.Close)
return &WeKnoraCloudEmbedder{
modelName: "test-embedding",
modelID: "test-model-id",
appID: "test-app-id",
apiKey: "test-api-key",
baseURL: server.URL,
client: server.Client(),
}
}
func TestWeKnoraCloudBatchEmbedPreservesInputOrder(t *testing.T) {
embedder := newWeKnoraCloudEmbedderTestServer(t, `{
"data": [
{"index": 1, "embedding": [0.3, 0.4]},
{"index": 0, "embedding": [0.1, 0.2]}
]
}`)
got, err := embedder.BatchEmbed(context.Background(), []string{"first", "second"})
if err != nil {
t.Fatalf("BatchEmbed returned error: %v", err)
}
want := [][]float32{{0.1, 0.2}, {0.3, 0.4}}
if !reflect.DeepEqual(got, want) {
t.Fatalf("BatchEmbed result = %v, want %v", got, want)
}
}
func TestWeKnoraCloudBatchEmbedRejectsMalformedResponse(t *testing.T) {
tests := []struct {
name string
response string
texts []string
wantErrMsg string
}{
{
name: "negative index",
response: `{"data": [{"index": -1, "embedding": [0.1, 0.2]}]}`,
texts: []string{"first"},
wantErrMsg: "response index -1 out of range",
},
{
name: "index above input range",
response: `{"data": [{"index": 1, "embedding": [0.1, 0.2]}]}`,
texts: []string{"first"},
wantErrMsg: "response index 1 out of range",
},
{
name: "duplicate index",
response: `{"data": [{"index": 0, "embedding": [0.1]}, {"index": 0, "embedding": [0.2]}]}`,
texts: []string{"first"},
wantErrMsg: "duplicate response index 0",
},
{
name: "missing result",
response: `{"data": [{"index": 0, "embedding": [0.1, 0.2]}]}`,
texts: []string{"first", "second"},
wantErrMsg: "missing embedding for input index 1",
},
}
for _, tt := range tests {
t.Run(tt.name, func(t *testing.T) {
embedder := newWeKnoraCloudEmbedderTestServer(t, tt.response)
_, err := embedder.BatchEmbed(context.Background(), tt.texts)
if err == nil {
t.Fatalf("BatchEmbed returned nil error, want error containing %q", tt.wantErrMsg)
}
if !strings.Contains(err.Error(), tt.wantErrMsg) {
t.Fatalf("BatchEmbed error = %q, want it to contain %q", err, tt.wantErrMsg)
}
})
}
}