Files
llmgw/internal/router/kb_embed_test.go
T
2026-10-29 21:56:05 +01:00

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