llmgw/internal/router/kb_embed.go

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
}