llmgw/internal/router/kb.go

package router

import (
	"context"
	"encoding/json"
	"errors"
	"fmt"
	"net/http"
	"strings"
	"time"
	"unicode"
	"unicode/utf8"

	"golang.org/x/text/unicode/norm"

	"gitea.mynodes.xyz/ilker/llmgw/internal/auth"
	"gitea.mynodes.xyz/ilker/llmgw/internal/guard"
	storepkg "gitea.mynodes.xyz/ilker/llmgw/internal/store"
)

// Knowledge bases on the client listener (BACKLOG F2): what llmgw-kb, run by
// a person on their own machine with their own key, reads and writes. The
// documents are read, cut and embedded there; what arrives here is a
// document's name, digest and chunks with their vectors.
//
// Which collections a caller may see is decided from the key's place in the
// tree and nothing the request says, beyond naming a project for a personal
// key. A collection the caller may not read is answered exactly as one that
// does not exist.

// projectHeader names the project a personal key acts in; the person must be
// a member of it.
const projectHeader = "X-Llmgw-Project"

// Bounds on one document, so one request cannot hold the control plane's
// memory or a transaction for long.
const (
	kbMaxBody       = 64 << 20
	kbMaxChunks     = 20000
	kbMaxChunkRunes = 8000
	kbMaxNameRunes  = 1024
)

// KBCollectionView is a collection as a caller sees it.
type KBCollectionView struct {
	ID   string `json:"id"`
	Name string `json:"name"`
	// Owner is the owning scope as scope:ID; OwnerName how people know it,
	// "Acme / hr".
	Owner      string `json:"owner"`
	OwnerName  string `json:"owner_name,omitempty"`
	EmbedModel string `json:"embed_model"`
	Dims       int    `json:"dims"`
	// Writable is whether this caller may add and remove documents.
	Writable bool `json:"writable"`
}

// KBDocumentView is a document as the manifest lists it.
type KBDocumentView struct {
	Name    string    `json:"name"`
	Digest  string    `json:"digest"`
	Bytes   int64     `json:"bytes"`
	Chunks  int       `json:"chunks"`
	AddedAt time.Time `json:"added_at"`
}

// KBManifest is a collection and its documents: what a sync compares the
// files on disk with.
type KBManifest struct {
	Collection KBCollectionView `json:"collection"`
	Documents  []KBDocumentView `json:"documents"`
}

// KBPutDocument is one document, read and embedded by the client.
type KBPutDocument struct {
	Name   string `json:"name"`
	Digest string `json:"digest"`
	Bytes  int64  `json:"bytes"`
	Chunks []struct {
		Text      string    `json:"text"`
		Embedding []float32 `json:"embedding"`
	} `json:"chunks"`
}

// KBPutResult is what adding a document did.
type KBPutResult struct {
	Name    string `json:"name"`
	Chunks  int    `json:"chunks"`
	Changed bool   `json:"changed"`
}

func (r *Router) registerKB(mux *http.ServeMux) {
	mux.HandleFunc("GET /v1/kb/collections", r.handleKBCollections)
	mux.HandleFunc("GET /v1/kb/collections/{id}/manifest", r.handleKBManifest)
	mux.HandleFunc("PUT /v1/kb/collections/{id}/documents", r.handleKBPut)
	mux.HandleFunc("DELETE /v1/kb/collections/{id}/documents", r.handleKBDelete)
}

// kbCaller is the caller's place in the tree, for deciding what it may read.
// A personal key has none unless it names a project it is a member of.
func (r *Router) kbCaller(w http.ResponseWriter, req *http.Request) (auth.Principal, storepkg.Effective, bool) {
	p, ok := auth.FromContext(req.Context())
	if !ok || r.store == nil {
		writeError(w, http.StatusNotImplemented, "invalid_request_error",
			"knowledge bases need authentication and the store, which this gateway runs without")
		return p, storepkg.Effective{}, false
	}
	if err := r.store.KnowledgeAvailable(); err != nil {
		writeError(w, http.StatusNotImplemented, "invalid_request_error", err.Error())
		return p, storepkg.Effective{}, false
	}
	e, err := r.kbEffective(req.Context(), p, req.Header.Get(projectHeader))
	if err != nil {
		var ke *kbError
		if errors.As(err, &ke) {
			writeError(w, ke.status, ke.kind, ke.msg)
		} else {
			writeError(w, http.StatusInternalServerError, "server_error", err.Error())
		}
		return p, e, false
	}
	return p, e, true
}

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,
		OwnerName: r.store.OwnerName(ctx, c.OwnerScope, c.OwnerID), EmbedModel: c.EmbedModel, Dims: c.Dims, Writable: writable}
}

