Co-Authored-By: Claude Opus 5.5 (1M context) <noreply@anthropic.com>
48 lines
1.6 KiB
Go
48 lines
1.6 KiB
Go
package router
|
|
|
|
import (
|
|
"net/http"
|
|
"net/http/httptest"
|
|
"strings"
|
|
"testing"
|
|
|
|
"gitea.mynodes.xyz/ilker/llmgw/internal/store"
|
|
)
|
|
|
|
// A document's text is embedded through the gateway's own API as the
|
|
// caller, in batches, each vector in its chunk's place.
|
|
func TestTheGatewayEmbedsADocumentSentAsText(t *testing.T) {
|
|
e := newFakeEngine(t, "qwen2.5-14b-instruct")
|
|
rt, secret := pricedRouter(t, e.srv.URL)
|
|
_ = rt.Handler() // mounts the API the embedding goes through
|
|
texts := make([]string, kbEmbedBatch*2+5)
|
|
for i := range texts {
|
|
texts[i] = strings.Repeat("söz ", i+1)
|
|
}
|
|
req := httptest.NewRequestWithContext(t.Context(), http.MethodPut, "/v1/kb/collections/x/documents", nil)
|
|
req.Header.Set("Authorization", "Bearer "+secret)
|
|
c := store.KBCollection{Name: "x", EmbedModel: "qwen2.5-14b-instruct", Dims: 3}
|
|
vecs, err := rt.kbEmbed(req, c, texts)
|
|
if err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
for i, v := range vecs {
|
|
// The fake engine's vector is the text's bytes, its words, and 1.
|
|
if len(v) != 3 || int(v[1]) != i+1 || int(v[0]) != len(texts[i]) {
|
|
t.Fatalf("chunk %d got %v", i, v)
|
|
}
|
|
}
|
|
|
|
// A model of other dimensions is refused, saying so.
|
|
c.Dims = 4
|
|
if _, err := rt.kbEmbed(req, c, texts[:2]); err == nil || !strings.Contains(err.Error(), "dimensions") {
|
|
t.Errorf("4 dimensions: %v", err)
|
|
}
|
|
// Without the caller's key, nothing is embedded.
|
|
bare := httptest.NewRequestWithContext(t.Context(), http.MethodPut, "/v1/kb/collections/x/documents", nil)
|
|
c.Dims = 3
|
|
if _, err := rt.kbEmbed(bare, c, texts[:1]); err == nil {
|
|
t.Error("embedded without a key")
|
|
}
|
|
}
|