llmgw/internal/control/adminapi/kb.go

package adminapi

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

	"gitea.mynodes.xyz/ilker/llmgw/internal/store"
)

// Knowledge bases, administered (BACKLOG F): collections made, granted,
// emptied and removed here; documents put into them by people with
// llmgw-kb, through the gateway.

// KBStore is the part of the store knowledge bases need.
type KBStore interface {
	KnowledgeAvailable() error
	Collections(ctx context.Context) ([]store.KBCollection, error)
	Collection(ctx context.Context, id string) (store.KBCollection, error)
	CreateCollection(ctx context.Context, c store.KBCollection) (store.KBCollection, error)
	DeleteCollection(ctx context.Context, id string) error
	Grant(ctx context.Context, g store.KBGrant) error
	Revoke(ctx context.Context, scope store.Scope, scopeID, collection string) (bool, error)
	Grants(ctx context.Context, collection string) ([]store.KBGrant, error)
	Documents(ctx context.Context, collection string) ([]store.KBDocument, error)
	DeleteDocument(ctx context.Context, collection, name string) (bool, error)
	SearchChunks(ctx context.Context, q store.KBSearch) ([]store.KBHit, error)
	SetKBDefault(ctx context.Context, d store.KBDefault) error
	CollectionChunks(ctx context.Context, c store.KBCollection, limit int) ([]store.KBHit, error)
	KBDefaults(ctx context.Context) ([]store.KBDefault, error)
}

// KBAdmin is what the knowledge-base endpoints need: the store, and a way to
// have a model embed a text (to learn its dimension, and to try a search).
type KBAdmin struct {
	Store KBStore
	Embed func(ctx context.Context, model, text string) ([]float32, error)
}

// SetKB enables the knowledge-base endpoints.
func (a *API) SetKB(k *KBAdmin) { a.kb = k }

// KBView is every collection, or why there can be none.
type KBView struct {
	Available bool `json:"available"`
	// Reason is why not, when not.
	Reason      string             `json:"reason,omitempty"`
	Collections []KBCollectionView `json:"collections"`
	// Defaults are the collections parts of the tree answer from when a
	// request names none.
	Defaults []KBDefaultView `json:"defaults"`
}

// KBDefaultView is one part of the tree's default collections.
type KBDefaultView struct {
	Target string `json:"target"`
	Name   string `json:"name,omitempty"`
	// Collections are ORG/NAME, in order.
	Collections []string  `json:"collections"`
	At          time.Time `json:"at"`
	By          string    `json:"by,omitempty"`
}

// KBDefaultRequest sets the collections a part of the tree answers from
// when a request names none; none clears it.
type KBDefaultRequest struct {
	// Target is org:, team:, project: (by path or ID) or key:KEY-ID.
	Target string `json:"target"`
	// Collections are IDs, names when unique, or ORG/NAME.
	Collections []string `json:"collections"`
}

// KBCollectionView is a collection as an administrator sees it.
type KBCollectionView struct {
	ID         string `json:"id"`
	Name       string `json:"name"`
	Org        string `json:"org"`
	OrgName    string `json:"org_name,omitempty"`
	Owner      string `json:"owner"`
	OwnerName  string `json:"owner_name,omitempty"`
	EmbedModel string `json:"embed_model"`
	Dims       int    `json:"dims"`
	Documents  int    `json:"documents"`
	Chunks     int    `json:"chunks"`
	// Grants are who else may read or write it.
	Grants    []KBGrantView `json:"grants"`
	CreatedAt time.Time     `json:"created_at"`
	CreatedBy string        `json:"created_by,omitempty"`
}

// KBGrantView is one grant.
type KBGrantView struct {
	Target string    `json:"target"`
	Name   string    `json:"name,omitempty"`
	Perm   string    `json:"perm"`
	At     time.Time `json:"at"`
	By     string    `json:"by,omitempty"`
}

// KBCreateRequest makes a collection.
type KBCreateRequest struct {
	// Owner is org:NAME, team:ORG/TEAM, project:ORG/TEAM/PROJECT (or the ID
	// after the colon), or user:USER-ID.
	Owner string `json:"owner"`
	// Org is the organization a person's collection belongs to; only for
	// user: owners.
	Org  string `json:"org,omitempty"`
	Name string `json:"name"`
	// Model is the catalog model that embeds its documents.
	Model string `json:"model"`
	// Dims is the model's embedding dimension; asked of the model when 0.
	Dims int `json:"dims,omitempty"`
}

