llmgw/internal/kb/kbsync/api.go

package kbsync

import (
	"bytes"
	"context"
	"crypto/tls"
	"crypto/x509"
	"encoding/json"
	"errors"
	"fmt"
	"io"
	"net/http"
	"net/url"
	"os"
	"strconv"
	"strings"
	"time"
)

// Client speaks to a gateway's knowledge-base and embeddings endpoints with
// a person's key.
type Client struct {
	// Base is the gateway's address, "https://llm.example.com".
	Base string
	Key  string
	// Project is the project a personal key acts in.
	Project string
	HTTP    *http.Client
	// Wait is how long a request waits out a model loading or a busy
	// gateway (429/503 with Retry-After) before giving up; 0 is 15 minutes.
	Wait time.Duration
	// Waiting, when set, is told each time a request waits, and why.
	Waiting func(why string, d time.Duration)
}

// NewHTTP is an HTTP client trusting the system's certificates, and those in
// caFile when one is given.
func NewHTTP(caFile string) (*http.Client, error) {
	tr := http.DefaultTransport.(*http.Transport).Clone()
	if caFile != "" {
		pem, err := os.ReadFile(caFile)
		if err != nil {
			return nil, err
		}
		pool, err := x509.SystemCertPool()
		if err != nil || pool == nil {
			pool = x509.NewCertPool()
		}
		if !pool.AppendCertsFromPEM(pem) {
			return nil, fmt.Errorf("no certificate in %s", caFile)
		}
		tr.TLSClientConfig = &tls.Config{RootCAs: pool, MinVersion: tls.VersionTLS12}
	}
	return &http.Client{Transport: tr, Timeout: 10 * time.Minute}, nil
}

// Collection is a collection as the gateway shows it to this key.
type Collection struct {
	ID         string `json:"id"`
	Name       string `json:"name"`
	Owner      string `json:"owner"`
	EmbedModel string `json:"embed_model"`
	Dims       int    `json:"dims"`
	Writable   bool   `json:"writable"`
}

// RemoteDoc is a document the collection holds.
type RemoteDoc struct {
	Name    string    `json:"name"`
	Digest  string    `json:"digest"`
	Bytes   int64     `json:"bytes"`
	Chunks  int       `json:"chunks"`
	AddedAt time.Time `json:"added_at"`
}

// Manifest is a collection and its documents.
type Manifest struct {
	Collection Collection  `json:"collection"`
	Documents  []RemoteDoc `json:"documents"`
}

// ChunkOut is one chunk as sent.
type ChunkOut struct {
	Text      string    `json:"text"`
	Embedding []float32 `json:"embedding"`
}

// DocOut is one document as sent.
type DocOut struct {
	Name   string     `json:"name"`
	Digest string     `json:"digest"`
	Bytes  int64      `json:"bytes"`
	Chunks []ChunkOut `json:"chunks"`
}

// APIError is an answer that was not a success.
type APIError struct {
	Status  int
	Message string
}

func (e *APIError) Error() string { return fmt.Sprintf("%d: %s", e.Status, e.Message) }

// do sends a request and decodes a JSON answer into out, waiting out answers
// that say to retry.
func (c *Client) do(ctx context.Context, method, path string, body any, out any) error {
	var payload []byte
	if body != nil {
		var err error
		if payload, err = json.Marshal(body); err != nil {
			return err
		}
	}
	wait := c.Wait
	if wait <= 0 {
		wait = 15 * time.Minute
	}
	deadline := time.Now().Add(wait)
	for {
		req, err := http.NewRequestWithContext(ctx, method, strings.TrimRight(c.Base, "/")+path, bytes.NewReader(payload))
		if err != nil {
			return err
		}
		req.Header.Set("Authorization", "Bearer "+c.Key)
		req.Header.Set("Content-Type", "application/json")
		req.Header.Set("User-Agent", "llmgw-kb")
		if c.Project != "" {
			req.Header.Set("X-Llmgw-Project", c.Project)
		}
		res, err := c.HTTP.Do(req)
		if err != nil {
			return err
		}
		data, err := io.ReadAll(io.LimitReader(res.Body, 256<<20))
		res.Body.Close()
		if err != nil {
			return err
		}
		if res.StatusCode < 300 {
			if out == nil {
				return nil
			}
			return json.Unmarshal(data, out)
		}
		retry := res.Header.Get("Retry-After")
		if retry != "" && (res.StatusCode == http.StatusTooManyRequests || res.StatusCode >= 500) && time.Now().Before(deadline) {
			secs, _ := strconv.Atoi(retry)
			d := time.Duration(max(secs, 1)) * time.Second
			if c.Waiting != nil {
				c.Waiting(errorMessage(data, res.Status), d)
			}
			select {
			case <-ctx.Done():
				return ctx.Err()
			case <-time.After(d):
			}
			continue
		}
		return &APIError{Status: res.StatusCode, Message: errorMessage(data, res.Status)}
	}
}