// kbCollection finds a collection by id or name among those the caller may
// read (or write), answering 404 for one it may not.
func (r *Router) kbCollection(w http.ResponseWriter, req *http.Request, e storepkg.Effective, write bool) (storepkg.KBCollection, bool) {
	want := req.PathValue("id")
	cs, err := r.store.CollectionsFor(req.Context(), e, write)
	if err != nil {
		writeError(w, http.StatusInternalServerError, "server_error", err.Error())
		return storepkg.KBCollection{}, false
	}
	for _, c := range cs {
		if c.ID == want {
			return c, true
		}
	}
	for _, c := range cs {
		if c.Name == want {
			return c, true
		}
	}
	msg := "no collection " + want
	if e.OrgID == "" {
		msg += " (a personal key names its project in " + projectHeader + ")"
	}
	writeError(w, http.StatusNotFound, "not_found_error", msg)
	return storepkg.KBCollection{}, false
}

func (r *Router) handleKBCollections(w http.ResponseWriter, req *http.Request) {
	_, e, ok := r.kbCaller(w, req)
	if !ok {
		return
	}
	read, err := r.store.CollectionsFor(req.Context(), e, false)
	if err != nil {
		writeError(w, http.StatusInternalServerError, "server_error", err.Error())
		return
	}
	write, err := r.store.CollectionsFor(req.Context(), e, true)
	if err != nil {
		writeError(w, http.StatusInternalServerError, "server_error", err.Error())
		return
	}
	writable := map[string]bool{}
	for _, c := range write {
		writable[c.ID] = true
	}
	out := []KBCollectionView{}
	for _, c := range read {
		out = append(out, r.kbView(req.Context(), c, writable[c.ID]))
	}
	writeJSON(w, http.StatusOK, map[string]any{"object": "list", "data": out})
}

func (r *Router) handleKBManifest(w http.ResponseWriter, req *http.Request) {
	_, e, ok := r.kbCaller(w, req)
	if !ok {
		return
	}
	c, ok := r.kbCollection(w, req, e, false)
	if !ok {
		return
	}
	docs, err := r.store.Documents(req.Context(), c.ID)
	if err != nil {
		writeError(w, http.StatusInternalServerError, "server_error", err.Error())
		return
	}
	writable, _ := r.store.CollectionsFor(req.Context(), e, true)
	m := KBManifest{Collection: r.kbView(req.Context(), c, containsCollection(writable, c.ID)), Documents: []KBDocumentView{}}
	for _, d := range docs {
		m.Documents = append(m.Documents, KBDocumentView{Name: d.Name, Digest: d.Digest, Bytes: d.Bytes,
			Chunks: d.Chunks, AddedAt: d.AddedAt})
	}
	writeJSON(w, http.StatusOK, m)
}

func containsCollection(cs []storepkg.KBCollection, id string) bool {
	for _, c := range cs {
		if c.ID == id {
			return true
		}
	}
	return false
}

// kbWriter is a caller about to change a collection: one it may write, and a
// caller whose policy does not forbid keeping what it sends.
func (r *Router) kbWriter(w http.ResponseWriter, req *http.Request) (auth.Principal, storepkg.KBCollection, bool) {
	p, e, ok := r.kbCaller(w, req)
	if !ok {
		return p, storepkg.KBCollection{}, false
	}
	if r.noStore(guard.Subject{KeyID: p.KeyID, UserID: p.UserID, ProjectID: e.ProjectID,
		TeamID: e.TeamID, OrgID: e.OrgID, Path: req.URL.Path}) {
		writeError(w, http.StatusForbidden, "permission_error",
			"the gateway's policy for this caller keeps nothing, and a knowledge base is documents kept")
		return p, storepkg.KBCollection{}, false
	}
	c, ok := r.kbCollection(w, req, e, true)
	return p, c, ok
}

// validDocName checks a document's name: what the client sends is a path
// relative to what it synced, with forward slashes.
func validDocName(n string) (string, error) {
	if !utf8.ValidString(n) {
		return "", errors.New("the name is not UTF-8")
	}
	n = norm.NFC.String(strings.TrimSpace(n))
	switch {
	case n == "":
		return "", errors.New("a document needs a name")
	case utf8.RuneCountInString(n) > kbMaxNameRunes:
		return "", fmt.Errorf("a name has at most %d characters", kbMaxNameRunes)
	case strings.ContainsFunc(n, unicode.IsControl):
		return "", errors.New("the name holds a control character")
	}
	return n, nil
}