// KBGrantRequest lets a scope read or write a collection, or with DELETE
// takes that away.
type KBGrantRequest struct {
	// Collection is its ID, its name when that is unique, or ORG/NAME.
	Collection string `json:"collection"`
	// Target is org:, team:, project:, key: or user: as for guard policies.
	Target string `json:"target"`
	// Perm is read or write; ignored by DELETE.
	Perm string `json:"perm,omitempty"`
}

// KBDocumentView is one document in a collection.
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"`
	AddedBy string    `json:"added_by,omitempty"`
}

// KBSearchRequest tries a search, as an administrator: no tenancy applies.
type KBSearchRequest struct {
	Collection string `json:"collection"`
	Query      string `json:"query"`
	K          int    `json:"k,omitempty"`
}

// KBHitView is one chunk a search found.
type KBHitView struct {
	Document string  `json:"document"`
	Ord      int     `json:"ord"`
	Score    float64 `json:"score"`
	Text     string  `json:"text"`
}

func (a *API) registerKB(mux Mux) {
	mux.HandleFunc("GET /api/v1/kb", a.kbList)
	mux.HandleFunc("POST /api/v1/kb/collections", a.kbCreate)
	mux.HandleFunc("DELETE /api/v1/kb/collections/{id}", a.kbDelete)
	mux.HandleFunc("GET /api/v1/kb/collections/{id}/documents", a.kbDocuments)
	mux.HandleFunc("DELETE /api/v1/kb/collections/{id}/documents", a.kbDeleteDocument)
	mux.HandleFunc("POST /api/v1/kb/grants", a.kbGrant(true))
	mux.HandleFunc("DELETE /api/v1/kb/grants", a.kbGrant(false))
	mux.HandleFunc("POST /api/v1/kb/search", a.kbSearch)
	mux.HandleFunc("POST /api/v1/kb/defaults", a.kbSetDefault)
	mux.HandleFunc("POST /api/v1/kb/collections/{id}/dataset", a.kbDataset)
}

// kbReady answers for a gateway without knowledge bases, and says whether
// it did.
func (a *API) kbReady(w http.ResponseWriter) bool {
	if a.kb == nil {
		writeError(w, http.StatusNotImplemented, "knowledge bases need the store, which this control plane runs without")
		return false
	}
	if err := a.kb.Store.KnowledgeAvailable(); err != nil {
		writeError(w, http.StatusNotImplemented, err.Error())
		return false
	}
	return true
}

func (a *API) names(ctx context.Context) map[string]string {
	if a.quotas == nil {
		return map[string]string{}
	}
	n, err := a.scopeNames(ctx)
	if err != nil {
		return map[string]string{}
	}
	return n
}

func (a *API) kbList(w http.ResponseWriter, r *http.Request) {
	if a.kb == nil {
		writeJSON(w, http.StatusOK, KBView{Reason: "knowledge bases need the store", Collections: []KBCollectionView{}, Defaults: []KBDefaultView{}})
		return
	}
	if err := a.kb.Store.KnowledgeAvailable(); err != nil {
		writeJSON(w, http.StatusOK, KBView{Reason: err.Error(), Collections: []KBCollectionView{}, Defaults: []KBDefaultView{}})
		return
	}
	cs, err := a.kb.Store.Collections(r.Context())
	if err != nil {
		writeError(w, http.StatusInternalServerError, err.Error())
		return
	}
	grants, err := a.kb.Store.Grants(r.Context(), "")
	if err != nil {
		writeError(w, http.StatusInternalServerError, err.Error())
		return
	}
	names := a.names(r.Context())
	v := KBView{Available: true, Collections: []KBCollectionView{}, Defaults: []KBDefaultView{}}
	for _, c := range cs {
		cv := KBCollectionView{ID: c.ID, Name: c.Name, Org: c.OrgID, OrgName: names["org:"+c.OrgID],
			Owner: string(c.OwnerScope) + ":" + c.OwnerID, OwnerName: names[string(c.OwnerScope)+":"+c.OwnerID],
			EmbedModel: c.EmbedModel, Dims: c.Dims, Grants: []KBGrantView{}, CreatedAt: c.CreatedAt, CreatedBy: c.CreatedBy}
		if docs, err := a.kb.Store.Documents(r.Context(), c.ID); err == nil {
			cv.Documents = len(docs)
			for _, d := range docs {
				cv.Chunks += d.Chunks
			}
		}
		for _, g := range grants {
			if g.Collection == c.ID {
				key := string(g.Scope) + ":" + g.ScopeID
				cv.Grants = append(cv.Grants, KBGrantView{Target: key, Name: names[key], Perm: string(g.Perm), At: g.AddedAt, By: g.AddedBy})
			}
		}
		v.Collections = append(v.Collections, cv)
	}
	label := map[string]string{}
	for _, c := range cs {
		label[c.ID] = or(names["org:"+c.OrgID], c.OrgID) + "/" + c.Name
	}
	defaults, err := a.kb.Store.KBDefaults(r.Context())
	if err != nil {
		writeError(w, http.StatusInternalServerError, err.Error())
		return
	}
	for _, d := range defaults {
		key := string(d.Scope) + ":" + d.ScopeID
		dv := KBDefaultView{Target: key, Name: names[key], Collections: []string{}, At: d.AddedAt, By: d.AddedBy}
		for _, id := range d.Collections {
			dv.Collections = append(dv.Collections, or(label[id], id))
		}
		v.Defaults = append(v.Defaults, dv)
	}
	writeJSON(w, http.StatusOK, v)
}