// errorMessage is the message in an OpenAI-shaped error, or the body.
func errorMessage(data []byte, status string) string {
	var e struct {
		Error struct {
			Message string `json:"message"`
		} `json:"error"`
	}
	if json.Unmarshal(data, &e) == nil && e.Error.Message != "" {
		return e.Error.Message
	}
	if s := strings.TrimSpace(string(data)); s != "" && len(s) < 500 {
		return s
	}
	return status
}

// Collections lists what this key may read.
func (c *Client) Collections(ctx context.Context) ([]Collection, error) {
	var out struct {
		Data []Collection `json:"data"`
	}
	err := c.do(ctx, http.MethodGet, "/v1/kb/collections", nil, &out)
	return out.Data, err
}

// Manifest reads a collection's documents, by its name or id.
func (c *Client) Manifest(ctx context.Context, collection string) (Manifest, error) {
	var m Manifest
	err := c.do(ctx, http.MethodGet, "/v1/kb/collections/"+url.PathEscape(collection)+"/manifest", nil, &m)
	return m, err
}

// Put sends a document; changed is false when the collection already had it.
func (c *Client) Put(ctx context.Context, collection string, d DocOut) (bool, error) {
	var out struct {
		Changed bool `json:"changed"`
	}
	err := c.do(ctx, http.MethodPut, "/v1/kb/collections/"+url.PathEscape(collection)+"/documents", d, &out)
	return out.Changed, err
}

// Delete removes a document; one already gone is not an error.
func (c *Client) Delete(ctx context.Context, collection, name string) error {
	err := c.do(ctx, http.MethodDelete, "/v1/kb/collections/"+url.PathEscape(collection)+"/documents?name="+url.QueryEscape(name), nil, nil)
	var ae *APIError
	if errors.As(err, &ae) && ae.Status == http.StatusNotFound && strings.HasPrefix(ae.Message, "no document") {
		return nil
	}
	return err
}

// embedBatch is how many chunks go in one embeddings request.
const embedBatch = 32

// Embed has the gateway embed texts with a model, in batches, in order.
func (c *Client) Embed(ctx context.Context, model string, texts []string) ([][]float32, error) {
	out := make([][]float32, 0, len(texts))
	for start := 0; start < len(texts); start += embedBatch {
		end := min(start+embedBatch, len(texts))
		var res struct {
			Data []struct {
				Index     int       `json:"index"`
				Embedding []float32 `json:"embedding"`
			} `json:"data"`
		}
		if err := c.do(ctx, http.MethodPost, "/v1/embeddings",
			map[string]any{"model": model, "input": texts[start:end], "encoding_format": "float"}, &res); err != nil {
			return nil, err
		}
		if len(res.Data) != end-start {
			return nil, fmt.Errorf("the gateway embedded %d of %d texts", len(res.Data), end-start)
		}
		batch := make([][]float32, end-start)
		for _, d := range res.Data {
			if d.Index < 0 || d.Index >= len(batch) {
				return nil, fmt.Errorf("an embedding came back with index %d", d.Index)
			}
			batch[d.Index] = d.Embedding
		}
		out = append(out, batch...)
	}
	return out, nil
}