The gateway embeds the documents llmgw-kb sends as text

Co-Authored-By: Claude Opus 5.5 (1M context) <noreply@anthropic.com>
This commit is contained in:
Ilker ManapandClaude Opus 5.5 committed 2026-10-29 21:56:05 +01:00
1 parent 01ad8c9207
commit cacaa35c3f
11 files changed
+350 -27

No files matched your search

+10
View File
@@ -25,6 +25,16 @@ is only released once its **exit criterion** is met, not when its code lands.
Until v1.0.0 the minor version is the breaking one: configuration formats and Go Until v1.0.0 the minor version is the breaking one: configuration formats and Go
APIs change between minors without a deprecation period. APIs change between minors without a deprecation period.
## Unreleased
### Added
- **The gateway embeds the documents `llmgw-kb` sends.** A document may be sent as its passages' text alone; the
gateway embeds them through its own `/v1/embeddings` as the caller, 32 at a time and four batches at once,
waiting out a busy engine rather than failing the document. About a kilobyte a passage leaves the machine instead
of a vector down and up again, some thirty times more. `llmgw-kb` does so with a gateway that says it can
(`server_embeds` in a collection's answer); `--embed-here` keeps the old way.
## v0.15.3 — 2026-10-29 ## v0.15.3 — 2026-10-29
A model stopped and kept, or removed with its files; the second mirror fixed where it was tried. A model stopped and kept, or removed with its files; the second mirror fixed where it was tried.
+2
View File
@@ -331,6 +331,8 @@ hidden files, version control folders and Office lock files are not.`,
f.BoolVar(&o.Delete, "delete", false, "remove documents this machine sent whose files are gone") f.BoolVar(&o.Delete, "delete", false, "remove documents this machine sent whose files are gone")
f.BoolVar(&o.Retry, "retry", false, "read again files skipped before") f.BoolVar(&o.Retry, "retry", false, "read again files skipped before")
f.IntVar(&o.Workers, "workers", 4, "files read and sent at once") f.IntVar(&o.Workers, "workers", 4, "files read and sent at once")
f.BoolVar(&o.EmbedHere, "embed-here", false,
"embed on this machine and send the vectors, even to a gateway that embeds text it is sent")
f.StringVar(&ocr, "ocr-languages", "tur+eng", "Tesseract's languages for images and scanned pages") f.StringVar(&ocr, "ocr-languages", "tur+eng", "Tesseract's languages for images and scanned pages")
f.Int64Var(&maxMiB, "max-size", 256, "skip files larger than this many MiB") f.Int64Var(&maxMiB, "max-size", 256, "skip files larger than this many MiB")
f.BoolVarP(&verbose, "verbose", "v", false, "list unchanged files too") f.BoolVarP(&verbose, "verbose", "v", false, "list unchanged files too")
+4
View File
@@ -2890,6 +2890,10 @@
"writable": { "writable": {
"type": "boolean", "type": "boolean",
"description": "Writable is whether this caller may add and remove documents." "description": "Writable is whether this caller may add and remove documents."
},
"server_embeds": {
"type": "boolean",
"description": "ServerEmbeds says a document may be sent as text alone: the gateway embeds its chunks, as the caller, with the collection's model."
} }
} }
}, },
+16 -14
View File
@@ -3,9 +3,9 @@
A knowledge base is a collection of a company's documents that the cluster A knowledge base is a collection of a company's documents that the cluster
can answer from. People put their documents into it from their own machines can answer from. People put their documents into it from their own machines
with `llmgw-kb`. The tool reads the files on that machine, cuts them into with `llmgw-kb`. The tool reads the files on that machine, cuts them into
passages ("chunks"), has the gateway embed each passage with the person's own passages ("chunks") and sends their text; the gateway embeds each passage
key, and sends the text and the vectors. The files themselves never leave with the person's own key and keeps the text and the vectors. The files
the machine. themselves never leave the machine.
## Why retrieval, and not training ## Why retrieval, and not training
@@ -159,10 +159,12 @@ llmgw ctl config set router.admission.bands '[{name: gold, priority: 300, satura
llmgw keys create import --project Acme/docs --band bulk llmgw keys create import --project Acme/docs --band bulk
``` ```
The people's questions then go first, and the import waits behind them. A The people's questions then go first, and the import waits behind them. The
request that waits as long as its band allows is answered 504 with gateway embeds a document's passages 32 at a time, four batches at once,
`Retry-After`, and `llmgw-kb` sends it again rather than counting the enough to fill an embedding replica's slots; a batch told to wait (`Retry-After`)
passage as failed. It gives up after 15 minutes. waits up to five minutes before the document is answered the same, and
`llmgw-kb` sends it again rather than counting it as failed. It gives up
after 15 minutes.
## For everyone: `llmgw-kb` ## For everyone: `llmgw-kb`
@@ -225,6 +227,7 @@ not changed.
| `--include`, `--exclude` | Glob patterns. `*.pdf` matches a file name; `reports/**/*.xlsx` matches a path. | | `--include`, `--exclude` | Glob patterns. `*.pdf` matches a file name; `reports/**/*.xlsx` matches a path. |
| `--hidden` | Reads hidden files and folders too (left out by default). | | `--hidden` | Reads hidden files and folders too (left out by default). |
| `-v` | Lists unchanged files too. | | `-v` | Lists unchanged files too. |
| `--embed-here` | Embeds the passages from this machine and sends the vectors. With a gateway that does not embed them itself it does so anyway. |
Documents are named by their path under the folder, with forward slashes, Documents are named by their path under the folder, with forward slashes,
on every system. With several folders, each folder's own name comes first. on every system. With several folders, each folder's own name comes first.
@@ -259,13 +262,12 @@ scanned page with no OCR.
### What leaves your machine ### What leaves your machine
For each new or changed document, two things leave the machine, both over For each new or changed document, one request leaves the machine, over the
the gateway's TLS connection: gateway's TLS connection: its name, its file's SHA-256 digest and size, and
its passages' text. The gateway embeds the passages through `/v1/embeddings`
- **The passages, to be embedded.** They go through `/v1/embeddings` like as you, so your key's limits, budget and guardrails apply as if you had
any request of yours. Your key's limits, budget and guardrails apply. asked. That is about a kilobyte a passage; embedding here (`--embed-here`)
- **The document itself.** Its name, its file's SHA-256 digest and size, brings each vector down and sends it up again, some thirty times more.
and the passages with their vectors.
The file itself is never sent. `--dry-run` shows exactly what would be. The file itself is never sent. `--dry-run` shows exactly what would be.
+5 -2
View File
@@ -63,6 +63,8 @@ type Collection struct {
EmbedModel string `json:"embed_model"` EmbedModel string `json:"embed_model"`
Dims int `json:"dims"` Dims int `json:"dims"`
Writable bool `json:"writable"` Writable bool `json:"writable"`
// ServerEmbeds is a gateway that embeds a document sent as text alone.
ServerEmbeds bool `json:"server_embeds"`
} }
// RemoteDoc is a document the collection holds. // RemoteDoc is a document the collection holds.
@@ -82,8 +84,9 @@ type Manifest struct {
// ChunkOut is one chunk as sent. // ChunkOut is one chunk as sent.
type ChunkOut struct { type ChunkOut struct {
Text string `json:"text"` Text string `json:"text"`
Embedding []float32 `json:"embedding"` // Embedding is empty when the gateway embeds the chunk.
Embedding []float32 `json:"embedding,omitempty"`
} }
// DocOut is one document as sent. // DocOut is one document as sent.
+26 -7
View File
@@ -46,8 +46,11 @@ type Options struct {
Retry bool Retry bool
// Workers is how many files are read and sent at once; 0 is 4. // Workers is how many files are read and sent at once; 0 is 4.
Workers int Workers int
Extract extract.Options // EmbedHere embeds the chunks on this machine even when the gateway
Chunk chunk.Options // would: a vector up and down for each, rather than the text alone.
EmbedHere bool
Extract extract.Options
Chunk chunk.Options
} }
// Event is something a sync did to one document, for showing progress. // Event is something a sync did to one document, for showing progress.
@@ -359,8 +362,12 @@ func (s *syncer) send(ctx context.Context, f file, name, digest, text string) {
texts[i] = c.Text() texts[i] = c.Text()
size += int64(len(texts[i])) size += int64(len(texts[i]))
} }
// About ten bytes of JSON a dimension. // Sent as text alone when the gateway embeds it; else about ten bytes
size += int64(len(cs) * s.coll.Dims * 10) // of JSON a dimension with each chunk.
serverEmbeds := s.coll.ServerEmbeds && !s.o.EmbedHere
if !serverEmbeds {
size += int64(len(cs) * s.coll.Dims * 10)
}
if size > maxDocJSON { if size > maxDocJSON {
s.fail(ctx, f, name, digest, fmt.Errorf("%d MiB to send is over the gateway's limit for one document; split the file", size>>20)) s.fail(ctx, f, name, digest, fmt.Errorf("%d MiB to send is over the gateway's limit for one document; split the file", size>>20))
return return
@@ -369,12 +376,19 @@ func (s *syncer) send(ctx context.Context, f file, name, digest, text string) {
s.emit(Event{Kind: "would-send", Name: name, Chunks: len(cs), Bytes: size}) s.emit(Event{Kind: "would-send", Name: name, Chunks: len(cs), Bytes: size})
return return
} }
out := DocOut{Name: name, Digest: digest, Bytes: f.size, Chunks: make([]ChunkOut, len(cs))}
if serverEmbeds {
for i := range cs {
out.Chunks[i] = ChunkOut{Text: texts[i]}
}
s.put(ctx, f, name, digest, out, len(cs), size)
return
}
vecs, err := s.c.Embed(ctx, s.coll.EmbedModel, texts) vecs, err := s.c.Embed(ctx, s.coll.EmbedModel, texts)
if err != nil { if err != nil {
s.fail(ctx, f, name, digest, fmt.Errorf("embedding with %s: %w", s.coll.EmbedModel, err)) s.fail(ctx, f, name, digest, fmt.Errorf("embedding with %s: %w", s.coll.EmbedModel, err))
return return
} }
out := DocOut{Name: name, Digest: digest, Bytes: f.size, Chunks: make([]ChunkOut, len(cs))}
for i := range cs { for i := range cs {
if len(vecs[i]) != s.coll.Dims { if len(vecs[i]) != s.coll.Dims {
s.fail(ctx, f, name, digest, fmt.Errorf("%s gave %d dimensions and the collection takes %d", s.fail(ctx, f, name, digest, fmt.Errorf("%s gave %d dimensions and the collection takes %d",
@@ -383,12 +397,17 @@ func (s *syncer) send(ctx context.Context, f file, name, digest, text string) {
} }
out.Chunks[i] = ChunkOut{Text: texts[i], Embedding: vecs[i]} out.Chunks[i] = ChunkOut{Text: texts[i], Embedding: vecs[i]}
} }
s.put(ctx, f, name, digest, out, len(cs), size)
}
// put sends a document and records it.
func (s *syncer) put(ctx context.Context, f file, name, digest string, out DocOut, chunks int, size int64) {
if _, err := s.c.Put(ctx, s.coll.ID, out); err != nil { if _, err := s.c.Put(ctx, s.coll.ID, out); err != nil {
s.fail(ctx, f, name, digest, err) s.fail(ctx, f, name, digest, err)
return return
} }
s.record(ctx, s.entry(f, name, digest, Sent, "", len(cs))) s.record(ctx, s.entry(f, name, digest, Sent, "", chunks))
s.emit(Event{Kind: "sent", Name: name, Chunks: len(cs), Bytes: size}) s.emit(Event{Kind: "sent", Name: name, Chunks: chunks, Bytes: size})
} }
func (s *syncer) fail(ctx context.Context, f file, name, digest string, err error) { func (s *syncer) fail(ctx context.Context, f file, name, digest string, err error) {
+29
View File
@@ -22,6 +22,10 @@ type fakeGateway struct {
puts int puts int
embedded int embedded int
busyFirst bool busyFirst bool
// serverEmbeds is a gateway that embeds text it is sent; vectors
// counts chunks that came with one.
serverEmbeds bool
vectors int
} }
func (g *fakeGateway) serve(t *testing.T) *Client { func (g *fakeGateway) serve(t *testing.T) *Client {
@@ -37,6 +41,7 @@ func (g *fakeGateway) serve(t *testing.T) *Client {
} }
m := Manifest{Collection: coll} m := Manifest{Collection: coll}
m.Collection.Writable = g.writable m.Collection.Writable = g.writable
m.Collection.ServerEmbeds = g.serverEmbeds
for _, d := range g.docs { for _, d := range g.docs {
m.Documents = append(m.Documents, d) m.Documents = append(m.Documents, d)
} }
@@ -48,6 +53,11 @@ func (g *fakeGateway) serve(t *testing.T) *Client {
g.mu.Lock() g.mu.Lock()
defer g.mu.Unlock() defer g.mu.Unlock()
g.puts++ g.puts++
for _, c := range d.Chunks {
if len(c.Embedding) > 0 {
g.vectors++
}
}
g.docs[d.Name] = RemoteDoc{Name: d.Name, Digest: d.Digest, Chunks: len(d.Chunks)} g.docs[d.Name] = RemoteDoc{Name: d.Name, Digest: d.Digest, Chunks: len(d.Chunks)}
_, _ = w.Write([]byte(`{"changed":true}`)) _, _ = w.Write([]byte(`{"changed":true}`))
}) })
@@ -270,3 +280,22 @@ func TestFiltersAndNames(t *testing.T) {
t.Errorf("an empty folder gave %v", err) t.Errorf("an empty folder gave %v", err)
} }
} }
func TestAGatewayThatEmbedsIsSentTextAlone(t *testing.T) {
g, c, l, dir := setup(t)
g.serverEmbeds = true
rep, k := run(t, c, l, Options{Collection: "ik", Roots: []string{dir}})
if len(k["sent"]) != 2 || rep.Failed != 0 {
t.Fatalf("%+v %v", rep, k)
}
if g.embedded != 0 || g.vectors != 0 {
t.Errorf("embedded here: %d asked, %d vectors sent", g.embedded, g.vectors)
}
// --embed-here sends the vectors all the same.
write(t, dir, "politika/yeni.md", "Yeni kural.\n")
run(t, c, l, Options{Collection: "ik", Roots: []string{dir}, EmbedHere: true})
if g.embedded == 0 || g.vectors == 0 {
t.Errorf("--embed-here: %d asked, %d vectors sent", g.embedded, g.vectors)
}
}
+44 -4
View File
@@ -53,6 +53,9 @@ type KBCollectionView struct {
Dims int `json:"dims"` Dims int `json:"dims"`
// Writable is whether this caller may add and remove documents. // Writable is whether this caller may add and remove documents.
Writable bool `json:"writable"` Writable bool `json:"writable"`
// ServerEmbeds says a document may be sent as text alone: the gateway
// embeds its chunks, as the caller, with the collection's model.
ServerEmbeds bool `json:"server_embeds"`
} }
// KBDocumentView is a document as the manifest lists it. // KBDocumentView is a document as the manifest lists it.
@@ -77,8 +80,10 @@ type KBPutDocument struct {
Digest string `json:"digest"` Digest string `json:"digest"`
Bytes int64 `json:"bytes"` Bytes int64 `json:"bytes"`
Chunks []struct { Chunks []struct {
Text string `json:"text"` Text string `json:"text"`
Embedding []float32 `json:"embedding"` // Embedding is the chunk's vector; with none on any chunk the
// gateway embeds them.
Embedding []float32 `json:"embedding,omitempty"`
} `json:"chunks"` } `json:"chunks"`
} }
@@ -124,7 +129,7 @@ func (r *Router) kbCaller(w http.ResponseWriter, req *http.Request) (auth.Princi
func (r *Router) kbView(ctx context.Context, c storepkg.KBCollection, writable bool) KBCollectionView { func (r *Router) kbView(ctx context.Context, c storepkg.KBCollection, writable bool) KBCollectionView {
return KBCollectionView{ID: c.ID, Name: c.Name, Owner: string(c.OwnerScope) + ":" + c.OwnerID, return KBCollectionView{ID: c.ID, Name: c.Name, Owner: string(c.OwnerScope) + ":" + c.OwnerID,
OwnerName: r.store.OwnerName(ctx, c.OwnerScope, c.OwnerID), EmbedModel: c.EmbedModel, Dims: c.Dims, Writable: writable} OwnerName: r.store.OwnerName(ctx, c.OwnerScope, c.OwnerID), EmbedModel: c.EmbedModel, Dims: c.Dims, Writable: writable, ServerEmbeds: true}
} }
// kbCollection finds a collection by id or name among those the caller may // kbCollection finds a collection by id or name among those the caller may
@@ -281,6 +286,19 @@ func (r *Router) handleKBPut(w http.ResponseWriter, req *http.Request) {
fmt.Sprintf("%d chunks: a document has at most %d", len(in.Chunks), kbMaxChunks)) fmt.Sprintf("%d chunks: a document has at most %d", len(in.Chunks), kbMaxChunks))
return return
} }
// Text alone: the gateway embeds it, as the caller, so their limits and
// usage count it as if they had asked.
embedded := 0
for _, ch := range in.Chunks {
if len(ch.Embedding) > 0 {
embedded++
}
}
if embedded > 0 && embedded < len(in.Chunks) {
writeError(w, http.StatusBadRequest, "invalid_request_error",
"some chunks have an embedding and some not: send every chunk's, or none for the gateway to embed")
return
}
chunks := make([]storepkg.KBChunk, len(in.Chunks)) chunks := make([]storepkg.KBChunk, len(in.Chunks))
for i, ch := range in.Chunks { for i, ch := range in.Chunks {
if !utf8.ValidString(ch.Text) || strings.TrimSpace(ch.Text) == "" { if !utf8.ValidString(ch.Text) || strings.TrimSpace(ch.Text) == "" {
@@ -292,7 +310,7 @@ func (r *Router) handleKBPut(w http.ResponseWriter, req *http.Request) {
fmt.Sprintf("chunk %d is over %d characters", i, kbMaxChunkRunes)) fmt.Sprintf("chunk %d is over %d characters", i, kbMaxChunkRunes))
return return
} }
if len(ch.Embedding) != c.Dims { if embedded > 0 && len(ch.Embedding) != c.Dims {
writeError(w, http.StatusBadRequest, "invalid_request_error", writeError(w, http.StatusBadRequest, "invalid_request_error",
fmt.Sprintf("chunk %d has %d dimensions; collection %s is embedded by %s in %d", i, fmt.Sprintf("chunk %d has %d dimensions; collection %s is embedded by %s in %d", i,
len(ch.Embedding), c.Name, c.EmbedModel, c.Dims)) len(ch.Embedding), c.Name, c.EmbedModel, c.Dims))
@@ -300,6 +318,28 @@ func (r *Router) handleKBPut(w http.ResponseWriter, req *http.Request) {
} }
chunks[i] = storepkg.KBChunk{Body: strings.ToValidUTF8(ch.Text, ""), Embedding: ch.Embedding} chunks[i] = storepkg.KBChunk{Body: strings.ToValidUTF8(ch.Text, ""), Embedding: ch.Embedding}
} }
if embedded == 0 {
texts := make([]string, len(chunks))
for i := range chunks {
texts[i] = chunks[i].Body
}
vecs, err := r.kbEmbed(req, c, texts)
if err != nil {
var ee *kbEmbedError
if errors.As(err, &ee) {
if ee.retryAfter != "" {
w.Header().Set("Retry-After", ee.retryAfter)
}
writeError(w, ee.status, "invalid_request_error", err.Error())
return
}
writeError(w, http.StatusBadGateway, "server_error", err.Error())
return
}
for i := range chunks {
chunks[i].Embedding = vecs[i]
}
}
d, changed, err := r.store.PutDocument(req.Context(), c, storepkg.KBDocument{Name: name, Digest: in.Digest, d, changed, err := r.store.PutDocument(req.Context(), c, storepkg.KBDocument{Name: name, Digest: in.Digest,
Bytes: in.Bytes, AddedBy: kbActor(p)}, chunks) Bytes: in.Bytes, AddedBy: kbActor(p)}, chunks)
if err != nil { if err != nil {
+134
View File
@@ -0,0 +1,134 @@
package router
import (
"context"
"encoding/json"
"fmt"
"net/http"
"strconv"
"sync"
"time"
storepkg "gitea.mynodes.xyz/ilker/llmgw/internal/store"
)
// A document sent as text alone is embedded here, through the gateway's own
// /v1/embeddings as the caller: their key, limits and usage, as if they had
// asked. The client then sends a kilobyte a chunk instead of the vector's
// twenty in each direction.
const (
// kbEmbedBatch is how many chunks go in one embeddings request, and
// kbEmbedAtOnce how many such requests are in flight for one document:
// enough to fill an embedding replica's slots.
kbEmbedBatch = 32
kbEmbedAtOnce = 4
// kbEmbedWait is how long one batch waits out answers that say to try
// again, before the document is refused with the same advice.
kbEmbedWait = 5 * time.Minute
)
// kbEmbedError is an embeddings answer that was not a success, passed on as
// the API gave it.
type kbEmbedError struct {
status int
retryAfter string
msg string
}
func (e *kbEmbedError) Error() string { return e.msg }
// kbEmbed embeds texts with the collection's model, as the request's caller.
func (r *Router) kbEmbed(req *http.Request, c storepkg.KBCollection, texts []string) ([][]float32, error) {
ctx, cancel := context.WithCancel(req.Context())
defer cancel()
out := make([][]float32, len(texts))
var (
mu sync.Mutex
first error
wg sync.WaitGroup
)
slots := make(chan struct{}, kbEmbedAtOnce)
for start := 0; start < len(texts); start += kbEmbedBatch {
end := min(start+kbEmbedBatch, len(texts))
wg.Add(1)
slots <- struct{}{}
go func() {
defer wg.Done()
defer func() { <-slots }()
vecs, err := r.kbEmbedBatch(ctx, req, c, texts[start:end])
mu.Lock()
defer mu.Unlock()
if err != nil {
if first == nil {
first = err
cancel()
}
return
}
copy(out[start:end], vecs)
}()
}
wg.Wait()
if first != nil {
return nil, first
}
return out, nil
}
// kbEmbedBatch embeds one batch, waiting out answers that say when to try
// again.
func (r *Router) kbEmbedBatch(ctx context.Context, req *http.Request, c storepkg.KBCollection, texts []string) ([][]float32, error) {
call := mediaCall{req: req.WithContext(ctx)}
deadline := time.Now().Add(kbEmbedWait)
for {
status, h, b, err := r.self(call, http.MethodPost, "/v1/embeddings",
map[string]any{"model": c.EmbedModel, "input": texts, "encoding_format": "float"})
if err != nil {
return nil, err
}
if status < 300 {
return kbVectors(b, len(texts), c)
}
retry := h.Get("Retry-After")
if retry == "" || (status != http.StatusTooManyRequests && status < 500) || time.Now().After(deadline) {
return nil, &kbEmbedError{status: status, retryAfter: retry,
msg: fmt.Sprintf("embedding with %s: %s", c.EmbedModel, errorMessage(b))}
}
secs, _ := strconv.Atoi(retry)
select {
case <-ctx.Done():
return nil, ctx.Err()
case <-time.After(time.Duration(max(secs, 1)) * time.Second):
}
}
}
// kbVectors reads an embeddings answer: one vector a text, of the
// collection's dimensions.
func kbVectors(b []byte, n int, c storepkg.KBCollection) ([][]float32, error) {
var res struct {
Data []struct {
Index int `json:"index"`
Embedding []float32 `json:"embedding"`
} `json:"data"`
}
if err := json.Unmarshal(b, &res); err != nil {
return nil, fmt.Errorf("embedding with %s: the answer is not JSON: %w", c.EmbedModel, err)
}
if len(res.Data) != n {
return nil, fmt.Errorf("embedding with %s: %d of %d texts came back", c.EmbedModel, len(res.Data), n)
}
out := make([][]float32, n)
for _, d := range res.Data {
if d.Index < 0 || d.Index >= n || out[d.Index] != nil {
return nil, fmt.Errorf("embedding with %s: an embedding came back with index %d", c.EmbedModel, d.Index)
}
if len(d.Embedding) != c.Dims {
return nil, &kbEmbedError{status: http.StatusBadRequest,
msg: fmt.Sprintf("%s gave %d dimensions and collection %s takes %d", c.EmbedModel, len(d.Embedding), c.Name, c.Dims)}
}
out[d.Index] = d.Embedding
}
return out, nil
}
+47
View File
@@ -0,0 +1,47 @@
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")
}
}
+33
View File
@@ -206,3 +206,36 @@ func TestAPersonalKeyReadsOnlyInTheProjectItNames(t *testing.T) {
t.Errorf("a project the person is not in: %d", code) t.Errorf("a project the person is not in: %d", code)
} }
} }
func TestADocumentSentAsTextIsEmbeddedByTheGateway(t *testing.T) {
rt := kbRouter(t)
h := rt.Handler()
ctx := t.Context()
s := rt.store
sfx := strings.ToLower(store.NewID(""))
org, _ := s.CreateOrg(ctx, "e-"+sfx)
team, _ := s.CreateTeam(ctx, org, "t")
prj, _ := s.CreateProject(ctx, team, "p")
// The fake engine embeds in three dimensions.
if _, err := s.CreateCollection(ctx, store.KBCollection{OwnerScope: store.ScopeProject, OwnerID: prj,
Name: "metin", EmbedModel: "qwen2.5-14b-instruct", Dims: 3}); err != nil {
t.Fatal(err)
}
key := kbKey(t, s, prj, "")
base := "/v1/kb/collections/metin/"
if code, _, body := call(t, h, key, http.MethodGet, base+"manifest", nil); !strings.Contains(body, `"server_embeds":true`) {
t.Errorf("manifest: %d %s", code, body)
}
textOnly := map[string]any{"name": "a.md", "digest": "d1", "bytes": 10, "chunks": []map[string]any{
{"text": "birinci parça"}, {"text": "ikinci parça burada"}}}
code, out, body := call(t, h, key, http.MethodPut, base+"documents", textOnly)
if code != http.StatusOK || out["chunks"] != float64(2) {
t.Fatalf("text only: %d %s", code, body)
}
// Some with a vector and some without is refused.
mixed := map[string]any{"name": "b.md", "digest": "d2", "bytes": 10, "chunks": []map[string]any{
{"text": "bir", "embedding": []float32{1, 0, 0}}, {"text": "iki"}}}
if code, _, body := call(t, h, key, http.MethodPut, base+"documents", mixed); code != http.StatusBadRequest {
t.Errorf("mixed: %d %s", code, body)
}
}