// kbTarget reads a grant's or an owner's target: org:, team: and project:
// by path or ID, key: and user: by ID.
func (a *API) kbTarget(ctx context.Context, target string) (store.Scope, string, error) {
	kind, what, ok := strings.Cut(strings.TrimSpace(target), ":")
	what = strings.TrimSpace(what)
	if !ok || what == "" {
		return "", "", fmt.Errorf("%q is not org:, team:, project:, key: or user: and what", target)
	}
	switch store.Scope(kind) {
	case store.ScopeKey, store.ScopeUser:
		return store.Scope(kind), what, nil
	case store.ScopeOrg, store.ScopeTeam, store.ScopeProject:
		if a.quotas == nil {
			return store.Scope(kind), what, nil
		}
		scope, id, _, err := a.resolveScope(ctx, target)
		return scope, id, err
	}
	return "", "", fmt.Errorf("%q: org, team, project, key or user", kind)
}

// kbFind reads a collection by ID, by its name when that is unique, or by
// ORG/NAME.
func (a *API) kbFind(ctx context.Context, ref string) (store.KBCollection, error) {
	ref = strings.TrimSpace(ref)
	if c, err := a.kb.Store.Collection(ctx, ref); err == nil {
		return c, nil
	}
	cs, err := a.kb.Store.Collections(ctx)
	if err != nil {
		return store.KBCollection{}, err
	}
	names := a.names(ctx)
	org, name, scoped := strings.Cut(ref, "/")
	var found []store.KBCollection
	for _, c := range cs {
		switch {
		case scoped && strings.EqualFold(c.Name, name) &&
			(c.OrgID == org || strings.EqualFold(names["org:"+c.OrgID], org)):
			found = append(found, c)
		case !scoped && strings.EqualFold(c.Name, ref):
			found = append(found, c)
		}
	}
	switch len(found) {
	case 0:
		return store.KBCollection{}, fmt.Errorf("no collection %q", ref)
	case 1:
		return found[0], nil
	}
	return store.KBCollection{}, fmt.Errorf("%q names a collection in more than one organization; write ORG/%s", ref, ref)
}

