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:
1 parent
01ad8c9207
commit
cacaa35c3f
11 files changed
+350
-27
No files matched your search
@@ -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.
|
||||||
|
|||||||
@@ -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")
|
||||||
|
|||||||
@@ -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
@@ -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.
|
||||||
|
|
||||||
|
|||||||
@@ -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.
|
||||||
|
|||||||
@@ -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) {
|
||||||
|
|||||||
@@ -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
@@ -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 {
|
||||||
|
|||||||
@@ -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
|
||||||
|
}
|
||||||
@@ -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")
|
||||||
|
}
|
||||||
|
}
|
||||||
@@ -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)
|
||||||
|
}
|
||||||
|
}
|
||||||
Reference in new issue
Block a user