llmgw/internal/control/adminapi/kb_test.go

package adminapi

import (
	"context"
	"encoding/json"
	"net/http"
	"net/url"
	"os"
	"path/filepath"
	"strings"
	"testing"

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

func kbModels() []catalog.Model {
	return append(testModels(), catalog.Model{Name: "bge-m3", Source: "/e",
		Capabilities: catalog.Capabilities{Embeddings: true},
		Profiles:     []catalog.Profile{{Name: "cpu", GPUClass: "h100-80g", TP: 1}}})
}

func TestKnowledgeBasesListSaysWhyThereAreNone(t *testing.T) {
	a, _ := setup(t, kbModels())
	s, err := store.OpenSQLite(context.Background(), filepath.Join(t.TempDir(), "s.db"))
	if err != nil {
		t.Fatal(err)
	}
	defer s.Close()
	a.SetKB(&KBAdmin{Store: s})
	var v KBView
	if code := call(t, a, "GET", "/api/v1/kb", "", &v); code != http.StatusOK || v.Available || !strings.Contains(v.Reason, "SQLite") {
		t.Errorf("got %d %+v", code, v)
	}
	if code := call(t, a, "POST", "/api/v1/kb/collections", `{"owner":"org:x","name":"n","model":"bge-m3"}`, nil); code != http.StatusNotImplemented {
		t.Errorf("creating on SQLite: %d", code)
	}
}

func TestKnowledgeBasesAdministered(t *testing.T) {
	dsn := os.Getenv("LLMGW_TEST_POSTGRES")
	if dsn == "" {
		t.Skip("set LLMGW_TEST_POSTGRES to a pgvector database to run this")
	}
	if u, err := url.Parse(dsn); err != nil || !strings.Contains(u.Path, "test") {
		t.Fatalf("refusing a database whose name does not say test")
	}
	ctx := context.Background()
	s, err := store.OpenPostgres(ctx, dsn)
	if err != nil {
		t.Fatal(err)
	}
	defer s.Close()
	if err := s.KnowledgeAvailable(); err != nil {
		t.Skip(err)
	}
	a, _ := setup(t, kbModels())
	a.SetQuotas(s)
	probes := 0
	a.SetKB(&KBAdmin{Store: s, Embed: func(_ context.Context, model, _ string) ([]float32, error) {
		probes++
		return make([]float32, 4), nil
	}})
	sfx := strings.ToLower(store.NewID(""))
	org, _ := s.CreateOrg(ctx, "Acme"+sfx)
	team, _ := s.CreateTeam(ctx, org, "hr")
	_, _ = s.CreateProject(ctx, team, "docs")
	other, _ := s.CreateOrg(ctx, "Other"+sfx)

	// Owner by path; the dimension asked of the model.
	var c KBCollectionView
	rec := func(method, path, body string, out any) int { return call(t, a, method, path, body, out) }
	body := `{"owner":"team:Acme` + sfx + `/hr","name":"politikalar","model":"bge-m3"}`
	if code := callCreated(t, a, body, &c); code != http.StatusCreated || c.Dims != 4 || probes != 1 || c.Org != org {
		t.Fatalf("create: %d %+v (probes %d)", code, c, probes)
	}
	// A model that does not embed is refused before anything is written.
	if code := rec("POST", "/api/v1/kb/collections", `{"owner":"org:`+org+`","name":"x","model":"llama"}`, nil); code != http.StatusBadRequest {
		t.Errorf("a chat model as the embedder: %d", code)
	}
	// Grants by path, never across organizations.
	if code := rec("POST", "/api/v1/kb/grants", `{"collection":"Acme`+sfx+`/politikalar","target":"project:Acme`+sfx+`/hr/docs","perm":"write"}`, nil); code != http.StatusOK {
		t.Errorf("grant: %d", code)
	}
	if code := rec("POST", "/api/v1/kb/grants", `{"collection":"`+c.ID+`","target":"org:`+other+`"}`, nil); code != http.StatusBadRequest {
		t.Errorf("a grant to another organization: %d", code)
	}
	var v KBView
	rec("GET", "/api/v1/kb", "", &v)
	var got *KBCollectionView
	for i := range v.Collections {
		if v.Collections[i].ID == c.ID {
			got = &v.Collections[i]
		}
	}
	if got == nil || len(got.Grants) != 1 || got.Grants[0].Perm != "write" || !strings.Contains(got.OwnerName, "hr") {
		t.Fatalf("list: %+v", got)
	}
	// By ORG/NAME too.
	if code := rec("DELETE", "/api/v1/kb/grants", `{"collection":"Acme`+sfx+`/politikalar","target":"project:Acme`+sfx+`/hr/docs"}`, nil); code != http.StatusOK {
		t.Errorf("revoke: %d", code)
	}
	if code := rec("DELETE", "/api/v1/kb/collections/"+c.ID, "", nil); code != http.StatusOK {
		t.Errorf("delete: %d", code)
	}
}

func callCreated(t *testing.T, a *API, body string, out any) int {
	t.Helper()
	mux := http.NewServeMux()
	a.Register(mux)
	r, _ := http.NewRequest("POST", "/api/v1/kb/collections", strings.NewReader(body))
	r.Header.Set("Content-Type", "application/json")
	w := &recorder{h: http.Header{}}
	mux.ServeHTTP(w, r)
	if w.code == http.StatusCreated {
		_ = json.Unmarshal(w.body, out)
	} else {
		t.Logf("create answered %d %s", w.code, w.body)
	}
	return w.code
}

type recorder struct {
	h    http.Header
	code int
	body []byte
}

func (r *recorder) Header() http.Header { return r.h }
func (r *recorder) Write(b []byte) (int, error) {
	if r.code == 0 {
		r.code = http.StatusOK
	}
	r.body = append(r.body, b...)
	return len(b), nil
}
func (r *recorder) WriteHeader(c int) { r.code = c }