func (a *API) kbCreate(w http.ResponseWriter, r *http.Request) {
	if !a.kbReady(w) {
		return
	}
	var in KBCreateRequest
	if err := json.NewDecoder(r.Body).Decode(&in); err != nil {
		writeError(w, http.StatusBadRequest, err.Error())
		return
	}
	scope, id, err := a.kbTarget(r.Context(), in.Owner)
	if err != nil {
		writeError(w, http.StatusBadRequest, err.Error())
		return
	}
	org := ""
	if scope == store.ScopeUser {
		if in.Org == "" {
			writeError(w, http.StatusBadRequest, "a person's collection names its organization (org)")
			return
		}
		if _, org, err = a.kbTarget(r.Context(), "org:"+in.Org); err != nil {
			writeError(w, http.StatusBadRequest, err.Error())
			return
		}
	}
	m, ok := a.ctl.Catalog().Resolve(in.Model)
	if !ok {
		writeError(w, http.StatusBadRequest, "no model "+in.Model+" in the catalog")
		return
	}
	if !m.Capabilities.Embeddings {
		writeError(w, http.StatusBadRequest, m.Name+" does not embed (capabilities.embeddings); a collection needs an embedding model")
		return
	}
	dims := in.Dims
	if dims == 0 {
		if a.kb.Embed == nil {
			writeError(w, http.StatusBadRequest, "give the model's dimension (dims): nothing here can ask it")
			return
		}
		// Asked of the model: a dimension typed by hand and wrong would be
		// found only when the first document is refused.
		ctx, cancel := context.WithTimeout(r.Context(), 10*time.Minute)
		defer cancel()
		v, err := a.kb.Embed(ctx, m.Name, "llmgw")
		if err != nil {
			writeError(w, http.StatusBadGateway, "asking "+m.Name+" for its dimension: "+err.Error()+
				" (is it running? or give dims)")
			return
		}
		dims = len(v)
	}
	c, err := a.kb.Store.CreateCollection(r.Context(), store.KBCollection{OwnerScope: scope, OwnerID: id, OrgID: org,
		Name: in.Name, EmbedModel: m.Name, Dims: dims, CreatedBy: actorOf(r)})
	if err != nil {
		code := http.StatusBadRequest
		if errors.Is(err, store.ErrNotFound) {
			code = http.StatusNotFound
		}
		writeError(w, code, err.Error())
		return
	}
	writeJSON(w, http.StatusCreated, KBCollectionView{ID: c.ID, Name: c.Name, Org: c.OrgID,
		Owner: string(c.OwnerScope) + ":" + c.OwnerID, EmbedModel: c.EmbedModel, Dims: c.Dims,
		Grants: []KBGrantView{}, CreatedAt: c.CreatedAt, CreatedBy: c.CreatedBy})
}

func (a *API) kbCollection(w http.ResponseWriter, r *http.Request) (store.KBCollection, bool) {
	if !a.kbReady(w) {
		return store.KBCollection{}, false
	}
	c, err := a.kbFind(r.Context(), r.PathValue("id"))
	if err != nil {
		writeError(w, http.StatusNotFound, err.Error())
		return store.KBCollection{}, false
	}
	return c, true
}

func (a *API) kbDelete(w http.ResponseWriter, r *http.Request) {
	c, ok := a.kbCollection(w, r)
	if !ok {
		return
	}
	if err := a.kb.Store.DeleteCollection(r.Context(), c.ID); err != nil {
		writeError(w, http.StatusInternalServerError, err.Error())
		return
	}
	writeJSON(w, http.StatusOK, map[string]any{"id": c.ID, "name": c.Name, "deleted": true})
}

func (a *API) kbDocuments(w http.ResponseWriter, r *http.Request) {
	c, ok := a.kbCollection(w, r)
	if !ok {
		return
	}
	docs, err := a.kb.Store.Documents(r.Context(), c.ID)
	if err != nil {
		writeError(w, http.StatusInternalServerError, err.Error())
		return
	}
	out := []KBDocumentView{}
	for _, d := range docs {
		out = append(out, KBDocumentView{Name: d.Name, Digest: d.Digest, Bytes: d.Bytes, Chunks: d.Chunks,
			AddedAt: d.AddedAt, AddedBy: d.AddedBy})
	}
	writeJSON(w, http.StatusOK, out)
}

func (a *API) kbDeleteDocument(w http.ResponseWriter, r *http.Request) {
	c, ok := a.kbCollection(w, r)
	if !ok {
		return
	}
	name := r.URL.Query().Get("name")
	gone, err := a.kb.Store.DeleteDocument(r.Context(), c.ID, name)
	if err != nil {
		writeError(w, http.StatusInternalServerError, err.Error())
		return
	}
	if !gone {
		writeError(w, http.StatusNotFound, "no document "+name+" in "+c.Name)
		return
	}
	writeJSON(w, http.StatusOK, map[string]any{"name": name, "deleted": true})
}