func (r *Router) handleKBPut(w http.ResponseWriter, req *http.Request) {
	p, c, ok := r.kbWriter(w, req)
	if !ok {
		return
	}
	var in KBPutDocument
	dec := json.NewDecoder(http.MaxBytesReader(w, req.Body, kbMaxBody))
	if err := dec.Decode(&in); err != nil {
		var tooBig *http.MaxBytesError
		if errors.As(err, &tooBig) {
			writeError(w, http.StatusRequestEntityTooLarge, "invalid_request_error",
				fmt.Sprintf("a document is at most %d MiB of JSON; split the file", kbMaxBody>>20))
			return
		}
		writeError(w, http.StatusBadRequest, "invalid_request_error", "the request is not valid JSON: "+err.Error())
		return
	}
	name, err := validDocName(in.Name)
	if err != nil {
		writeError(w, http.StatusBadRequest, "invalid_request_error", err.Error())
		return
	}
	switch {
	case in.Digest == "":
		writeError(w, http.StatusBadRequest, "invalid_request_error", "a document needs the digest of its file")
		return
	case len(in.Chunks) == 0:
		writeError(w, http.StatusBadRequest, "invalid_request_error", "a document needs at least one chunk")
		return
	case len(in.Chunks) > kbMaxChunks:
		writeError(w, http.StatusBadRequest, "invalid_request_error",
			fmt.Sprintf("%d chunks: a document has at most %d", len(in.Chunks), kbMaxChunks))
		return
	}
	chunks := make([]storepkg.KBChunk, len(in.Chunks))
	for i, ch := range in.Chunks {
		if !utf8.ValidString(ch.Text) || strings.TrimSpace(ch.Text) == "" {
			writeError(w, http.StatusBadRequest, "invalid_request_error", fmt.Sprintf("chunk %d has no UTF-8 text", i))
			return
		}
		if utf8.RuneCountInString(ch.Text) > kbMaxChunkRunes {
			writeError(w, http.StatusBadRequest, "invalid_request_error",
				fmt.Sprintf("chunk %d is over %d characters", i, kbMaxChunkRunes))
			return
		}
		if len(ch.Embedding) != c.Dims {
			writeError(w, http.StatusBadRequest, "invalid_request_error",
				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))
			return
		}
		chunks[i] = storepkg.KBChunk{Body: strings.ToValidUTF8(ch.Text, ""), Embedding: ch.Embedding}
	}
	d, changed, err := r.store.PutDocument(req.Context(), c, storepkg.KBDocument{Name: name, Digest: in.Digest,
		Bytes: in.Bytes, AddedBy: kbActor(p)}, chunks)
	if err != nil {
		writeError(w, http.StatusInternalServerError, "server_error", err.Error())
		return
	}
	if changed {
		r.kbAudit(req.Context(), p, "PUT /v1/kb/collections/{id}/documents", c, name, http.StatusOK)
	}
	writeJSON(w, http.StatusOK, KBPutResult{Name: name, Chunks: d.Chunks, Changed: changed})
}

func (r *Router) handleKBDelete(w http.ResponseWriter, req *http.Request) {
	p, c, ok := r.kbWriter(w, req)
	if !ok {
		return
	}
	name, err := validDocName(req.URL.Query().Get("name"))
	if err != nil {
		writeError(w, http.StatusBadRequest, "invalid_request_error", err.Error())
		return
	}
	gone, err := r.store.DeleteDocument(req.Context(), c.ID, name)
	if err != nil {
		writeError(w, http.StatusInternalServerError, "server_error", err.Error())
		return
	}
	if !gone {
		writeError(w, http.StatusNotFound, "not_found_error", "no document "+name+" in "+c.Name)
		return
	}
	r.kbAudit(req.Context(), p, "DELETE /v1/kb/collections/{id}/documents", c, name, http.StatusOK)
	writeJSON(w, http.StatusOK, map[string]any{"name": name, "deleted": true})
}

// kbActor is who changed a collection, as the audit log names people: the
// person, or the key when it belongs to no one.
func kbActor(p auth.Principal) string {
	if p.UserID != "" {
		return "user:" + p.UserID
	}
	return "key:" + p.KeyID
}

// kbAudit records a change to a collection. The document's name is kept; its
// text never is.
func (r *Router) kbAudit(ctx context.Context, p auth.Principal, action string, c storepkg.KBCollection, doc string, status int) {
	err := r.store.WriteAudit(context.WithoutCancel(ctx), storepkg.AuditEvent{Actor: kbActor(p), Source: "gateway",
		Action: action, Target: "collection:" + c.ID, Detail: c.Name + ": " + doc, Status: status})
	if err != nil {
		r.log.Warn("could not record a knowledge base change in the audit log", "collection", c.ID, "err", err)
	}
}