func (a *API) kbGrant(grant bool) http.HandlerFunc {
	return func(w http.ResponseWriter, r *http.Request) {
		if !a.kbReady(w) {
			return
		}
		var in KBGrantRequest
		if err := json.NewDecoder(r.Body).Decode(&in); err != nil {
			writeError(w, http.StatusBadRequest, err.Error())
			return
		}
		c, err := a.kbFind(r.Context(), in.Collection)
		if err != nil {
			writeError(w, http.StatusNotFound, err.Error())
			return
		}
		scope, id, err := a.kbTarget(r.Context(), in.Target)
		if err != nil {
			writeError(w, http.StatusBadRequest, err.Error())
			return
		}
		if !grant {
			gone, err := a.kb.Store.Revoke(r.Context(), scope, id, c.ID)
			if err != nil {
				writeError(w, http.StatusInternalServerError, err.Error())
				return
			}
			if !gone {
				writeError(w, http.StatusNotFound, "no such grant")
				return
			}
			writeJSON(w, http.StatusOK, KBGrantView{Target: string(scope) + ":" + id})
			return
		}
		perm := store.KBPerm(strings.ToLower(strings.TrimSpace(in.Perm)))
		if perm == "" {
			perm = store.KBRead
		}
		if err := a.kb.Store.Grant(r.Context(), store.KBGrant{Scope: scope, ScopeID: id, Collection: c.ID,
			Perm: perm, AddedBy: actorOf(r)}); err != nil {
			writeError(w, http.StatusBadRequest, err.Error())
			return
		}
		writeJSON(w, http.StatusOK, KBGrantView{Target: string(scope) + ":" + id, Perm: string(perm), At: time.Now(), By: actorOf(r)})
	}
}

func (a *API) kbSearch(w http.ResponseWriter, r *http.Request) {
	if !a.kbReady(w) {
		return
	}
	var in KBSearchRequest
	if err := json.NewDecoder(r.Body).Decode(&in); err != nil {
		writeError(w, http.StatusBadRequest, err.Error())
		return
	}
	c, err := a.kbFind(r.Context(), in.Collection)
	if err != nil {
		writeError(w, http.StatusNotFound, err.Error())
		return
	}
	k := in.K
	if k <= 0 || k > 50 {
		k = 5
	}
	q := store.KBSearch{OrgID: c.OrgID, Collections: []store.KBCollection{c}, Text: in.Query, K: k}
	// Said, not swallowed: a search that quietly lost its meaning half
	// looks like a collection that does not know the answer.
	if a.kb.Embed != nil {
		v, err := a.kb.Embed(r.Context(), c.EmbedModel, in.Query)
		if err != nil {
			writeError(w, http.StatusBadGateway, c.EmbedModel+" did not embed the question: "+err.Error())
			return
		}
		q.Embeddings = map[string][]float32{c.EmbedModel: v}
	}
	hits, err := a.kb.Store.SearchChunks(r.Context(), q)
	if err != nil {
		writeError(w, http.StatusInternalServerError, err.Error())
		return
	}
	out := []KBHitView{}
	for _, h := range hits {
		out = append(out, KBHitView{Document: h.Document, Ord: h.Ord, Score: h.Score, Text: h.Body})
	}
	writeJSON(w, http.StatusOK, out)
}

func (a *API) kbSetDefault(w http.ResponseWriter, r *http.Request) {
	if !a.kbReady(w) {
		return
	}
	var in KBDefaultRequest
	if err := json.NewDecoder(r.Body).Decode(&in); err != nil {
		writeError(w, http.StatusBadRequest, err.Error())
		return
	}
	scope, id, err := a.kbTarget(r.Context(), in.Target)
	if err != nil {
		writeError(w, http.StatusBadRequest, err.Error())
		return
	}
	d := store.KBDefault{Scope: scope, ScopeID: id, AddedBy: actorOf(r)}
	for _, ref := range in.Collections {
		c, err := a.kbFind(r.Context(), ref)
		if err != nil {
			writeError(w, http.StatusNotFound, err.Error())
			return
		}
		d.Collections = append(d.Collections, c.ID)
	}
	if err := a.kb.Store.SetKBDefault(r.Context(), d); err != nil {
		writeError(w, http.StatusBadRequest, err.Error())
		return
	}
	writeJSON(w, http.StatusOK, KBDefaultView{Target: string(scope) + ":" + id, Collections: in.Collections, At: time.Now(), By: actorOf(r)})
}